export rec2coeff, showpoly

title_color =   :underline
name_color  =   :light_blue

function show(io::IO,m::Measure)
    printstyled("Measure dλ(t)=w(t)dt\n";color=title_color)
    print("name:\t\t")
    printstyled("$(m.name)\n"; color=name_color)
    println("w(t):\t\t$(m.w)")
    println("dom:\t\t($(m.dom[1]), $(m.dom[2]))")
    print("symmetric:\t")
    m.symmetric ? col=:green : col=:red
    printstyled("$(m.symmetric)\n";color=col)
    if typeof(m.pars)!=Dict{Any,Any}
        print("pars:")
        for (i,d) in enumerate(m.pars)
            printstyled("\t\t$d\n")
        end
    end
end

function show(io::IO,m::MultiMeasure)
    p = length(m.name)
    printstyled("$p-variate measure dλ(t)= ∏_{i=1}^$p w_i(t)dt_i\n";color=title_color)
    print("name:")
    [ printstyled("\t\t$(m.name[i])\n"; color=name_color) for i = 1:p]
    println("w(t):\t\t∏_{i=1}^$p w_i(t)")
    print("dom:")
    [ printstyled("\t\t($(m.dom[i][1]), $(m.dom[i][2]))\n") for i=1:p ]
    print("symmetric:")
    for i=1:p
        m.symmetric[i] ? col=:green : col=:red
        i==1 ? printstyled("\t$(m.symmetric[i])\n";color=col) : printstyled("\t\t$(m.symmetric[i])\n";color=col)
    end
    if length(findall(x->typeof(x)!=Dict{Any,Any},m.pars))>=1 print("pars:")
        for i=1:p
            print("\t\t")
            if typeof(m.pars[i])!=Dict{Any,Any}
                for (i,d) in enumerate(m.pars[i])
                    i==1 ? print("$d") : print(", $d")
                end
                print("\n")
            else
                print("∅\n")
            end
        end
    end
end

function show(io::IO,op::OrthoPoly;showmeasure::Bool=true)
    printstyled("\nUnivariate orthogonal polynomials\n";color=title_color)
    print("name:"); printstyled("\t\t$(op.name)\n";color=name_color)
    print("degree:"); printstyled("\t\t$(op.deg)\n")
    print("#coeffs:"); printstyled("\t$(length(op.α))\n")
    length(op.α)<=7 ? n=length(op.α) : n=7
    print("α ="); printstyled("\t\t$(op.α[1:n])")
    n<length(op.α) ? print("...\n") : print("\n")
    print("β ="); printstyled("\t\t$(op.β[1:n])")
    n<length(op.α) ? print("...\n") : print("\n")
    if showmeasure
        print("\n")
        show(io,op.meas)
    end
end

function show(io::IO,q::Quad;showmeasure::Bool=true)
    printstyled("\nQuadrature rule\n";color=title_color)
    print("name:"); printstyled("\t\t$(q.name)\n";color=name_color)
    print("N:"); printstyled("\t\t$(q.Nquad)\n")
    q.Nquad<=7 ? n=q.Nquad : n=7
    print("nodes"); printstyled("\t\t$(q.nodes[1:n])")
    n<q.Nquad ? print("...\n") : print("\n")
    print("weights"); printstyled("\t\t$(q.weights[1:n])")
    n<q.Nquad ? print("...\n") : print("\n")
    if showmeasure
        print("\n")
        show(io,q.meas)
    end
end

function show(io::IO,opq::OrthoPolyQ;showmeasure::Bool=false)
    show(io,opq.op)
    show(io,opq.quad;showmeasure=showmeasure)
end

function show(io::IO,mop::MultiOrthoPoly;showmeasure::Bool=false)
    p = length(mop.name)
    printstyled("\n$p-variate orthogonal polynomials\n";color=title_color)
    print("name:");
    [ printstyled("\t\t$(mop.name[i])\n";color=name_color) for i=1:p ]
    print("deg:"); printstyled("\t\t$(mop.deg)\n")
    print("dim:"); printstyled("\t\t$(mop.dim)\n")
    print("ind:");
    mop.dim<=7 ? n=mop.dim : n=7
    # printstyled("\t$(mop.ind[1,:])\n")
    [ printstyled("\t\t$(mop.ind[i,:])\n") for i=1:n ]
    n<mop.dim ? print("\t\t...\n") : print("\n")

    print("\n")
    showmeasure ? show(io,mop.meas) : ()
end

function show(io::IO,t::Tensor)
    printstyled("\n$(t.dim)-dimensional tensor\n";color=title_color)
    print("dim:"); printstyled("\t\t$(t.dim)\n")
    print("nonzeros:"); printstyled("\t$(length(t.T.nzind))\n")
    show(io,t.op;showmeasure=false)
end

"""
```
rec2coeff(deg::Int,a::Vector{Float64},b::Vector{Float64})
rec2coeff(a,b) = rec2coeff(length(a),a,b)
```
Get the coefficients of the orthogonal polynomial of degree up to `deg` specified by its
recurrence coefficients `(a,b)`. The function returns the values ``c_i^{(k)}`` from
```math
p_k (t) = t^d + \\sum_{i=0}^{k-1} c_i t^i,
```
where ``k`` runs from `1` to `deg`.

The call `rec2coeff(a,b)` outputs all possible recurrence coefficients given `(a,b)`.
"""
function rec2coeff(deg::Int,a::Vector{Float64},b::Vector{Float64})
    @assert deg >= 1 "degree must be positive (you asked for degree = $deg)"
    @assert length(a) == length(b) && length(a) >= deg "incorrect number of recurrence coefficients"
    c = Vector{Vector{Float64}}(undef,deg)
    @inbounds c[1] = [ -a[1] ]
    deg == 1 && return c
    @inbounds c[2] = [ a[1]*a[2] - b[2], -a[1] - a[2] ]
    deg == 2 && return c
    for k in 2:deg-1
        c[k+1] = [ -a[k+1]*c[k][1] - b[k+1]*c[k-1][1] ]
        for i = 1:k-2
            push!(c[k+1],c[k][i] - a[k+1]*c[k][i+1] - b[k+1]*c[k-1][i+1])
        end
        push!(c[k+1],c[k][k-1] - a[k+1]*c[k][k] - b[k+1])
        push!(c[k+1],c[k][k] - a[k+1])
    end
    return c
end
rec2coeff(α,β) = rec2coeff(length(α),α,β)

"""
```
showpoly(coeffs::Vector{Float64};sym::String,digits::Integer)
showpoly(coeffs::Vector{Vector{Float64}};sym::String,digits::Integer)
```
Show the monic polynomial with coefficients `coeffs` in a human readable way.
They keyword `sym` sets the name of the variable, and `digits` controls the number of shown digits.
```@repl
julia> showpoly([1.2, 2.3, 3.4456])
x^3 + 3.45x^2 + 2.3x + 1.2
julia> showpoly([1.2, 2.3, 3.4456], sym="t", digits=2)
t^3 + 3.45t^2 + 2.3t + 1.2
```

Tailored to types from `PolyChaos.jl`
```
showpoly(d::Integer,α::Vector{Float64},β::Vector{Float64};upto::Bool,sym::String,digits::Integer)
showpoly(d::Integer,op::OrthoPoly;upto::Bool=true,sym::String,digits::Integer)
showpoly(d::Integer,opq::OrthoPolyQ;upto::Bool=true,sym::String,digits::Integer)
```
Show the monic orthogonal polynomials with recurrence coefficients `(α,β)` up to degree `d`.
Setting the keyword `upto` to false prints the monic polynomial of degree equal to `d`.
```@repl
julia> op = OrthoPoly("gaussian",5);
julia> showpoly(3,op)
1
x
x^2 - 1.0
x^3 - 3.0x

julia> showpoly(3,op; upto=false)
x^3 - 3.0x

julia> showpoly(3,op; sym="t")
1
t
t^2 - 1.0
t^3 - 3.0t
```

The following function calls show all orthogonal polynomials given `(α,β)`.
```
showpoly(α::Vector{Float64},β::Vector{Float64};sym::String,digits::Integer)
showpoly(op::Union{OrthoPoly,OrthoPolyQ};sym::String,digits::Integer)
```

```@repl
julia> showpoly(op; sym="y")
1
y
y^2 - 1.0
y^3 - 3.0y
y^4 - 6.0y^2 + 3.0
y^5 - 10.0y^3 + 15.0y
y^6 - 15.0y^4 + 45.0y^2 - 15.0
```

[Thanks @pfitzseb for providing this functionality.](https://discourse.julialang.org/t/how-to-define-verbose-output-for-a-polynomial/21317/5)

"""
function showpoly(coeffs::Vector{Float64};sym::String="x",digits::Integer=2)
    io = Base.stdout
    length(coeffs) > 1 ? print(io, sym*"^", length(coeffs)) : print(io, sym)
    for (i, c) in enumerate(reverse(coeffs))
        abs(round(c, digits=digits)) == 0. && continue
        ex = length(coeffs) - i
        print(io, ' ', c > 0 ? '+' : '-', ' ')
        print(io, abs(round(c, digits=digits)))
        ex > 0 && print(io, sym)
        ex > 1 && print(io, '^', ex)
    end
    print(io, '\n')
end

function showpoly(coeffs::Vector{Vector{Float64}};sym::String="x",digits::Integer=2)
    print("1\n")
    map(c->showpoly(c,sym=sym,digits=digits),coeffs)
    print()
end

function showpoly(d::Integer,α::Vector{Float64},β::Vector{Float64};upto::Bool=true,sym::String="x",digits::Integer=2)
    @assert d >= 0 "degree has to be non-negative."
    d == 0 && return print("1")
    cfs = upto ? rec2coeff(d,α,β) : rec2coeff(d,α,β)[end]
    showpoly(cfs,sym=sym,digits=digits)
end

showpoly(d::Integer,op::OrthoPoly;upto::Bool=true,sym::String="x",digits::Integer=2) = showpoly(d,op.α,op.β;upto=upto,sym=sym,digits=digits)
showpoly(d::Integer,opq::OrthoPolyQ;upto::Bool=true,sym::String="x",digits::Integer=2) = showpoly(d,opq.op;upto=upto,sym=sym,digits=digits)

showpoly(α::Vector{Float64},β::Vector{Float64};sym::String="x",digits::Integer=2) = showpoly(length(α),α,β;sym=sym,digits=digits)
showpoly(op::Union{OrthoPoly,OrthoPolyQ};sym::String="x",digits::Integer=2) = showpoly(size(coeffs(op),1),op;sym=sym,digits=digits)


using JuMP, Gurobi, LinearAlgebra, Suppressor, StatsBase, Compat, Mosek, DataFrames, CSV, Random
include("convex_penalties_edm.jl")
include("exact_edm.jl")

results_template = DataFrame(n=Int[], r=Real[], t_sdp=Real[], err_sdpdtr=Real[], nocorr_sdp=Int[], t_ex=Real[], err_exdtr=Real[], nocorr_ex=Int[], MIPGap=Real[], nodeCountEX=Int[], t_oa=Real[], err_oadtr=Real[], nocorr_oa=Int[], MIPGapOA=Real[], nodeCountOA=Int[], cutCountOA=Int[])

results_edmradio_ex=similar(results_template, 0)

# Gram matrix recovery test problem
r=2
d_radio=0.3 # You will need to vary this to reproduce all of the results in Table 6
n=30  # You will need to vary this to reproduce all of the results in Table 6
lambda=1.0*n^2
gamma=1.0/n

for ARG in ARGS #1-4
    d_radio=[0.1, 0.2, 0.3, 0.4][parse(Int, ARG)]

    xs=rand(n, r).-0.5 #generate n random points in r dimensional space uniformly distributed on [-0.5, 0.5]^r
    D=zeros(n,n) #the matrix we would like to recover.
    for i=1:n #in a loop is slightly slower, but unimportant
        for j=1:(i-1)
            D[i,j]=max(0.0, norm(xs[i,:].-xs[j,:])^2+0.01*(randn(1)[1])) #Corrupt the distance with some noise, in a symmetric fashion
            D[j,i]=D[i,j]
        end
    end

    # Get observed set of coordinates
    coords=[repeat(1:n, inner=size(1:n,1)) repeat(1:n, outer=size(1:n,1))];
    entries=Int[]
    for i=1:n^2
        if norm(D[coords[i,1], coords[i,2]])<=d_radio^2
            push!(entries, i) #Remark: D_ij=D_ji (including noise), so this ensures that if D_ij is in the radio range, so is D_ji; no need to account for symmetry directly
        end
    end

    numEntriesSampled=size(entries,1)

    # Solve the sdp relaxation
    t_sdp=@elapsed G_imputed, Y_relax=getMatrixTraceNorm_EDM_Gram(coords, entries, D, n, numEntriesSampled, d_radio, lambda, gamma, r);
    u,v,=svd(Y_relax)
    Y_ws=u[:,1:r]*u[:,1:r]'

    D_imputed=diag(G_imputed)*ones(n)'+ones(n)*diag(G_imputed)'.-2*G_imputed
    @show err_sdpDTR=norm(D_imputed.-D)/norm(D)
    @show sdp_corr=sum(abs.(D.-D_imputed).<1e-2)
    @show sdp_incorr =sum(abs.(D.-D_imputed).>=1e-2)

    # Commented out as the submitted version of the paper only includes multi tree
    t_oa=@elapsed G_imputed,nodecountOA, cutCountOA, miogapOA,=getMatrixFrobeniusNorm_cuttingplanesmultitree_exact_inout(coords, entries, D, n, numEntriesSampled, lambda, d_radio, r, gamma, Y_ws, Y_relax);
    D_imputed=diag(G_imputed)*ones(n)'+ones(n)*diag(G_imputed)'.-2*G_imputed
    @show err_oaDTR=norm(D_imputed.-D)/norm(D)
    @show oa_corr=sum(abs.(D.-D_imputed).<1e-2)
    @show oa_incorr =sum(abs.(D.-D_imputed).>=1e-2)

    # Commented out as the submitted version of the paper only includes multi tree
    t_oa=-1.0#@elapsed G_imputed,nodecountOA, cutCountOA, miogapOA,=getMatrixFrobeniusNorm_cuttingplanes_exact_inout(coords, entries, D, n, numEntriesSampled, lambda, d_radio, r, gamma, Y_ws, Y_relax);
    # D_imputed=diag(G_imputed)*ones(n)'+ones(n)*diag(G_imputed)'.-2*G_imputed
    @show err_oaDTR=-1.0#norm(D_imputed.-D)/norm(D)
    @show oa_corr=-1.0#sum(abs.(D.-D_imputed).<1e-2)
    @show oa_incorr =-1.0#sum(abs.(D.-D_imputed).>=1e-2)
    miogapOA=-1.0; nodeCountOA=-1; cutCountOA=-1

    t_ex=-1.0@elapsed G_imputed,miogap,nodeCountEX=getMatrixExact_EDM_Gram(coords, entries, D, n, r, numEntriesSampled, d_radio, lambda, gamma, D_imputed);
    # D_imputed=diag(G_imputed)*ones(n)'+ones(n)*diag(G_imputed)'.-2*G_imputed
    @show err_exDTR=-1.0#norm(D_imputed.-D)/norm(D)
    @show ex_corr=-1.0#sum(abs.(D.-D_imputed).<1e-2)
    @show ex_incorr =-1.0#sum(abs.(D.-D_imputed).>=1e-2)
    miogap=-1.0; nodeCountEX=-1

    push!(results_edmradio_ex, [n, r, t_sdp, err_sdpDTR, sdp_corr, t_ex, err_exDTR, ex_corr, miogap, nodeCountEX, t_oa, err_oaDTR, oa_corr, miogapOA, nodecountOA, cutCountOA])
