"""
requireents \\

For calculating the required entropies  \\
Arguments:  probs, dictionary of probabilities  \\
            calc, the required entropies  \\
Returns:    ents, dictionary calc => entropy  \\
"""

function requireents(probs::Dict{String, Array{Float64}}, calc::Array{String, 1}; base::Int=2)
  ents = Dict{String, Float64}()
  in("HX", calc) && (ents["HX"] = entropy(probs["PX"], base))
  in("HY", calc) && (ents["HY"] = entropy(probs["PY"], base))
  in("HXY", calc) && (ents["HXY"] = condentropy(probs["PY"], probs["PXY"], base=base))
  return ents
end


struct PolyharmonicSpline <: UQModel
    w::Array{Float64}
    v::Array{Float64}
    c::Array{Float64,2}
    k::Int64
    n::Array{Symbol}
    output::Symbol

    function PolyharmonicSpline(data::DataFrame, k::Int64, output::Symbol)
        f = data[:, output]

        centers = select(data, Not(output))
        n = propertynames(centers)
        centers = Matrix{Float64}(centers)

        dim = size(centers, 1)

        A = zeros(dim, dim)
        for i in 1:dim, j in 1:dim
            if i == j
                continue
            end
            r = sqrt(sum((centers[i, :] - centers[j, :]) .^ 2))
            A[i, j] = ϕ(r, k)
        end

        B = [ones(dim, 1) centers]

        M = [A B; transpose(B) zeros(size(B, 2), size(B, 2))]
        F = [f; zeros(size(B, 2), 1)]

        wv = M \ F

        w = wv[1:dim, :]
        v = wv[(dim + 1):size(wv, 1), :]

        return new(w, v, centers, k, n, output)
    end
end

function ϕ(r::Float64, k::Int64)
    if k % 2 != 0
        return r^k
    elseif r < 1
        return r^(k - 1) * log(r^r)
    else
        return r^k * log(r)
    end
end

function calc(ps::PolyharmonicSpline, x::Array{Float64,1})
    r = sqrt.(sum((ps.c .- transpose(x)) .^ 2; dims=2))
    f = sum(ϕ.(r, ps.k) .* ps.w)
    return f += (transpose(ps.v) * [1; x])[1]
end

function evaluate!(ps::PolyharmonicSpline, df::DataFrame)
    x = Matrix{Float64}(df[:, ps.n]) # convert to matrix and order variables by ps.n

    out = map(row -> calc(ps, convert(Array, row)), eachrow(x))
    return df[!, ps.output] = out
end


struct StatelessSoftmax{T} <: AbstractStatelessPolicy where {T}
    θ::T

    function StatelessSoftmax(::Type{T}, num_actions::Int) where {T}
        θ = zeros(T, num_actions)
        return new{typeof(θ)}(θ)
    end

    function StatelessSoftmax(num_actions::Int) where {T}
        return StatelessSoftmax(Float64, num_actions)
    end
end

function params(π::StatelessSoftmax)
    return (π.θ,)
end

function paramsvec(π::StatelessSoftmax)
    return copy(vec(π.θ))
end

function paramsfromvec!(π::StatelessSoftmax, θ)
    @. π.θ = θ
    return π
end

function rrule(::typeof(paramsfromvec!), π::StatelessSoftmax, θ)
    @. π.θ = θ
    function paramsfromvec_statelesssoftmax!(ȳ)
        dθ = ȳ.θ
        return NoTangent(), ȳ, dθ
    end
    return π, paramsfromvec_statelesssoftmax!
end

struct LinearSoftmax{T} <: AbstractPolicy where {T}
    θ::T

    function LinearSoftmax(::Type{T}, num_features::Int, num_actions::Int) where {T}
        θ = zeros(T, (num_features, num_actions))
        return new{typeof(θ)}(θ)
    end

    function LinearSoftmax(num_features::Int, num_actions::Int)
        return LinearSoftmax(Float64, num_features, num_actions)
    end
end

function params(π::LinearSoftmax)
    return (π.θ,)
end

function paramsvec(π::LinearSoftmax)
    return copy(vec(π.θ))
end

function paramsfromvec!(π::LinearSoftmax, θ)
    π.θ .= reshape(θ, size(π.θ))
    return π
end

function rrule(::typeof(paramsfromvec!), π::LinearSoftmax, θ)
    π.θ .= reshape(θ, size(π.θ))
    function paramsfromvec_softmax!(ȳ)
        dθ = vec(ȳ.θ)
        return NoTangent(), ȳ, dθ
    end
    return π, paramsfromvec_softmax!
end

struct SoftmaxBuffer{TA, TP, TS} <: Any
    action::TA
    p::TP
    ψ::TS

    function SoftmaxBuffer(π::StatelessSoftmax)
        T = eltype(π.θ)
        n = length(π.θ)
        p = zeros(T, n)
        ψ = (zero(π.θ),)
        action = zeros(Int, 1)
        return new{typeof(action),typeof(p), typeof(ψ)}(action, p, ψ)
    end

    function SoftmaxBuffer(π::LinearSoftmax)
        T = eltype(π.θ)
        n = size(π.θ, 2)
        p = zeros(T, n)
        ψ = (zero(π.θ),)
        action = zeros(Int, 1)
        return new{typeof(action),typeof(p),typeof(ψ)}(action, p, ψ)
    end
end

function softmax!(out::AbstractVector, x::AbstractVector)
    @. out = x
    softmax!(out)
    return nothing
end

function softmax!(x::AbstractVector)
    maxx = maximum(x)
    x .-= maxx
    @. x = exp(x)
    x ./= sum(x)
    return nothing
end

function rrule(::typeof(softmax!), x::AbstractVector)
    softmax!(x)
    function softmax_pullback(ȳ)
        b = dot(ȳ, x)
        gx = @. (ȳ - b) * x
        return NoTangent(), gx
    end
    return x, softmax_pullback
end

function softmax(x::Vector{T}) where {T}
    maxx = maximum(x)
    y = x .- maxx
    @. y = exp(y)
    tot = sum(y)
    z = @. y / tot
    return z
end

function rrule(::typeof(softmax), x::Vector{T}) where {T}
    y = softmax(x)
    function softmax_pullback(ȳ)
        b = dot(ȳ, y)
        gx = @. (ȳ - b) * y
        return NoTangent(), gx
    end
    return y, softmax_pullback
end

function sample_discrete(p)
    n = length(p)
    i = 1
    c = p[1]
    u = rand()
    while c < u && i < n
        c += p[i += 1]
    end

    return i
end

function (π::StatelessSoftmax)()
    p = softmax(π.θ)
    return Categorical(p)
end

function (π::StatelessSoftmax)(buff::SoftmaxBuffer)
    softmax!(buff.p, π.θ)
    return Categorical(buff.p)
end

function (π::StatelessSoftmax)(s)
    return π()
end

function (π::StatelessSoftmax)(buff::SoftmaxBuffer, s)
    return π(buff)
end

function logpdf(π::StatelessSoftmax, a)
    p = similar(π.θ)
    softmax!(p, π.θ)
    return log(p[a])
end

function rrule(::typeof(logpdf), π::StatelessSoftmax, a)
    p = similar(π.θ)
    softmax!(p, π.θ)
    logp = log(p[a])
    T = typeof(logp)
    function logpdf_stateless_softmax_pullback(Ȳ)
        p *= -T(1.0)
        p[a] += T(1.0)
        p *= Ȳ
        dθ = Tangent{typeof(π)}(θ=p)
        return NoTangent(), dθ, NoTangent()
    end
    return logp, logpdf_stateless_softmax_pullback
end

function logpdf!(buff::SoftmaxBuffer, π::StatelessSoftmax, a)
    softmax!(buff.p, π.θ)
    return log(buff.p[a])
end

function logpdf!(buff::SoftmaxBuffer, π::StatelessSoftmax, s, a)
    return logpdf!(buff, π, a)
end

function grad_logpdf!(ψ, π::StatelessSoftmax, a)
    softmax!(ψ[1], π.θ)
    logp = log(ψ[1][a])
    T = eltype(π.θ)
    @. ψ[1] *= -T(1.0)
    ψ[1][a] += T(1.0)
    return logp, ψ
end

function grad_logpdf!(buff::SoftmaxBuffer, π::StatelessSoftmax, a)
    return grad_logpdf!(buff.ψ, π, a)
end

function grad_logpdf!(buff::SoftmaxBuffer, π::StatelessSoftmax, s, a)
    return grad_logpdf!(buff.ψ, π, a)
end

function sample_with_trace!(ψ, action, π::StatelessSoftmax)
    softmax!(ψ[1], π.θ)
    a = sample_discrete(ψ[1])
    action[1] = a

    logp = log(ψ[1][a])
    T = typeof(logp)
    @. ψ[1] *= -T(1.0)
    ψ[1][a] += T(1.0)
    return a, logp, ψ
end

function sample_with_trace!(buff::SoftmaxBuffer, π::StatelessSoftmax)
    return sample_with_trace!(buff.ψ, buff.action, π)
end

function sample_with_trace!(ψ, action, π::StatelessSoftmax, s)
    return sample_with_trace!(ψ, action, π)
end

function sample_with_trace!(buff::SoftmaxBuffer, π::StatelessSoftmax, s)
    return sample_with_trace!(buff.ψ, buff.action, π)
end

function (π::LinearSoftmax)(s)
    T = eltype(π.θ)
    p = zeros(T, size(π.θ,2))
    mul!(p, π.θ', s)
    softmax!(p)
    return Categorical(p)
end

function rrule(π::LinearSoftmax, s)
    T = eltype(π.θ)
    p = zeros(T, size(π.θ,2))
    mul!(p, π.θ', s)
    softmax!(p)
    d = Categorical(p)
    function linear_softmax_dist_pullback(ȳ)
        b = dot(ȳ.p, p)
        dp = @. (ȳ.p - b) * p
        ψ = zero(π.θ)
        G = ψ'
        for i in 1:length(dp)
            @. G[i, :] = s * dp[i]
        end
        ds = zero(s)
        mul!(ds, π.θ, dp)
        dθ = Tangent{typeof(π)}(θ=ψ)
        return dθ, ds
    end
    return d, linear_softmax_dist_pullback
end

function (π::LinearSoftmax)(buff::SoftmaxBuffer, s)
    mul!(buff.p, π.θ', s)
    softmax!(buff.p)
    return Categorical(buff.p)
end

function logpdf(π::LinearSoftmax, s, a)
    return logpdf_softmax(π.θ, s, a)
end

function logpdf_softmax(θ, s, a)
    T = eltype(θ)
    p = zeros(T, size(θ,2))
    mul!(p, θ', s)
    softmax!(p)
    return log(p[a])
end

function rrule(::typeof(logpdf_softmax), θ, s, a)
    T = eltype(θ)
    p = zeros(T, size(θ,2))
    mul!(p, θ', s)
    softmax!(p)
    logp = log(p[a])

    function logpdf_linear_softmax_pullback(ȳ)
        p *= -T(1.0)
        p[a] += T(1.0)
        p *= ȳ
        ψ = zero(θ)
        G = ψ'
        for i in 1:length(p)
            @. G[i, :] = s * p[i]
        end
        ds = zero(s)
        mul!(ds, θ, p)
        return NoTangent(), ψ, ds, NoTangent()
    end
    return logp, logpdf_linear_softmax_pullback
end

function logpdf_softmax(θ, s::AbstractMatrix, a)
    T = eltype(θ)
    num_a = size(θ,2)
    n = size(s,2)
    p = zeros(T, (num_a, n))
    mul!(p, θ', s)
    logps = zeros(T, (1,n))
    for i in 1:n
        p_i = @view p[:, i]
        softmax!(p_i)
        logps[1,i] = log(p_i[a[i]])
    end
    return logps
end

function rrule(::typeof(logpdf_softmax), θ, s::AbstractMatrix, a)
    T = eltype(θ)
    num_a = size(θ,2)
    n = size(s,2)
    p = zeros(T, (num_a, n))
    mul!(p, θ', s)
    logps = zeros(T, (1,n))
    for i in 1:n
        p_i = @view p[:, i]
        softmax!(p_i)
        logps[1,i] = log(p_i[a[1,i]])
    end

    function logpdf_linear_softmax_batch_pullback(ȳ)
        ψ = zero(θ)
        G = ψ'
        for i in 1:n
            p[:, i] *= -one(T)
            p[a[1, i], i] += one(T)
            p[:, i] *= ȳ[1, i]
            for j in 1:num_a
                x = @view s[:, i]
                @. G[j, :] += x * p[j,i]
            end
        end
        ds = zero(s)
        mul!(ds, θ, p)
        return NoTangent(), ψ, ds, ZeroTangent()
    end
    return logps, logpdf_linear_softmax_batch_pullback
end

function logpdf!(buff::SoftmaxBuffer, π::LinearSoftmax, s, a)
    mul!(buff.p, π.θ', s)
    softmax!(buff.p)
    return log(buff.p[a])
end

function grad_logpdf!(ψ, π::LinearSoftmax, s, a)
    T = eltype(π.θ)
    p = zeros(T, size(π.θ,2))
    mul!(p, π.θ', s)
    softmax!(p)
    logpa = log(p[a])
    T = typeof(logpa)
    @. p *= -T(1.0)
    p[a] += T(1.0)
    G = ψ[1]'
    for i in 1:length(p)
        @. G[i, :] = s * p[i]
    end
    return logpa, ψ
end

function grad_logpdf!(buff::SoftmaxBuffer, π::LinearSoftmax, s, a)
    mul!(buff.p, π.θ', s)
    softmax!(buff.p)
    logpa = log(buff.p[a])
    T = typeof(logpa)
    @. buff.p *= -T(1.0)
    buff.p[a] += T(1.0)
    G = buff.ψ[1]'
    for i in 1:length(buff.p)
        @. G[i, :] = s * buff.p[i]
    end
    return logpa, buff.ψ
end

function sample_with_trace!(ψ, action, π::LinearSoftmax, s)
    T = eltype(π.θ)
    p = zeros(T, size(π.θ,2))
    mul!(p, π.θ', s)
    softmax!(p)
    action[1] = sample_discrete(p)
    logp = log(p[action[1]])
    T = typeof(logp)
    @. p *= -T(1.0)

    p[action[1]] += T(1.0)
    G = ψ[1]'
    for i in 1:length(p)
        @. G[i, :] = s * p[i]
    end
    return action[1], logp, ψ
end

function sample_with_trace!(buff::SoftmaxBuffer, π::LinearSoftmax, s)
    mul!(buff.p, π.θ', s)
    softmax!(buff.p)
    a = sample_discrete(buff.p)
    buff.action[1] = a
    logp = log(buff.p[a])
    T = typeof(logp)
    @. buff.p *= -T(1.0)

    buff.p[a] += T(1.0)
    G = buff.ψ[1]'
    for i in 1:length(buff.p)
        @. G[i, :] = s * buff.p[i]
    end
    return a, logp, buff.ψ
end


using Distributions
using Lora

# using Gadfly

nrats, ndays = 30, 5
nvars = 9

a_μ = BasicContUnvParameter(:a_μ, logtarget=(p::Float64, v::Vector) -> logpdf(Normal(0, 100), p), nkeys=nvars)

a_σ2 = BasicContUnvParameter(:a_σ2, logtarget=(p::Float64, v::Vector) -> logpdf(InverseGamma(0.01, 0.01), p), nkeys=nvars)

b_μ = BasicContUnvParameter(:b_μ, logtarget=(p::Float64, v::Vector) -> logpdf(Normal(0, 100), p), nkeys=nvars)

b_σ2 = BasicContUnvParameter(:b_σ2, logtarget=(p::Float64, v::Vector) -> logpdf(InverseGamma(0.01, 0.01), p), nkeys=nvars)

y_σ2 = BasicContUnvParameter(:y_σ2, logtarget=(p::Float64, v::Vector) -> logpdf(InverseGamma(0.01, 0.01), p), nkeys=nvars)

x = Data(:x)

a = BasicContMuvParameter(
  :a,
  logtarget=(p::Vector{Float64}, v::Vector) -> logpdf(MvNormal(v[1]*ones(nrats), sqrt(v[2])), p),
  nkeys=nvars
)

b = BasicContMuvParameter(
  :b,
  logtarget=(p::Vector{Float64}, v::Vector) -> logpdf(MvNormal(v[3]*ones(nrats), sqrt(v[4])), p),
  nkeys=nvars
)

ylogtarget(p::Vector{Float64}, v::Vector) =
  logpdf(MvNormal(repeat(v[7], inner=[ndays])+repeat(v[8], inner=[ndays]).*repeat(v[6], outer=[nrats]), sqrt(v[5])), p)

y = BasicContMuvParameter(:y, logtarget=ylogtarget, nkeys=nvars)

model = GenericModel([a_μ, a_σ2, b_μ, b_σ2, y_σ2, x, a, b, y], isindexed=false)

mcrange = BasicMCRange(nsteps=120000, burnin=100000)

X, = dataset("rats", "age")
X = vec(X)-mean(X)
Y, = dataset("rats", "weight")
v0 = Dict{Symbol, Any}(
  :a_μ=>150.,
  :a_σ2=>1.,
  :b_μ=>10.,
  :b_σ2=>1.,
  :y_σ2=>1.,
  :x=>X,
  :a=>fill(250., nrats),
  :b=>fill(6., nrats),
  :y=>vec(hcat([Y[i, :] for i in 1:30]...))
)

samplers = [
  [
    MH(1.), MH(x::Float64 -> Truncated(Normal(x, 1.), 0, Inf), symmetric=false),
    MH(1.), MH(x::Float64 -> Truncated(Normal(x, 1.), 0, Inf), symmetric=false),
    MH(x::Float64 -> Truncated(Normal(x, 1.), 0, Inf), symmetric=false)
  ];
  fill(MH(1*ones(nrats)), 2); MH(1*ones(nrats*ndays))
]

kk = [:a_μ, :a_σ2, :b_μ, :b_σ2, :y_σ2, :a, :b, :y]

samplers = Dict(zip(kk, samplers))

basicmcjobs =
  Dict([(k, BasicMCJob(
    model,
    samplers[k],
    BasicMCRange(nsteps=1, burnin=0),
    v0,
    pindex=model[k].index,
    jointtarget=true,
    resetpstate=false
  )) for k in kk])

tuners = Dict(zip(kk, fill(RobertsRosenthalMCTuner(0.44), 8)))

gibbsjob = GibbsJob(model, basicmcjobs, mcrange, v0, jointtarget=true, tuner=tuners, verbose=true)

@time run(gibbsjob)

output(gibbsjob)

chains = Dict(gibbsjob)

mean(chains[:b_μ].value)

# plot(x=1:length(chains[:b_μ].value), y=chains[:b_μ].value, Geom.line)


using Gaugefields

function heatbathtest_4D(NX,NY,NZ,NT,β,NC)
    Dim = 4
    Nwing = 1

    U = Initialize_Gaugefields(NC,Nwing,NX,NY,NZ,NT,condition = "cold")
    println(typeof(U))

    gauge_action = GaugeAction(U)
    plaqloop = make_loops_fromname("plaquette",Dim=Dim)
    append!(plaqloop,plaqloop')
    βinp = β/2
    push!(gauge_action,βinp,plaqloop)

    rectloop = make_loops_fromname("rectangular",Dim=Dim)
    append!(rectloop,rectloop')
    βinp = β/2
    push!(gauge_action,βinp,rectloop)
    hnew = Heatbath_update(U,gauge_action)

    show(gauge_action)

    temp1 = similar(U[1])
    temp2 = similar(U[1])
    temp3 = similar(U[1])

    comb = 6
    factor = 1/(comb*U[1].NV*U[1].NC)
    @time plaq_t = calculate_Plaquette(U,temp1,temp2)*factor
    println("plaq_t = $plaq_t")
    poly = calculate_Polyakov_loop(U,temp1,temp2)
    println("polyakov loop = $(real(poly)) $(imag(poly))")

    numhb = 1000
    for itrj = 1:numhb

        heatbath!(U,hnew)

        plaq_t = calculate_Plaquette(U,temp1,temp2)*factor
        poly = calculate_Polyakov_loop(U,temp1,temp2)

        if itrj % 40 == 0
            println("$itrj plaq_t = $plaq_t")
            println("$itrj polyakov loop = $(real(poly)) $(imag(poly))")
        end
    end

    close(fp)
    return plaq_t

end

NX = 4
NY = 4
NZ = 4
NT = 4
NC = 3
β = 5.7
heatbathtest_4D(NX,NY,NZ,NT,β,NC)


"""
Filter function for polynomial chaos

- @args p: λ, Δ, ℓ, op, ..., mode

"""
function Base.filter!(u::T, p...) where {T<:AbstractArray{<:AbstractFloat,1}}

    q0 = eachindex(u) |> first
    q1 = eachindex(u) |> last

    λ = p[1]
    mode = p[end]

    if mode == :l2
        if length(p) > 2
            Δ = p[2]
            op = p[3]
            filter_l2!(u, op, λ, Δ)
        else
            filter_l2!(u, λ)
        end
    elseif mode == :l1
        for i = q0+1:q1
            sc = 1.0 - 5.0 * λ * i * (i - 1) * ℓ[i] / abs(u[i])
            if sc < 0.0
                sc = 0.0
            end
            u[i] *= sc
        end
    elseif mode == :lasso
        nr = length(u)
        _λ = abs(u[end]) / (nr * (nr - 1) * ℓ[end])
        for i = q0+1:q1
            sc = 1.0 - _λ * i * (i - 1) * ℓ[i] / abs(u[i] + 1e-8)
            if sc < 0.0
                sc = 0.0
            end
            u[i] *= sc
        end
    elseif mode == :dev
        nr = length(u)
        for i = q0+1:q1
            _λ = λ * abs(u[i] / (u[1] + 1e-10)) / (nr * (nr - 1) * ℓ[i])
            u[i] /= (1.0 + _λ * i^2 * (i - 1)^2)
        end
    else
        throw("unavailable filter mode")
    end

end

function Base.filter!(u::AbstractArray{<:AbstractFloat,2}, dim::Integer, p...)

    if dim == 1
        for j in axes(u, 2)
            _u = @view u[:, j]
            filter!(_u, p...)
        end
    elseif dim == 2
        for i in axes(u, 1)
            _u = @view u[i, :]
            filter!(_u, p...)
        end
    end

end

function Base.filter!(u::AbstractArray{<:AbstractFloat,3}, dim::Integer, p...)

    if dim == 1
        for k in axes(u, 3), j in axes(u, 2)
            _u = @view u[:, j, k]
            filter!(_u, p...)
        end
    elseif dim == 2
        for k in axes(u, 3), i in axes(u, 1)
            _u = @view u[i, :, k]
            filter!(_u, p...)
        end
    elseif dim == 3
        for j in axes(u, 2), i in axes(u, 1)
            _u = @view u[i, j, :]
            filter!(_u, p...)
        end
    end

end

function filter_l2!(u::T, λ) where {T<:AbstractArray{<:AbstractFloat,1}}

    q0 = eachindex(u) |> first
    q1 = eachindex(u) |> last

    for i = q0+1:q1
        u[i] /= (1.0 + λ * i^2 * (i - 1)^2)
    end

    return nothing

end

function filter_l2!(
    u::T1,
    op::T2,
    λ,
    Δ,
) where {T1<:AbstractArray{<:AbstractFloat,1},T2<:AbstractOrthoPoly}

    q0 = eachindex(u) |> first
    q1 = eachindex(u) |> last

    uRan = evaluatePCE(u, op.quad.nodes, op)
    δ = 0.5 * (maximum(uRan) - minimum(uRan))
    _λ = λ * (exp(δ / Δ) - 1.0)
    #_λ = λ * δ / Δ

    for i = q0+1:q1
        u[i] /= (1.0 + _λ * i^2 * (i - 1)^2)
    end

    return nothing

end


Base.@kwdef struct DensityAnalysis{D, K, B}
    datalimits::D=automatic
    npoints::Int=200
    kernel::K=automatic
    bandwidth::B=automatic
end

# Work around lack of length 1 tuple method
# TODO: also add weights support here
_kde(data::NTuple{1, Any}; kwargs...) = kde(data...; kwargs...)
_kde(data::Tuple; kwargs...) = kde(data; kwargs...)

function compute_datalimits(positional, datalimits)
    return if datalimits === automatic
        map(v -> mapreduce(extrema, extend_extrema, v), Tuple(positional))
    else
        datalimits
    end
end

applydatalimits(f::Function, d) = map(f, d)
applydatalimits(limits::Union{AbstractArray, Tuple}, _) = limits

function _density(data...; datalimits, npoints, kwargs...)
    k = _kde(data; kwargs...)
    es = applydatalimits(datalimits, data)
    rgs = map(e -> range(e...; length=npoints), es)
    res = pdf(k, rgs...)
    return (rgs..., res)
end

function (d::DensityAnalysis)(input::ProcessedLayer)
    datalimits = compute_datalimits(input.positional, d.datalimits)
    options = valid_options(; datalimits, d.npoints, d.kernel, d.bandwidth)
    output = map(input) do p, n
        return _density(p...; pairs(n)..., pairs(options)...), (;)
    end
    N = length(input.positional)
    labels = set(input.labels, N+1 => "pdf")
    plottypes = [LinesFill, Heatmap, Volume]
    default_plottype = plottypes[N]
    plottype = Makie.plottype(input.plottype, default_plottype)
    return ProcessedLayer(output; plottype, labels)
end

"""
    density(; datalimits=automatic, kernel=automatic, bandwidth=automatic, npoints=200)

Fit a kernel density estimation of `data`.
Here, `datalimits` specifies the range for which the density should be calculated,
and `kernel` and `bandwidth` are forwarded to `KernelDensity.kde`.
`npoints` is the number of points used by Makie to draw the line
"""
density(; options...) = transformation(DensityAnalysis(; options...))


factor_ab(x) = logpdf(MvNormal(PDMat([1.0 0.5; 0.5 1.0])), x)
factor_bc(x) = logpdf(MvNormal([2.0, 2.0], PDMat([2.0 -1.0; -1.0 2.0])), x)

struct FactorProposal
    addresses :: Vector{String}
    last :: Vector{String}
    std :: Float64
end
FactorProposal(addresses) = FactorProposal(addresses, [addresses[1]], 0.5)
function (fp::FactorProposal)(x::Dict)
    new = deepcopy(x)
    address = rand(fp.addresses)
    new[address] = randn() * fp.std + new[address]
    fp.last[1] = address
    new
end
function (fp::FactorProposal)(x_prime::Dict, x::Dict)
    address = fp.last[1]
    logpdf(Normal(x[address], fp.std), x_prime[address]) - log(length(fp.addresses))
end

#=
@testset "make field and sample" begin
    factors = Dict(["a", "b"] => factor_ab, ["b", "c"] => factor_bc)
    evidence = Dict("b" => 3.0)
    #evidence = Dict()
    field = RandomField(factors, evidence)
    addresses = ["a", "c"]
    proposal = FactorProposal(addresses)
    initial_values = Dict("a" => 0.0, "b" => 3.0, "c" => 0.0)
    samples = mh(field, [proposal], initial_values; burn=1000, thin=100, num_iterations=11000)
    for address in ["a", "b", "c"]
        plot_marginal(samples, address, "out", "factor-marginal-$address.png")
    end
end
=#

@testset "using generative field" begin
    factors = Dict(["a", "b"] => factor_ab, ["b", "c"] => factor_bc)
    field = RandomField(factors)
    addresses = ["a", "b", "c"]
    proposal = FactorProposal(addresses)
    init = Dict("a" => 0.0, "b" => 0.0, "c" => 0.0)
    gf = GenerativeField(field, proposal, init)

    # generative fields can be used in generic stochastic functions
    t = trace()
    sample(t, "field", gf)
    @info t

    # note that we can pass upstream samped values as evidence
    t = trace()
    upstream_value = sample(t, "a", Normal())
    field = RandomField(factors, Dict("a" => upstream_value))
    proposal = FactorProposal(["b", "c"])
    init = Dict("a" => 0.0, "b" => 0.0, "c" => 0.0)
    gf = GenerativeField(field, proposal, init)
    sample(t, "field", gf)
    @info t
end


@testset "dataitem.jl" begin
    @testset "Standardization of the Variable type" begin

        v = Variable("x", 100 .+ 5*randn(100))

        vstd = Extremes.standardize(v)

        m = mean(vstd.value)
        s = std(vstd.value)

        @test m ≈ 0.0 atol = sqrt(eps())
        @test s ≈ 1.0 atol = sqrt(eps())
        @test vstd.offset ≈ 100 atol = 5
        @test vstd.scale ≈ 5 atol = 1

        ṽ = Extremes.reconstruct(vstd)

        @test ṽ.name == v.name
        @test ṽ.value ≈ v.value atol = sqrt(eps())

        v = Variable("x", [-1; 0 ; 1])
        vstd = Extremes.standardize(v)

        @test vstd.offset == 0 # integer comparison
        @test vstd.scale == 1 # integer comparison

        v = Variable("x", ones(10))
        vstd = Extremes.standardize(v)

        @test vstd.scale == 1 # integer comparison

        vstd = VariableStd("x", [-1; 0 ; 1])
        @test vstd.offset == 0 # integer comparison
        @test vstd.scale == 1 # integer comparison

        @test_throws AssertionError Extremes.VariableStd("x", [-2; 0 ; 2])
        @test_throws AssertionError Extremes.VariableStd("x", [0; 1 ; 2])

        # Test on the transform function
        y = randn(3)
        x₁ = Extremes.standardize(Variable("x₁", [-1; 0; 1]))
        x₂ = Extremes.standardize(Variable("x₂", [0; 1; 2]))
        x₃ = Extremes.standardize(Variable("x₃", [0; 2; 4]))

        model = BlockMaxima(Variable("y", y), locationcov=[x₁; x₂], logscalecov = [x₃])

        θ̃ = collect(1.0:6.0)

        fm_std = MaximumLikelihoodEVA(model, θ̃)

        fm = Extremes.transform(fm_std)

        θ̂ = fm.θ̂

        @test θ̂[1] ≈ θ̃[1] - θ̃[2]*0/1 - θ̃[3]*1/1
        @test θ̂[2] ≈ θ̃[2]/1
        @test θ̂[3] ≈ θ̃[3]/1
        @test θ̂[4] ≈ θ̃[4] - θ̃[5]*2/2
        @test θ̂[5] ≈ θ̃[5]/2
        @test θ̂[6] ≈ θ̃[6]

    end

end


using StatsBase, DataFrames, Flux
include("2048_fn.jl")

using Flux
model = Flux.Chain(
  Dense(14*16+4, 128, relu),
  Dense(128, 64, relu),
  Dense(64, 32, relu),
  Dense(32, 1, relu));

loss(x,y) = sum((model(x) .- y).^2)
opt = ADAM(params(model))

@time simulate_one_game_and_train!(model, opt, 0.9)

init_grid = init_game()
@time (init_grid, fnl, seq, cart, tf, reward) = simulate_game!(init_grid, model, 0.5)

# score every move using untrained model
grid_try = copy(init_grid)
initial_model = [Tracker.data(model(grid_move_to_feature(sim_seq(grid_try, seq[1:i+1], cart[1:i+1], tf[1:i+1]), seq[i])))[1] for i=1:length(seq)-1]

# now fit the model once
x = generate_x_based_on_episode(init_grid, seq, cart, tf);
y = discounted_rewards(reward, 0.99)[1:end-1];
xy = [(x[:,j], y[j]) for j in 1:size(x)[2]]
Flux.@epochs 100 Flux.train!(loss, xy, opt)

updated_model = [Tracker.data(model(grid_move_to_feature(sim_seq(grid_try, seq[1:i+1], cart[1:i+1], tf[1:i+1]), seq[i])))[1] for i=1:length(seq)-1]
using Plots
plot(cumsum(reward))
plot!(updated_model)
plot!(y)


using EmpiricalRiskMinimization
srand(123)

# data
n = 2000; k = 30;
d = k + 1;
X = randn(n, k); theta = randn(d);
y = [ones(n) X] * theta + 0.5 * randn(n);

# model
M = Model(X, y, embedall=true)
train(M);
status(M);
theta = thetaopt(M);


# Initialising simulation conditions

using Distributions: Normal
using Random: rand!
using LinearAlgebra: normalize, dot, norm

# Constants (SI units)
const kB = 1.38064852e-23 #Boltzmann
const c = 3e8 # Speed of light (m/s)
const ε₀ = 8.854e-12 # Vacuum permittivity
const h̄ = 6.626e-34 / (2 * π) # Reduced Planck
const g = 9.81 # Gravity
const a₀ = 5.28e-11 # Bohr radius

# ATOM INITIALISATION

# Define a neutral atom species
struct AtomSpecies
    m :: Float64    # Atomic mass (kg)
    aₛ :: Float64   # Scattering length (m)
    σ :: Float64    # Scattering cross-section (m^2)
    α :: Float64    # Real part of polarizability
end

const Rb87 = AtomSpecies(1.454660e-25, 98a₀, 8*pi * ( 98a₀ )^2,
                         1.1e-38)

# Dipole trap constant
kappa(species::AtomSpecies) = species.α / (2 * ε₀ * c)

# Initialise N (test) particles with position and velocity components
# distributed uniformly in the range [-size, size] and [-v, v], respectively.
function uniform_cloud(N, size, v)
    positions = 2.0 * size * (rand(Float64, 3, N) .- 0.5)
    velocities = 2.0 * v * (rand(Float64, 3, N) .- 0.5)
    return positions, velocities
end

# Initialise Boltzmann-distributed velocities at temperature
# T for N particles, each of mass m.
boltzmann_velocities(N, m, T) = reshape(rand( Normal(0, sqrt(kB * T / m)), 3*N), 3, N)

# Initialise Boltzmann-distributed positions in a harmonic trap at
# temperature T. Generates N particles, each of mass m.
function harmonic_boltzmann_positions(N, m, T, ωx, ωy, ωz)
    coeff = sqrt(kB * T / m)
    positions = zeros(Float64, 3, N)
    rand!( Normal(0, coeff / ωx), view(positions, 1, :) )
    rand!( Normal(0, coeff / ωy), view(positions, 2, :) )
    rand!( Normal(0, coeff / ωz), view(positions, 3, :) )
    return positions
end

# FIELD INITIALISATION

# Convert values to time-parametrised functions
function time_parametrize(value)
    f(_) = value
    f() = value
    return f
end

# Allow functions to be parametrised (for convenience)
function time_parametrize(func::Function)
    f(t) = func(t)
    function f()
        @warn "Time-dependent function given no argument. Assuming t = 0."
        return f(0)
    end
    return f
end

# Parametrise multiple functions/values at once (for convenience)
time_parametrize(args...) = [time_parametrize(arg) for arg in args]

# Harmonic potential (time-dependent)
struct HarmonicField
    # Time-parametrised trap frequencies
    ωx :: Function
    ωy :: Function
    ωz :: Function
    centre :: Function
    # Initialise
    HarmonicField(ωx, ωy, ωz, centre = [0.0,0,0]) = (
        new( time_parametrize(ωx, ωy, ωz, centre)... )
    )
end

# Gaussian beam (time-dependent)
struct GaussianBeam
    focus :: Function
    direction :: Function
    power :: Function
    waist :: Function
    wavelength :: Function

    function GaussianBeam(args...)
        focus, direction, power, waist, wavelength = time_parametrize(args...)
        unit_dir(t) = normalize(direction(t))
        return new(focus, unit_dir, power, waist, wavelength)
    end
end

# E.g. gravity
struct UniformField
    strength :: Function # Vector acceleration (m/s^2)
    origin :: Function
    UniformField(str, origin = [0,0,0]) = new( time_parametrize(str, origin)... )
end

const gravity = UniformField([0, -g, 0])

# Acceleration (constant field)
function acceleration(field::UniformField, positions, species, t, output)
    # 'species' parameter required to conform to correct call signature.
    Nt = size(positions, 2)
    strx, stry, strz = field.strength(t)
    Threads.@threads for atom in 1:Nt
        output[1, atom] += strx
        output[2, atom] += stry
        output[3, atom] += strz
    end
    return nothing
end

# Potential (constant field)
function potential(field::UniformField, positions, species, t, output)
    # 'species' parameter required to conform to correct call signature.
    Nt = size(positions, 2)
    strx, stry, strz = field.strength(t)
    ox, oy, oz = field.origin(t)
    m = species.m
    Threads.@threads for atom in 1:Nt
        px, py, pz = view(positions, :, atom)
        output[atom] -= m * (strx * (px - ox) + stry * (py - oy) + strz * (pz - oz))
    end
    return nothing
end

# Overload accel/potential by generating new output array when none provided.
acceleration(f, p, s, t) = acceleration(f, p, s, t, zeros(Float64, size(p)))
potential(f, p, s, t) = potential(f, p, s, t, zeros(Float64, size(p, 2)))
# Initialise a potential or acceleration function to
potential(f) = (args...) -> potential(f, args...) #pos, spec, t, out
acceleration(f) = (args...) -> acceleration(f, args...)

# Acceleration (harmonic field)
function acceleration(field::HarmonicField, positions, species, t, output)
    # 'species' parameter required to conform to correct call signature.
    Nt = size(positions, 2)
    cx, cy, cz = field.centre(t)
    ωx2, ωy2, ωz2 = field.ωx(t)^2, field.ωy(t)^2, field.ωz(t)^2
    Threads.@threads for atom in 1:Nt
        output[1, atom] -= ωx2 * (positions[1, atom] - cx)
        output[2, atom] -= ωy2 * (positions[2, atom] - cy)
        output[3, atom] -= ωz2 * (positions[3, atom] - cz)
    end
    return nothing
end

# Potential (harmonic field)
function potential(field::HarmonicField, positions, species, t, output)
    Nt = size(positions, 2)
    m = species.m
    cx, cy, cz = field.centre(t)
    ωx2, ωy2, ωz2 = field.ωx(t)^2, field.ωy(t)^2, field.ωz(t)^2
    Threads.@threads for atom in 1:Nt
        x, y, z = view(positions, :, atom)
        output[atom] += 0.5 * m * (ωx2 * (x-cx)^2 + ωy2 * (y-cy)^2 + ωz2 * (z-cz)^2)
    end
    return nothing
end

# Get parameters of a Gaussian beam
parameters(b::GaussianBeam, t::Float64) = ( b.focus(t), b.direction(t), b.power(t),
                                   b.waist(t), b.wavelength(t) )

# Acceleration (Gaussian beam)
function acceleration(field::GaussianBeam, positions, species, t, output)
    foc, dir, P, w₀, _ = parameters(field, t)
    κ = kappa(species)
    N = size(positions, 2)

    fx, fy, fz = foc
    ux, uy, uz = dir

    r²coeff = - 2 / w₀^2
    acoeff = 8 * P * κ / (π * species.m * w₀^4)
    Threads.@threads for atom in 1:N
        px, py, pz = view(positions, :, atom)
        dx = px - fx # Displacement from beam focus
        dy = py - fy
        dz = pz - fz
        z = dx*ux + dy*uy + dz*uz # Axial coordinate
        rx = dx - (z * ux) # Radial vector
        ry = dy - (z * uy)
        rz = dz - (z * uz)
        r² = rx^2 + ry^2 + rz^2
        a = (-acoeff * exp(r²coeff * r²))
        output[1, atom] += a*rx
        output[2, atom] += a*ry
        output[3, atom] += a*rz
    end
    return nothing
end

# Potential (Gaussian beam)
function potential(field::GaussianBeam, positions, species, t, output)
    foc, dir, P, w₀, _ = parameters(field, t)

    fx, fy, fz = foc
    ux, uy, uz = dir

    Nt = size(positions, 2)
    κ = kappa(species)
    r²coeff = - 2 / w₀^2
    Ucoeff = - κ * 2 * P / (π * w₀^2)
    Threads.@threads for atom in 1:Nt
        px, py, pz = view(positions, :, atom)
        dx = px - fx
        dy = py - fy
        dz = pz - fz
        z = dx*ux + dy*uy + dz*uz
        r² = (dx^2 + dy^2 + dz^2) - z^2
        output[atom] += Ucoeff * exp(r²coeff * r²) - Ucoeff
    end
    return nothing
end

# EVAPORATION PROBABILITIES

# An exponential ramp
exponential_ramp(start, stop, τ) = (t) -> stop + (start - stop) * exp(- t / τ)

# A linear ramp
linear_ramp(start, stop, τ) = (t) -> start + (stop - start) * (t / τ)

# Perfect energy-based removal
function energy_evap(εₜ, positions, velocities, conditions,
        potential, t, out = zeros(Float64, size(positions, 2)))

    potential(positions, conditions.species, t, out)
    #out initially contains pe
    m = conditions.species.m
    N = size(positions, 2)
    for atom in 1:N
        #vx, vy, vz = view(velocities, :, atom)
        #ke = 0.5 * m * (vx^2 + vy^2 + vz^2)
        if out[atom] > εₜ
            out[atom] = 1 # Replace PE with evaporation probability
        else
            out[atom] = 0
        end
    end
    return out
end

# Overload for correct type signature and time-parametrisation
function energy_evap(depth, potential)
    εₜ = time_parametrize(depth)
    return (pos, vel, cond, t, o) -> energy_evap(εₜ(t), pos, vel, cond, potential, t, o)
end

# Removal when beyond ellipsoid
function ellipsoid_evap(ωx, ωy, ωz, T, p, centre,
                        species, positions, out = zeros(Float64, size(positions, 2)))

    N = size(positions, 2)

    m = species.m
    σx² = kB * T / (m * ωx^2)
    σy² = kB * T / (m * ωy^2)
    σz² = kB * T / (m * ωz^2)

    bound = -2 * log(p) # Boundary condition
    cx, cy, cz = centre

    for atom in 1:N
        x = positions[1, atom] - cx
        y = positions[2, atom] - cy
        z = positions[3, atom] - cz
        if (x^2 / σx² + y^2 / σy² + z^2 / σz²) > bound
            out[atom] = 1.0
        else
            out[atom] = 0.0
        end
    end

    return out
end

# Overload for correct type signature and time-parametrisation
function ellipsoid_evap(ωx, ωy, ωz, T, p, centre = [0.0,0,0])
    ω_x, ω_y, ω_z, T = time_parametrize(ωx, ωy, ωz, T)
    centre = time_parametrize(centre)
    function evap(pos, vel, cond, t, o)
        return ellipsoid_evap(ω_x(t), ω_y(t), ω_z(t), T(t), p, centre(t), cond.species, pos, o)
    end
    return evap
end

# No evaporation (evap. probability 0)
no_evap(pos, args...) = zeros(Float64, size(pos, 2))

# SIMULATION CONDITIONS
struct SimulationConditions
    species :: AtomSpecies

    F :: Float64 # Initial F-scale
    positions :: Matrix{Float64} # Initial test-particle positions
    velocities :: Matrix{Float64} # Initial test-particle velocities

    acceleration :: Function # Time-parametrised, vectorised acceleration func.
    potential :: Function # Potential (also parametrised/vectorised)

    threebodyloss :: Float64 # 0 # 3-body loss constant (m^6/s, NOT cm^6/s)
    evaporate :: Function # no_evap # Particle loss-function
    τ_bg :: Float64 # Inf # Background loss time constant

    # Initialisation
    function SimulationConditions(spec, F, pos, vel, accel, potential;
            K = 0, evap = no_evap, τ_bg = Inf) # Optional arguments

        # Check that array sizes are 3 x Nt
        num_components, Nt = size(pos)
        size(vel) != (num_components, Nt) && throw(DimensionMismatch(
            "Velocity and position arrays must have the same size.")
        )
        (num_components != 3) && throw(ArgumentError(
            "Position and velocity must have three components.")
        )

        return new(spec, F, pos, vel, accel, potential, K, evap, τ_bg)
    end
end


include("ReflectBlock.jl")

num_grover_step(prob::Real) = Int(round(pi/4/sqrt(prob))) - 1

function GroverSearch(oracle, num_bit::Int; psi = uniform_state(num_bit))
    uni_reflect = Reflect(uniform_state(num_bit))
    # solve a real search problem
    num_iter = num_grover_step(1.0/(1<<num_bit))

    for i in 1:num_iter
        psi = uni_reflect(oracle(psi))
    end
    psi
end


abstract type Node end

struct LeafNode <: Node
    value
end

struct CategoricalNode <: Node
    attribute
    children  :: Dict{Any, Node}
    majority  :: LeafNode
end

struct NumericNode <: Node
    attribute
    split     :: Number
    left      :: Node
    right     :: Node
end

# Calculate Impurity
impurity(df, current, clf) = begin
    n = nrow(df)

    if eltype(df[!, current]) <: Number
        df = sort!(df, current)

        countl, countr = Counter(df[1:1, clf.topredict]), Counter(df[2:end, clf.topredict])
        minimp, splitat = weightedsum(map(clf.criterion, [countl, countr]), [1/n, 1-1/n]), 1

        for i = 2:(n-1)
            x = df[i, clf.topredict]
            countl[x] += 1
            countr[x] -= 1
            imp = weightedsum(map(clf.criterion, [countl, countr]), [i/n, 1-i/n])

            if imp <= minimp
                minimp, splitat = imp, i
            end
        end
        splitval = (df[splitat, current] + df[splitat+1, current])/2
        minimp, (splitat, splitval)
    else
        sdfs = values(groupby(df, current))
        imps = [(nrow(sdf)/n)*clf.criterion(sdf[!, clf.topredict]) for sdf in sdfs]
        sum(imps), nothing
    end
end

# Building a Tree
Node(df, attributes, depth, clf) = begin
    if (length(Set(df[:, clf.topredict])) == 1 || depth >= clf.maxdepth  ||
        length(attributes) == 0  || nrow(df) <= 2)
        return LeafNode(findmajority(df[:, clf.topredict]))
    end

    minimp, mincurrent, minsplit = Inf, nothing, nothing
    for current in clf.attrfilter(attributes)
        imp, split = impurity(df, current, clf)
        if imp <= minimp
            minimp, mincurrent, minsplit = imp, current, split
        end
    end
    current, split = mincurrent, minsplit

    if eltype(df[!, current]) <: Number
        df = sort!(df, current)
        splitat, splitval = split
        dfl, dfr = df[1:splitat, :], df[splitat:end, :]
        left  = Node(dfl, setdiff(attributes, [current]), depth+1, clf)
        right = Node(dfr, setdiff(attributes, [current]), depth+1, clf)
        NumericNode(current, splitval, left, right)
    else
        children = Dict()
        for ((c,), sdf) in pairs(groupby(df, current))
            children[c] = Node(sdf, setdiff(attributes, [current]), depth+1, clf)
        end
        majority = LeafNode(findmajority(df[:, clf.topredict]))
        CategoricalNode(current, children, majority)
    end
end

# Predicting
predict(node::Node, df::DataFrame) = [predict(node, x) for x in eachrow(df)]

predict(node::LeafNode, x::DataFrameRow) = node.value

predict(node::CategoricalNode, x::DataFrameRow) = begin
    category = x[node.attribute]
    switchto = get(node.children, category, node.majority)
    predict(switchto, x)
end

predict(node::NumericNode, x::DataFrameRow) = begin
    if x[node.attribute] <= node.split
        predict(node.left, x)
    else
        predict(node.right, x)
    end
end


using PointProcessDistance
using Test
using Distances

@testset "PointProcessDistance.jl" begin
    n = 100
    x = rand(2, n)
    dist = PeriodicEuclidean(ones(2))

    rs = 0.:0.01:sqrt(2) / 2
    k = K(rs, x; dist)
    d = D(rs, x, n - 1; dist)
    t = T(rs, x; dist)
    h = H(rs, x; dist)

end


@testset "Functions" begin
    Base.info("Testing functions...")
    inputvec = [-5.0, 0.0, 5.0]
    @test sigmoid(inputvec[2]) == 0.5
    @test sigmoid.(inputvec) == map(x -> 1.0 / (1.0 + exp(-x)), inputvec)
    @test relu(inputvec[1]) == zero(inputvec[1])
    @test relu.(inputvec) == [0.0, 0.0, inputvec[3]]

    inputvec_softmax = [0.0, 0.5, 1.0]
    @test softmax(inputvec_softmax) == [0.36787944117144233, 0.6065306597126334, 1.0] ./ 1.9744101008840758
    @test onehot([1,3,2], 1:3) == [[1,0,0] [0,0,1] [0,1,0]]
    @test crossentropyerror([1,0,0], [1,0,0]) == -log(1.0 + 1.0e-7)
    @test crossentropyerror([0.6,0.4,0], [0,0,1]) == -log(1.0e-7)
    Base.info("Done testing functions.")
end


module Exercise1

export rowmatmul!, colmatmul!, rowmatmuldot!, colmatmuldot!, turbomul!

using LoopVectorization
import LinearAlgebra: dot

"""
   rowmatmul!(C, A, B)

Set C = A*B, looping over rows first
"""
function rowmatmul!(C,A,B)
  for i in 1:size(C,1)
    for j in 1:size(C,2)
      C[i,j] = zero(eltype(C))
      for k in 1:size(A,2)
        C[i,j] += A[i,k]*B[k,j]
      end
    end
  end
end

function colmatmul!(C,A,B)
  for j in 1:size(C,1)
    for i in 1:size(C,2)
      C[i,j] = zero(eltype(C))
      for k in 1:size(A,2)
        C[i,j] += A[i,k]*B[k,j]
      end
    end
  end
end

function rowmatmuldot!(C,A,B)
  error("not implemented")
end

function colmatmuldot!(C,A,B)
  error("not implemented")
end

function turbomul!(C, A, B)
  @turbo for n ∈ indices((C,B), 2), m ∈ indices((C,A), 1)
      Cmn = zero(eltype(C))
      for k ∈ indices((A,B), (2,1))
          Cmn += A[m,k] * B[k,n]
      end
      C[m,n] = Cmn
  end
end

end


export xlogx, hb, hxi, g_p

# Entropy log computation

function xlogx(p::Number)
    if p == 0
        0
    else
        p * log2(p)
    end
end

function xlogx(p)
    ret = zero(p)
    nonzero = p .!= 0
    ret[nonzero] = p[nonzero] .* log2(p[nonzero])
    return ret
end

function hb(t, p)
    return -xlogx(p) - xlogx(t-p)
end

function hb(p)
    return hb(1, p)
end

function hxi(p::Real)
    #return [hb(2*p) ones(length(p),1) hb(2*p)+1]
    PrimalEntropy([hb(2*p); 1; hb(2*p)+1; hb(1/2+p); 2*hb(1/2,p)-2*p; hb(1/2,p)+1/2; hb(2*p)+1; hb(p); hb(2*p)+2*p; hb(1/2,p)+1/2; hb(2*p)+1; hb(1/2,p)+1/2; hb(2*p)+1; hb(2*p)+1; hb(2*p)+1], 1)
    #PrimalEntropy([hb(2*p); ones(length(p),1); hb(2*p)+1; hb(1/2+p); 2*hb(1/2,p)-2*p; hb(1/2,p)+1/2; hb(2*p)+1; hb(p); hb(2*p)+2*p; hb(1/2,p)+1/2; hb(2*p)+1; hb(1/2,p)+1/2; hb(2*p)+1; hb(2*p)+1; hb(2*p)+1], 1)
end

function g_p(p)
    hxi(p) + hb(p) * matusrentropy(1,14) + (1 + 2*p - hb(2*p)/2) * (matusrentropy(1,23) + matusrentropy(2,4))
end


"""
    ThreeStateMorse()

Three state morse potential referred to as Model IA here:
[J. Chem. Phys. 150, 244102 (2019)](https://doi.org/10.1063/1.5096276)

Models IB and IC retain the same functional form and need only a change of parameters.
"""
Parameters.@with_kw struct ThreeStateMorse <: DiabaticModel
    d1::Float64 = 0.02
    d2::Float64 = 0.02
    d3::Float64 = 0.003

    α1::Float64 = 0.4
    α2::Float64 = 0.65
    α3::Float64 = 0.65

    r1::Float64 = 4.0
    r2::Float64 = 4.5
    r3::Float64 = 6.0

    c1::Float64 = 0.02
    c2::Float64 = 0.0
    c3::Float64 = 0.02

    a12::Float64 = 0.005
    a13::Float64 = 0.005
    a23::Float64 = 0.0

    α12::Float64 = 32.0
    α13::Float64 = 32.0
    α23::Float64 = 0.0

    r12::Float64 = 3.40
    r13::Float64 = 4.97
    r23::Float64 = 0.0
end

NQCModels.ndofs(::ThreeStateMorse) = 1
NQCModels.nstates(::ThreeStateMorse) = 3

function NQCModels.potential(model::ThreeStateMorse, R::Real)
    V_ii(x, d, α, r, c) = d * (1 - exp(-α*(x-r)))^2 + c
    V_ij(x, a, α, r) = a * exp(-α*(x-r)^2)

    V11 = V_ii(R, model.d1, model.α1, model.r1, model.c1)
    V22 = V_ii(R, model.d2, model.α2, model.r2, model.c2)
    V33 = V_ii(R, model.d3, model.α3, model.r3, model.c3)

    V12 = V_ij(R, model.a12, model.α12, model.r12)
    V13 = V_ij(R, model.a13, model.α13, model.r13)
    V23 = V_ij(R, model.a23, model.α23, model.r23)

    return Hermitian(SMatrix{3,3}(V11, V12, V13, V12, V22, V23, V13, V23, V33))
end

function NQCModels.derivative(model::ThreeStateMorse, R::Real)

    function D_ii(x, d, α, r)
        ex = exp(-α*(x-r))
        return 2 * d * α * (ex - ex^2)
    end

    D_ij(x, a, α, r) = -2 * a * α * (x-r) * exp(-α*(x-r)^2)

    D11 = D_ii(R, model.d1, model.α1, model.r1)
    D22 = D_ii(R, model.d2, model.α2, model.r2)
    D33 = D_ii(R, model.d3, model.α3, model.r3)

    D12 = D_ij(R, model.a12, model.α12, model.r12)
    D13 = D_ij(R, model.a13, model.α13, model.r13)
    D23 = D_ij(R, model.a23, model.α23, model.r23)

    return Hermitian(SMatrix{3,3}(D11, D12, D13, D12, D22, D23, D13, D23, D33))

end


function calc_grad( pw::PWGrid, Potentials, Focc, psi::Array{ComplexF64,2} )

    Ngwx    = size(psi)[1]
    Nstates = size(psi)[2]
    Ω = pw.Ω
    Ns = pw.Ns
    #
    grad = zeros( ComplexF64, Ngwx, Nstates )

    H_psi = op_H( pw, Potentials, psi )
    for i = 1:Nstates
        grad[:,i] = H_psi[:,i]
        for j = 1:Nstates
            grad[:,i] = grad[:,i] - dot( psi[:,j], H_psi[:,i] ) * psi[:,j]
        end
        grad[:,i] = Focc[i]*grad[:,i]
    end
    return grad

end


module TestUnitsAndUncertainites

using TinyModia
using DifferentialEquations
using ModiaPlot
using Measurements
using Unitful

TestModel = Model(
    T = (0.2 ± 0.02)u"s",
    u0 = 8.5u"kg" ± 0.85u"kg",
    k = (1 ± 0.1)u"kg^-1",
    m = (10 ± 1)u"kg",
    F = (20 ± 2)u"N",

    x = Var(init = (5.5 ± 0.55)u"kg"),
    v = Var(init = (1 ± 0.1)u"m/s"),
    r = Var(init = 1u"m/s"),
    equations = :[
    T*der(x) + x = u0
    y = x
    z = x*k

    F = m * a
    der(v) = a
    der(r) = v ]
)

model = @instantiateModel(TestModel, FloatType = Measurement{Float64})

simulate!(model, Tsit5(), stopTime = 1.0)

plot(model, ["T", "x", "der(x)", "y", "a", "v", "r"])

end


@doc raw"""
    NegBinomialLikelihood(r::Real)

## Arguments
- `r::Real` number of failures until the experiment is stopped

---

[Negative Binomial likelihood](https://en.wikipedia.org/wiki/Negative_binomial_distribution) with number of failures `r`
```math
    p(y|r, f) = {y + r - 1 \choose y} (1 - \sigma(f))^r \sigma(f)^y,
```
if ``r\in \mathbb{N}`` or
```math
    p(y|r, f) = \frac{\Gamma(y + r)}{\Gamma(y + 1)\Gamma(r)} (1 - \sigma(f))^r \sigma(f)^y,
```
if ``r\in\mathbb{R}``.
Where ``\sigma`` is the logistic function
"""
struct NegBinomialLikelihood{L,Tr} <: AbstractLikelihood
    invlink::L
    r::Tr
end

NegBinomialLikelihood(r) = NegBinomialLikelihood(LogisticLink(), r)

(l::NegBinomialLikelihood)(f::Real) = NegativeBinomial(l.r, l.invlink(f))

implemented(::NegBinomialLikelihood, ::Union{<:AnalyticVI,<:GibbsSampling}) = true

function (l::NegBinomialLikelihood)(y::Real, f::Real)
    return pdf(l(f), y)
end

function Distributions.loglikelihood(l::NegBinomialLikelihood, y::Real, f::Real)
    return logpdf(l(f), y)
end

function Base.show(io::IO, l::NegBinomialLikelihood)
    return print(io, "Negative Binomial Likelihood (r = $(l.r))")
end

function compute_proba(
    l::NegBinomialLikelihood, μ::AbstractVector{<:Real}, σ²::AbstractVector{<:Real}
)
    N = length(μ)
    T = eltype(μ)
    pred = zeros(T, N)
    sig_pred = zeros(T, N)
    for i in 1:N
        x = pred_nodes .* sqrt(max(σ²[i], zero(T))) .+ μ[i]
        pred[i] = dot(pred_weights, l.invlink.(x))
        sig_pred[i] = dot(pred_weights, l.invlink.(x) .^ 2) - pred[i]^2
    end
    return pred, sig_pred
end

## Local Updates ##
function init_local_vars(::NegBinomialLikelihood, batchsize::Int, T::DataType=Float64)
    return (; c=rand(T, batchsize), θ=zeros(T, batchsize))
end

function local_updates!(
    local_vars,
    l::NegBinomialLikelihood,
    y::AbstractVector,
    μ::AbstractVector,
    diagΣ::AbstractVector,
)
    map!(sqrt_expec_square, local_vars.c, μ, diagΣ) # √E[f^2]
    map!(local_vars.θ, local_vars.c, y) do c, y
        (l.r + y) * tanh(c / 2) / c
    end # E[ω]
    return local_vars
end

function sample_local!(
    local_vars, l::NegBinomialLikelihood, y::AbstractVector, f::AbstractVector
)
    map!(local_vars.θ, f, y) do f, y
        rand(PolyaGamma(y + Int(l.r), abs(f)))
    end
    return local_vars
end

## Global Updates ##

@inline function ∇E_μ(l::NegBinomialLikelihood, ::AOptimizer, y::AbstractVector, state)
    return ((y .- l.r) / 2,)
end
@inline function ∇E_Σ(::NegBinomialLikelihood, ::AOptimizer, y::AbstractVector, state)
    return (state.θ / 2,)
end

## ELBO Section ##

AugmentedKL(l::NegBinomialLikelihood, state, y) = PolyaGammaKL(l, state, y)

function logabsbinomial(n, k)
    return log(binomial(n, k))
end

function negbin_logconst(y, r::Real)
    return loggamma.(y .+ r) - loggamma.(y .+ 1) .- loggamma(r)
end

function negbin_logconst(y, r::Int)
    return logabsbinomial.(y .+ (r - 1), y)
end

function expec_loglikelihood(
    l::NegBinomialLikelihood,
    ::AnalyticVI,
    y::AbstractVector,
    μ::AbstractVector,
    diag_cov::AbstractVector,
    state,
)
    tot = Zygote.@ignore(sum(negbin_logconst(y, l.r))) - log(2.0) * sum(y .+ l.r)
    tot += dot(μ, (y .- l.r)) / 2 - dot(state.θ, μ) / 2 - dot(state.θ, diag_cov) / 2
    return tot
end

function PolyaGammaKL(l::NegBinomialLikelihood, state, y)
    return PolyaGammaKL(y .+ l.r, state.c, state.θ)
end


"""
```julia
LogisticLikelihood()
```

Bernoulli likelihood with a logistic link for the Bernoulli likelihood
```math
    p(y|f) = \\sigma(yf) = \\frac{1}{1+\\exp(-yf)},
```
(for more info see : [wiki page](https://en.wikipedia.org/wiki/Logistic_function))

---

For the analytic version the likelihood, it is augmented via:
```math
    p(y|f,ω) = exp(0.5(yf - (yf)^2 ω))
```
where ``ω ~ PG(ω | 1, 0)``, and `PG` is the Polya-Gamma distribution
See paper : [Efficient Gaussian Process Classification Using Polya-Gamma Data Augmentation](https://arxiv.org/abs/1802.06383)
"""
struct LogisticLikelihood{T<:Real} <: ClassificationLikelihood{T}
    c::AbstractVector{T}
    θ::AbstractVector{T}
    function LogisticLikelihood{T}() where {T<:Real}
        new{T}()
    end
    function LogisticLikelihood{T}(c::AbstractVector{<:Real},θ::AbstractVector{<:Real}) where {T<:Real}
        new{T}(c,θ)
    end
end

function LogisticLikelihood()
    LogisticLikelihood{Float64}()
end

function init_likelihood(likelihood::LogisticLikelihood{T},inference::Inference{T},nLatent::Int,nSamplesUsed::Int,nFeatures::Int) where T
    if inference isa AnalyticVI || inference isa GibbsSampling
        LogisticLikelihood{T}(abs.(rand(T,nSamplesUsed)),zeros(T,nSamplesUsed))
    else
        LogisticLikelihood{T}()
    end
end

function pdf(l::LogisticLikelihood,y::Real,f::Real)
    logistic(y*f)
end

function logpdf(l::LogisticLikelihood,y::T,f::T) where {T<:Real}
    -log(one(T)+exp(-y*f))
end

function Base.show(io::IO,model::LogisticLikelihood{T}) where T
    print(io,"Bernoulli Likelihood with Logistic Link")
end

function compute_proba(l::LogisticLikelihood{T},μ::AbstractVector{T},σ²::AbstractVector{T}) where {T<:Real}
    N = length(μ)
    pred = zeros(T,N)
    σ²_pred = zeros(T,N)
    for i in 1:N
        x = pred_nodes.*sqrt(max(σ²[i],zero(T))).+μ[i]
        pred[i] = dot(pred_weights,logistic.(x))
        σ²_pred[i] = max(dot(pred_weights,logistic.(x).^2)-pred[i]^2,zero(T))
    end
    return pred, σ²_pred
end

### Local Updates Section ###

function local_updates!(l::LogisticLikelihood{T},y::AbstractVector,μ::AbstractVector,diag_cov::AbstractVector) where {T}
    l.c .= sqrt.(diag_cov+abs2.(μ))
    l.θ .= 0.5*tanh.(0.5.*l.c)./l.c
end

function sample_local!(l::LogisticLikelihood,y::AbstractVector,f::AbstractVector) where {T}
    pg = PolyaGammaDist()
    set_ω!(l,draw.([pg],[1.0],f))
end

### Natural Gradient Section ###

@inline ∇E_μ(l::LogisticLikelihood,::AOptimizer,y::AbstractVector) where {T} = (0.5*y,)
@inline ∇E_Σ(l::LogisticLikelihood,::AOptimizer,y::AbstractVector) where {T} = (0.5*l.θ,)

### ELBO Section ###

function expec_log_likelihood(l::LogisticLikelihood{T},i::AnalyticVI,y::AbstractVector,μ::AbstractVector,diag_cov::AbstractVector) where {T}
    tot = -(0.5*length(y)*logtwo)
    tot += 0.5.*(dot(μ,y)-dot(l.θ,diag_cov)-dot(l.θ,μ))
    return tot
end

AugmentedKL(l::LogisticLikelihood,::AbstractVector) = PolyaGammaKL(l)

function PolyaGammaKL(l::LogisticLikelihood{T}) where {T}
    sum(broadcast(PolyaGammaKL,ones(T,length(l.c)),l.c,l.θ))
end

### Gradient Section ###

@inline grad_log_pdf(::LogisticLikelihood{T},y::Real,f::Real) where {T} = y*logistic(-y*f)

@inline hessian_log_pdf(::LogisticLikelihood{T},y::Real,f::Real) where {T<:Real} = -exp(y*f)/logistic(-y*f)^2


# Primal infeasible test problems
# x >= 0, elements in A >=0, elements in b <0 and s in R+

nn = 1
rng = Random.MersenneTwister(74747)

@testset "Primal infeasible QP problems - Testset 1" begin
  for iii = 1:nn

    # choose size of problem
    n = rand(rng, 5:50)
    m = 2 * n
    A = sprand(rng, m, n, 0.8)
    b = -rand(rng, m)
    A = [A; -sparse(1.0I, n, n)]
    b = [b; zeros(n)]

    # create dual feasible problem
    P = generate_pos_def_matrix(n, rng)
    ytrue = rand(rng, m + n, 1)
    xtrue = rand(rng, n, 1)
    q = (-P * xtrue -  A' * ytrue)[:]

    constraint = COSMO.Constraint(-A, b, COSMO.Nonnegatives)
    ra = 0.

    settings = COSMO.Settings(max_iter=10000, eps_abs = 1e-5, eps_rel=1e-5)
    model = COSMO.Model()
    assemble!(model, P, q, constraint, settings = settings)
    res1 = COSMO.optimize!(model);

     @test res1.status == :Primal_infeasible
  end
end
nothing


function autotransition_init(K::Integer, order::Integer, partition::String="")
        a = rand(Dirichlet(ones(K)/K)) #uninformative prior on initial state probabilities
        A = strong_autotrans_matrix(K)
        no_emission_symbols = Int(4^(order+1)) #alphabet size for the order
        emission_dists = [generate_emission_dist(no_emission_symbols) for i in 1:K]
        #generate the BHMM with the appropriate transition matrix and emissions distributions
        hmm = BHMM(a, A, emission_dists, partition)
end
            #function to construct BHMM transition matrix with strong priors on auto-transition
            function strong_autotrans_matrix(states::Integer, prior_dope::AbstractFloat=(states*250.0), prior_background::AbstractFloat=.1)
                transition_matrix=zeros(states,states)
                for k in 1:states
                    dirichlet_params = fill(prior_background, states)
                    dirichlet_params[k] = prior_dope
                    transition_matrix[k,:] = rand(Dirichlet(dirichlet_params))
                end
                return transition_matrix
            end

            #function to construct BHMM state emission distribution from uninformative dirichlet over the alphabet size
            function generate_emission_dist(no_emission_symbols, prior=Dirichlet(ones(no_emission_symbols)/no_emission_symbols))
                return Categorical(rand(prior))
            end


module Periodogram

export periodogram

""" Returns a maximum-likelihood estimate of the amplitude of a
sinusoidal signal in the series of `xs`, sampled at times `ts`, for
each of the frequencies `fs`. """
function periodogram(ts, xs, fs)
    amps = zeros(fs)

    for i in eachindex(amps)
        cs = cos(2*pi*fs[i]*ts)
        ss = sin(2*pi*fs[i]*ts)
        os = ones(ts)

        A = cat(2, cs, ss, os)
        a,b,m = A\xs

        amps[i] = sqrt(a*a + b*b)
    end

    amps
end

end


module DiscreteHMM

using SPPL

@sppl (debug) function sppl_model(n, categories,
        transition_matrix, obs_matrix)
    latent = array(n)
    X = array(n)

    # Sample initial point.
    latent[0] ~ Choice(zip(categories, [.3, .3, .4]))
    Switch(latent[0], categories) do z
        index = parse(Int, z)
        X[0] ~  Choice(zip(categories, obs_matrix[index]))
    end

    # Sample remaining points.
    for i in 1 : n
        Switch(latent[i-1], categories) do z
            index = parse(Int, z)
            SPPL.Sequence(
                latent[i] ~ Choice(zip(categories, transition_matrix[index])),
                Switch(latent[i], categories) do z
                    index = parse(Int, z)
                    X[i] ~ Choice(zip(categories, obs_matrix[index]))
                end)
        end
    end
end

transition_matrix = [
                     [.1, .4, .5],
                     [.2, .1, .7],
                     [.6, .1, .3]
                    ]

m = sppl_model(5, ["1", "2", "3"],
               transition_matrix, transition_matrix)

end # module


abstract type AbstractNewtonSolver{T} <: NonlinearSolver{T} end

@define newton_solver_variables begin
    x::Vector{T}
    J::Matrix{T}

    x₀::Vector{T}
    x₁::Vector{T}
    y₀::Vector{T}
    y₁::Vector{T}
    δx::Vector{T}
    δy::Vector{T}

    F!::FT
    Jparams::TJ

    linear::TL

    params::NonlinearSolverParameters{T}
    status::NonlinearSolverStatus{T}
end

function setInitialConditions!(s::AbstractNewtonSolver{T}, x₀::Vector{T}) where {T}
    s.x[:] = x₀
end

function computeJacobian(s::AbstractNewtonSolver)
    computeJacobian(s.x, s.J, s.Jparams)
end

function check_jacobian(s::AbstractNewtonSolver)
    println("Condition Number of Jacobian: ", cond(s.J))
    println("Determinant of Jacobian:      ", det(s.J))
    println("minimum(|Jacobian|):          ", minimum(abs.(s.J)))
    println("maximum(|Jacobian|):          ", maximum(abs.(s.J)))
    println()
end

function print_jacobian(s::AbstractNewtonSolver)
    display(s.J)
    println()
end


module essentials
	using Gaston
	set(showable="svg")
	using Yao
	using StatsBase: Histogram, fit
	using BitBasis: BitStr
	Θ = LinRange(0, 2π, 100) # 50
	Φ = LinRange(0, π, 20)
	r = 1
	x = [r*cos(θ)*sin(ϕ) for θ in Θ, ϕ in Φ]
	y = [r*sin(θ)*sin(ϕ) for θ in Θ, ϕ in Φ]
	z = [r*cos(ϕ)        for θ in Θ, ϕ in Φ]
	function plot_blochsph(a::Matrix{ComplexF64})
		surf(x, y, z, lc = :turquoise,
			Axes(view = "equal xyz",
				pm3d = "depthorder",
				palette = :gray,
				style = "fill transparent solid 0.1", tics=:false, border=:false, colorbox=:false), w=:pm3d)
		surf!([0,0], [-1,1], [0,0], lc=:red)
		surf!([-1,1], [0,0], [0,0], lc=:brown)
		surf!([0,0], [0,0], [-1,1], lc=:orange)
		surf!([0,0,0,1], [0,0,-1.2,0], [1.2,-1.2,0,0.2], supp=["|0〉","|1〉","x","y"], w="labels")
		p1 = 2 * acos(abs(a[1]))
		p2 = angle(a[2]) - angle(a[1])
		x1 = r*sin(p1)*cos(p2)
		y1 = r*sin(p2)*sin(p1)
		z1 = r*cos(p1)
		surf!([0], [0], [0], supp=[[y1] [-x1] [z1]], w=:vector, lw=5, lc="'turquoise'")
	end

	function plot_hist(x::Array{BitStr{n,Int},1}) where n
		set(preamble="set xtics font ',$(n<=3 ? 15 : 15/(2^(n-3)))'")
		hist = fit(Histogram, Int.(x), 0:2^n)
		bar(hist.edges[1][1:end-1], hist.weights, fc="'dark-red'", Axes(title = :Histogram, xtics = (0:(2^n-1), "|" .* string.(0:(2^n-1), base=2, pad=n) .* "〉"), yrange = "[0:]"))
	end
end


using JuMP, EAGO, Gurobi, CSV, DataFrames, Ipopt
#m = Model(with_optimizer(Ipopt.Optimizer))

m = Model(with_optimizer(EAGO.Optimizer, relaxed_optimizer = Gurobi.Optimizer(OutputFlag=0)))

x_Idx = Any[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]
@variable(m, x[x_Idx])
JuMP.set_lower_bound(x[21], 0.0)
JuMP.set_lower_bound(x[1], -0.5)
JuMP.set_upper_bound(x[1], 0.5)
JuMP.set_lower_bound(x[2], 5.4)
JuMP.set_upper_bound(x[2], 6.4)
JuMP.set_lower_bound(x[3], 0.4)
JuMP.set_upper_bound(x[3], 1.4)
JuMP.set_lower_bound(x[4], 4.9)
JuMP.set_upper_bound(x[4], 5.9)
JuMP.set_lower_bound(x[5], 1.3)
JuMP.set_upper_bound(x[5], 2.3)
JuMP.set_lower_bound(x[6], 3.9)
JuMP.set_upper_bound(x[6], 4.9)
JuMP.set_lower_bound(x[7], 2.1)
JuMP.set_upper_bound(x[7], 3.1)
JuMP.set_lower_bound(x[8], 4.1)
JuMP.set_upper_bound(x[8], 5.1)
JuMP.set_lower_bound(x[9], 2.8)
JuMP.set_upper_bound(x[9], 3.8)
JuMP.set_lower_bound(x[10], 3.0)
JuMP.set_upper_bound(x[10], 4.0)
JuMP.set_lower_bound(x[11], 3.9)
JuMP.set_upper_bound(x[11], 4.9)
JuMP.set_lower_bound(x[12], 3.2)
JuMP.set_upper_bound(x[12], 4.2)
JuMP.set_lower_bound(x[13], 4.7)
JuMP.set_upper_bound(x[13], 5.7)
JuMP.set_lower_bound(x[14], 2.3)
JuMP.set_upper_bound(x[14], 3.3)
JuMP.set_lower_bound(x[15], 5.6)
JuMP.set_upper_bound(x[15], 6.6)
JuMP.set_lower_bound(x[16], 2.3)
JuMP.set_upper_bound(x[16], 3.3)
JuMP.set_lower_bound(x[17], 6.0)
JuMP.set_upper_bound(x[17], 7.0)
JuMP.set_lower_bound(x[18], 1.9)
JuMP.set_upper_bound(x[18], 2.9)
JuMP.set_lower_bound(x[19], 6.9)
JuMP.set_upper_bound(x[19], 7.9)
JuMP.set_lower_bound(x[20], 1.0)
JuMP.set_upper_bound(x[20], 2.0)
JuMP.set_upper_bound(x[21], 10.0)
JuMP.set_lower_bound(x[22], -2.0)
JuMP.set_upper_bound(x[22], 2.0)
JuMP.set_lower_bound(x[23], -2.0)
JuMP.set_upper_bound(x[23], 2.0)
JuMP.set_lower_bound(x[24], -2.0)
JuMP.set_upper_bound(x[24], 2.0)

# ----- Constraints ----- #
@NLconstraint(m, e2, x[22]*x[1]+ (x[1])^2*x[23]+ (x[1])^3*x[24]-x[2]+x[21] == 0.0)
@NLconstraint(m, e3, x[22]*x[3]+ (x[3])^2*x[23]+ (x[3])^3*x[24]-x[4]+x[21] == 0.0)
@NLconstraint(m, e4, x[22]*x[5]+ (x[5])^2*x[23]+ (x[5])^3*x[24]-x[6]+x[21] == 0.0)
@NLconstraint(m, e5, x[22]*x[7]+ (x[7])^2*x[23]+ (x[7])^3*x[24]-x[8]+x[21] == 0.0)
@NLconstraint(m, e6, x[22]*x[9]+ (x[9])^2*x[23]+ (x[9])^3*x[24]-x[10]+x[21] == 0.0)
@NLconstraint(m, e7, x[22]*x[11]+ (x[11])^2*x[23]+ (x[11])^3*x[24]-x[12]+x[21] == 0.0)
@NLconstraint(m, e8, x[22]*x[13]+ (x[13])^2*x[23]+ (x[13])^3*x[24]-x[14]+x[21] == 0.0)
@NLconstraint(m, e9, x[22]*x[15]+ (x[15])^2*x[23]+ (x[15])^3*x[24]-x[16]+x[21] == 0.0)
@NLconstraint(m, e10, x[22]*x[17]+ (x[17])^2*x[23]+ (x[17])^3*x[24]-x[18]+x[21] == 0.0)
@NLconstraint(m, e11, x[22]*x[19]+ (x[19])^2*x[23]+ (x[19])^3*x[24]-x[20]+x[21] == 0.0)

# ----- Objective ----- #
@NLobjective(m, Min, ( (x[1])^2+ ((-5.9)+x[2])^2+ ((-0.9)+x[3])^2+ ((-5.4)+x[4])^2+ ((-1.8)+x[5])^2+ ((-4.4)+x[6])^2+ ((-2.6)+x[7])^2+ ((-4.6)+x[8])^2+ ((-3.3)+x[9])^2+ ((-3.5)+x[10])^2+ ((-4.4)+x[11])^2+ ((-3.7)+x[12])^2+ ((-5.2)+x[13])^2+ ((-2.8)+x[14])^2+ ((-6.1)+x[15])^2+ ((-2.8)+x[16])^2+ ((-6.5)+x[17])^2+ ((-2.4)+x[18])^2+ ((-7.4)+x[19])^2+ ((-1.5)+x[20])^2))

JuMP.optimize!(m)

backend_opt = backend(m).optimizer.model.optimizer
last_relaxed_opt = backend_opt.relaxed_optimizer
run_time = backend_opt._run_time
println("run time: $run_time")


import LinearAlgebra: BlasInt
const LA = LinearAlgebra

@enum Factorization begin
    QR_FACT
    LU_FACT
end
struct MatrixWorkspace{T} <: AbstractMatrix{Complex{T}}
    A::Matrix{Complex{T}} # Matrix
    d::Vector{T} # Inverse of scaling factors
    # factorizations
    fact::Base.RefValue{Factorization}
    factorized::Base.RefValue{Bool}
    lu::Union{Nothing,LA.LU{Complex{T},Matrix{Complex{T}}}} # LU Factorization of D * J
    qr::LA.QRPivoted{Complex{T},Matrix{Complex{T}}} # Pivoted QR Factorization of D * J
    # qr worksapce
    qr_perm::Vector{Int} # work copy of the permutation in the qr
    qr_work::Vector{Complex{T}} # preallocated memory for pivoted qr
    qr_rwork::Vector{T} # preallocated memory for pivoted qr
    qr_ormqr_work::Vector{Complex{T}}
    # overdetermined
    qr_rhs::Vector{Complex{T}}
    # iterative refinement
    ir_r::Vector{Complex{T}}
    ir_δx::Vector{Complex{T}}
    residual_abs::Vector{T}
    # cond estimators
    lu_cond_work::Vector{Complex{T}}
    lu_cond_rwork::Vector{T}
    qr_cond_work::Vector{ComplexF64}
    qr_cond_rwork::Vector{T}
    # inf_norm_est
    inf_norm_est_work::Vector{Complex{T}}
    inf_norm_est_rwork::Vector{T}
    row_scaling::Vector{T}
end

MatrixWorkspace(Â::AbstractMatrix) =
    MatrixWorkspace(Â, promote_type(real(eltype(Â)), Float64))

function MatrixWorkspace(Â::AbstractMatrix, ::Type{T}) where {T}
    A = Matrix{Complex{T}}(Â)
    m, n = size(A)
    d = ones(m)
    fact = m == n ? LU_FACT : QR_FACT
    # lu
    lu = m == n ? LA.lu(A; check = false) : nothing
    # qr
    qr = LA.qr(A, Val(true))
    qr_perm = zeros(Int, length(qr.p))
    qr_work = Vector{Complex{T}}(undef, 1)
    qr_ormqr_work = Vector{Complex{T}}(undef, 1)
    qr_rwork = Vector{T}(undef, 2n)

    ir_r = Vector{Complex{T}}(undef, m)
    ir_δx = Vector{Complex{T}}(undef, n)
    residual_abs = Vector{T}(undef, m)
    qr_rhs = copy(ir_r)
    lu_cond_work = Vector{Complex{T}}(undef, 2n)
    lu_cond_rwork = Vector{T}(undef, 2n)
    qr_cond_work = Vector{Complex{T}}(undef, 2n)
    qr_cond_rwork = Vector{T}(undef, n)
    inf_norm_est_work = Vector{Complex{T}}(undef, n)
    inf_norm_est_rwork = Vector{T}(undef, n)
    row_scaling = ones(m)
    MatrixWorkspace(
        A,
        d,
        Ref(fact),
        Ref(true),
        lu,
        qr,
        qr_perm,
        qr_work,
        qr_rwork,
        qr_ormqr_work,
        qr_rhs,
        ir_r,
        ir_δx,
        residual_abs,
        lu_cond_work,
        lu_cond_rwork,
        qr_cond_work,
        qr_cond_rwork,
        inf_norm_est_work,
        inf_norm_est_rwork,
        row_scaling,
    )
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)

"""
    updated!(MW::MatrixWorkspace)

Indicate that the matrix `MW` got updated.
"""
updated!(MW::MatrixWorkspace) = (MW.factorized[] = false; MW)

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

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

"""
    factorization!(MW::MatrixWorkspace, fact::Factorization)

Update the used factorization in `MW` to fact.
"""
function factorization!(MW::MatrixWorkspace, fact::Factorization)
    if fact == LU_FACT && MW.lu === nothing
        throw(ArgumentError("Cannot set `LU_FACT` for non-square matrix."))
    end
    MW.factorized[] = MW.fact[] == fact
    MW.fact[] = fact
    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},
    b::Union{AbstractVector,Nothing} = nothing,
    ipiv::Union{Vector{I},Nothing} = nothing,
    ::Val{Pivot} = Val(true),
) where {T,I<:Integer,Pivot}
    m, n = size(A)
    minmn = min(m, n)
    # LU Factorization
    @inbounds begin
        for k = 1:minmn
            # find index max
            kp = k
            if Pivot
                amax = zero(real(T))
                for i = k:m
                    if T <: Complex
                        absi = abs2(A[i, k])
                    else
                        absi = abs(A[i, k])
                    end
                    if absi > amax
                        kp = i
                        amax = absi
                    end
                end
            end
            if !(ipiv isa Nothing)
                ipiv[k] = kp
            end
            if !iszero(A[kp, k])
                if k != kp
                    # Interchange
                    for i = 1:n
                        tmp = A[k, i]
                        A[k, i] = A[kp, i]
                        A[kp, i] = tmp
                    end
                    if !(b isa Nothing)
                        b[k], b[kp] = b[kp], b[k]
                    end
                end
                # Scale first column
                # Akkinv = @fastmath inv(A[k,k])
                Akkinv = inv(A[k, k])
                for i = k+1:m
                    A[i, k] *= Akkinv
                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

#############
# Custom QR #
#############
"""
    geqp3!(A::AbstractMatrix{ComplexF64},
        jpvt::AbstractVector{BlasInt},
        tau::AbstractVector{ComplexF64},
        work::Vector{ComplexF64},
        rwork=Vector{Float64}(undef, 2*size(A,2)))

Version of `LAPACK.geqp3!` which also accepts the work buffers.
"""
function geqp3!(
    A::AbstractMatrix{ComplexF64},
    jpvt::AbstractVector{BlasInt},
    tau::AbstractVector{ComplexF64},
    work::Vector{ComplexF64},
    rwork = Vector{Float64}(undef, 2 * size(A, 2)),
)
    m, n = size(A)
    lda = stride(A, 2)
    lda == 0 && return nothing
    jpvt .= BlasInt(0)
    lwork = BlasInt(-1)
    info = Ref{BlasInt}()
    for i = 1:2
        ccall(
            (LA.LAPACK.@blasfunc(zgeqp3_), LA.LAPACK.liblapack),
            Cvoid,
            (
                Ref{BlasInt},
                Ref{BlasInt},
                Ptr{ComplexF64},
                Ref{BlasInt},
                Ptr{BlasInt},
                Ptr{ComplexF64},
                Ptr{ComplexF64},
                Ref{BlasInt},
                Ptr{Float64},
                Ptr{BlasInt},
            ),
            m,
            n,
            A,
            lda,
            jpvt,
            tau,
            work,
            lwork,
            rwork,
            info,
        )
        LA.LAPACK.chklapackerror(info[])
        if i == 1
            lwork = BlasInt(real(work[1]))
            resize!(work, lwork)
        end
    end
    nothing
end

function factorize!(WS::MatrixWorkspace)
    if WS.fact[] == LU_FACT
        @inbounds copyto!(WS.lu.factors, WS.A)
        lu!(WS.lu.factors, nothing, WS.lu.ipiv)
    elseif WS.fact[] == QR_FACT
        @inbounds copyto!(WS.qr.factors, WS.A)
        geqp3!(WS.qr.factors, WS.qr.jpvt, WS.qr.τ, WS.qr_work, WS.qr_rwork)
    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
@inline 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;
    singular_exception::Bool = false,
)
    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])
        xj = x[j] = 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; check::Bool = false)
    x === b || copyto!(x, b)
    _ipiv!(LU, x)
    ldiv_unit_lower!(LU.factors, x)
    ldiv_upper!(LU.factors, x; singular_exception = check)
    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

"""
    ormqr!(side, trans, A, tau, C, work)

Version of `LAPACK.ormqr!` which also accepts the work buffer.
Computes `Q * C` (`trans = N`), `transpose(Q) * C` (`trans = T`), `adjoint(Q) * C`
(`trans = C`) for `side = L` or the equivalent right-sided multiplication
for `side = R` using `Q` from a `QR` factorization of `A` computed using
`geqrf!`. `C` is overwritten.
"""
function ormqr!(
    side::AbstractChar,
    trans::AbstractChar,
    A::Matrix{ComplexF64},
    tau::Vector{ComplexF64},
    C::Vector{ComplexF64},
    work::Vector{ComplexF64},
)
    m, n = (size(C, 1), 1)
    mA = size(A, 1)
    k = length(tau)
    lwork = BlasInt(-1)
    info = Ref{BlasInt}()
    for i = 1:2  # first call returns lwork as work[1]
        ccall(
            (LA.LAPACK.@blasfunc(zunmqr_), LA.LAPACK.liblapack),
            Cvoid,
            (
                Ref{UInt8},
                Ref{UInt8},
                Ref{BlasInt},
                Ref{BlasInt},
                Ref{BlasInt},
                Ptr{ComplexF64},
                Ref{BlasInt},
                Ptr{ComplexF64},
                Ptr{ComplexF64},
                Ref{BlasInt},
                Ptr{ComplexF64},
                Ref{BlasInt},
                Ptr{BlasInt},
            ),
            side,
            trans,
            m,
            n,
            k,
            A,
            max(1, stride(A, 2)),
            tau,
            C,
            max(1, stride(C, 2)),
            work,
            lwork,
            info,
        )
        LA.LAPACK.chklapackerror(info[])
        if i == 1
            lwork = BlasInt(real(work[1]))
            resize!(work, lwork)
        end
    end
    C
end

# This is an adoption of the Julia implementation of ldiv! for QRPivoted
function qr_ldiv!(
    x,
    A::LA.QRPivoted{ComplexF64},
    b::Vector{ComplexF64},
    perm::Vector{Int},
    ormqr_work::Vector{ComplexF64},
      #corank::Int,
)
    mA, nA = size(A.factors)
    nr = min(mA, nA)
    nrhs = length(b)
    nr == 0 && return x

    # The following is equivalent to LA.lmul!(LA.adjoint(A.Q), b) but
    # we can also pass the preallocated work vector
    ormqr!('L', 'C', A.factors, A.τ, b, ormqr_work)
    ldiv_upper!(A.factors, b; singular_exception = false)
    @inbounds for i = 1:nr
        x[i] = b[i]
        perm[i] = A.p[i]
    end
    Base.invpermute!!(x, perm)
    # The following would cover the case of rank deficient qr.
    # This is currently not used but we keep it here since it will be used later.
    # if corank == nr
    #     fill!(x, zero(eltype(x)))
    #     return x
    # end
    # rank = nr - corank
    # if rank ≠ nA || rank ≠ nr
    #     C, τ = LA.LAPACK.tzrzf!(A.factors[1:rank,:])
    #     LA.ldiv!(LA.UpperTriangular(C[1:rank,1:rank]),view(LA.lmul!(LA.adjoint(A.Q), b), 1:rank))
    #     for i in 1:rank
    #         x[i] = b[i]
    #     end
    #     for i in rank+1:nA
    #         x[i] = zero(eltype(x))
    #     end
    #     LA.LAPACK.ormrz!('L', eltype(x)<:Complex ? 'C' : 'T', C, τ, reshape(x, length(x), 1))
    #     invpermute!(x, A.p)
    # end
    return x
end

function LA.ldiv!(x::AbstractVector, WS::MatrixWorkspace, b::AbstractVector)
    WS.factorized[] || factorize!(WS)
    if WS.fact[] == LU_FACT
        lu = WS.lu
        if lu !== nothing
            lu_ldiv!(x, lu, b)
        end
    elseif WS.fact[] == QR_FACT
        copyto!(WS.qr_rhs, b)
    # if issquare(WS.A)
            # # we applied row scaling for square matrices
            # @inbounds for i in eachindex(Jac.b)
            #     Jac.b[i] = b[i] / Jac.D[i]
            # end
    #
        qr_ldiv!(x, WS.qr, WS.qr_rhs, WS.qr_perm, WS.qr_ormqr_work)
    end
    x
end

##########################
## Iterative refinement ##
##########################
"""
    residual!(u, A, x, b, ::Type{T}=eltype(u))

Compute the residual `Ax-b` in precision `T` and store in `u`.
"""
function residual!(
    u::AbstractVector{S},
    A::AbstractMatrix{S},
    x::AbstractVector{S},
    b::AbstractVector{S},
    ::Type{T} = eltype(u),
) where {S,T}
    @boundscheck size(A, 1) == length(b) && size(A, 2) == length(x)
    m, n = size(A)
    @inbounds for i = 1:m
        dot = A[i, 1] * convert(T, x[1])
        for j = 2:n
            dot = muladd(A[i, j], convert(T, x[j]), dot)
        end
        u[i] = dot - b[i]
    end
    u
end

"""
    iterative_refinement_step!([x,] MW::MatrixWorkspace, x̂, b, norm::AbstractNorm=InfNorm(), T=eltype(x̂))

Apply one step of iterative refinement where the residual is computed with precision `T`.
Stores the result in `x` and returns the norm of the update `δx`. If `x` is not provided
`x̂` is updated inplace.
"""
function iterative_refinement_step!(
    WS::MatrixWorkspace,
    x̂,
    b,
    norm::AbstractNorm = InfNorm(),
    ::Type{T} = eltype(x̂),
) where {T}
    iterative_refinement_step!(x̂, WS, x̂, b, norm, T)
end
function iterative_refinement_step!(
    x,
    WS::MatrixWorkspace,
    x̂,
    b,
    norm::AbstractNorm = InfNorm(),
    ::Type{T} = eltype(x̂),
) where {T}
    residual!(WS.ir_r, WS.A, x̂, b, T)
    δx = LA.ldiv!(WS.ir_δx, WS, WS.ir_r)
    for i in eachindex(x)
        x[i] = x̂[i] - δx[i]
    end
    return norm(δx)
end

#########################
## Condition estimates ##
#########################
"""
    rcond(WS::MatrixWorkspace)

Computes the reciprocal condition number of the matrix `A` in `WS` wrt the ∞-norm.
If `A` is a non-square matrix, the reciprocal condition number of `R` in the
QR decomposition is computed.
"""
function rcond(WS::MatrixWorkspace)
    WS.factorized[] || factorize!(WS)
    if WS.fact[] == LU_FACT
        lu = WS.lu
        if lu === nothing
            NaN
        else
            rcond(lu, LA.opnorm(WS.A, Inf), WS.lu_cond_work, WS.lu_cond_rwork)
        end
    else # WS.fact[] == QR_FACT
        rcond(WS.qr, WS.qr_cond_work, WS.qr_cond_rwork)
    end
end
function rcond(LU::LA.LU, anorm, work, rwork)
    A = LU.factors
    n = size(A, 1)
    rcond = Ref{Float64}()
    info = Ref{Int64}()
    lda = max(1, stride(A, 2))
    ccall(
        (LA.LAPACK.@blasfunc(zgecon_), LinearAlgebra.LAPACK.liblapack),
        Cvoid,
        (
            Ref{UInt8},
            Ref{Int64},
            Ptr{ComplexF64},
            Ref{Int64},
            Ref{Float64},
            Ref{Float64},
            Ptr{ComplexF64},
            Ptr{Float64},
            Ptr{Int64},
        ),
        'I',
        n,
        A,
        lda,
        anorm,
        rcond,
        work,
        rwork,
        info,
    )
    rcond[]
end

function rcond(qr::LA.QRPivoted, work, rwork)
    # use trcon to estimate cond of R
    m, n = size(qr.factors)
    rcond = Ref{Float64}(1)
    info = Ref{BlasInt}()
    ccall(
        (LA.LAPACK.@blasfunc(ztrcon_), LA.LAPACK.liblapack),
        Cvoid,
        (
            Ref{UInt8},
            Ref{UInt8},
            Ref{UInt8},
            Ref{BlasInt},
            Ptr{ComplexF64},
            Ref{BlasInt},
            Ref{Float64},
            Ptr{ComplexF64},
            Ptr{Float64},
            Ptr{BlasInt},
        ),
        'I',
        'U',
        'N',
        n,
        qr.factors,
        m,
        rcond,
        work,
        rwork,
        info,
    )
    if m == n
        # Q has condition number 1 in the 2 norm
        # Some experiments indicate that in the Inf norm this is roughly
        # 0.8√n
        rcond[] / (0.8 * √(n))
    else
        rcond[]
    end
end

"""
    inf_norm_est(WS::MatrixWorkspace, g::Union{Nothing,Vector{<:Real}}=nothing, d::Union{Nothing,Vector{<:Real}}=nothing)
    inf_norm_est(lu::LA.LU, g::Union{Nothing,Vector{<:Real}}=nothing, d::Union{Nothing,Vector{<:Real}}=nothing)
    inf_norm_est(lu::LA.LU, g::Union{Nothing,Vector{<:Real}}, d::Union{Nothing,Vector{<:Real}}, work{<:Complex}, rwork::Vector{<:Real})

Estimation of the infinity norm of `diag(d)⁻¹A⁻¹g` where `g` and `d` are optional positive vectors
and `A=LU`. If `g` is `nothing` the infinity norm of `A⁻¹` is estimated.
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 inf_norm_est(
    lu::LA.LU,
    g::Union{Nothing,Vector{<:Real}} = nothing,
    d::Union{Nothing,Vector{<:Real}} = nothing,
)
    n = size(lu.factors, 2)
    work = Vector{eltype(lu.factors)}(undef, n)
    rwork = Vector{real(eltype(lu.factors))}(undef, n)
    inf_norm_est(lu, g, d, work, rwork)
end
function inf_norm_est(
    WS::MatrixWorkspace,
    g::Union{Nothing,Vector{<:Real}} = nothing,
    d::Union{Nothing,Vector{<:Real}} = nothing,
)
    WS.fact[] == LU_FACT || factorization!(WS, LU_FACT)
    WS.factorized[] || factorize!(WS)
    inf_norm_est(WS.lu, g, d, WS.inf_norm_est_work, WS.inf_norm_est_rwork)
end
function inf_norm_est(
    lu::LA.LU,
    g::Union{Nothing,Vector{<:Real}},
    d::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 !== nothing
        x ./= d
    end
    lu_ldiv_adj!(y, lu, x)
    if g !== nothing
        y .*= g
    end
    γ = sum(abs, y)
    y ./= abs.(y)
    if g !== nothing
        y .*= g
    end
    lu_ldiv!(z, lu, ξ)
    if d !== nothing
        x .= real.(z) ./ d
    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 !== nothing
            x ./= d
        end
        lu_ldiv_adj!(y, lu, x)
        if g !== nothing
            y .*= g
        end
        γ̄ = γ
        γ = sum(abs, y)
        γ ≤ γ̄ && break
        ξ .= y ./ abs.(y)
        if g !== nothing
            ξ .*= g
        end
        lu_ldiv!(z, lu, ξ)
        if d !== nothing
            x .= real.(z) ./ d
        else
            x .= real.(z)
        end
        k += 1
        if x[j] == LA.norm(x, Inf) || k > 5
            break
        end
    end
    γ
end

"""
    row_scaling!(WS::MatrixWorkspace,
        norm::Union{InfNorm, WeightedNorm{<:InfNorm}},
        residual::Union{Nothing,AbstractVector{<:Real}} = nothing)

Compute a suitable scaling `D` of the rows of `WS.A` and store it in `WS.row_scaling`.
The returned scaling has the property that the from `norm` induced operator for `D⁻¹ * WS.A`
is 1.
The vector `residual` is a an optional vector containing the absolute values of the limiting
residual of Newton's method. The scaling is performaned in such a way that `norm(residual)`
is *not* increase. Unless the from `norm` induced operator norm is less than 1.
Then the residual is increased by the minimal amount necessary.
"""
function row_scaling!(
    WS::MatrixWorkspace{T},
    norm::Union{InfNorm,WeightedNorm{<:InfNorm}},
    r::Union{Nothing,AbstractVector{<:Real}} = nothing,
) where {T}
    row_scaling!(WS.row_scaling, WS, norm, r)
end
function row_scaling!(
    d::AbstractVector,
    WS::MatrixWorkspace{T},
    norm::Union{InfNorm,WeightedNorm{<:InfNorm}},
    r::Union{Nothing,AbstractVector{<:Real}} = nothing,
) where {T}
    d .= 1e-3
    m = length(d)
    if isa(norm, WeightedNorm)
        for j = 1:size(WS.A, 2), i = 1:m
            d[i] += abs(WS.A[i, j]) * norm[j]
        end
    else
        for j = 1:size(WS.A, 2), i = 1:m
            d[i] += abs(WS.A[i, j])
        end
    end

    if size(WS, 1) == 1
        d[1] = max(1, d[1])
        return d
    end

    if r !== nothing
        d̂ = min(maximum(d), 1.0)
        # compute the norm of r under the scaling of all rows whose norm is larger
        # than d̂, i.e. all rows that either decrease the value of |rᵢ| or
        # or increase it by the minimal abount to still have opnorm 1.
        norm_scaled_r = r[1] / max(d[1], d̂)
        for i = 2:m
            norm_scaled_r = max(norm_scaled_r, r[i] / max(d[i], d̂))
        end
        for i = 1:length(d)
            r̂ᵢ = r[i] / norm_scaled_r
            if isfinite(r̂ᵢ)
                d[i] = max(d[i], r̂ᵢ)
            end
        end
    else
        # Don't scale zero rows -> set dᵢ = 1 if too small
        d_max_eps = sqrt(eps())
        for i = 1:m
            if d[i] < d_max_eps
                d[i] = one(T)
            end
        end
    end
    d
end

"""
    scaled_cond!(WS::MatrixWorkspace, norm, residual)

Perform first a row scaling as described by [`row_scaling!`](@ref) and then
compute the condition number ``κ`` of the row-equilibrated matrix `WS.A`.
This is the best possible condition number under row scaling (e.g. [Higham02, Thm 7.5]).

[Higham02]: Higham, Nicholas J. Accuracy and stability of numerical algorithms. Vol. 80. Siam, 200
"""
function scaled_cond!(
    WS::MatrixWorkspace,
    norm::Union{<:WeightedNorm{<:InfNorm},InfNorm},
    r::Union{Nothing,AbstractVector{<:Real}} = nothing,
)
    d = row_scaling!(WS, norm, r)
    # row scaling makes ||A||_∞ = 1 so we don't this here
    if (norm isa WeightedNorm)
        cond_est = inf_norm_est(WS, d, weights(norm))
    else
        cond_est = inf_norm_est(WS, d)
    end
    cond_est
end

#####################
## JacobianMonitor ##
#####################
struct JacobianMonitor{T}
    J::MatrixWorkspace{T}
    cond::Base.RefValue{Float64}
    forward_err::Base.RefValue{Float64}
    # stats
    factorizations::Base.RefValue{Int}
    ldivs::Base.RefValue{Int}
end
function JacobianMonitor(A::AbstractMatrix)
    J = MatrixWorkspace(A)
    cond = Ref(1.0)
    forward_err = Ref(0.0)
    JacobianMonitor(J, cond, forward_err, Ref(0), Ref(0))
end

updated!(JM::JacobianMonitor) = updated!(JM.J)
jacobian(JM::JacobianMonitor) = JM.J

function Base.show(io::IO, JM::JacobianMonitor{T}) where {T}
    println(io, "JacobianMonitor{$T}:")
    println(io, " • cond → ", round(JM.cond[], sigdigits = 5))
    println(io, " • forward_err → ", round(JM.forward_err[], sigdigits = 5))
    println(io, " • # factorizations → ", JM.factorizations[])
    println(io, " • # ldivs → ", JM.ldivs[])
end

"""
    init!(JM::JacobianMonitor)

(Re-)initialize the `JacobianMonitor`.
"""
function init!(JM::JacobianMonitor; keep_stats::Bool = false)
    JM.cond[] = 1.0
    JM.forward_err[] = 0.0
    JM.factorizations[] = 0.0
    JM.ldivs[] = 0.0
    JM
end

"""
    forward_err(JM::JacobianMonitor)

Return the last estimate of the forward error as computed by [`forward_err`](@ref) or [`strong_forward_err`](@ref).
"""
forward_err(JM::JacobianMonitor) = JM.forward_err[]

"""
    forward_err!(x̂::AbstractVector, JM::JacobianMonitor, b::AbstractVector, norm::AbstractNorm)

Compute an estimate of the forward error ||x-x̂||/||x|| and update the value in `JM`.
"""
function forward_err!(
    JM::JacobianMonitor,
    x̂::AbstractVector,
    b::AbstractVector,
    norm::AbstractNorm,
    T = eltype(x̂),
)
    forward_err!(x̂, JM, x̂, b, norm, T)
end
function forward_err!(
    x::AbstractVector,
    JM::JacobianMonitor,
    x̂::AbstractVector,
    b::AbstractVector,
    norm::AbstractNorm,
    T = eltype(x̂),
)
    norm_x̂ = norm(x̂)
    JM.forward_err[] = iterative_refinement_step!(x, JM.J, x̂, b, norm, T) / norm_x̂
end

"""
    strong_forward_err!(x̂::AbstractVector, JM::JacobianMonitor, b::AbstractVector, norm::AbstractNorm)

Compute a more robust estimate of the forward error ||x-x̂||/||x|| using eq. (2.14) in [Demmel, Section 2.4.4]
and update the value in `JM`.
If `jacobian(JM).fact[] == QR_FACT` then this falls back to [`forward_err!`](@ref).

[Demmel]: Demmel, James W. Applied numerical linear algebra. Vol. 56. Siam, 1997.
"""
function strong_forward_err!(
    JM::JacobianMonitor,
    x̂::AbstractVector,
    b::AbstractVector,
    norm::InfNorm,
    T = eltype(x̂),
)
    strong_forward_err!(x̂, JM, x̂, b, norm, T)
end
function strong_forward_err!(
    x::AbstractVector,
    JM::JacobianMonitor,
    x̂::AbstractVector,
    b::AbstractVector,
    norm::Union{InfNorm,WeightedNorm{InfNorm}},
    T = eltype(x̂),
)
    if jacobian(JM).fact[] == QR_FACT
        return JM.forward_err[] = forward_err!(x, JM, x̂, b, norm)
    end
    norm_x̂ = norm(x̂)
    WS = jacobian(JM)
    residual!(WS.ir_r, WS.A, x̂, b, T)
    WS.residual_abs .= abs.(b)
    for j = 1:size(WS.A, 2)
        ax̂ⱼ = abs(x̂[j])
        for i = 1:size(WS.A, 1)
            WS.residual_abs[i] += abs(WS.A[i, j]) * ax̂ⱼ
        end
    end
    WS.residual_abs .*= eps()
    WS.residual_abs .+= abs.(WS.ir_r)

    if norm isa WeightedNorm
        JM.forward_err[] = inf_norm_est(WS, WS.residual_abs, weights(norm)) / norm_x̂
    else
        JM.forward_err[] = inf_norm_est(WS, WS.residual_abs) / norm_x̂
    end
end

"""
    cond!(JM::JacobianMonitor, norm, residual)

Compute the condition number ``κ`` of the Jacobian ``J`` with respect to the
infinity norm and store it. If ``J`` is a square matrix the condition number of the
row-equilibrated Jacobian is computed. See [`scaled_cond!`](@ref) for details.
"""
function cond!(
    JM::JacobianMonitor,
    norm::AbstractNorm = InfNorm(),
    r::Union{Nothing,AbstractVector{<:Real}} = nothing,
)
    if jacobian(JM).fact[] == LU_FACT
        JM.cond[] = scaled_cond!(jacobian(JM), norm, r)
    else
        # TODO: The weighted norm is NOT considered currently.
        JM.cond[] = inv(rcond(jacobian(JM)))
    end
end

"""
    apply_row_scaling!(r, JM::JacobianMonitor)

Apply the computed row scaling used in [`cond!`](@ref) resp. [`scaled_cond!`](@ref) to `r`.
"""
apply_row_scaling!(r::AbstractVector, JM::JacobianMonitor) = (r ./= JM.J.row_scaling; r)

"""
    cond(JM::JacobianMonitor)

Returns the with [`cond!`](@ref) computed condition number.
"""
LA.cond(JM::JacobianMonitor) = JM.cond[]

"""
    enum JacobianMonitorUpdates

## Cases

* `JAC_MONITOR_UPDATE_NOTHING`: Do nothing
* `JAC_MONITOR_UPDATE_FERR`: Update the forward error estimate
* `JAC_MONITOR_UPDATE_COND`: Update the condition number estimate
* `JAC_MONITOR_UPDATE_ALL`: Update the forward error and condition number estimate
"""
@enum JacobianMonitorUpdates begin
    JAC_MONITOR_UPDATE_NOTHING
    JAC_MONITOR_UPDATE_FERR
    JAC_MONITOR_UPDATE_COND
    JAC_MONITOR_UPDATE_ALL
end

"""
    ldiv!(x̂::AbstractVector,
                  JM::JacobianMonitor,
                  b::AbstractVector,
                  norm::AbstractNorm,
                  update::JacobianMonitorUpdates=JAC_MONITOR_UPDATE_NOTHING)

solve the linear system `jacobian(JM)x̂=b`. `update` controls the computation of additional
informations, see [`JacobianMonitorUpdates`](@ref).
"""
function LA.ldiv!(
    x̂::AbstractVector,
    JM::JacobianMonitor,
    b::AbstractVector,
    norm::AbstractNorm = InfNorm(),
    update::JacobianMonitorUpdates = JAC_MONITOR_UPDATE_NOTHING,
)
    # stats update
    JM.factorizations[] += !jacobian(JM).factorized[]
    JM.ldivs[] += 1

    LA.ldiv!(x̂, jacobian(JM), b)
    if update == JAC_MONITOR_UPDATE_FERR || update == JAC_MONITOR_UPDATE_ALL
        forward_err!(JM, x̂, b, norm)
    end
    if update == JAC_MONITOR_UPDATE_COND || update == JAC_MONITOR_UPDATE_ALL
        cond!(JM)
    end
    x̂
end

#########################
## Hermite Normal Form ##
#########################

"""
    hnf(A, T=elytpe(A))

Compute the hermite normal form `H` of `A` by overwriting `A` and the corresponding transformation
matrix `U` using precision T. This is `A*U == H` and `H` is a lower triangular matrix.

The implementation follows the algorithm described in [1].

[1] Kannan, Ravindran, and Achim Bachem. "Polynomial algorithms for computing the Smith and Hermite normal forms of an integer matrix." SIAM Journal on Computing 8.4 (1979): 499-507.
"""
function hnf(A, T = eltype(A))
    H = similar(A, T)
    H .= A
    U = similar(A, T)
    hnf!(H, U)
    H, U
end

# use checked arithmethic
⊡(x, y) = Base.checked_mul(x, y)
⊞(x, y) = Base.checked_add(x, y)

"""
    hnf!(H, U, A)

Inplace version of [hnf](@ref).

    hnf!(A, U)

Inplace version of [hnf](@ref) overwriting `A` with `H`.
"""
hnf!(H, U, A) = hnf!(copyto!(H, A), U)
function hnf!(A, U)
    n = size(A, 1)
    U .= 0
    @inbounds for i = 1:n
        U[i, i] = one(eltype(U))
    end
    @inbounds for i = 1:(n-1)
        ii = i ⊞ 1
        for j = 1:i
            if !iszero(A[j, j]) || !iszero(A[j, ii])
                # 4.1
                r, p, q = gcdx(A[j, j], A[j, ii])
                # 4.2
                d_j = -A[j, ii] ÷ r
                d_ii = A[j, j] ÷ r
                for k = 1:n
                    a_kj, a_kii = A[k, j], A[k, ii]
                    A[k, j] = a_kj ⊡ p ⊞ a_kii ⊡ q
                    A[k, ii] = a_kj ⊡ d_j ⊞ a_kii ⊡ d_ii

                    u_kj, u_kii = U[k, j], U[k, ii]
                    U[k, j] = u_kj ⊡ p ⊞ u_kii ⊡ q
                    U[k, ii] = u_kj ⊡ d_j ⊞ u_kii ⊡ d_ii
                end
            end
            # 4.3
            if j > 1
                reduce_off_diagonal!(A, U, j)
            end
        end
        #5
        reduce_off_diagonal!(A, U, ii)
    end
    # Special case for 1 × 1 matrices to guarantee that the diagonal is positive
    # This comes up in polyhedral homotopy for univariate polynomials
    if n == 1
        U[1, 1] = flipsign(U[1, 1], A[1, 1])
        A[1, 1] = flipsign(A[1, 1], A[1, 1])
    end

    nothing
end

@inline function reduce_off_diagonal!(A, U, k)
    n = size(A, 1)
    @inbounds if A[k, k] < 0
        for i = 1:n
            A[i, k] = -A[i, k]
            U[i, k] = -U[i, k]
        end
    end
    @inbounds for z = 1:(k-1)
        if !iszero(A[z, z])
            r = -ceil(eltype(A), A[k, z] / A[z, z])
            for i = 1:n
                U[i, z] = U[i, z] ⊞ r ⊡ U[i, k]
                A[i, z] = A[i, z] ⊞ r ⊡ A[i, k]
            end
        end
    end
    nothing
end


include("euler/euler.jl")
using .Primes: prime_numbers

function compute(n::Int)::Int
    primes = prime_numbers(10 ^ n ÷ n)
    consecutive = 0
    i = 10 ^ (n - 1)
    while true
        if i ∈ primes
            consecutive = 0
        else
            divisors = 0
            t = i
            for prime ∈ primes
                if t % prime == 0
                    divisors += 1
                    t ÷= prime
                end
                if prime > t || t == 1 || divisors == n
                    break
                end
            end
            if divisors == n
                consecutive += 1
            else
                consecutive = 0
            end
        end
        if consecutive == n return i - n + 1 end
        i += 1
    end
end