end
CSV.write("oavsgurobi_edmgram.csv", results_edmradio_ex, append=true)

Euclidean Distance Matrix Completion/createTable6.jl


"""
    gbₕ_free(Tₐ,Tₗ,d,Dₕ₀)
    gbₕ_free(Tₐ,Tₗ,d)

Leaf boundary layer conductance for heat under **free** convection (m s-1).

# Arguments

- `Tₐ` (°C): air temperature
- `Tₗ` (°C): leaf temperature
- `d` (m): characteristic dimension, *e.g.* leaf width (see eq. 10.9 from Monteith and Unsworth, 2013).
- `Dₕ₀ = 21.5e-6`: molecular diffusivity for heat at base temperature. Use value from
[`Constants`](@ref) if not provided.

# Note

`R` and `Dₕ₀` can be found using [`Constants`](@ref). To transform in ``mol\\ m^{-2}\\ s^{-1}``,
use [`ms_to_mol`](@ref).

# References

Leuning, R., F. M. Kelliher, DGG de Pury, et E.-D. SCHULZE. 1995. « Leaf nitrogen,
photosynthesis, conductance and transpiration: scaling from leaves to canopies ». Plant,
Cell & Environment 18 (10): 1183‑1200.

Monteith, John, et Mike Unsworth. 2013. Principles of environmental physics: plants,
animals, and the atmosphere. Academic Press. Paragraph 10.1.3, eq. 10.9.
"""
function gbₕ_free(Tₐ,Tₗ,d,Dₕ₀ = Constants().Dₕ₀)
    zeroT = zero(Tₐ) # make it type stable

    if abs(Tₗ-Tₐ) > zeroT
        Gr = 1.58e8 * d^3.0 * abs(Tₗ-Tₐ) # Grashof number (Monteith and Unsworth, 2013)
        # !Note: Leuning et al. (1995) use 1.6e8 (eq. E4).
        # Leuning et al. (1995) eq. E3:
        Gbₕ_free = 0.5 * get_Dₕ(Tₐ,Dₕ₀) * (Gr^0.25) / d
    else
        Gbₕ_free = zeroT
    end

    return Gbₕ_free
end

"""
    gbₕ_forced(Wind,d)

Boundary layer conductance for heat under **forced** convection (m s-1). See eq. E1 from
Leuning et al. (1995) for more details.

# Arguments

- `Wind` (m s-1): wind speed
- `d` (m): characteristic dimension, *e.g.* leaf width (see eq. 10.9 from Monteith and Unsworth, 2013).

# Notes

`d` is the minimal dimension of the surface of an object in contact with the air.

# References

Leuning, R., F. M. Kelliher, DGG de Pury, et E.-D. SCHULZE. 1995. « Leaf nitrogen,
photosynthesis, conductance and transpiration: scaling from leaves to canopies ». Plant,
Cell & Environment 18 (10): 1183‑1200.
"""
function gbₕ_forced(Wind,d)
    0.003 * sqrt(Wind / d)
end

"""
    get_Dₕ(T,Dₕ₀)
    get_Dₕ(T)

Dₕ -molecular diffusivity for heat at base temperature- from Dₕ₀ (corrected by temperature).
See Monteith and Unsworth (2013, eq. 3.10).

# Arguments

- `Tₐ` (°C): temperature
- `Dₕ₀`: molecular diffusivity for heat at base temperature. Use value from [`Constants`](@ref)
if not provided.

# References

Monteith, John, et Mike Unsworth. 2013. Principles of environmental physics: plants,
animals, and the atmosphere. Academic Press. Paragraph 10.1.3.
"""
function get_Dₕ(T,Dₕ₀ = Constants().Dₕ₀)
    Dₕ₀ * (1 + 0.007*T)
end


densities_T = [25.0, 38.0, 1.0] # TCR, pMHC, Ac
densities_TC1 = [15.0, 38.0, 10.0, 1.0] # TCR, pMHC, CD4, Ac
densities_TC2 = [8.0,  38.0, 17.0, 1.0]
densities_TC3 = [3.0,  38.0, 22.0, 1.0]

rates_T = [-3.37638358011457
            -0.302013187183543]
rates_TC1 = [-1.72689171836867
            -0.646695212785203]
rates_TC2 = [-0.644293816490510
            0.617796010276307]
rates_TC3 = [0.891017383250648
            2.16934160248305]

sd_data = CSV.read("C:/Users/stravaglino3/Downloads/julia_code/n_stds_for_julia.csv")
nT_sds = sd_data[:,1]
nTC1_sds = Array(sd_data[:,2:3])
nTC2_sds = Array(sd_data[:,4:5])
nTC3_sds = Array(sd_data[:,6:7])

scatter(n_se_tcr[:,1],n_se_tcr[:,2],yerror=n_se_tcr[:,3],markercolor=:black,markersize=5)
scatter!(n_se_1[:,1],n_se_1[:,2],yerror=n_se_1[:,3],markercolor=:blue,markersize=5)
scatter!(n_se_2[:,1],n_se_2[:,2],yerror=n_se_2[:,3],markercolor=:red,markersize=5)
h = scatter!(n_se_3[:,1],n_se_3[:,2],yerror=n_se_3[:,3],markercolor=:cyan,markersize=5,legend=false)

tspan = (0.0,16.0)
t = collect(0.01:0.01:16)

gr()
# bimolecular TCR-pMHC bonds
u₀ = [0.0]
f = (du,u,p,t) -> bimolecular!(du,u,p,t,densities_T,[])
n_T = Array(solve_model(rates_T))'[2:end]
n_T_sd = n_T .± nT_sds

# TCR-pMHC-CD4 (1)
u₀ = [0.0, 0.0]
f = (du,u,p,t) -> trimolecular!(du,u,p,t,densities_TC1,10 .^ rates_T)
n_TC1 = Array(solve_model(rates_TC1))'[2:end,:]
n_TC1_sd = sum(n_TC1 .± nTC1_sds, dims=2)

# TCR-pMHC-CD4 (2)
f = (du,u,p,t) -> trimolecular!(du,u,p,t,densities_TC2,10 .^ rates_T)
n_TC2 = Array(solve_model(rates_TC2))'[2:end,:]
n_TC2_sd = sum(n_TC2 .± nTC2_sds, dims=2)

# TCR-pMHC-CD4 (3)
f = (du,u,p,t) -> trimolecular!(du,u,p,t,densities_TC3,10 .^ rates_T)
n_TC3 = Array(solve_model(rates_TC3))'[2:end,:]
n_TC3_sd = sum(n_TC3 .± nTC3_sds, dims=2)

plot!(t, n_T_sd, errorbars=:ribbon, alpha=0.1,msc=:black)
plot!(t, n_T, lw=3, lc=:black)
plot!(t,n_TC1_sd, errorbars=:ribbon, alpha=0.1, msc=:blue)
plot!(t, sum(n_TC1,dims=2), lw=3, lc=:blue)
plot!(t,n_TC2_sd, errorbars=:ribbon, alpha=0.1, msc=:red)
plot!(t, sum(n_TC2,dims=2), lw=3, lc=:red)
plot!(t,n_TC3_sd, errorbars=:ribbon, alpha=0.1, msc=:cyan)
plot!(t, sum(n_TC3,dims=2), lw=3, lc=:cyan)

xlabel!("Contact Time tc")
ylabel!("<n>")

#------------------------------------------------------------------------------
# CD4 fitting

cd4_1 = CSV.read("cd4_data_from_muaz.csv")
t_cd41, Pa_cd41, n_se_cd41 = clean_data(cd4_1)
n_cd41 = Pa2n(Pa_cd41)

f = (du,u,p,t) -> bimolecular!(du,u,p,t,dens,cons) # enclose constants

u₀ = [0.0]          # Initial condition
tspan = (0.0,8.0)  # Simulation interval
pinit = [1e-4,0.5]      # equation parameter. p = [k₁, k₋₁]
densities_C = [1700., 850., 1.] #CD4, pMHC, A
cons = []

rates_C,hes_C,l_C,boot_C = optimization(bimolecular!,n_cd41,densities_C',[],model_opt2,t_cd41)

# Setup the ODE problem, then solve
# plotly()
t = collect(0.01:0.01:16)
tspan = (0.0,16.0)
f = (du,u,p,t) -> bimolecular!(du,u,p,t,densities_C,cons)
prob = ODEProblem(f,u₀,tspan,10 .^ rates_C)
sol_C = Array(solve(prob,Vern7(),saveat=t,abstol=1e-10,reltol=1e-10))'
plot(t,sol_C,lc=:orange,lw=3)
scatter!(n_se_cd41[:,1],n_se_cd41[:,2],yerror=n_se_cd41[:,3],markercolor=:orange)

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

densities_T =   [25.0, 38.0, 1.0]       # TCR, pMHC, Ac
densities_TC1 = [15.0, 38.0, 10.0, 1.0] # TCR, pMHC, CD4, Ac
densities_TC2 = [8.0,  38.0, 17.0, 1.0]
densities_TC3 = [3.0,  38.0, 22.0, 1.0]
densities_C =   [1700., 850., 1.]       # CD4, pMHC, A

n_T_norm = n_T ./25. .*15.
n_TC1_norm = sum(n_TC1,dims=2)
n_TC2_norm = sum(n_TC2,dims=2) ./8. .*15.# ./17. .*10.
n_TC3_norm = sum(n_TC3,dims=2) ./3. .*15.# ./22. .*10.
n_C_norm = sol_C ./1700. .*10. ./850. .*38

plotly()
t = collect(0.01:0.01:16)
plot(t,n_T_norm,lw=3,lc=:black,label="TCR")
plot!(t,n_TC1_norm,lw=3,lc=:blue,label="TCR 15, CD4 10")
plot!(t,n_TC2_norm,lw=3,lc=:red,label="TCR 8, CD4 17")
plot!(t,n_TC3_norm,lw=3,lc=:cyan,label="TCR 3, CD4 22")
plot!(t,n_C_norm,lw=3,lc=:orange,yscale=:linear,label="CD4",legend=:bottomright)
# yaxis!("log<n>")
yaxis!("<n>")
xlabel!("Contact Time tc")
# savefig("n_log_scaled")

coop_indx_TC1 = cooperativity_index.(n_T_norm,n_C_norm,n_TC1_norm)
coop_indx_TC2 = cooperativity_index.(n_T_norm,n_C_norm,n_TC2_norm)
coop_indx_TC3 = cooperativity_index.(n_T_norm,n_C_norm,n_TC3_norm)
plot(t,coop_indx_TC1 .* 100,lw=3,lc=:blue,label="Coop TCR 15, CD4 10")
plot!(t,coop_indx_TC2 .* 100,lw=3,lc=:red,label="Coop TCR 8, CD4 17")
plot!(t,coop_indx_TC3 .* 100,lw=3,lc=:cyan,label="Coop TCR 3, CD4 22")
xlabel!("Contact Time tc")
ylabel!("Cooperative Index (%)")


using SimpleSATSolver
using SatisfiabilityInterface
using Test

function is_sat(p::SATProblem)
    status, results = solve(p, SimpleSAT())

    return status == :sat && satisfies(p, results)
end

function is_unsat(p::SATProblem)
    status, results = solve(p, SimpleSAT())

    return status == :unsat
end

is_sat(v::Vector{Vector{Int}}) = is_sat(SATProblem(v))
is_unsat(v::Vector{Vector{Int}}) = is_unsat(SATProblem(v))

@testset "SimpleSATSolver.jl" begin
    @test is_sat([ [2] ])
    @test is_sat([ [1, -5, 4], [-1, 5, 3, 4], [-3, -4] ])
    @test is_sat([[1, 2], [-1, 3], [-3, -2], [-1, 3], [-1, -3]])

    @test is_unsat([ [1], [-1] ])

    @test is_sat([ [1, 2], [-1, 3], [-3, -2], [-1, 3], [-1, -3] ])  # forces backtracking

    # 2-colouring of linear graph with 3 vertices:
    @test is_sat( [[1, 2], [-1, -2], [3, 4], [-3, -4], [5, 6], [-5, -6], [-1, -3], [-2, -4], [-3, -5], [-4, -6]] )

    # 3-colouring of ring with 11 vertices:
    @test is_sat( [[1, 2, 3], [-1, -2], [-1, -3], [-2, -3], [4, 5, 6], [-4, -5], [-4, -6], [-5, -6], [7, 8, 9], [-7, -8], [-7, -9], [-8, -9], [10, 11, 12], [-10, -11], [-10, -12], [-11, -12], [13, 14, 15], [-13, -14], [-13, -15], [-14, -15], [16, 17, 18], [-16, -17], [-16, -18], [-17, -18], [19, 20, 21], [-19, -20], [-19, -21], [-20, -21], [22, 23, 24], [-22, -23], [-22, -24], [-23, -24], [25, 26, 27], [-25, -26], [-25, -27], [-26, -27], [28, 29, 30], [-28, -29], [-28, -30], [-29, -30], [31, 32, 33], [-31, -32], [-31, -33], [-32, -33], [-1, -4], [-2, -5], [-3, -6], [-4, -7], [-5, -8], [-6, -9], [-7, -10], [-8, -11], [-9, -12], [-10, -13], [-11, -14], [-12, -15], [-13, -16], [-14, -17], [-15, -18], [-16, -19], [-17, -20], [-18, -21], [-19, -22], [-20, -23], [-21, -24], [-22, -25], [-23, -26], [-24, -27], [-25, -28], [-26, -29], [-27, -30], [-28, -31], [-29, -32], [-30, -33], [-31, -1], [-32, -2], [-33, -3]])

    @testset "Load CNF file" begin
        p = read_cnf("factorisation.cnf")
        @test is_sat(p)
    end

end


import GreyDecision.Random
import GreyDecision.Utility
import GreyDecision.GreyNumbers: GreyNumber

@testset "Random Grey Numbers" begin

    @testset "Single Random GreyNumber" begin
        g = rand(GreyNumber)
        @test first(g) >= 0.0
        @test last(g) <= 1.0
    end

    @testset "Typed Random GreyNumber" begin
        g = rand(GreyNumber{Float64})
        @test first(g) >= 0.0
        @test last(g) <= 1.0

        h = rand(GreyNumber{UInt8})
        @test first(h) < 256
        @test last(h) >= 0
    end

    @testset "Array of random grey numbers" begin
        v = rand(GreyNumber{Float64}, 10)

        @test length(v) == 10
        @test typeof(v[1]) == GreyNumber{Float64}
        @test first(v[1]) >= 0.0
        @test last(v[1]) <= 1.0
    end

    @testset "Matrix of random grey numbers" begin
        v = rand(GreyNumber{Float64}, 10, 20)

        @test size(v) == (10, 20)
        @test typeof(v[1, 1]) == GreyNumber{Float64}
        @test first(v[1, 1]) >= 0.0
        @test last(v[1, 2]) <= 1.0
        @test first(v[2, 1]) >= 0.0
        @test last(v[2, 2]) <= 1.0
    end

end


using Flux.Losses: crossentropy, binarycrossentropy, logitbinarycrossentropy

@testset "Losses" begin

x = [1.,2.,3.]
cx = gpu(x)
@test crossentropy(x,x) ≈ crossentropy(cx,cx)
@test crossentropy(x,x, agg=identity) ≈ crossentropy(cx,cx, agg=identity) |> cpu
@test crossentropy(x,x, agg=x->mean([1.0;2.0;3.0].*x)) ≈ crossentropy(cx,cx, agg=x->mean(gpu([1.0;2.0;3.0]).*x))

x = [-1.1491, 0.8619, 0.3127]
y = [1, 1, 0.]
@test binarycrossentropy(σ.(x), y) ≈ binarycrossentropy(gpu(σ.(x)), gpu(y))
@test logitbinarycrossentropy(x, y) ≈ logitbinarycrossentropy(gpu(x), gpu(y))

@testset "GPU grad tests" begin
  x = rand(Float32, 3,3)
  y = rand(Float32, 3,3)

  for loss in ALL_LOSSES
    gpu_gradtest(loss, x, y)
  end
end

end #testset


@testset "Multiplication" begin
    v = [1,2,3]
    tv = v.'
    mat = diagm([1,2,3])

    @test (tv*v) === 14
    @test (tv*mat)::TransposedVector == [1 4 9]
    @test_throws Exception [1]*reshape([1],(1,1)) # no longer permitted
    @test_throws Exception tv*tv
    @test (v*tv)::Matrix == [1 2 3; 2 4 6; 3 6 9]
    @test_throws Exception v*v # Was previously a missing method error, now an error message
    @test_throws Exception mat*tv

    @test_throws Exception tv*v.'
    @test (tv*mat.')::TransposedVector == [1 4 9]
    @test_throws Exception [1]*reshape([1],(1,1)).' # no longer permitted
    @test tv*tv.' === 14
    @test_throws Exception v*tv.'
    @test (v*v.')::Matrix == [1 2 3; 2 4 6; 3 6 9]
    @test (mat*tv.')::Vector == [1,4,9]

    @test (tv.'*v.')::Matrix == [1 2 3; 2 4 6; 3 6 9]
    @test_throws Exception tv.'*mat.'
    @test (v.'*mat.')::TransposedVector == [1 4 9]
    @test_throws Exception tv.'*tv.'
    @test v.'*tv.' === 14
    @test_throws Exception v.'*v.'
    @test (mat.'*tv.')::Vector == [1,4,9]

    @test_throws Exception tv.'*v
    @test_throws Exception tv.'*mat
    @test (v.'*mat)::TransposedVector == [1 4 9]
    @test (tv.'*tv)::Matrix == [1 2 3; 2 4 6; 3 6 9]
    @test_throws Exception v.'*tv
    @test v.'*v === 14
    @test_throws Exception mat.'*tv

    z = [1+im,2,3]
    cz = z'
    mat = diagm([1+im,2,3])

    @test cz*z === 15 + 0im

    @test_throws Exception cz*z'
    @test (cz*mat')::TransposedVector == [-2im 4 9]
    @test_throws Exception [1]*reshape([1],(1,1))' # no longer permitted
    @test cz*cz' === 15 + 0im
    @test_throws Exception z*vz'
    @test (z*z')::Matrix == [2 2+2im 3+3im; 2-2im 4 6; 3-3im 6 9]
    @test (mat*cz')::Vector == [2im,4,9]

    @test (cz'*z')::Matrix == [2 2+2im 3+3im; 2-2im 4 6; 3-3im 6 9]
    @test_throws Exception cz'*mat'
    @test (z'*mat')::TransposedVector == [-2im 4 9]
    @test_throws Exception cz'*cz'
    @test z'*cz' === 15 + 0im
    @test_throws Exception z'*z'
    @test (mat'*cz')::Vector == [2,4,9]

    @test_throws Exception cz'*z
    @test_throws Exception cz'*mat
    @test (z'*mat)::TransposedVector == [2 4 9]
    @test (cz'*cz)::Matrix == [2 2+2im 3+3im; 2-2im 4 6; 3-3im 6 9]
    @test_throws Exception z'*cz
    @test z'*z === 15 + 0im
    @test_throws Exception mat'*cz

end


module EnsembleKalmanInversions

using Distributions
using OffsetArrays
using ProgressBars
using Random
using Printf
using LinearAlgebra
using Suppressor: @suppress
using Statistics
using EnsembleKalmanProcesses.EnsembleKalmanProcessModule
using EnsembleKalmanProcesses.ParameterDistributionStorage

using EnsembleKalmanProcesses.EnsembleKalmanProcessModule: sample_distribution

using ..InverseProblems: n_ensemble, observation_map, forward_map, tupify_parameters

#####
##### Priors
#####

function lognormal_with_mean_std(mean, std)
    k = std^2 / mean^2 + 1
    μ = log(mean / sqrt(k))
    σ = sqrt(log(k))
    return LogNormal(μ, σ)
end

struct ConstrainedNormal{FT}
    # θ is the original constrained paramter, θ̃ is the unconstrained parameter ~ N(μ, σ)
    # θ = lower_bound + (upper_bound - lower_bound)/（1 + exp(θ̃)）
    μ :: FT
    σ :: FT
    lower_bound :: FT
    upper_bound :: FT
end

# Scaling factor to give the parameter a magnitude of one
sf(prior) = 1 / abs(prior.μ)

# Model priors are sometimes constrained; EKI deals with unconstrained, Normal priors.
convert_prior(prior::LogNormal) = Normal(sf(prior) * prior.μ, sf(prior) * prior.σ)
convert_prior(prior::Normal) = Normal(sf(prior) * prior.μ, sf(prior) * prior.σ)
convert_prior(prior::ConstrainedNormal) = Normal(prior.μ, prior.σ)

# Convert parameters to unconstrained for EKI
forward_parameter_transform(prior::LogNormal, parameter) = log(parameter^sf(prior))
forward_parameter_transform(prior::Normal, parameter) = parameter * sf(prior)
forward_parameter_transform(cn::ConstrainedNormal, parameter) =
    log((cn.upper_bound - parameter) / (cn.upper_bound - cn.lower_bound))

# Convert parameters from unconstrained (EKI) to constrained
inverse_parameter_transform(prior::LogNormal, parameter) = exp(parameter / sf(prior))
inverse_parameter_transform(prior::Normal, parameter) = parameter / sf(prior)
inverse_parameter_transform(cn::ConstrainedNormal, parameter) =
    cn.lower_bound + (cn.upper_bound - cn.lower_bound) / (1 + exp(parameter))

# Convenience vectorized version
inverse_parameter_transform(priors::NamedTuple, parameters::Vector) =
    NamedTuple(name => inverse_parameter_transform(priors[name], parameters[i])
               for (i, name) in enumerate(keys(priors)))

# Convert covariance from unconstrained (EKI) to constrained
inverse_covariance_transform(::Tuple{Vararg{LogNormal}}, parameters, covariance) =
    Diagonal(exp.(parameters)) * covariance * Diagonal(exp.(parameters))

inverse_covariance_transform(::Tuple{Vararg{Normal}}, parameters, covariance) = covariance

function inverse_covariance_transform(cn::Tuple{Vararg{ConstrainedNormal}}, parameters, covariance)
    upper_bound = [cn[i].upper_bound for i = 1:length(cn)]
    lower_bound = [cn[i].lower_bound for i = 1:length(cn)]
    dT = Diagonal(@. -(upper_bound - lower_bound) * exp(parameters) / (1 + exp(parameters))^2)
    return dT * covariance * dT'
end

mutable struct EnsembleKalmanInversion{I, P, E, M, O, F, S, R}
    inverse_problem :: I
    parameter_distribution :: P
    ensemble_kalman_process :: E
    mapped_observations :: M
    noise_covariance :: O
    inverting_forward_map :: F
    iteration :: Int
    iteration_summaries :: S
    resampler :: R
end

"""
    parameter_ensemble(eki::EnsembleKalmanInversion)

Return a `Vector` of parameter sets (in physical / constrained space) for each ensemble member.
"""
function parameter_ensemble(eki::EnsembleKalmanInversion)
    priors = eki.inverse_problem.free_parameters.priors
    return parameter_ensemble(eki.ensemble_kalman_process, priors)
end

function parameter_ensemble(ensemble_kalman_process, priors)
    unconstrained_parameters = get_u_final(ensemble_kalman_process) # (N_params, ensemble_size) array
    ensemble_size = size(unconstrained_parameters, 2)
    return [inverse_parameter_transform(priors, unconstrained_parameters[:, n]) for n = 1:ensemble_size]
end

Base.show(io::IO, eki::EnsembleKalmanInversion) =
    print(io, "EnsembleKalmanInversion", '\n',
              "├── inverse_problem: ", typeof(eki.inverse_problem).name.wrapper, '\n',
              "├── parameter_distribution: ", typeof(eki.parameter_distribution).name.wrapper, '\n',
              "├── ensemble_kalman_process: ", typeof(eki.ensemble_kalman_process), '\n',
              "├── mapped_observations: ", summary(eki.mapped_observations), '\n',
              "├── noise_covariance: ", summary(eki.noise_covariance), '\n',
              "├── inverting_forward_map: ", typeof(eki.inverting_forward_map).name.wrapper, '\n',
              "├── iteration: $(eki.iteration)", '\n',
              "└── resampler: $(typeof(eki.resampler))")

construct_noise_covariance(noise_covariance::AbstractMatrix, y) = noise_covariance

function construct_noise_covariance(noise_covariance::Number, y)
    # Independent noise for synthetic observations
    n_obs = length(y)
    return noise_covariance * Matrix(I, n_obs, n_obs)
end

"""
    EnsembleKalmanInversion(inverse_problem; noise_covariance=1e-2)

 Return an object that interfaces with [EnsembleKalmanProcesses.jl](https://github.com/CliMA/EnsembleKalmanProcesses.jl)
and uses Ensemble Kalman Inversion to iteratively "solve" the inverse problem:

```math
y = G(θ) + η,
```

for the parameters ``θ``, where ``y`` is a "normalized" vector of observations,
``G(θ)`` is a forward map that predicts the observations, and ``η ∼ N(0, Γ_y)`` is zero-mean
random noise with covariance matrix ``Γ_y`` representing uncertainty in the observations.

By "solve", we mean that the iteration finds the parameter values ``θ`` that minimizes the
distance between ``y`` and ``G(θ)``.

The "forward map output" `G` can have many interpretations. The specific statistics that `G` computes
have to be selected for each use case to provide a concise summary of the complex model solution that
contains the values that we would most like to match to the corresponding truth values `y`. For example,
in the context of an ocean-surface boundary layer parametrization, this summary could be a vector of
concatenated `u`, `v`, `b`, `e` profiles at all or some time steps of the CATKE solution.

(For more details on the Ensemble Kalman Inversion algorithm refer to the
[EnsembleKalmanProcesses.jl Documentation](https://clima.github.io/EnsembleKalmanProcesses.jl/stable/ensemble_kalman_inversion/).)

Arguments
=========

- `inverse_problem :: InverseProblem`: Represents an inverse problem representing the comparison between
                                       synthetic observations generated by
                                       [Oceananigans.jl](https://clima.github.io/OceananigansDocumentation/stable/)
                                       and model predictions, also generated by Oceananigans.jl.

- `noise_covariance` (`AbstractMatrix` or `Number`): normalized covariance representing observational
                                                     uncertainty. If `noise_covariance isa Number` then
                                                     it's converted to an identity matrix scaled by
                                                     `noise_covariance`.

- `resampler`: controls resampling procedure when the forward map contains NaNs. See `NaNResampler`.
"""
function EnsembleKalmanInversion(inverse_problem; noise_covariance=1e-2, resampler=NaNResampler())

    free_parameters = inverse_problem.free_parameters
    original_priors = free_parameters.priors

    transformed_priors = [Parameterized(convert_prior(prior)) for prior in original_priors]
    no_constraints = [[no_constraint()] for _ in transformed_priors]

    parameter_distribution = ParameterDistribution(transformed_priors,
                                                   no_constraints,
                                                   collect(string.(free_parameters.names)))

    ek_process = Inversion()
    initial_ensemble = sample_distribution(parameter_distribution, n_ensemble(inverse_problem))

    # Build EKP-friendly observations "y" and the covariance matrix of observational uncertainty "Γy"
    y = dropdims(observation_map(inverse_problem), dims=2) # length(forward_map_output) column vector
    Γy = construct_noise_covariance(noise_covariance, y)

    # The closure G(θ) maps (N_params, ensemble_size) array to (length(forward_map_output), ensemble_size)
    function inverting_forward_map(θ)
        θ = parameter_ensemble(ensemble_kalman_process, original_priors)

        return forward_map(inverse_problem, θ)
    end

    ensemble_kalman_process = EnsembleKalmanProcess(initial_ensemble, y, Γy, ek_process)

    eki = EnsembleKalmanInversion(inverse_problem,
                                  parameter_distribution,
                                  ensemble_kalman_process,
                                  y,
                                  Γy,
                                  inverting_forward_map,
                                  0,
                                  OffsetArray([], -1),
                                  resampler)

    return eki
end

"""
    UnscentedKalmanInversion(inverse_problem, prior_mean, prior_cov;
                             noise_covariance = 1e-2, α_reg = 1, update_freq = 0)

Return an object that interfaces with [EnsembleKalmanProcesses.jl](https://github.com/CliMA/EnsembleKalmanProcesses.jl)
and uses Unscented Kalman Inversion to iteratively "solve" the inverse problem:

```math
y = G(θ) + η,
```

for the parameters ``θ``, where ``y`` is a "normalized" vector of observations,
``G(θ)`` is a forward map that predicts the observations, and ``η ∼ N(0, Γ_y)`` is zero-mean
random noise with covariance matrix ``Γ_y`` representing uncertainty in the observations.

By "solve", we mean that the iteration finds the parameter values ``θ`` that minimizes the
distance between ``y`` and ``G(θ)``.

(For more details on the Unscented Kalman Inversion algorithm refer to the
[EnsembleKalmanProcesses.jl Documentation](https://clima.github.io/EnsembleKalmanProcesses.jl/stable/unscented_kalman_inversion/).)

Arguments
=========

- `inverse_problem :: InverseProblem`: an inverse problem representing the comparison between
                                       synthetic observations generated by
                                       [Oceananigans.jl](https://clima.github.io/OceananigansDocumentation/stable/)
                                       and model predictions also generated by Oceananigans.jl.

- `prior_mean :: Vector{Float64}`: prior mean

- `prior_cov :: Matrix{Float64}`: prior covariance

- `noise_covariance :: Float64`: observation error covariance

- `α_reg :: Float64`: regularization parameter toward the prior mean (0 < `α_reg` ≤ 1);
                      default `α_reg=1` implies no regularization

- `update_freq :: IT`: set to 0 when the inverse problem is not identifiable (default), namely the
                       inverse problem has multiple solutions, the covariance matrix will represent
                       only the sensitivity of the parameters, instead of posterior covariance information;
                       set to 1 (or anything > 0) when the inverse problem is identifiable, and
                       the covariance matrix will converge to a good approximation of the
                       posterior covariance with an uninformative prior
"""
function UnscentedKalmanInversion(inverse_problem, prior_mean, prior_cov;
                                  noise_covariance = 1e-2, α_reg = 1, update_freq = 0, resampler = nothing)

    free_parameters = inverse_problem.free_parameters
    original_priors = free_parameters.priors

    transformed_priors = [Parameterized(convert_prior(prior)) for prior in original_priors]
    no_constraints = [[no_constraint()] for _ in transformed_priors]
    parameter_distribution = ParameterDistribution(transformed_priors, no_constraints, collect(string.(free_parameters.names)))

    # Build EKP-friendly observations "y" and the covariance matrix of observational uncertainty "Γy"
    y = dropdims(observation_map(inverse_problem), dims=2) # length(forward_map_output) column vector
    Γy = construct_noise_covariance(noise_covariance, y)

    # The closure G(θ) maps (N_params, ensemble_size) array to (length(forward_map_output), ensemble_size)
    function G(θ)
        batch_size = size(θ, 2)
        inverted_parameters = [inverse_parameter_transform.(values(original_priors), θ[:, i]) for i = 1:batch_size]
        return forward_map(inverse_problem, inverted_parameters)
    end

    ensemble_kalman_process = EnsembleKalmanProcess(y, Γy, Unscented(prior_mean, prior_cov, α_reg, update_freq))

    eki = EnsembleKalmanInversion(inverse_problem,
                                  parameter_distribution,
                                  ensemble_kalman_process,
                                  y,
                                  Γy,
                                  G,
                                  0,
                                  OffsetArray([], -1),
                                  resampler)

    return eki
end

"""
    UnscentedKalmanInversionPostprocess(eki)

Returns
=======

- `mean :: Matrix{Float64}`: `N_iterations` × `N_parameters` mean matrix
- `cov :: Vector{Matrix{Float64}}`: `N_iterations` vector of `N_parameters` × `N_parameters` covariance matrix
- `std :: Matrix{Float64}`: `N_iterations` × `N_parameters` standard deviation matrix
- `err :: Vector{Float64}`: `N_iterations` error array
"""
function UnscentedKalmanInversionPostprocess(eki)
    original_priors = eki.inverse_problem.free_parameters.priors
    θ_mean_raw = hcat(eki.ensemble_kalman_process.process.u_mean...)
    θθ_cov_raw = eki.ensemble_kalman_process.process.uu_cov

    θ_mean = similar(θ_mean_raw)
    θθ_cov = similar(θθ_cov_raw)
    θθ_std_arr = similar(θ_mean_raw)

    for i = 1:size(θ_mean, 2) # number of iterations
        θ_mean[:, i] = inverse_parameter_transform.(values(original_priors), θ_mean_raw[:, i])
        θθ_cov[i] = inverse_covariance_transform(values(original_priors), θ_mean_raw[:, i], θθ_cov_raw[i])

        for j = 1:size(θ_mean, 1) # number of parameters
            θθ_std_arr[j, i] = sqrt(θθ_cov[i][j, j])
        end
    end

    return θ_mean, θθ_cov, θθ_std_arr, eki.ensemble_kalman_process.err
end

"""
    struct IterationSummary{P, M, C, V, E}

Container with information about each iteration of the Ensemble Kalman Process.
"""
struct IterationSummary{P, M, C, V, E}
    parameters :: P # constrained
    ensemble_mean :: M # constrained
    ensemble_cov :: C # constrained
    ensemble_var :: V
    mean_square_errors :: E
end

"""
    IterationSummary(eki, parameters, forward_map)

Return the summary for Ensemble Kalman Process `eki` with free `parameters` and `forward_map`.
"""
function IterationSummary(eki, parameters, forward_map)
    N_observations, N_ensemble = size(forward_map)
    original_priors = eki.inverse_problem.free_parameters.priors

    ensemble_mean = mean(parameters, dims=2)
    constrained_ensemble_mean = inverse_parameter_transform.(values(original_priors), ensemble_mean)
    constrained_ensemble_mean = tupify_parameters(eki.inverse_problem, constrained_ensemble_mean)

    ensemble_covariance = cov(parameters, dims=2)
    constrained_ensemble_covariance = inverse_covariance_transform(values(original_priors), parameters, ensemble_covariance)
    constrained_ensemble_variance = tupify_parameters(eki.inverse_problem, diag(constrained_ensemble_covariance))

    constrained_parameters = inverse_parameter_transform.(values(original_priors), parameters)

    constrained_parameters = [tupify_parameters(eki.inverse_problem, constrained_parameters[:, i])
                              for i = 1:size(constrained_parameters, 2)]

    mean_square_errors = [
        mapreduce((x, y) -> (x - y)^2, +, eki.mapped_observations, view(forward_map, :, m)) / N_observations
        for m = 1:N_ensemble
    ]

    return IterationSummary(constrained_parameters,
                            constrained_ensemble_mean,
                            constrained_ensemble_covariance,
                            constrained_ensemble_variance,
                            mean_square_errors)
end

function IterationSummary(eki, parameters)
    original_priors = eki.inverse_problem.free_parameters.priors

    ensemble_mean = mean(parameters, dims=2)
    constrained_ensemble_mean = inverse_parameter_transform.(values(original_priors), ensemble_mean)
    constrained_ensemble_mean = tupify_parameters(eki.inverse_problem, constrained_ensemble_mean)

    ensemble_covariance = cov(parameters, dims=2)
    constrained_ensemble_covariance = inverse_covariance_transform(values(original_priors), parameters, ensemble_covariance)
    constrained_ensemble_variance = tupify_parameters(eki.inverse_problem, diag(constrained_ensemble_covariance))

    constrained_parameters = inverse_parameter_transform.(values(original_priors), parameters)

    constrained_parameters = [tupify_parameters(eki.inverse_problem, constrained_parameters[:, i])
                              for i = 1:size(constrained_parameters, 2)]

    return IterationSummary(constrained_parameters,
                            constrained_ensemble_mean,
                            constrained_ensemble_covariance,
                            constrained_ensemble_variance,
                            fill(NaN, size(constrained_parameters, 2)))
end

"""
    IterationSummary(eki)

Return the summary for Ensemble Kalman Process `eki` before any iteration.
"""
function IterationSummary(eki)
    parameters = get_u_final(eki.ensemble_kalman_process) # (N_params, ensemble_size) array
    return IterationSummary(eki, parameters)
end

function Base.show(io::IO, is::IterationSummary)
    print(io, "IterationSummary(ensemble = ", length(is.mean_square_errors), ")", '\n',
              "                      ", param_str.(keys(is.ensemble_mean))..., '\n',
              "       ensemble_mean: ", param_str.(values(is.ensemble_mean))..., '\n',
              "   ensemble_variance: ", param_str.(values(is.ensemble_var))..., '\n',
              particle_str.(1:length(is.parameters), is.mean_square_errors, is.parameters)...)
    return nothing
end

quick_summary(iter, is) = println("Iter $iter ", is.ensemble_mean)

function param_str(p::Symbol)
    p_str = string(p)
    length(p_str) > 9 && (p_str = p_str[1:9])
    return @sprintf("% 9s | ", p_str)
end

param_str(p::Number) = @sprintf("%1.3e | ", p)

particle_str(particle, error, parameters) =
    @sprintf("% 7s particle % 3d: ", " ", particle) *
    string(param_str.(values(parameters))...) *
    @sprintf("error = %.3e", error) * "\n"

"""
    sample(eki, θ, G, n)

Generate `n` new particles sampled from a multivariate Normal distribution parameterized
by the ensemble mean and covariance computed based on the `N_θ` × `N_ensemble` ensemble
array `θ`, under the condition that all `n` particles lead to forward map outputs that
are "stable" (don't include `NaNs`). `G` is the inverting forward map computed on
ensemble `θ`.

Return an `N_θ` × `n` array of new particles, along with the inverting forward
map output corresponding to the new particles.
"""
function sample(eki, θ, G, n)

    n_params, ens_size = size(θ)
    G_length = size(G, 1)

    ens_dist = eki.resampler.distribution(θ, G)

    found_θ = zeros((n_params, 0))
    found_G = zeros((G_length, 0))

    while size(found_θ, 2) < n
        θ_sample = rand(ens_dist, ens_size)
        G_sample = eki.inverting_forward_map(θ_sample)

        nan_values = nan_cols(G_sample)
        success_columns = findall(Bool.(1 .- nan_values))

        found_θ = hcat(found_θ, θ_sample[:, success_columns])
        found_G = hcat(found_G, G_sample[:, success_columns])
    end

    return found_θ[:, 1:n], found_G[:, 1:n]
end

nan_cols(G) = vec(mapslices(any, isnan.(G); dims=1))

"""
    iterate!(eki::EnsembleKalmanInversion; iterations=1)

Iterate the ensemble Kalman inversion problem `eki` forward by `iterations`.
"""
function iterate!(eki::EnsembleKalmanInversion; iterations = 1, show_progress = true)

    eki.iteration == 0 && (iterations += 1; eki.iteration -= 1)

    iterator = show_progress ? ProgressBar(1:iterations) : 1:iterations

    for _ in iterator

        eki.iteration += 1

        θ = get_u_final(eki.ensemble_kalman_process) # (N_params, ensemble_size) array
        G = eki.inverting_forward_map(θ) # (len(G), ensemble_size)

        eki.iteration != 0 && update_ensemble!(eki.ensemble_kalman_process, G)

        resample!(eki.resampler, G, θ, eki)

        # Save the parameter values and mean square error between forward map
        # and observations at the current iteration
        summary = IterationSummary(eki, θ, G)
        push!(eki.iteration_summaries, summary)
    end

    # Return ensemble mean (best guess for optimal parameters)
    best_parameters = eki.iteration_summaries[end].ensemble_mean

    return tupify_parameters(eki.inverse_problem, best_parameters)
end

#####
##### Resampling
#####

abstract type EnsembleDistribution end

function ensemble_dist(θ)
    μ = [mean(θ, dims=2)...]
    Σ = cov(θ, dims=2)
    return MvNormal(μ, Σ)
end

struct FullEnsembleDistribution <: EnsembleDistribution end

(::FullEnsembleDistribution)(θ, G) = ensemble_dist(θ)

struct SuccessfulEnsembleDistribution <: EnsembleDistribution end

(::SuccessfulEnsembleDistribution)(θ, G) = ensemble_dist(θ[:, findall(Bool.(1 .- nan_cols(G)))])

resample!(::Nothing, args...) = nothing

struct NaNResampler{D}
    abort_fraction :: Float64
    distribution :: D
end

NaNResampler(; abort_fraction=0.0, distribution=FullEnsembleDistribution()) = NaNResampler(abort_fraction, distribution)

function resample!(resampler::NaNResampler, G, θ, eki)

    # `ensemble_size` vector of bits indicating, for each ensemble member, whether the forward map contained `NaN`s
    nan_values = nan_cols(G)
    nan_columns = findall(Bool.(nan_values)) # indices of columns (particles) with `NaN`s
    nan_count = length(nan_columns)
    nan_fraction = nan_count / size(θ, 2)

    if nan_fraction > resampler.abort_fraction
        error("The forward map for $(100nan_fraction)% of particles included NaNs. Consider \n" *
              "    1. Increasing the `abort_fraction` for `resampler = NaNResampler(abort_fraction=0.0)` \n" *
              "        when constructing `EnsembleKalmanInversion`; \n" *
              "    2. Reducing the time-step for `InverseProblem.simulation`, \n" *
              "    3. Evolving `InverseProblem.simulation` for less time \n" *
              "    4. Narrowing `FreeParameters` priors.")
    end

    if nan_fraction > 0
        @warn "The forward map for $nan_count particles ($(100nan_fraction)%) included NaNs. Resampling
                $nan_count particles from a multivariate Normal distribution parameterized by the
                ensemble mean and covariance."

        found_θ, found_G = sample(eki, θ, G, nan_count)
        view(θ, :, nan_columns) .= found_θ
        view(G, :, nan_columns) .= found_G

        new_process = EnsembleKalmanProcess(θ,
                                            eki.mapped_observations,
                                            eki.noise_covariance,
                                            eki.ensemble_kalman_process.process)

        eki.ensemble_kalman_process = new_process
    end

    return G
end

end # module


"An ordinary differential equation"
immutable ODE
  var::DReal.Ex{Float64}         # independent variable
  rhs::DReal.Ex{Float64}        # f(var)
end

"A system of differential equations"
immutable ODESystem
  name::ASCIIString
  odes::Vector{ODE}
end

ODESystem(odes::Vector{ODE}) = ODESystem(genvar("flow_"), odes)

to_c{T}(exs::Vector{Ex{T}}) = [ex.e for ex in exs]

"Returns constraint that [x_1 ...] is the integral of the system of differential
equations from 0 to time_0 with initial conditions x_0 and p_0"
function integral(
  sys::ODESystem,         # a system of odes
  varst::Vector{Ex{Float64}},     # variables which will be constrained to integral output
  t0::Ex{Float64},        # lower bound of integration
  tend::Ex{Float64},      # upper bound of integration
  vars0::Vector{Ex{Float64}};     # Initial conditions
  ctx::Context=global_context())

  nvars = length(varst)
  @assert nvars == length(vars0)
  @assert nvars == length(sys.odes)
  odevars = [ode.var for ode in sys.odes]
  oderhses = [ode.rhs for ode in sys.odes]
  opensmt_define_ode(ctx.ctx, sys.name, to_c(odevars), to_c(oderhses), Cuint(length(odevars)))
  Ex{Bool}(opensmt_mk_integral(ctx.ctx, to_c(varst), t0.e, tend.e, to_c(vars0), Cuint(nvars), sys.name), no_vars())
end

function integral(
  odes::Vector{ODE},      # a system of odes
  varst::Vector{Ex{Float64}},     # variables which will be constrained to integral output
  t0::Ex{Float64},        # lower bound of integration
  tend::Ex{Float64},      # upper bound of integration
  vars0::Vector{Ex{Float64}};     # Initial conditions
  ctx::Context=global_context())

  sys = ODESystem(odes)
  integral(sys, varst, t0, tend, vars0; ctx=ctx)
end


boundingbox(d::Vector{T}) where {T <: Number} = minimum(d)..maximum(d)
boundingbox(d::Set{T}) where {T<:Number} = minimum(d)..maximum(d)

"Return the bounding box of the union of two or more bounding boxes."
unionbox(d::Domain) = d
unionbox(d1::Domain, d2::Domain) = unionbox(promote_domains(d1, d2)...)
unionbox(d1::Domain, d2::Domain, domains...) =
    unionbox(unionbox(d1,d2), domains...)

unionbox(d1::Domain{T}, d2::Domain{T}) where {T} = unionbox1(d1, d2)
unionbox1(d1, d2) = unionbox2(d1, d2)
unionbox2(d1, d2) = FullSpace{eltype(d1)}()
unionbox1(d1::EmptySpace, d2) = d2
unionbox1(d1::FullSpace, d2) = d1
unionbox2(d1, d2::EmptySpace) = d1
unionbox2(d1, d2::FullSpace) = d2

unionbox(d1::D, d2::D) where {D<:FixedInterval} = d1

function unionbox(d1::AbstractInterval{T}, d2::AbstractInterval{T}) where {T}
    a, b = endpoints(d1)
    c, d = endpoints(d2)
    A = min(a, c)
    B = max(b, d)
    isinf(A) && isinf(B) ? FullSpace{T}() : A..B
end

unionbox(d1::HyperRectangle{T}, d2::HyperRectangle{T}) where {T} =
    Rectangle{T}(map(unionbox, components(d1), components(d2)))

"Return the bounding box of the intersection of two or more bounding boxes."
intersectbox(d::Domain) = d
intersectbox(d1::Domain, d2::Domain) = intersectbox(promote_domains(d1, d2)...)
intersectbox(d1::Domain, d2::Domain, domains...) =
    intersectbox(intersectbox(d1,d2), domains...)

intersectbox(d1::Domain{T}, d2::Domain{T}) where {T} = intersectbox1(d1, d2)
intersectbox1(d1, d2) = intersectbox2(d1, d2)
intersectbox2(d1, d2) = FullSpace{eltype(d1)}()
intersectbox1(d1::EmptySpace, d2) = d1
intersectbox1(d1::FullSpace, d2) = d2
intersectbox2(d1, d2::EmptySpace) = d2
intersectbox2(d1, d2::FullSpace) = d1

intersectbox(d1::D, d2::D) where {D<:FixedInterval} = d1

intersectbox(d1::AbstractInterval{T}, d2::AbstractInterval{T}) where {T} =
    intersectdomain(d1, d2)

function intersectbox(d1::HyperRectangle{T}, d2::HyperRectangle{T}) where {T}
    d = Rectangle{T}(map(intersectbox, components(d1), components(d2)))
    isempty(d) ? EmptySpace{T}() : d
end

boundingbox(d::AbstractMappedDomain) = map_boundingbox(boundingbox(superdomain(d)), forward_map(d))

function map_boundingbox(box::Interval, fmap)
    l,r = (infimum(box),supremum(box))
    ml = fmap(l); mr = fmap(r)
    min(ml,mr)..max(ml,mr)
end

# This is a best effort implementation, it could be wrong for some maps,
# because we only map the corners. Hence we restrict to affine maps here.
map_boundingbox(box::HyperRectangle, fmap::AbstractAffineMap) =
    map_boundingbox_generic(box, fmap)

# This is a best effort implementation. It could be wrong for some maps,
# because we only map the corners.
function map_boundingbox_generic(box::HyperRectangle{T}, fmap) where {T}
    mapped_corners = map(fmap, corners(box))
    left = [minimum(x[j] for x in mapped_corners) for j in 1:dimension(box)]
    right = [maximum(x[j] for x in mapped_corners) for j in 1:dimension(box)]
    Rectangle{T}(left, right)
end


include("../examples/lorenz63.jl")
include("../src/clvs.jl")
#include("../src/lss.jl")
include("../src/adjoint_lss.jl")
using Test
using DiffEqSensitivity, OrdinaryDiffEq, Zygote
function test_dlorenz63()
    s = [10., 28., 8/3]
    m = 10
    u0 = rand(3,m)
    u_trj = lorenz63(u0, s, 1)[2,:,:]
    eps = 1.e-4
    du_x = (lorenz63(u_trj .+ eps.*[1; 0.; 0.],
    					s, 1)[2,:,:] -
    			lorenz63(u_trj .- eps.*[1; 0.; 0.],
    					s, 1)[2,:,:])./(2*eps)
    du_y = (lorenz63(u_trj .+ eps.*[0; 1.; 0.],
    					s, 1)[2,:,:] -
    			lorenz63(u_trj .- eps.*[0; 1.; 0.],
    					s, 1)[2,:,:])./(2*eps)
    du_z = (lorenz63(u_trj .+ eps.*[0; 0.; 1.],
    					s, 1)[2,:,:] -
    			lorenz63(u_trj .- eps.*[0; 0.; 1.],
    					s, 1)[2,:,:])./(2*eps)
    du_fd = reshape(collect([du_x; du_y; du_z]),
    					3, 3, m)
    du = dlorenz63(u_trj', s)
    @test all(isapprox.(du_fd - du, 0.,atol=1.e-8)) == 1
end
function test_les()
    s = [10., 28., 8/3]
    m = 1
    n = 20000
    u0 = rand(3,m)
    u_trj = lorenz63(u0, s, n)[:,:,1]
    du_trj = dlorenz63(u_trj, s)
    les, clv_trj = clvs(du_trj,3)
    println(les/0.005)
    @test all(isapprox.(les/0.005, [0.91, 0.,
    	-14.572], rtol=1.e-1)) == 1
end
function test_lss()
    s = [10., 28., 8/3]
    m = 1
    n = 3000
    n_runup = 5000
    u0 = rand(3,m)
    u_init = lorenz63(u0, s, n)[end,:,:]
    d = 3
    d_u = 2
    dJ = zeros(1,d,n+1)
    dJ[1,3,:] .= 1.
    n_samples = 100
    dJds = zeros(1,n_samples)
    vsh = zeros(d,n+1,n_samples)
    for i=1:n_samples
    	println("Starting LSS, sample ", i)
    	u_trj = lorenz63(u_init, s, n)[:,:,1]
    	du_trj = dlorenz63(u_trj, s)
        X = perturbation(u_trj,s) #ith col in T_{u_{i+1}} M
    	f = vectorField(u_trj,s)
    	#f = zeros(d,n+1)
    	y, xi = lss(du_trj, X, f, s, d_u)
		dJds[:,i] = compute_sens(y, xi, dJ, f)
    	@show dJds[i]
    	vsh[:,:,i] = y
    	u_init .= reshape(u_trj[end,:],3,1)
    end
    @test isapprox((sum(dJds)/n_samples),1.0,rtol=0.1)
    return vsh, dJds
end
function test_condition_number()
    s = [10., 28., 8/3]
    m = 1
    n_runup = 5000
    u0 = rand(3,m)
    u_init = lorenz63(u0, s, n_runup)[end,:,:]
    d = 3
    d_u = 2
    n_samples = 30
    dJds = zeros(n_samples)
    condnum = zeros(n_samples)
    n_arr = LinRange(500, 5000, n_samples)
    for i=1:n_samples
    	println("Starting LSS, sample ", i)
    	n = floor(Int64,n_arr[i])
    	u_trj = lorenz63(u_init, s, n)[:,:,1]
    	du_trj = dlorenz63(u_trj, s)
        X = perturbation(u_trj,s) #ith col in T_{u_{i+1}} M
    	f = zeros(d,n+1)
    	J = u_trj[:,3]
    	dJ = zeros(d,n+1)
    	dJ[3,:] .= 1.

    	y, dJds[i], c = lss(u_trj, du_trj, X, f, J,
    					 dJ, s, d_u)
    	condnum[i] = c
    	println(dJds[i])
    	u_init .= reshape(u_trj[end,:],3,1)
    end
    @test isapprox((sum(dJds)/n_samples),1.0,rtol=0.1)
    return dJds, condnum
end
function test_ensemble_sensitivity()
#let
    s = [10., 28., 8/3]
	du0, ds = zeros(3), zeros(3)
	n_samples = 1000
	#prob = ODEProblem(lorenz63_ad, u0, (0.,1.), s)
    for i =1:n_samples
		u0 = rand(3)
		println("sample number, ", i)
		#sol = solve(prob,Tsit5(),saveat=0.005,sensealg=QuadratureAdjoint())
    	du01, ds1 = Zygote.gradient(obj_fun, u0, s)
		@show du01, ds1
		du0 .+= du01/n_samples
		ds .+= ds1/n_samples
	end
	@test isapprox(ds[2], 1., rtol=0.1)
	return du0, ds
end
function test_tangent_ad()
	s = [10., 28., 8/3]
	du0, ds = zeros(3), zeros(3)
	n_samples = 1000
	#prob = ODEProblem(lorenz63_ad, u0, (0.,1.), s)
    for i =1:n_samples
		u0 = rand(3)
		println("sample number, ", i)
		#sol = solve(prob,Tsit5(),saveat=0.005,sensealg=QuadratureAdjoint())
    	du01, ds1 = Zygote.gradient(obj_fun, u0, s)
		@show du01, ds1
		du0 .+= du01/n_samples
		ds .+= ds1/n_samples
	end
	return du0, ds
end
#Tangent NILSS
function obj_fun_Q(v0,u0)
	sol=	solve(u0 + eps*v0, s, (0.,1.))
	return sol[:,end]
end
# To get Q, differentiate obj_fun_Q wrt v0
function obj_fun_v(u0, v0)
	sol=	solve(u0 + eps*v0, s + eps*ds0, (0.,1.))
	return sol[:,end]
end
# To get v, differentiate obj_fun_v wrt ds0

#Adjoint NILSS
function obj_fun_vT(w0,u0)
	sol = solve(u0, s, (0.,1.))
	sol[:, end] .+= eps*w0
	J = sum(sol[3,:])/n
	return J
end
#Derivative of obj_fun_vT wrt w0 to get vT
function obj_fun_QT(w0,u0)
	sol = solve(u0, s, (0.,1.))
	sol[:, end] .+= eps*w0
	return sol[:,end]
end
#Derivative of obj_fun_QT wrt w0 to get QT
function test_adjoint_lss()
    s = [10., 28., 8/3]
    m = 1
    n = 3000
    n_runup = 5000
    u0 = rand(3,m)
    u_init = lorenz63(u0, s, n)[end,:,:]
    d = 3
    d_u = 2
    dJ = zeros(d,n+1)
    dJ[3,:] .= 1.0/(n+1)
    n_samples = 100
    dJds = zeros(1,n_samples)
    vsh = zeros(d,n+1,n_samples)
    for i=1:n_samples
    	println("Starting LSS, sample ", i)
    	u_trj = lorenz63(u_init, s, n)[:,:,1]
		u_trj = reverse(u_trj, dims=1)
		duT_trj = permutedims(dlorenz63(u_trj, s),[2,1,3])
        X = perturbation(u_trj,s) #ith col in T_{u_{i+1}} M
    	g = vectorField(u_trj,s)
		u_next = lorenz63(reshape(u_trj[1,:],d,:),s,1)[end,:,1]
		g = [vectorField(reshape(u_next,:,d),s) g[:,2:n+1]]
		X = reshape(X, 1, d, n+1)
    	f = zeros(d,n+1)
    	J = u_trj[:,3]
    	y, xi = lss(duT_trj, dJ, f, s, 2, g)
    	vsh[:,:,i] = y
    	u_init .= reshape(u_trj[1,:],3,1)
		dJds[:,i] = compute_sens(y, xi, X, f)
		@show dJds[:,i]
    end
    #@test isapprox((sum(dJds)/n_samples),1.0,rtol=0.1)
    return vsh, dJds
end


using FocusedBlindDecon
using Test
using LinearAlgebra
using Misfits

include("adj_tests.jl")
include("fbd_tests.jl")


module Variables

include("./constants.jl")
using JuMP
using .Constants

export init_variables

function init_variables(m)

    #Dimensions of box
    @variable(m, lob_length_of_box[i=1:length(boxes)] == boxes[i][1])
    @variable(m, wob_width_of_box[i=1:length(boxes)] == boxes[i][2])
    @variable(m, hob_height_of_box[i=1:length(boxes)] == boxes[i][3])

    #Front-left coordinate of each box
    @variable(m, flbx_front_left_bot_x[i=1:length(boxes)] >= 0)
    @variable(m, flby_front_left_bot_y[i=1:length(boxes)] >= 0)
    @variable(m, flbz_front_left_bot_z[i=1:length(boxes)] >= 0)

    #Check if length (pi) is parallel to X, Y or Z-axis
    @variable(m, lpx_length_of_box_parallel_to_x[i=1:length(boxes)], Bin)
    @variable(m, lpy_length_of_box_parallel_to_y[i=1:length(boxes)], Bin)
    @variable(m, lpz_length_of_box_parallel_to_z[i=1:length(boxes)], Bin)

    #Check if width (qi) is parallel to X, Y or Z-axis
    @variable(m, wpx_width_of_box_parallel_to_x[i=1:length(boxes)], Bin)
    @variable(m, wpy_width_of_box_parallel_to_y[i=1:length(boxes)], Bin)
    @variable(m, wpz_width_of_box_parallel_to_z[i=1:length(boxes)], Bin)

    #Check if height (ri) is parallel to X, Y or Z-axis
    @variable(m, hpx_height_of_box_parallel_to_x[i=1:length(boxes)], Bin)
    @variable(m, hpy_height_of_box_parallel_to_y[i=1:length(boxes)], Bin)
    @variable(m, hpz_height_of_box_parallel_to_z[i=1:length(boxes)], Bin)

    # box_j_x_axis_cover = lpx_length_of_box_parallel_to_x[j]*lob_length_of_box[j] + wpx_width_of_box_parallel_to_x[j]*wob_width_of_box[j]
    # box_j_y_axis_cover = lpy_length_of_box_parallel_to_y[j]*lob_length_of_box[j] + wpy_width_of_box_parallel_to_y[j]*wob_width_of_box[j]
    @expressions(m, begin
        box_x_axis_cover[i=1:length(boxes)],
        (
        m[:lpx_length_of_box_parallel_to_x][i]*fix_value(m[:lob_length_of_box][i]) +
        m[:wpx_width_of_box_parallel_to_x][i]*fix_value(m[:wob_width_of_box][i])
        )

        box_y_axis_cover[i=1:length(boxes)],
        (
        m[:lpy_length_of_box_parallel_to_y][i]*fix_value(m[:lob_length_of_box][i]) +
        m[:wpy_width_of_box_parallel_to_y][i]*fix_value(m[:wob_width_of_box][i])
        )
    end)

    #Make objective available for the whole program. Later used in objective function.
    @expression(
        m, heights[i=1:length(boxes)], m[:flbz_front_left_bot_z][i] + fix_value(m[:hob_height_of_box][i])
    )
    @variable(m, H >= maxHeight >= 0)
    @constraint(m, m[:maxHeight] .>= m[:heights])

    return m
end

end


"""
    rescale_hidden!(rbm, λ::AbstractArray)

For continuous hidden units with a scale parameter,
scales parameters such that hidden unit activations
are divided by `λ`. For other hidden units does
nothing. The resulting RBM is equivalent to the
original one.
"""
function rescale_hidden!(rbm::RBM, λ::AbstractArray)
    @assert size(rbm.hidden) == size(λ)
    if rescale_activations!(rbm.hidden, λ)
        rbm.w .*= reshape(λ, map(one, size(rbm.visible))..., size(rbm.hidden)...)
    end
    return rbm
end

"""
    rescale_activations!(layer, λ::AbstractArray)

For continuous layers with scale parameters, re-parameterizes
such that unit activations are divided by `λ`, and returns `true`.
For other layers just returns `false`.
"""
rescale_activations!(layer::Union{Binary,Spin,Potts}, λ::AbstractArray) = false

#= Note that λ < 0 can lead to trouble, e.g. for ReLU which
must have positive activations. So we dissallow it below. =#

function rescale_activations!(layer::Union{Gaussian,ReLU}, λ::AbstractArray)
    @assert size(layer) == size(λ)
    @assert all(λ .> 0)
    layer.θ .*= λ
    layer.γ .*= λ.^2
    return true
end

function rescale_activations!(layer::dReLU, λ::AbstractArray)
    @assert size(layer) == size(λ)
    @assert all(λ .> 0)
    layer.θp .*= λ
    layer.θn .*= λ
    layer.γp .*= λ.^2
    layer.γn .*= λ.^2
    return true
end

function rescale_activations!(layer::Union{pReLU,xReLU}, λ::AbstractArray)
    @assert size(layer) == size(λ)
    @assert all(λ .> 0) # makes life simpler
    layer.θ .*= λ
    layer.Δ .*= λ
    layer.γ .*= λ.^2
    return true
end


# 2D FWI gradient test with 4 sources
# The receiver positions and the source wavelets are the same for each of the four experiments.
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January 2017
#
# Mathias Louboutin, mlouboutin3@gatech.edu
# Updated July 2020

parsed_args = parse_commandline()

nlayer = parsed_args["nlayer"]
tti = parsed_args["tti"]
fs =  parsed_args["fs"]

### Model
model, model0, dm = setup_model(parsed_args["tti"], 4)
q, srcGeometry, recGeometry, info = setup_geom(model)
dt = srcGeometry.dt[1]

###################################################################################################

@testset "TWRI gradient test w.r.t m with $(nlayer) layers and tti $(tti) and freesurface $(fs)" begin
	optw = TWRIOptions(grad_corr=false, comp_alpha=false, weight_fun=nothing, eps=0, params=:m)
	# Gradient test
	h = 5f-2
	maxiter = 5
	err1 = zeros(maxiter)
	err2 = zeros(maxiter)
	h_all = zeros(maxiter)
	modelH = deepcopy(model0)

	# Observed data
	opt = Options(sum_padding=true, free_surface=parsed_args["fs"])
	F = judiModeling(info, model, srcGeometry, recGeometry; options=opt)
	F0 = judiModeling(info, model0, srcGeometry, recGeometry; options=opt)
	d = F*q
	d0 = F0*q
	y = 2.5f0*(d0 - d)

	# FWI gradient and function value for m0
	Jm0, gradm = twri_objective(model0, q, d, y; options=opt, optionswri=optw)

	dJ = dot(gradm, dm)
	@printf("Perturbation size is %2.2e and reference objective function is %2.4e \n", dJ, Jm0)

	for j=1:maxiter
		# FWI gradient and function falue for m0 + h*dm
		modelH.m = model0.m + h*dm
		Jm, _ = twri_objective(modelH, q, d, y; options=opt, optionswri=optw)

		# Check convergence
		err1[j] = abs(Jm - Jm0)
		err2[j] = abs(Jm - Jm0 - h*dJ)
		j == 1 ? prev = 1 : prev = j - 1
		@printf("h = %2.2e, phi = %2.4e, e1 = %2.2e, rate = %2.2e", h, Jm, err1[j], err1[prev]/err1[j])
		@printf(", e2  = %2.2e, rate = %2.2e \n", err2[j], err2[prev]/err2[j])
		h_all[j] = h
		h = h * .8f0
	end
	# CHeck convergence rates
	rate_1 = sum(err1[1:end-1]./err1[2:end])/(maxiter - 1)
	rate_2 = sum(err2[1:end-1]./err2[2:end])/(maxiter - 1)
	@test isapprox(rate_1, 1.25f0; rtol=5f-2)
	@test isapprox(rate_2, 1.5625f0; rtol=5f-2)
end

@testset "TWRI gradient test w.r.t y with $(nlayer) layers and tti $(tti) and freesurface $(fs)" begin
	optw = TWRIOptions(grad_corr=false, comp_alpha=false, weight_fun=nothing, eps=0, params=:y)
	# Gradient test
	h = 5f-2
	maxiter = 5
	err1 = zeros(maxiter)
	err2 = zeros(maxiter)
	h_all = zeros(maxiter)
	modelH = deepcopy(model0)

	# Observed data
	opt = Options(sum_padding=true, free_surface=parsed_args["fs"])
	F = judiModeling(info, model, srcGeometry, recGeometry; options=opt)
	F0 = judiModeling(info, model0, srcGeometry, recGeometry; options=opt)
	d = F*q
	d0 = F0*q
	y = .1f0 * (d0 - d)
	dy = .05f0*(d0 - d)

	# FWI gradient and function value for m0
	Jm0, grady = twri_objective(model0, q, d, y; options=opt, optionswri=optw)

	dJ = dot(grady, dy)
	@printf("Perturbation size is %2.4e and reference objective function is %2.4e \n", dJ, Jm0)

	for j=1:maxiter
		# FWI gradient and function falue for m0 + h*dm
		yloc = y + h * dy
		Jm, _ = twri_objective(model0, q, d, yloc;options=opt, optionswri=optw)

		# Check convergence
		err1[j] = abs(Jm - Jm0)
		err2[j] = abs(Jm - Jm0 - h*dJ)
		j == 1 ? prev = 1 : prev = j - 1
		@printf("h = %2.2e, phi = %2.4e, e1 = %2.2e, rate = %2.2e", h, Jm, err1[j], err1[prev]/err1[j])
		@printf(", e2  = %2.2e, rate = %2.2e \n", err2[j], err2[prev]/err2[j])
		h_all[j] = h
		h = h * .8f0
	end
	# CHeck convergence rates
	rate_1 = sum(err1[1:end-1]./err1[2:end])/(maxiter - 1)
	rate_2 = sum(err2[1:end-1]./err2[2:end])/(maxiter - 1)
	@test isapprox(rate_1, 1.25f0; rtol=5f-2)
	@test isapprox(rate_2, 1.5625f0; rtol=5f-2)
end


using GalerkinSparseGrids
using Base.Test
using Cubature
using ODE

tests = [
	"elementary.jl",
	"hier_DG.jl",
	"vhier_DG.jl",
	"differentiation.jl",
	"solvers.jl"
]

for filename in tests
	name = first(splitext(filename))
	@testset "$name" begin
		include(filename)
	end
end


#= Glacier mass balance toy model

Test with ideal data of a hybrid MB model based on neural networks
for the ablation and accumulation components that respect a classic
temperature-index model, mixed with model interpretation using
SINDy (Brunton et al., 2016).

=#

## Environment and packages
cd(@__DIR__)
using Pkg; Pkg.activate("../."); Pkg.instantiate()
using Infiltrator
using Debugger

using Flux, DiffEqFlux
using Flux: @epochs
using Zygote
using Plots
gr()
using Base: @kwdef
using Statistics
using DataDrivenDiffEq
using ModelingToolkit
using LinearAlgebra
# Set a random seed for reproduceable behaviour
using Random

#Random.seed!(2345)

@kwdef mutable struct Hyperparameters
    batchsize::Int = 500    # batch size
    η::Float64 = 0.1   # learning rate
    epochs::Int = 500        # number of epochs
    use_cuda::Bool = true   # use gpu (if cuda available)
end

###############################################
############  FUNCTIONS   #####################
###############################################

# Fake quadratic response to snowfall and temperature
α(P) = P.^(2)
μ(T) = T.^(2)

# Temperature-index equation for point glacier mass balance
function MB(P, T, Up, Ut)
    T_melt = 0
    MB = Up(P) - Ut(max.(T.-T_melt, 0))
    return MB
end

# Toy function to create artificial reference data
function toy_MB(α, μ, P, T)
    T_melt = 0.0
    MB = α(P) - μ(max.(T.-T_melt, 0.0))
    return MB
end

function train()

    # We determine the loss function
    function loss(batch)
        l, l_acc, l_abl = 0.0f0, 0.0f0, 0.0f0
        num = 0
        for (x, y) in batch

            # Make NN predictions
            p_batch = x[1,:]'
            t_batch = x[2,:]'
            pdd_batch = max.(t_batch.-0, 0)
            Ŷp = Up(p_batch)
            Ŷt = Ut(pdd_batch)

            # We evaluate the MB as the combination of Accumulation - Ablation
            w_pc=1000
            l_MB = sqrt(Flux.Losses.mse(MB(p_batch, t_batch, Up, Ut), y; agg=mean))
            l_range_acc = sum((max.((Ŷp/p_batch).-110, 0)).*w_pc)
            l_range_abl = sum((max.((Ŷt/pdd_batch).-110, 0)).*w_pc)

            #l += l_MB
            l += l_MB + l_range_acc + l_range_abl
            l_acc += l_range_acc
            l_abl += l_range_abl
            num +=  size(x, 2)

            # println("Accumulation loss: ", l_range_acc)
            # println("Ablation loss: ", l_range_abl)

        end

        return l/num
    end

    # Container to track the losses
    losses = Float32[]

    # Callback to show the loss during training
    callback(l) = begin
        push!(losses, l)
        if length(losses)%50==0
            println("Current loss after $(length(losses)) iterations: $(losses[end])")
        end
        false
    end

    function hybrid_train!(loss, ps_Up, ps_Ut, data, opt)
        # Retrieve model parameters
        ps_Up = Params(ps_Up)
        ps_Ut = Params(ps_Ut)

        for batch in data
        # back is a method that computes the product of the gradient so far with its argument.
        train_loss_Up, back_Up = Zygote.pullback(() -> loss(batch), ps_Up)
        train_loss_Ut, back_Ut = Zygote.pullback(() -> loss(batch), ps_Ut)
        # Callback to track the training
        callback(train_loss_Up)
        # Apply back() to the correct type of 1.0 to get the gradient of loss.
        gs_Up = back_Up(one(train_loss_Up))
        gs_Ut = back_Ut(one(train_loss_Ut))
        # Insert what ever code you want here that needs gradient.
        # E.g. logging with TensorBoardLogger.jl as histogram so you can see if it is becoming huge.
        Flux.update!(opt, ps_Up, gs_Up)
        Flux.update!(opt, ps_Ut, gs_Ut)
        # Here you might like to check validation set accuracy, and break out to do early stopping.
        end
    end

    #########################################
    ##########################################

    ######### Define the network  ############
    # We determine the hyperameters for the training
    hyparams = Hyperparameters()

    # Leaky ReLu as activation function
    leakyrelu(x, a=0.01) = max(a*x, x)
    relu_acc(x) = min(max(0, x), 30)
    relu_abl(x) = min(max(0, x), 35)

    # Define the networks 1->5->5->5->1
    global Up = Chain(
        Dense(1,10,initb = Flux.zeros),
        BatchNorm(10, leakyrelu),
        Dense(10,10,initb = Flux.zeros),
        BatchNorm(10, leakyrelu),
        Dense(10,5,initb = Flux.zeros),
        BatchNorm(5, leakyrelu),
        Dense(5,1, relu, initb = Flux.zeros)
    )

    global Ut = Chain(
        Dense(1,10,initb = Flux.zeros),
        BatchNorm(10, leakyrelu),
        Dense(10,10,initb = Flux.zeros),
        BatchNorm(10, leakyrelu),
        Dense(10,5,initb = Flux.zeros),
        BatchNorm(5, leakyrelu),
        Dense(5,1, relu, initb = Flux.zeros)
    )

    # We define an optimizer
    #opt = RMSProp(hyparams.η, 0.95)
    opt = ADAM(hyparams.η)

    # We get the model parameters to be trained
    ps_Up = Flux.params(Up)
    ps_Ut = Flux.params(Ut)

    sqnorm(x) = sum(abs2, x)

    #######  We generate toy data to train the model  ########
    global snowfall_toy = rand(0.0f0:100.0f0, 500)
    global temperature_toy = 100f0*sin.((1:500)./20)
    forcings = hcat(snowfall_toy, temperature_toy)

    MB_toy = toy_MB(α, μ, snowfall_toy, temperature_toy)

    # Plot the toy dataset
    l1 = @layout [a b; c; d]
    p1₁ = plot(1:500, snowfall_toy, label="Snowfall", color="midnightblue")
    p1₂ = plot(1:500, temperature_toy, label="Temperature", color="darkred")
    hline!(p1₂, [0], c="black", label="")
    p1₃ = plot(1:500, MB_toy, label="Mass balance")
    hline!(p1₃, [0], c="black", label="")

    X = vcat(snowfall_toy', temperature_toy')
    Y = collect(MB_toy')

    # Very important:
    # Shuffle batches and use batchsize similar to dataset size for easy training
    data = Flux.Data.DataLoader((X, Y), batchsize=hyparams.batchsize, (X, Y), shuffle=true)

    # We train the mass balance hybrid model
    @epochs hyparams.epochs hybrid_train!(loss, ps_Up, ps_Ut, data, opt)

    # Perform MB simulations with trained model
    global T_melt = 0
    global snow_norm = X[1,:]
    global temp_norm = X[2,:]
    global PDD_toy = max.(temp_norm.-T_melt, 0)
    global Ŷ_MB = MB(snow_norm', temp_norm', Up, Ut)
    global Ŷ_acc = Up(snow_norm')
    global Ŷ_abl = Ut(PDD_toy')

    # Plot comparison with toy data
    p1₄ = plot(1:500, Ŷ_MB', label="Simulated mass balance")
    hline!(p1₄, [0], c="black", label="")
    p1 = plot(p1₁,p1₂,p1₃, p1₄,layout=l1)

    l2 = @layout [a b;c d]
    p2₁ = plot(1:500, α(snowfall_toy')', label="α")
    plot!(1:500, Ŷ_acc', label="Up")
    hline!(p2₁, [0], c="black", label="")
    p2₂ = plot(1:500, α(snowfall_toy')'./snowfall_toy, label="α/P")
    plot!(1:500, Ŷ_acc'./snowfall_toy, label="Up/P")
    hline!(p2₂, [0], c="black", label="")

    p2₃ = plot(1:500, μ(PDD_toy')', label="μ", lw=2)
    plot!(1:500, Ŷ_abl', label="Ut", lw=2)
    hline!(p2₃, [0], c="black", label="")
    p2₄ = plot(1:500, μ(PDD_toy')'./PDD_toy, label="μ/T", lw=2)
    plot!(1:500, Ŷ_abl'./PDD_toy, label="Ut/T", lw=2)
    hline!(p2₄, [0], c="black", label="")

    p2 = plot(p2₁, p2₂, p2₃, p2₄, layout=l2)

    if(!ispath(pwd(),"plots"))
        mkdir("plots")
    end

    display(p1)
    display(p2)
    savefig(p1,joinpath(pwd(),"plots","hybrid_MB_model.png"))
    savefig(p2,joinpath(pwd(),"plots","hybrid_acc_abl.png"))

    ### Let's take a look at the raw learnt dynamics
    X_raw = 1:150
    Ŷp_raw = Up(X_raw')
    Yp_raw = α(X_raw')
    Ŷt_raw = Ut(X_raw')
    Yt_raw = μ(X_raw')

    l3 = @layout [a b]
    p3₁ = plot(X_raw, Ŷp_raw', lw=3, xlabel="Snowfall (P)", ylabel="Accumulation (α(P))",label="Up (inferred function)")
    plot!(X_raw, Yp_raw', l2=3, lw=3, label="α (true function")
    p3₂ = plot(X_raw, Ŷt_raw', lw=3, xlabel="Temperature (T)", ylabel="Ablation (μ(T))",label="Ut (inferred function)")
    plot!(X_raw, Yt_raw', l2=3, lw=3, label="μ (true function")
    p3 = plot(p3₁, p3₂, layout=l3)
    display(p3)
    savefig(p3,joinpath(pwd(),"plots","recovered_functions.png"))
end

###################################################################
########################  SINDy    ################################
###################################################################

function sindy()

    ## Symbolic regression via sparse regression ( SINDy based )
    # Create a Basis
    Xp = Float64.(copy(reshape(snowfall_toy, (1,length((snowfall_toy))))))
    Xt = max.(temperature_toy.-T_melt, 0)
    @variables X_P[1:1]

    # Generate the basis functions, multivariate polynomials up to deg 5
    b = [polynomial_basis(X_P, 5);]
    basis_p = Basis(b, X_P)

    # Create an optimizer for the SINDy problem
    λ, ρ = 9999e-1, 5.0
    opt = ADMM(λ, ρ)

    #opt = STRRidge(0.1)# Create the thresholds which should be used in the search process

    # opt = SR3(9e-1, 1.0)
    # λs = exp10.(-5:0.1:-1)

    # Test on ideal data
    println("Accumulation:")
    println("SINDy on ideal data \n")
    Ψ_p_true = SINDy(Xp, α(snow_norm'), basis_p, opt, maxiter = 1000) # Succeed
    #Ψ_true = SINDy(Float32.(reshape(X_raw, (1, length(X_raw)))), Float32.(Yp_raw), basis_p, opt, maxiter = 100, normalize = true)
    println(Ψ_p_true)
    print_equations(Ψ_p_true, show_parameter=true)
    p̂_p_true = parameters(Ψ_p_true)
    println("Ψ_true Parameter guess : $(p̂_p_true) \n")

    println("SINDy on synthetic range predicted data \n")
    syn_range = 1:100
    Ψ_p_range = SINDy(Float64.(syn_range), Float64.(Up(syn_range')), basis_p, opt, p=[0], maxiter = 20000, denoise=true, normalize=true)
    println(Ψ_p_range)
    print_equations(Ψ_p_range)
    p̂_p_range = parameters(Ψ_p_range)
    println("Ψ First parameter guess : $(p̂_p_range)")

    println("SINDy on predicted data \n")
    Ψ_p = SINDy(Xp, Float64.(Ŷ_acc), basis_p, opt, p=[0], maxiter = 200000, denoise=false, normalize=true, convergence_error = 1e-8)
    println(Ψ_p)
    print_equations(Ψ_p)
    p̂_p = parameters(Ψ_p)
    println("Ψ First parameter guess : $(p̂_p) \n \n")

    println("Ablation: ")
    println("SINDy on ideal data \n")
    Ψ_t_true = SINDy(Xt, μ(temp_norm'), basis_p, opt, maxiter = 1000) # Succeed
    #Ψ_true = SINDy(Float32.(reshape(X_raw, (1, length(X_raw)))), Float32.(Yp_raw), basis_p, opt, maxiter = 100, normalize = true)
    println(Ψ_t_true)
    print_equations(Ψ_t_true, show_parameter=true)
    p̂_t_true = parameters(Ψ_t_true)
    println("Ψ_true Parameter guess : $(p̂_t_true) \n")

    println("SINDy on synthetic range predicted data \n")
    syn_range = 1:100
    Ψ_t_range = SINDy(Float64.(syn_range), Float64.(Ut(syn_range')), basis_p, opt, p=[0], maxiter = 20000, denoise=true, normalize=true)
    println(Ψ_t_range)
    print_equations(Ψ_t_range)
    p̂_t_range = parameters(Ψ_t_range)
    println("Ψ First parameter guess : $(p̂_t_range)")

    println("SINDy on predicted data \n")
    Ψ_t = SINDy(Xt, Float64.(Ŷ_abl), basis_p, opt, p=[0], maxiter = 200000, denoise=false, normalize=true, convergence_error = 1e-8)
    println(Ψ_t)
    print_equations(Ψ_t)
    p̂_t = parameters(Ψ_t)
    println("Ψ First parameter guess : $(p̂_t) \n")
end

###############################################################
###########################  MAIN #############################
###############################################################

train()

#sindy()


using SimpleProbabilitySets
using Base.Test, Distributions

@testset "SimpleProbabilitySets" begin
    d = Categorical([0.1, 0.2, 0.7])
    d2 = Categorical([0.4, 0.4, 0.2])
    pb = PBox(d, d2)
    @test cdfs(pb)[2] == [0.4, 0.8, 1.0]
    @test pints(pb)[1] ≈ [0.1, 0.0, 0.2] atol = 1e-6

    pl = [0.1, 0.2, 0.5]
    pu = [0.4, 0.4, 0.7]
    uncsmall = 0.05
    unclarge = 0.6
    pi = PInterval(pl, pu)
    pi2 = PInterval(pl, uncsmall)
    pi3 = PInterval(pl, unclarge)
    @test pi.plower == [0.1, 0.2, 0.5]
    @test pi2.plower ≈ [0.05, 0.15, 0.45] atol = 1e-6
    @test pi3.pupper ≈ [0.7, 0.8, 1.0] atol = 1e-6

    p = psample(pi)
    @test sum(p) ≈ 1.0 atol = 1e-9
    @test all(pi.plower .< p .< pi.pupper)
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    𝓒(x)

Return closure of `x`.
"""
𝓒(x) = x ./ sum(x)

"""
    Composition(partscomps)
    Composition(parts, comps)
    Composition(part₁=comp₁, part₂=part₂, ...)
    Composition(comps)
    Composition(comp₁, comp₂, ...)

A D-part composition as defined by Aitchison 1986.

## Examples

A 2-part composition with parts `a=0.1` and `b=0.8`:

```julia
julia> Composition(a=0.2, b=0.8)
julia> Composition((a=0.2, b=0.8))
julia> Composition((:a, :b), (0.2, 0.8))
```

When the names of the parts are not specified, the
constructor uses default names `part1`, `part2`,
..., `partD`:

```
julia> Composition(0.1, 0.8)
julia> Composition((0.1, 0.8))
```
"""
struct Composition{D,PARTS}
  data::NamedTuple{PARTS,NTuple{D,Union{Float64,Missing}}}
end

Composition(data::NamedTuple) =
  Composition{length(data),keys(data)}(data)

Composition(; data...) = Composition((; data...))

Composition(parts::NTuple, comps) =
  Composition((; zip(parts, Tuple(comps))...))

Composition(comps) =
  Composition(ntuple(i->Symbol("w$i"), length(comps)), comps)

Composition(comp::Real, comps...) = Composition((comp, comps...))

"""
    parts(c)

Parts in the composition `c`.
"""
parts(::Composition{D,PARTS}) where {D,PARTS} = PARTS

"""
    components(c)

Components in the composition `c`.
"""
components(c::Composition) = getfield(c, :data) |> values |> SVector

Base.getproperty(c::Composition, p::Symbol) = getproperty(getfield(c, :data), p)

# -------------
# VECTOR SPACE
# -------------

+(c₁::Composition{D,PARTS}, c₂::Composition{D,PARTS}) where {D,PARTS} =
  Composition(PARTS, 𝓒(components(c₁) .* components(c₂)))

-(c::Composition) = Composition(parts(c), 𝓒(1 ./ components(c)))

-(c₁::Composition, c₂::Composition) = c₁ + -c₂

*(λ::Real, c::Composition) = Composition(parts(c), 𝓒(components(c).^λ))

==(c₁::Composition, c₂::Composition) =
  parts(c₁) == parts(c₂) && 𝓒(components(c₁)) ≈ 𝓒(components(c₂))

⋅(c₁::Composition{D}, c₂::Composition{D}) where {D} = begin
  x, y = components(c₁), components(c₂)
  sum(log(x[i]/x[j])*log(y[i]/y[j]) for j=1:D for i=j+1:D) / D
end

norm(c::Composition) = √(c⋅c)

distance(c₁::Composition, c₂::Composition) = norm(c₁ - c₂)

# --------------------
# RANDOM COMPOSITIONS
# --------------------

"""
    rand(Composition{D}, n=1)

Generates `D`-part composition at random according to a
balanced Dirichlet distribution.
"""
function rand(rng::Random.AbstractRNG,
              ::Random.SamplerType{<:Composition{D}}) where {D}
  α = fill(1.0, D)
  d = Dirichlet(α)
  x = rand(rng, d)
  Composition(x)
end

# -----------
# IO METHODS
# -----------

function Base.show(io::IO, c::Composition)
  w = [(@sprintf "%.03f" w) for w in components(c)]
  show(io, join(w, " : "))
end

function Base.show(io::IO, mime::MIME"text/plain",
                   c::Composition{D,PARTS}) where {D,PARTS}
  w = components(c)
  x = Vector{Float64}()
  p = Vector{Symbol}()
  m = Vector{Symbol}()
  for i in 1:D
    if ismissing(w[i])
      push!(m, PARTS[i])
    else
      push!(p, PARTS[i])
      push!(x, w[i])
    end
  end
  plt = barplot(p, x, title="$D-part composition")
  isempty(m) || annotate!(plt, :t, "missing: $(join(m,", "))")
  show(io, mime, plt)
end


# * Derivatives

function diffop!(dest::AbstractMatrix,
                 L::BSplineOrRestricted, R::BSplineOrRestricted,
                 a, b, op=I)
    assert_compatible_bases(L,R)

    k = max(order(L),order(R))
    (bandwidth(dest,1) ≤ k-1) && (bandwidth(dest,2) ≤ k-1) ||
        throw(DimensionMismatch("Insufficient bandwidths of destination matrix"))

    # ∂' = -∂ ⇒ for weak Laplacians (in case of Dirichlet0 boundary
    # conditions), we negate the left basis.
    χ = (iseven(a) ? 1 : -1)*basis_functions(L, a)
    ξ = basis_functions(R, b)

    overlap_matrix!(dest, χ, op*ξ, weights(parent(R)))

    dest
end

diffop!(dest::BandedMatrix, B::BSplineOrRestricted, o) =
    diffop!(dest, B, B, o)

@materialize function *(Ac::AdjointBasisOrRestricted{<:BSpline},
                        D::Derivative,
                        B::BasisOrRestricted{<:BSpline})
    T -> begin
        Matrix(undef, parent(Ac), B, T)
    end
    dest::AbstractMatrix{T} -> begin
        diffop!(dest, parent(Ac), B, 0, 1)
    end
end

@materialize function *(Ac::AdjointBasisOrRestricted{<:BSpline},
                        Dc::QuasiAdjoint{<:Any,<:Derivative},
                        D::Derivative,
                        B::BasisOrRestricted{<:BSpline})
    T -> begin
        Matrix(undef, parent(Ac), B, T)
    end
    dest::AbstractMatrix{T} -> begin
        diffop!(dest, parent(Ac), B, 1, 1)
    end
end

@materialize function *(Ac::AdjointBasisOrRestricted{<:BSpline},
                        Dc::QuasiAdjoint{<:Any,<:Derivative},
                        O::QuasiDiagonal,
                        D::Derivative,
                        B::BasisOrRestricted{<:BSpline})
    T -> begin
        Matrix(undef, parent(Ac), B, T)
    end
    dest::AbstractMatrix{T} -> begin
        # Evaluate the quasi-diagonal operator O on the quadrature roots
        # of B.
        op = Diagonal(getindex.(Ref(O.diag), locs(parent(B))))
        diffop!(dest, parent(Ac), B, 1, 1, op)
    end
end

@materialize function *(Ac::AdjointBasisOrRestricted{<:BSpline},
                        Dc::QuasiAdjoint{<:Any,<:Derivative},
                        O::QuasiDiagonal,
                        P::QuasiDiagonal,
                        D::Derivative,
                        B::BasisOrRestricted{<:BSpline})
    T -> begin
        Matrix(undef, parent(Ac), B, T)
    end
    dest::AbstractMatrix{T} -> begin
        # Evaluate the quasi-diagonal operators O & P on the
        # quadrature roots of B.
        op = Diagonal(getindex.(Ref(O.diag), locs(parent(B))) .*
                      getindex.(Ref(P.diag), locs(parent(B))))
        diffop!(dest, parent(Ac), B, 1, 1, op)
    end
end


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

export determinize_fsa
"""
  `determinize_fsa(fsa)`

Returns a deterministic finite state acceptor. `fsa` must be an acceptor ([`is_acceptor`](@ref)).

Requires the WFST to be defined in a left distributive and weakly divisible semiring.

See [M. Mohri, F. Pereira, Fernando, M. Riley, Michael "Speech Recognition with Weighted Finite-State Transducers" in Springer Handb. Speech Process. 2008](https://cs.nyu.edu/~mohri/pub/hbka.pdf) for details.
"""
function determinize_fsa(fst::WFST{W,A,D,I,O}) where {W,A,D,I,O}
  if !isleft(W)
    throw(ErrorException("Semiring $W is not left distributive, fst cannot be determinized."))
  end
  if !isweaklydivisible(W)
    throw(ErrorException("Semiring $W is not weakly divisible, fst cannot be determinized."))
  end
  dfst = WFST(get_isym(fst), get_osym(fst); W = W)
  i = get_initial(fst)
  w = get_initial_weight(fst,i)
  initial!(dfst,get_initial(fst))

  S = Queue{Set{Tuple{D,W}}}()
  qq = Set([(i,w)])
  cnt = 1
  qq2state = Dict( qq => cnt ) # new states
  enqueue!(S,qq)
  while !isempty(S)
    pp = dequeue!(S)
    lab2weight = Dict{D,W}() # set in line 8 of Fig.10 in [1]
    lab_nextstate2weight = Dict{D,Dict{D,W}}() # set in line 9 of Fig.10 in [1]

    for (p,v) in pp
      for arc in fst[p]
        lab, w, n = get_ilabel(arc), get_weight(arc), get_nextstate(arc)
        vw = v*w
        if get_ilabel(arc) in keys(lab2weight)
          lab2weight[lab] +=  vw
        else
          lab2weight[lab]  =  vw
        end
        if lab in keys(lab_nextstate2weight)
          if n in keys(lab_nextstate2weight[lab])
            lab_nextstate2weight[lab][n] += vw
          else
            lab_nextstate2weight[lab][n] = vw
          end
        else
          lab_nextstate2weight[lab] = Dict(n => vw)
        end
      end
    end

    for lab in keys(lab2weight)
      qq = Set{Tuple{D,W}}()
      wp = lab2weight[lab]
      for n in keys(lab_nextstate2weight[lab])
        vw = lab_nextstate2weight[lab][n]
        push!(qq,(n, one(W) / wp * vw))
      end
      if !(qq in keys(qq2state))
        cnt += 1
        qq2state[qq] = cnt
        src, dest = qq2state[pp], cnt
        if max(src,dest) - length(dfst) > 0
          add_states!(dfst, max(src,dest) - length(dfst) )  # expand states
        end
        push!(dfst[src],Arc{W,D}(lab,lab,wp,dest))
        rho = zero(W)
        any_final = false
        for (n,v) in qq
          if isfinal(fst,n)
            any_final = true
            rho += v*get_final_weight(fst,n)
          end
        end
        if any_final final!(dfst,cnt,rho) end
        enqueue!(S,qq)
      else
        src, dest = qq2state[pp], qq2state[qq]
        push!(dfst[src],Arc{W,D}(lab,lab,wp,dest))
      end
    end
  end
  return dfst
end


using COSMO, Test, LinearAlgebra, Random, SparseArrays
rng = Random.MersenneTwister(41)

@testset "Interface" begin
# set!() -  to directly set the problem data without building constraints first

P = [4. 1;1 2]
q = [1; 1.]
A = [1. 1;1 0; 0 1]
l = [1.;0;0]
u = [1.;0.7;0.7]
Aa = [A; -A]
b = [u; -l]
convex_sets = [COSMO.Nonnegatives(3); COSMO.Nonnegatives(3)]

model = COSMO.Model()
COSMO.set!(model, P, q, Aa, b, convex_sets )
res = COSMO.optimize!(model)
@test isapprox(res.obj_val, 1.88, atol = 1e-3)

qf = [1; 1; 1]
Pf = zeros(1, 1)
Af = [1 2; 1 2]
bf = [1; 2]
model = COSMO.Model()
@test_throws DimensionMismatch COSMO.set!(model, P, qf, Aa, b, convex_sets)
@test_throws DimensionMismatch COSMO.set!(model, Pf, q, Aa, b, convex_sets)
@test_throws DimensionMismatch COSMO.set!(model, P, q, Af, b, convex_sets)
@test_throws DimensionMismatch COSMO.set!(model, P, q, Aa, bf, convex_sets)

# assemble! function test different combinations of P an q

#P not sparse, q dense
P = rand(rng, 2, 2)
q = rand(rng, 2)
A = rand(rng, 1, 2)
b = 1.
constraint = COSMO.Constraint(A, b, COSMO.Nonnegatives)
model = COSMO.Model();
COSMO.assemble!(model, P, q, constraint)
@test model.p.P == P
@test model.p.q == q
@test model.p.A == -A
@test model.p.b == [b]

# sparse input data
COSMO.empty_model!(model)
COSMO.assemble!(model, sparse(P), sparse(q), constraint)
@test model.p.P == P
@test model.p.q == q

A = rand(rng, 5, 1)
b = rand(rng, 5)
constraint = COSMO.Constraint(A, b, COSMO.Nonnegatives)

# P number, q vector
COSMO.empty_model!(model)
COSMO.assemble!(model, 1., [1.], constraint)
@test model.p.P == hcat([1.])
@test model.p.q == [1.]

# P vector, q number
COSMO.empty_model!(model)
COSMO.assemble!(model, [1.], 1., constraint)
@test model.p.P == hcat([1.])
@test model.p.q == [1.]

# P number, q number
COSMO.empty_model!(model)
COSMO.assemble!(model, 1., 1., constraint, settings = COSMO.Settings(verbose = true))
@test model.p.P == hcat([1.])
@test model.p.q == [1.]

# P vector, q matrix
COSMO.empty_model!(model)
COSMO.assemble!(model, [1.], hcat([1.]), constraint)
@test model.p.P == hcat([1.])
@test model.p.q == [1.]

# Test case where dimension of an A in one constraint is inconsistent with P
A = 1.
b = 0.
P = sparse(1.0I, 2, 2)
q = rand(rng, 2)
constraint = COSMO.Constraint(A, b, COSMO.Nonnegatives)
model = COSMO.Model();
@test_throws DimensionMismatch COSMO.assemble!(model, P, q, [constraint])

end


include("../oracles/assignment_oracle.jl")

using Test

@testset "Assignment Oracle" begin
    C = [2.0 6; 3 5]
    c = reshape(C, 4)
    (z, w) = assignment_oracle(c)
    W = reshape(w, (2,2))
    @test z ≈ 7.0
    @test W ≈ [1 0; 0 1]

    c_ham = assignment_hamming_reduce(c)
    (z2, w2) = assignment_oracle(c_ham)
    W2 = reshape(w2, (2,2))
    @test z2 ≈ 0
    @test W2 ≈ W
end