text:
# This script gets run once, on a developer's machine, to generate artifacts
using Pkg.Artifacts
using Downloads

function generate_artifacts(ver="latest",repo="https://github.com/JuliaPlots/PlotlyJS.jl")
    artifacts_toml = joinpath(dirname(@__DIR__), "Artifacts.toml")

    # if Artifacts.toml does not exist we also do not have to remove it
    isfile(artifacts_toml) && rm(artifacts_toml)

    plotschema_url = "https://api.plot.ly/v2/plot-schema?sha1"
    plotly_url = "https://cdn.plot.ly/plotly-$ver.min.js"

    plotlyartifacts_hash = create_artifact() do artifact_dir
        Downloads.download(plotschema_url, joinpath(artifact_dir, "plot-schema.json"))
        Downloads.download(plotly_url, joinpath(artifact_dir, "plotly.min.js"))
    end

    tarball_hash = archive_artifact(plotlyartifacts_hash, "plotly-artifacts-$ver.tar.gz")

    bind_artifact!(artifacts_toml, "plotly-artifacts", plotlyartifacts_hash; download_info = [
        ("$repo/releases/download/plotly-artifacts-$ver/plotly-artifacts-$ver.tar.gz", tarball_hash)
    ])
end


text:
export Descent, Momentum, AdaGrad, Adam

const ε = 1e-8

"""
    Descent(η=0.1)
Basic gradient descent optimizer with learning rate `η`.

# Parameters
- learning rate : `η`

# Example

"""
struct Descent
    eta::Float64
end

Descent(;η::Float64 = 0.1) = Descent(η)

apply!(opt::Descent, x::Array, Δ::Array) = opt.eta .* Δ

"""
    Momentum(η=0.01, α=0.9, velocity)
Momentum gradient descent optimizer with learning rate `η` and parameter of velocity `α`.

# Parameters
- learning rate : `η`
- parameter of velocity : `α`

# Example

"""
struct Momentum
    eta::Float64
    alpha::Float64
    velocity::Dict
end

Momentum(η = 0.01, α = 0.9) = Momentum(η, α, Dict())

function apply!(opt::Momentum, x::Array, Δ::Array)
    η, α = opt.eta, opt.alpha
    v = get!(() -> zero(x), opt.velocity, x)::typeof(x) #get!() function returns a view.
    @. v = (α*v - η*Δ)
    @. Δ = -v
end

"""
    AdaGrad(η = 0.01)
Gradient descent optimizer with learning rate attenuation.

# Parameters
- η : initial learning rate

# Examples

"""
struct AdaGrad
    eta::Float64
    h::Dict
end

AdaGrad(η = 0.01) = AdaGrad(η, Dict())

function apply!(opt::AdaGrad, x::Array, Δ::Array)
    η = opt.eta
    h = get!(() -> zero(x), opt.h, x)::typeof(x)
    @. h += Δ^2
    @. Δ *= η / (√h + ε)
    return Δ
end

"""
    Adam(η=0.01, β=(0.9, 0.99))
Gradient descent adaptive moment estimation optimizer.

# Parameters
- η : learning rate
- β : Decay of momentums

# Examples

"""
struct Adam
    eta::Float64
    beta::Tuple{Float64, Float64}
    recode::Dict
end

Adam(η = 0.01, β = (0.9, 0.99)) = Adam(η, β, Dict())

function apply!(opt::Adam, x::Array, Δ::Array)
    η, β = opt.eta, opt.beta
    mt, vt , βp = get!(opt.recode, x) do
        (zero(x), zero(x), Float64[β[1], β[2]])
    end::Tuple{typeof(x),typeof(x),Vector{Float64}}
    @. mt = β[1] * mt + (1 - β[1]) * Δ
    @. vt = β[2] * vt + (1 - β[2]) * Δ^2
    @. Δ =  mt / (1 - βp[1]) / (√(vt / (1 - βp[2])) + ε) * η
    βp .= βp .* β
    return Δ
end

ext: jl


text:
function optimize_controls!(
        m::ClimateModel;
        obj_option = "temp", temp_goal = 2.0, budget=10., expenditure = 0.5,
        max_deployment = Dict("mitigate"=>1., "remove"=>1., "geoeng"=>1., "adapt"=>0.4),
        max_slope = Dict("mitigate"=>1. /40., "remove"=>1. /40., "geoeng"=>1. /20., "adapt"=>0.),
        max_update = Dict("mitigate"=>nothing, "remove"=>nothing, "geoeng"=>nothing, "adapt"=>0.1),
        temp_final = nothing,
        delay_deployment = Dict(
            "mitigate"=>0,
            "remove"=>10,
            "geoeng"=>30,
            "adapt"=>0
        ),
        cost_exponent = 2.,
        mitigation_penetration = nothing,
        print_status = false, print_statistics = false, print_raw_status = true,
    )

    # Translate JuMP printing options
    if print_status
        if print_statistics
            bool_str = "yes"
        else
            bool_str = "no"
        end
        print_int = 1
    else
        print_int = 0
        bool_str = "no"
    end

    # Shorthands
    tarr = t(m)
    Earr = E(m)
    τ = τd(m)
    dt = m.domain.dt
    t0 = tarr[1]
    tp = m.domain.present_year
    q = m.economics.baseline_emissions
    qGtCO2 = ppm_to_GtCO2(q)
    N = length(tarr)

    # Set default temperature goal for end year
    if isnothing(temp_final)
        temp_final = temp_goal
    elseif temp_final >= temp_goal
        temp_final = temp_goal
    end

    # Set defaults for start_deployment
    start_deployment = Dict()
    for (key, item) in delay_deployment
        start_deployment[key] = t0 + delay_deployment[key]
    end

    max_slope_update = Dict()
    for (key, item) in max_update
        if isnothing(item)
            max_slope_update[key] = 0
            max_update[key] = Inf
        else
            max_slope_update[key] = 1
        end
    end

    if typeof(cost_exponent) in [Int64, Float64]
        cost_exponent = Dict(
            "mitigate"=>cost_exponent,
            "remove"=>cost_exponent,
            "geoeng"=>cost_exponent,
            "adapt"=>cost_exponent
        )
    end

    model_optimizer = Model(optimizer_with_attributes(Ipopt.Optimizer,
        "acceptable_tol" => 1.e-8, "max_iter" => Int64(1e8),
        "acceptable_constr_viol_tol" => 1.e-3, "constr_viol_tol" => 1.e-4,
        "print_frequency_iter" => 50,  "print_timing_statistics" => bool_str,
        "print_level" => print_int,
    ))

    function fM_JuMP(α)
        if α <= 0.
            return 100.
        else
            base_cost = α ^ cost_exponent["mitigate"]
            if isnothing(mitigation_penetration)
                return base_cost
            else
                penetration_factor = (
                    1. /(1. - exp( - (1. - α) / (1. - mitigation_penetration)))
                )
                return base_cost * penetration_factor
            end
        end
    end
    register(model_optimizer, :fM_JuMP, 1, fM_JuMP, autodiff=true)

    function fA_JuMP(α)
        if α <= 0.
            return 1000.
        else
            return α ^ cost_exponent["adapt"]
        end
    end
    register(model_optimizer, :fA_JuMP, 1, fA_JuMP, autodiff=true)

    function fR_JuMP(α)
        if α <= 0.
            return 1000.
        else
            return α ^ cost_exponent["remove"]
        end
    end
    register(model_optimizer, :fR_JuMP, 1, fR_JuMP, autodiff=true)

    function fG_JuMP(α)
        if α <= 0.
            return 1000.
        else
            return α ^ cost_exponent["geoeng"]
        end
    end
    register(model_optimizer, :fG_JuMP, 1, fG_JuMP, autodiff=true)

    function log_JuMP(x)
        if x <= 0.
            return -1000.0
        else
            return log(x)
        end
    end
    register(model_optimizer, :log_JuMP, 1, log_JuMP, autodiff=true)

    function discounting_JuMP(t)
        if t < tp
            return 0.
        else
            return (
                (1. + m.economics.ρ) ^
                (-(t - tp))
            )
        end
    end
    register(model_optimizer, :discounting_JuMP, 1, discounting_JuMP, autodiff=true)

    # constraints on control variables
    @variables(model_optimizer, begin
            0. <= M[1:N] <= max_deployment["mitigate"]  # emissions reductions
            0. <= R[1:N] <= max_deployment["remove"]  # negative emissions
            0. <= G[1:N] <= max_deployment["geoeng"]  # geoengineering
            0. <= A[1:N] <= max_deployment["adapt"]  # adapt
    end)

    control_vars = Dict(
        "mitigate" => M,
        "remove" => R,
        "geoeng" => G,
        "adapt" => A
    )
    controls = Dict(
        "mitigate" => m.controls.mitigate,
        "remove" => m.controls.remove,
        "geoeng" => m.controls.geoeng,
        "adapt" => m.controls.adapt
    )

    domain_idx = (tarr .> tp) # don't update past or present
    if tarr[1] == tp
        domain_idx = (tarr .>= tp) # unless present is also first timestep
    end

    M₀ = m.economics.mitigate_init
    R₀ = m.economics.remove_init
    G₀ = m.economics.geoeng_init
    A₀ = m.economics.adapt_init

    control_inits = Dict(
        "mitigate" => M₀,
        "remove" => R₀,
        "geoeng" => G₀,
        "adapt" => A₀
    )

    for (key, control) in control_vars
        if control_inits[key] !== nothing
            fix(control_vars[key][1], control_inits[key]; force = true)
            Nstart = 2
        else
            Nstart = 1
        end

        for idx in Nstart:N
            if idx <= length(tarr[.~domain_idx])
                fix(control_vars[key][idx], controls[key][idx]; force = true)
            else
                if tarr[idx] < start_deployment[key]
                    if control_inits[key] !== nothing
                        fix(control_vars[key][idx], control_inits[key]; force = true)
                    else
                        fix(control_vars[key][idx], 0.; force = true)
                    end
                end
            end
        end
    end

    # add integral function as a new variable defined by first order finite differences
    @variable(model_optimizer, cumsum_qMR[1:N]);
    for i in 1:N-1
        @constraint(
            model_optimizer, cumsum_qMR[i+1] - cumsum_qMR[i] ==
            (dt * (m.physics.r * (q[i+1] * (1. - M[i+1]) - q[1] * R[i+1])))
        )
    end
    @constraint(
        model_optimizer, cumsum_qMR[1] == (dt * (m.physics.r * (q[1] * (1. - M[1]) - q[1] * R[1])))
    );

    # add temperature kernel as new variable defined by first order finite difference
    @variable(model_optimizer, cumsum_KFdt[1:N]);
    for i in 1:N-1
        @NLconstraint(
            model_optimizer, cumsum_KFdt[i+1] - cumsum_KFdt[i] ==
            (
                dt *
                exp( (tarr[i+1] - (t0 - dt)) / τ ) * (
                    m.physics.a * log_JuMP(
                        (m.physics.c0 + cumsum_qMR[i+1]) /
                        (m.physics.c0)
                    ) - m.economics.Finf*G[i+1] )
            )
        )
    end
    @NLconstraint(
        model_optimizer, cumsum_KFdt[1] ==
        (
            dt *
            exp( dt / τ ) * (
                m.physics.a * log_JuMP(
                    (m.physics.c0 + cumsum_qMR[1]) /
                    (m.physics.c0)
                ) - m.economics.Finf*G[1] )
         )
    );

    # Add constraint of rate of changes
    present_idx = findmin(abs.(tarr .- tp))[2]

    @variables(model_optimizer, begin
            -max_slope["mitigate"] <= dMdt[present_idx+max_slope_update["mitigate"]:N-1] <= max_slope["mitigate"]
            -max_slope["remove"] <= dRdt[present_idx+max_slope_update["remove"]:N-1] <= max_slope["remove"]
            -max_slope["geoeng"] <= dGdt[present_idx+max_slope_update["geoeng"]:N-1] <= max_slope["geoeng"]
            -max_slope["adapt"] <= dAdt[present_idx+max_slope_update["adapt"]:N-1] <= max_slope["adapt"]
    end);
    for i in present_idx+max_slope_update["mitigate"]:N-1
        @constraint(model_optimizer, dMdt[i] == (M[i+1] - M[i]) / dt)
    end
    for i in present_idx+max_slope_update["remove"]:N-1
        @constraint(model_optimizer, dRdt[i] == (R[i+1] - R[i]) / dt)
    end
    for i in present_idx+max_slope_update["geoeng"]:N-1
        @constraint(model_optimizer, dGdt[i] == (G[i+1] - G[i]) / dt)
    end
    for i in present_idx+max_slope_update["adapt"]:N-1
        @constraint(model_optimizer, dAdt[i] == (A[i+1] - A[i]) / dt)
    end

    if present_idx > 1
        @variables(model_optimizer, begin
            -max_update["mitigate"] <= dM <= max_update["mitigate"]
            -max_update["remove"] <= dR <= max_update["remove"]
            -max_update["geoeng"] <= dG <= max_update["geoeng"]
            -max_update["adapt"] <= dA <= max_update["adapt"]
        end);
        @constraint(model_optimizer, dM == (M[present_idx+1] - M[present_idx]))
        @constraint(model_optimizer, dR == (R[present_idx+1] - R[present_idx]))
        @constraint(model_optimizer, dG == (G[present_idx+1] - G[present_idx]))
        @constraint(model_optimizer, dA == (A[present_idx+1] - A[present_idx]))
    end

    if obj_option == "net_benefit"
        # in practice we solve the equivalent problem of minimizing the net cost (- net benefit)
        @NLobjective(model_optimizer, Min,
            sum(
                (
                    (1 - A[i]) * m.economics.β *
                    Earr[i] *
                    ((m.physics.T0 +
                        (
                             (m.physics.a * log_JuMP(
                                        (m.physics.c0 + cumsum_qMR[i]) /
                                        (m.physics.c0)
                                    ) - m.economics.Finf*G[i]
                            ) +
                            m.physics.κ /
                            (τ * m.physics.B) *
                            exp( - (tarr[i] - (t0 - dt)) / τ ) *
                            cumsum_KFdt[i]
                        ) / (m.physics.B + m.physics.κ)
                    )
                    )^2 +
                    m.economics.mitigate_cost * qGtCO2[i] *
                    fM_JuMP(M[i]) +
                    m.economics.adapt_cost * fA_JuMP(A[i]) +
                    m.economics.remove_cost * fR_JuMP(R[i]) +
                    m.economics.geoeng_cost * Earr[i] *
                    fG_JuMP(G[i])
                ) *
                discounting_JuMP(tarr[i]) *
                dt
            for i=1:N)
        )

    elseif obj_option == "temp"
        @NLobjective(model_optimizer, Min,
            sum(
                (
                    m.economics.mitigate_cost * qGtCO2[i] *
                    fM_JuMP(M[i]) +
                    m.economics.adapt_cost * fA_JuMP(A[i]) +
                    m.economics.remove_cost * fR_JuMP(R[i]) +
                    m.economics.geoeng_cost * Earr[i] *
                    fG_JuMP(G[i])
                ) *
                discounting_JuMP(tarr[i]) *
                dt
            for i=1:N)
        )

        for i in 1:N-1
            @NLconstraint(model_optimizer,
                (1 - A[i]) * m.economics.β *
                Earr[i] *
                ((m.physics.T0 +
                    (
                        (m.physics.a * log_JuMP(
                                    (m.physics.c0 + cumsum_qMR[i]) /
                                    (m.physics.c0)
                                ) - 8.5*G[i]
                        ) +
                        m.physics.κ /
                        (τ * m.physics.B) *
                        exp( - (tarr[i] - (t0 - dt)) / τ ) *
                        cumsum_KFdt[i]
                    ) / (m.physics.B + m.physics.κ)
                )
                )^2 <= (
                    m.economics.β *
                    Earr[i] *
                    temp_goal^2
                )
            )
        end
        i=N
        @NLconstraint(model_optimizer,
            (1 - A[i]) * m.economics.β *
            Earr[i] *
            ((m.physics.T0 +
                (
                    (m.physics.a * log_JuMP(
                                (m.physics.c0 + cumsum_qMR[i]) /
                                (m.physics.c0)
                            ) - 8.5*G[i]
                    ) +
                    m.physics.κ /
                    (τ * m.physics.B) *
                    exp( - (tarr[i] - (t0 - dt)) / τ ) *
                    cumsum_KFdt[i]
                ) / (m.physics.B + m.physics.κ)
            )
            )^2 <= (
                m.economics.β *
                Earr[i] *
                temp_final^2
            )
        )

    elseif obj_option == "budget"
        @NLobjective(model_optimizer, Min,
            sum(
                (1 - A[i]) * m.economics.β *
                Earr[i] *
                ((m.physics.T0 +
                    (
                        (m.physics.a * log_JuMP(
                                    (m.physics.c0 + cumsum_qMR[i]) /
                                    (m.physics.c0)
                                ) - 8.5*G[i]
                        ) +
                        m.physics.κ /
                        (τ * m.physics.B) *
                        exp( - (tarr[i] - (t0 - dt)) / τ ) *
                        cumsum_KFdt[i]
                    ) / (m.physics.B + m.physics.κ)
                )
                )^2 *
                discounting_JuMP(t[i]) *
                dt
            for i=1:N)
        )

        @NLconstraint(model_optimizer,
            sum(
                (
                    m.economics.mitigate_cost * qGtCO2[i] *
                    fM_JuMP(M[i]) +
                    m.economics.adapt_cost * fA_JuMP(A[i]) +
                    m.economics.remove_cost * fR_JuMP(R[i]) +
                    m.economics.geoeng_cost * Earr[i] *
                    fG_JuMP(G[i])
                ) *
                discounting_JuMP(t[i]) *
                dt
            for i=1:N) <= budget
        )

    elseif obj_option == "expenditure"
        @NLobjective(model_optimizer, Min,
            sum(
                (1 - A[i]) * m.economics.β *
                Earr[i] *
                ((m.physics.T0 +
                    (
                        (m.physics.a * log_JuMP(
                                    (m.physics.c0 + cumsum_qMR[i]) /
                                    (m.physics.c0)
                                ) - 8.5*G[i]
                        ) +
                        m.physics.κ /
                        (τ * m.physics.B) *
                        exp( - (tarr[i] - (t0 - dt)) / τ ) *
                        cumsum_KFdt[i]
                    ) / (m.physics.B + m.physics.κ)
                )
                )^2 *
                discounting_JuMP(t[i]) *
                dt
            for i=1:N)
        )

        for i in 1:N
            @NLconstraint(model_optimizer,
                (
                    m.economics.mitigate_cost * qGtCO2[i] *
                    fM_JuMP(M[i]) +
                    m.economics.adapt_cost * fA_JuMP(A[i]) +
                    m.economics.remove_cost * fR_JuMP(R[i]) +
                    m.economics.geoeng_cost * Earr[i] *
                    fG_JuMP(G[i])
                ) <= expenditure * Earr[i]
            )
        end
    end

    optimize!(model_optimizer)

    if print_raw_status
        print(raw_status(model_optimizer), "\n")
    end

    mitigate_values = value.(M)[domain_idx]
    mitigate_values[q[domain_idx].==0.] .= 0.
    getfield(m.controls, :mitigate)[domain_idx] = mitigate_values
    getfield(m.controls, :remove)[domain_idx] = value.(R)[domain_idx]
    getfield(m.controls, :geoeng)[domain_idx] = value.(G)[domain_idx]
    getfield(m.controls, :adapt)[domain_idx] = value.(A)[domain_idx]

    return model_optimizer
end


text:
"""
Tools for the exploration of chaos and nonlinear dynamics
"""
module ChaosTools

using Reexport
@reexport using DynamicalSystemsBase

using DynamicalSystemsBase: DS, DDS, CDS
using DynamicalSystemsBase: MDI, TDI
using DynamicalSystemsBase: stateeltype
using DiffEqBase: AbstractODEIntegrator, u_modified!

include("lyapunovs.jl")

include("orbitdiagrams/discrete_diagram.jl")
include("orbitdiagrams/poincare.jl")

include("dimensions/entropies.jl")
include("dimensions/dims.jl")

include("nlts.jl")

include("periodicity/periodic_points.jl")
include("periodicity/period.jl")

include("chaos_detection.jl")

include("partially_predictable.jl")

# Ugly methods that shouldn't exist:
include("ugliness.jl")

end # module


text:
module Input

using ..TypeSynonyms: RationalComplex

struct FiberData
    E        :: Vector{Float64}
    μ        :: Vector{Float64}
    G        :: Vector{Float64}
    radii    :: Vector{Float64}
    r_inner  :: Float64
    center   :: RationalComplex
    n_terms  :: Int64
end

struct CohesiveData
    E  :: Float64
    μ  :: Float64
    G  :: Float64
end

struct CellData
    l1       :: Float64
    l2       :: Float64
    γ        :: Float64
    fibers   :: Vector{FiberData}
    cohesive :: CohesiveData
end

end # module Input


text:
# scale columns of 2d Array by first row elements

calc(::BM_TEST, ::FUNT, x) = calc(BM_JBase1(), FUNT(), x)
calc(::BM_JBase1, ::FUNT, x) = begin
    ix = 1
    nx = size(x,1)
    @inbounds for j = 1:size(x,2)
        vx = view(x, ix:ix+nx-1)
        scale!(vx, x[ix])
        ix += nx
    end
    return x
end
calc(::BM_JBase2, ::FUNT, x) = begin
    ix = 1
    nx = size(x,1)
    @inbounds for j = 1:size(x,2)
        vx = view(x, ix:ix+nx-1)
        scale!(x[ix], vx)
        ix += nx
    end
    return x
end
calc(::BM_BLAS, ::FUNT, x) = begin
    ix = 1
    nx = size(x,1)
    @inbounds for j = 1:size(x,2)
        BLAS.scal!(nx, x[ix], pointer(x,ix), 1)
        ix += nx
    end
    return x
end
calc(::BM_Broadcast, ::FUNT, x) = begin
    ix = 1
    nx = size(x,1)
    @inbounds for j = 1:size(x,2)
        vx = view(x, ix:ix+nx-1)
        broadcast!(*, vx, vx, x[ix])
        ix += nx
    end
    return x
end
calc(::BM_Forloop, ::FUNT, x) = scaleloop1_2d(x)
function scaleloop1_2d(x::Array)
    ix = 1
    nx = size(x,1)
    @inbounds for j = 1:size(x,2)
        v = x[ix]
        for i = ix:ix+nx-1
            x[i] *= v
        end
        ix += nx
    end
    return x
end
calc(::BM_FAO, ::FUNT, x) = begin
    ix = 1
    nx = size(x,1)
    @inbounds for i = 1:size(x,2)
        unsafe_fast_scale!(x, ix, x[ix], nx)
        ix += nx
    end
    return x
end

Base.start{T}(p::BenchCase{T,FUNT}, n::Int) =
    (rand(n,16), )

function Base.run{Op,FUNT}(p::BenchCase{Op,FUNT}, n::Int, s)  # bench type, config, start value
    calc(Op(), FUNT(), s...)
    return nothing
end


text:
@testset "linter" begin

import Dolo

# an example model shouldn't raise any error
models_dir = joinpath(Dolo.pkg_path,"examples","models")
for fname in (readdir(models_dir))
    filename = joinpath(models_dir, fname)
    if isfile(filename)
        println(filename)
        ch = Dolo.lint(filename)
        @assert length(ch) == 0
    end
end

# for now, we check that problematic models don't make linter failxx
errors_dir = joinpath(Dolo.pkg_path,"linter","models")
for fname in (readdir(errors_dir))
    filename = joinpath(errors_dir, fname)
    if isfile(filename)
        ok = true
        try
            ch = Dolo.lint(filename)
        catch err
            ok = false
            println("Uncaught error while parsing")
            println(err)
        end
        @assert ok
    end
end

end


text:
# Indepndent component factorization (ICF) based on Independent component analysis (ICA)
# We will compare FCF with ICF

"""
    icf(Z; opt = "orth")

Apply independent component analysis to Z, return estimated
mixing matrix and independent components

# Arguments
-   `Z`: an array of vectors of the same length, represent the
       realizations of mixed independent signals
-   `obj`: the type of loss function, valid options are "kurt" and "ent"
-   `opt`: string, type of optimization, valid option: "orth", "sphe".
        The "sphe" is only designed for "kurt" loss function

# Outputs
-   `Aest`: estimated mixing matrix of size `s`-by-`s`, where `s` = size(Z,1)
-   `Xest`: estimated independent component `Xest = pinv(Aest)*Z`
"""
function icf(Z::Array{Array{D, 1}, 1}; obj = "kurt", opt = "orth") where D <: Number
    # get number of the components and the dimension of each components
    s = size(Z,1)
    T = size(Z[1],1)

    # center every components
    Zbar = mean.(Z)
    Zc = Z .- [Zbar[i]*ones(T) for i in 1:s]

    # compute the covariance matrix
    # whiten the Zc
    C = [Zc[i]'*Zc[j]/T for i = 1:s, j = 1:s]
    eigcov = eigen(C);
    Λ, U = eigcov.values, eigcov.vectors
    Λ[Λ.<0] .= 0 # keep non-negative eigs
    Σ = Diagonal(Λ).^0.5
    Y = U*pinv(Σ)*U'*Zc

    if obj == "kurt"
        # find the orthogonal matrix minimizer neg_abs_sum_kurt

        # optimization using optim
        if opt == "orth"
            # loss function
            F = W -> neg_abs_sum_kurt(W'*Y)

            # gradient of loss function
            grad_F = W -> grad_neg_abs_sum_kurt(W, Y)

            # using Stiefel manifold
            W = OptOrtho(F, grad_F, s)
        elseif opt == "sphe"
            # using Sphere manifold and find one column at a time
            # loss function
            F = W -> neg_abs_sum_kurt(W[:,:]'*Y)

            # gradient of loss function
            grad_F = W -> grad_neg_abs_sum_kurt(W, Y)

            # first column
            W = reshape(OptSphere(F, grad_F, s), :, 1)

            # find 2nd to s-1th column one at a time
            for i = 2: s - 1
                pY = (I - W*W')*Y

                # updated loss function
                F = W -> neg_abs_sum_kurt(W[:,:]'*pY)

                # updated gradient of loss function
                grad_F = W -> grad_neg_abs_sum_kurt(W, pY)

                # find i th column, attach it to W
                Wi = OptSphere(F, grad_F, s)
                W = (qr(hcat(W, Wi)).Q)[:, 1:i]
            end

            # find the s-th column by take qr decomposition
            W = Matrix(qr(hcat(W, randn(s))).Q)
        end
    elseif obj == "ent"
        F = W -> ent_sum(W'*Y)

        grad_F = W -> grad_ent_sum(W, Y)

        # using Stiefel manifold
        W = OptOrtho(F, grad_F, s)
    end

    Aest = U*Σ*U'*W
    xest = pinv(Aest) * Z

    # order Xest by the descending order of absolute free kurtosis if obj = "kurt"
    if obj == "kurt"
        order = sortperm(xest, by = x -> -abs(kurt(x)))
    elseif obj == "ent"
        order = sortperm(xest, by = x -> -negent(x))
    end

    xest = xest[order]
    Aest = Aest[:,order]

    return Aest, xest
end

"""
    neg_abs_sum_kurt(Z)

Calculate the sum of negative absolute value of kurtosis of rows of Z
"""
function neg_abs_sum_kurt(Z)
    return -sum(abs.(kurt.(Z)))
end

"""
    kurt(z)

This function returns the kurtosis of z

# Arguments
-   `z`: an vector regarded as independent realizations of a random variables

# Outputs
-   `k`: a scalar, kurtosis of z
"""
function kurt(z)
    #--------------------------------------------------------------------------
    # Syntax:       k = kurt(z)
    #
    # Input:        z: an vector regarded as independent realizations of a random variables
    #
    # Outputs:      k: a scalar, kurtosis of z
    #
    # Description:  This function returns the kurtosis of z
    #
    # Authors:      Hao Wu
    #               lingluanwh@gmail.com
    #
    # Date:         Jan 20, 2019
    #--------------------------------------------------------------------------
    z = z .- mean(z)
    k = mean(z.^4) - 3*mean(z.^2)^2
    return k
end

"""
grad_neg_abs_sum_kurt(W, Z)

This function calcualtes the gradient of neg_abs_sum_kurt(W'*Z) w.r.t W

# Arguments
-   `W`: a matrix such that size(W, 1) = size(Z, 1)
-   `Z`: an array of vectors of the same length, regarded as the
       realizations of mixed independent signals

# Outputs
-   `grad`: the gradient of neg_abs_sum_kurt(W'*Z) w.r.t W
"""
function grad_neg_abs_sum_kurt(W, Z)
    # get number of components and dimension
    s = size(Z, 1)

    # center rows of Z
    Z = [Z[i] .- mean(Z[i]) for i = 1: s]

    # initalize grad matrix
    grad = complex(zeros(size(W)))

    # Compute the gradient
    X = (W[:,:])'*Z

    for l = 1: size(grad , 2)
        sl = sign(kurt(X[l]))
        for k = 1: size(grad, 1)
            grad[k, l] = -sl * (4 * mean(Z[k].*(X[l].^3)) - 12 * mean(X[l].^2)*mean(Z[k].*X[l]))
        end
    end

    return grad
end

"""
negent(z)

This function returns the negentropy of z

# Arguments
-   `z`: an vector regarded as independent realizations of a random variables

# Outputs
-   `k`: a scalar, negentropy of z
"""
function negent(z)
    # approximate the entropy of z by cumulants
    z = z .- mean(z); # center z
    return kurt(z)^2/48 + mean(z.^3)^2/12
end

"""
    ent_sum(Z)

Calculate the negative um of negentropy of rows of Z
"""
function ent_sum(Z)
    return -sum(negent.(Z))
end

"""
grad_ent_sum(W, Z)

This function calcualtes the gradient of ent_sum(W'*Z) w.r.t W

# Arguments
-   `W`: a matrix such that size(W, 1) = size(Z, 1)
-   `Z`: an array of vectors of the same length, regarded as the
       realizations of mixed independent signals

# Outputs
-   `grad`: the gradient of ent_sum(W'*Z) w.r.t W
"""
function grad_ent_sum(W, Z)
    # get number of components and dimension
    s = size(Z, 1)

    # center rows of Z
    Z = [Z[i] .- mean(Z[i]) for i = 1: s]

    # initalize grad matrix
    grad = complex(zeros(size(W)))

    # Compute the gradient
    X = (W[:,:])'*Z

    for l = 1: size(grad , 2)
        for k = 1: size(grad, 1)
            grad[k, l] =  -1 * kurt(X[l])/24 * (4 * mean(Z[k].*(X[l].^3)) - 12 * mean(X[l].^2)*mean(Z[k].*X[l]))
            grad[k, l] -= mean(X[l].^3)/6 * (3 * mean(Z[k] .* (X[l].^2)))
        end
    end

    return grad
end


text:
using JuMP, Ipopt, GLPK, Gurobi
using DualDecomposition, DRMSMIP
using CSV, SparseArrays

const DD = DualDecomposition
const parallel = DD.parallel

function run_main(foo)
    nStage = parse(Int, foo[1])
    nScenario = parse(Int, foo[2])
    radius = parse(Float64, foo[3])
    filename = foo[4]

    bus_data = CSV.File("./TEP/6-bus/bus_data.csv" ; datarow=2, delim=',')
    branch_data = CSV.File("./TEP/6-bus/branch_data.csv" ; datarow=2, delim=',');

    @assert nScenario <= 36

    nBus = length(bus_data)
    nBranch = length(branch_data)

    discount = 0.95
    LC = 1000 # load curtailment penalty

    generator_efficiency = 0.9
    pump_efficiency = 0.9
    evaporation_remain = 0.9
    r_init = zeros(nBus)
    r_init[3] = 50
    h_max = zeros(nBus)
    h_max[3] = 200

    MAX_LINE_ADDED = 5
    Buses = 1:nBus
    Branches = 1:nBranch

    solMethod = BM.ProximalMethod
    #solMethod = BM.TrustRegionMethod

    from = branch_data.From
    to = branch_data.To
    n0 = branch_data.n0_ij
    gamma = 100 ./ branch_data.Reactance
    fbar = branch_data.barf_ij
    cost = float(branch_data.Cost)

    genMax = bus_data.Gen_Max
    genLevel = bus_data.Gen_Level
    load = bus_data.Load

    inflow = Dict()
    outflow = Dict()
    for bus in Buses
        temp = []
        temp2 = []
        for branch in Branches
            if from[branch] == bus
                push!(temp, branch)
            end
            if to[branch] == bus
                push!(temp2, branch)
            end
        end
        outflow[bus] = temp
        inflow[bus] = temp2
    end
    #@show inflow
    #@show outflow

    function create_tree()::DRMSMIP.DR_Tree
        D_samp = generate_sample(float(bus_data.Load))
        set = DRMSMIP.WassersteinSet(D_samp, radius, DRMSMIP.norm_L1)
        cost_coeff = Float64[]
        for branch in Branches
            for line in 1:MAX_LINE_ADDED
                if line > n0[branch]
                    push!(cost_coeff, cost[branch])# cost for alpha_ijk
                else
                    push!(cost_coeff, 0.0)
                end
            end
        end
        tree = DRMSMIP.DR_Tree([0.0], set, cost_coeff)
        add_nodes!(tree, 1, 1)
        return tree
    end

    function generate_sample(π::Array{Float64})::Array{DRMSMIP.Sample}
        ret = Array{DRMSMIP.Sample}(undef, nScenario)
        multiplier = range(1.0,stop=1.2,length=nScenario)
        for ii in 1:nScenario
            ξ = π * multiplier[ii]
            ret[ii] = DRMSMIP.Sample(ξ, 1/nScenario)
        end
        return ret
    end

    function add_nodes!(tree::DRMSMIP.DR_Tree, id::Int, stage::Int)
        if stage < nStage-1
            for ii in 1:nScenario
                demand = tree.nodes[id].set.samples[ii].ξ
                D_samp = generate_sample(demand)
                set = DRMSMIP.WassersteinSet(D_samp, radius, DRMSMIP.norm_L1)
                cost_coeff = Float64[]
                for branch in Branches
                    for line in 1:MAX_LINE_ADDED
                        push!(cost_coeff, cost[branch] * discount ^ stage )# cost for alpha_ijk
                    end
                end
                for branch in Branches
                    for line in 1:MAX_LINE_ADDED
                        push!(cost_coeff, -cost[branch] * discount ^ stage )# cost for alpha0_ijk
                    end
                end
                for bus in Buses
                    push!(cost_coeff, LC * discount ^ stage)# cost for DC_i
                end
                for bus in Buses
                    push!(cost_coeff, 0.0)# cost for beta_i
                end
                DD.addchild!(tree, id, demand, set, cost_coeff)
                childid = length(tree.nodes)
                add_nodes!(tree, childid, stage+1)
            end
        elseif stage == nStage-1
            for ii in 1:nScenario
                demand = tree.nodes[id].set.samples[ii].ξ
                cost_coeff = Float64[]
                # for branch in Branches
                #     for line in 1:MAX_LINE_ADDED
                #         push!(cost_coeff, cost[branch] * discount ^ stage )# cost for alpha_ijk
                #     end
                # end
                # for branch in Branches
                #     for line in 1:MAX_LINE_ADDED
                #         push!(cost_coeff, -cost[branch] * discount ^ stage )# cost for alpha0_ijk
                #     end
                # end
                for bus in Buses
                    push!(cost_coeff, LC * discount ^ stage)# cost for DC_i
                end
                DD.addchild!(tree, id, demand, nothing, cost_coeff)
            end
        end
    end

    function create_scenario_model(tree::DRMSMIP.DR_Tree, id::Int)
        hist = DD.get_history(tree, id)
        #m = Model(GLPK.Optimizer)
        m = Model(Gurobi.Optimizer)
        JuMP.set_optimizer_attribute(m, "OutputFlag", 0)

        @variable(m, alpha[1:nStage,Branches,1:MAX_LINE_ADDED], Bin)
        #@variable(m, 1 >= alpha[1:nStage,Branches,1:MAX_LINE_ADDED] >= 0)
        @variable(m, alpha0[2:nStage,Branches,1:MAX_LINE_ADDED])
        @variable(m, generation[1:nStage,Buses]>=0)
        @variable(m, dc[1:nStage,Buses]>=0)
        @variable(m, π/2 >= theta[1:nStage,Buses] >= -π/2)
        @variable(m, flow[1:nStage,Branches,1:MAX_LINE_ADDED])
        @variable(m, beta[1:nStage,Buses]>=0)
        @variable(m, bG[2:nStage,Buses]>=0)
        @variable(m, bP[2:nStage,Buses]>=0)

        a0 = zeros(nBranch, MAX_LINE_ADDED)
        for branch in Branches, k in 1: MAX_LINE_ADDED
            if k <= n0[branch]
                a0[branch,k] = 1
            end
        end

        @constraints(m, begin
                        [branch in Branches,k in 1:n0[branch]], alpha[1, branch, k] == 1
                        [branch in Branches,k in 2:MAX_LINE_ADDED], alpha[1, branch, k] <= alpha[1, branch, k-1]
                        [bus in Buses], beta[1, bus] == r_init[bus]
                    end)

        for stage = 2:nStage
            demand = tree.nodes[hist[stage]].ξ

            for bus in Buses
                @constraints(m, begin
                    sum(flow[stage, branch, k] for branch in inflow[bus], k in 1:MAX_LINE_ADDED)-sum(flow[stage, branch, k] for branch in outflow[bus], k in 1:MAX_LINE_ADDED) + generation[stage, bus] + dc[stage, bus] + generator_efficiency*bG[stage,bus] - bP[stage,bus]== demand[bus]
                end)
            end

            @constraints(m, begin

                [branch in Branches], sum(flow[stage,branch, k] for k in 1:n0[branch]) - gamma[branch]*n0[branch]*(theta[stage,from[branch]] - theta[stage,to[branch]]) == 0

                [branch in Branches,k in n0[branch]+1:MAX_LINE_ADDED], flow[stage,branch, k] - gamma[branch] * (theta[stage,from[branch]] - theta[stage,to[branch]]) - gamma[branch]*pi * (1-alpha0[stage,branch, k]) <= 0
                [branch in Branches,k in n0[branch]+1:MAX_LINE_ADDED], flow[stage,branch, k] - gamma[branch] * (theta[stage,from[branch]] - theta[stage,to[branch]]) + gamma[branch]*pi * (1-alpha0[stage,branch, k]) >= 0

                [branch in Branches,k in 1:MAX_LINE_ADDED], flow[stage,branch, k] <= fbar[branch] * alpha0[stage,branch, k]
                [branch in Branches,k in 1:MAX_LINE_ADDED], flow[stage,branch, k] >= -fbar[branch] * alpha0[stage,branch, k]

                [bus in Buses], generation[stage,bus] <= genMax[bus]
                [bus in Buses], dc[stage,bus] <= demand[bus]

                [branch in Branches,k in 1:MAX_LINE_ADDED], alpha0[stage,branch, k] == alpha[stage-1,branch, k]
                #[bus in Buses], demand[stage,bus] == bus_data[bus,4]
            end)

            if stage < nStage
                @constraints(m, begin
                    [branch in Branches,k in 1:MAX_LINE_ADDED], alpha[stage,branch, k] >= alpha0[stage,branch, k]
                    [branch in Branches,k in 2:MAX_LINE_ADDED], alpha[stage,branch, k] <= alpha[stage,branch, k-1]
                end)
            end

            @constraints(m, begin
                [bus in Buses], beta[stage,bus] == evaporation_remain * beta[stage-1,bus] + pump_efficiency * bP[stage,bus] - bG[stage,bus]
                [bus in Buses], beta[stage,bus] <= h_max[bus]
                [bus in Buses], bG[stage,bus] <= evaporation_remain * beta[stage-1,bus]
            end)
        end

        @objective(m, Min, 0 )

        # open("model"*string(id)*".lp", "w") do f
        #     print(f, m)
        # end
        return m
    end

    function dual_decomp(tree::DRMSMIP.DR_Tree, nodelist, leafdict)
        # Create DualDecomposition instance.
        #algo = DRMSMIP.DRMS_LagrangeDual(tree, BM.TrustRegionMethod)
        algo = DRMSMIP.DRMS_LagrangeDual(tree, solMethod)

        ##
        parallel.partition(length(nodelist[nStage]))
        ##

        # Add Lagrange dual problem for each scenario s.
        models = Dict{Int,JuMP.Model}(nodelist[nStage][s] => create_scenario_model(tree,nodelist[nStage][s]) for s in parallel.getpartition())
        for s in parallel.getpartition()
            id = nodelist[nStage][s]
            DD.add_block_model!(algo, s, models[id])
        end

        coupling_variables = create_coupling_variables(tree, models, nodelist, leafdict)

        # Set nonanticipativity variables as an array of symbols.
        DD.set_coupling_variables!(algo, coupling_variables)

        bundle_init = DRMSMIP.initialize_bundle(tree, algo)
        return algo, bundle_init
    end

    function create_coupling_variables(tree::DRMSMIP.DR_Tree, models, nodelist, leafdict)
        coupling_variables = Vector{DD.CouplingVariableRef}()
        for s in parallel.getpartition()
            id = nodelist[nStage][s]
            model = models[id]
            hist = DD.get_history(tree, id)

            for stage in 1:nStage-1
                root = hist[stage]
                var_id = 1
                alpha_ref = model[:alpha]
                for branch in Branches
                    for k in 1:MAX_LINE_ADDED
                        push!(coupling_variables, DD.CouplingVariableRef(leafdict[id], [root, var_id], alpha_ref[stage, branch, k]))
                        var_id += 1
                    end
                end
                if stage >= 2
                    alpha0_ref = model[:alpha0]
                    for branch in Branches
                        for k in 1:MAX_LINE_ADDED
                            push!(coupling_variables, DD.CouplingVariableRef(leafdict[id], [root, var_id], alpha0_ref[stage, branch, k]))
                            var_id += 1
                        end
                    end
                    dc_ref = model[:dc]
                    for bus in Buses
                        push!(coupling_variables, DD.CouplingVariableRef(leafdict[id], [root, var_id], dc_ref[stage, bus]))
                        var_id += 1
                    end
                    beta_ref = model[:beta]
                    for bus in Buses
                        push!(coupling_variables, DD.CouplingVariableRef(leafdict[id], [root, var_id], beta_ref[stage, bus]))
                        var_id += 1
                    end
                end
            end
            # dummy coupling variables
            var_id = 1
            dc_ref = model[:dc]
            for bus in Buses
                push!(coupling_variables, DD.CouplingVariableRef(s, [id, var_id], dc_ref[nStage, bus]))
                var_id += 1
            end
        end
        return coupling_variables
    end

    function solve_DD!(algo::DRMSMIP.DRMS_LagrangeDual, bundle_init::Array{Float64,1}, fathom_val::Union{Float64,Nothing}=nothing)
        # Solve the problem with the solver; this solver is for the underlying bundle method.
        # DD.run!(algo, optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0), bundle_init)
        DD.run!(algo, optimizer_with_attributes(Gurobi.Optimizer, "OutputFlag" => 0), bundle_init, fathom_val)
    end

    function set_to_0!(m::JuMP.Model,stage::Int,branch::Int,k::Int)
        alpha = m[:alpha]
        x = alpha[stage,branch,k]
        set_upper_bound(x, 0)
        set_lower_bound(x, 0)
    end

    function set_to_1!(m::JuMP.Model,stage::Int,branch::Int,k::Int)
        alpha = m[:alpha]
        x = alpha[stage,branch,k]
        set_upper_bound(x, 1)
        set_lower_bound(x, 1)
    end

    function unset_bound!(m::JuMP.Model,stage::Int,branch::Int,k::Int)
        alpha = m[:alpha]
        x = alpha[stage,branch,k]
        set_upper_bound(x, 1)
        set_lower_bound(x, 0)
    end

    function set_beta!(m::JuMP.Model,stage::Int,bus::Int,val::Float64)
        beta = m[:beta]
        x = beta[stage,bus]
        set_upper_bound(x, val)
        set_lower_bound(x, val)
    end

    function branch_and_bound(tree::DRMSMIP.DR_Tree, time_limit::Float64)

        nodelist = DD.get_stage_id(tree)
        leafdict = DD.leaf2block(nodelist[nStage])
        main_prob = make_main_prob(tree,nodelist)

        parallel.init()
        algo, bundle_init = dual_decomp(tree,nodelist,leafdict)
        feasible_prob = get_problem_copy(algo, nodelist) # list of subproblems
        incumbent = [[algo,-Inf]]
        final_solution = nothing
        tot_time = nothing
        z_hi = nothing
        z_lo = nothing
        if parallel.is_root()
            time_start = time()
            z_hi = Inf
            z_lo = -Inf
            while true
                tot_time = time() - time_start
                if tot_time > time_limit
                    println("maximum time limit reached")
                    parallel.bcast(true)
                    break
                elseif length(incumbent) == 0
                    #parallel.finalize()
                    #return final_solution, feasible_prob, main_prob, nodelist
                    parallel.bcast(true)
                    break
                end
                parallel.bcast(false)
                current_node = pop!(incumbent)
                current = current_node[1]
                solve_DD!(current,bundle_init,z_hi)

                # bounding
                bounding = current.block_model.dual_bound < z_hi
                parallel.bcast(bounding)
                if bounding
                    alpha_identical, alpha, beta = solution_identical(tree,current,nodelist)
                    #println(alpha)
                    #readline()
                    parallel.bcast((alpha_identical, alpha, beta))
                    if alpha_identical
                        println("alpha identical!")
                        # get feasible objective value and replace z_hi if it is smaller than z_hi
                        z = feasible_objective(tree, feasible_prob, main_prob, alpha, beta, nodelist)
                        z_dl = current.block_model.dual_bound
                        if z < z_hi
                            z_hi = z
                            z_lo = z_dl
                            final_solution = alpha, beta
                            # delete all problems with dual bound greater than z_hi
                            for p in length(incumbent):-1:1
                                lb = incumbent[p][2]
                                if lb > z_hi
                                    deleteat!(incumbent,p)
                                    parallel.bcast(p)
                                end
                            end
                        end
                        parallel.bcast(0)
                    else
                        println("alpha not identical!")
                        # use heuristic to get feasible objective value and replace z_hi if it is smaller than z_hi
                        h_alpha, h_beta = heuristic_solution(tree,alpha,beta)
                        parallel.bcast((h_alpha, h_beta))
                        z = feasible_objective(tree, feasible_prob, main_prob, h_alpha, h_beta, nodelist)
                        z_dl = current.block_model.dual_bound
                        if z < z_hi
                            z_hi = z
                            z_lo = z_dl
                            final_solution = h_alpha, h_beta
                            # delete all problems with dual bound greater than z_hi
                            for p in length(incumbent):-1:1
                                lb = incumbent[p][2]
                                if lb > z_hi
                                    deleteat!(incumbent,p)
                                    parallel.bcast(p)
                                end
                            end
                        end
                        parallel.bcast(0)

                        #branching
                        branching = z > z_dl * (1+1e-6)
                        parallel.bcast(branching)
                        if branching

                            parallel.bcast(z_dl)

                            # select one variable
                            # add bound constraints
                            next0, next1 = branch_variable(tree, current, alpha, nodelist, leafdict)

                            # push to incumbent
                            # always try increasing the lines
                            push!(incumbent, [next0,z_dl])
                            push!(incumbent, [next1,z_dl])
                        end
                    end
                    println("ub: ", z_hi, " lb: ", z_lo)
                end
                println("incumbent: ",length(incumbent))
            end
        else
            while true
                termination = parallel.bcast(nothing)
                if termination
                    break
                end
                current_node = pop!(incumbent)
                current = current_node[1]
                solve_DD!(current,bundle_init,z_hi)

                bounding = parallel.bcast(nothing)
                if bounding
                    solution_identical(tree,current,nodelist)
                    alpha_identical, alpha, beta = parallel.bcast(nothing)

                    if alpha_identical
                        feasible_objective(tree, feasible_prob, main_prob, alpha, beta, nodelist)
                        delid = parallel.bcast(nothing)
                        while delid != 0
                            deleteat!(incumbent,p)
                            delid = parallel.bcast(nothing)
                        end
                    else
                        h_alpha, h_beta = parallel.bcast(nothing)
                        feasible_objective(tree, feasible_prob, main_prob, h_alpha, h_beta, nodelist)
                        delid = parallel.bcast(nothing)
                        while delid != 0
                            deleteat!(incumbent,delid)
                            delid = parallel.bcast(nothing)
                        end
                        branching = parallel.bcast(nothing)
                        if branching
                            z_lb = parallel.bcast(nothing)
                            next0, next1 = branch_variable(tree, current, alpha, nodelist, leafdict)

                            # push to incumbent
                            # always try increasing the lines
                            push!(incumbent, [next0,z_lb])
                            push!(incumbent, [next1,z_lb])
                        end
                    end
                end
            end
        end

        dual_decomp_results(tree, feasible_prob, main_prob, final_solution, nodelist, tot_time, z_hi, z_lo)

        parallel.finalize()
    end

    function get_problem_copy(LD::DRMSMIP.DRMS_LagrangeDual, nodelist)::Dict{Int,JuMP.Model}
        feasible_prob = Dict()
        for s in parallel.getpartition()
            m = copy(LD.block_model.model[s])
            set_optimizer(m, Gurobi.Optimizer)
            JuMP.set_optimizer_attribute(m, "OutputFlag", 0)

            alpha = m[:alpha]
            alpha0 = m[:alpha0]
            dc = m[:dc]
            for branch in Branches
                for k in n0[branch]+1:MAX_LINE_ADDED
                    set_objective_coefficient(m, alpha[1,branch,k], cost[branch])
                end
            end
            for stage in 2:nStage-1
                for branch in Branches
                    for k in 1:MAX_LINE_ADDED
                        set_objective_coefficient(m, alpha[stage,branch,k], cost[branch] * discount ^ (stage-1))
                        set_objective_coefficient(m, alpha0[stage,branch,k], -cost[branch] * discount ^ (stage-1))
                    end
                end
            end
            for stage in 2:nStage
                for bus in Buses
                    set_objective_coefficient(m, dc[stage,bus], LC * discount ^ (stage-1))
                end
            end
            feasible_prob[s] = m
        end
        return feasible_prob
    end

    function make_main_prob(tree::DRMSMIP.DR_Tree, nodelist::Array{Array{Int}})
        K = tree.K
        model = JuMP.Model(Gurobi.Optimizer)
        JuMP.set_optimizer_attribute(model, "OutputFlag", 0)

        @variable(model, P[2:length(tree.nodes)] >= 0)
        @variable(model, w[k=1:K-1, id=nodelist[k+1], s=1:tree.nodes[DD.get_parent(tree,id)].set.N] >= 0)

        DRMSMIP.con_E!(tree, model, nodelist)
        DRMSMIP.con_P!(tree, model, nodelist)
        DRMSMIP.con_M!(tree, model, nodelist)
        DRMSMIP.con_N!(tree, model, nodelist)

        @objective(model, Max, 0)
        return model
    end

    function solution_identical(tree::DRMSMIP.DR_Tree, LD::DRMSMIP.DRMS_LagrangeDual, nodelist)
        n = parallel.sum(DD.num_coupling_variables(LD.block_model))
        solutions = Dict{Int,SparseVector{Float64}}()
        for (s,m) in DD.block_model(LD)
            solutions[s] = sparsevec(Dict{Int,Float64}(), n)
        end
        for var in DD.coupling_variables(LD)
            # @assert has_block_model(LD, var.key.block_id)
            solutions[var.key.block_id][DD.index_of_λ(LD, var)] = JuMP.value(var.ref)
        end
        solutions_combined = parallel.combine_dict(solutions)

        if parallel.is_root()
            alpha_average = zeros(length(tree.nodes), nBranch, MAX_LINE_ADDED)
            beta_holder = zeros(length(tree.nodes), nBus)
            for s in 1:length(nodelist[nStage])
                leaf = nodelist[nStage][s]
                hist = DD.get_history(tree, leaf)
                for stage in 1:nStage - 1
                    node = hist[stage]
                    for branch in Branches
                        for k in 1:MAX_LINE_ADDED
                            var_id = (branch-1)*MAX_LINE_ADDED+k
                            i = LD.var_to_index[s,[node,var_id]]
                            alpha_average[node,branch,k] += solutions_combined[s][i]
                        end
                    end
                    if stage>=2
                        for bus in Buses
                            var_id = nBranch*MAX_LINE_ADDED*2+nBus+bus
                            i = LD.var_to_index[s,[node,var_id]]
                            beta_holder[node,bus] = max(beta_holder[node,bus], solutions_combined[s][i])
                        end
                    else
                        for bus in Buses
                            beta_holder[node,bus] = r_init[bus]
                        end
                    end
                end
            end
            alpha_identical = true
            for node in 1:length(tree.nodes)
                stage = DD.get_stage(tree,node)
                if stage < nStage
                    for branch in Branches
                        for k in 1:MAX_LINE_ADDED
                            alpha_average[node,branch,k] = alpha_average[node,branch,k] / nScenario^(nStage-stage)
                            if alpha_average[node,branch,k] < 1e-9
                                alpha_average[node,branch,k] = 0
                            elseif alpha_average[node,branch,k] > 1-1e-9
                                alpha_average[node,branch,k] = 1
                            else
                                alpha_identical = false
                            end
                        end
                    end
                end
            end
            return alpha_identical, alpha_average, beta_holder
        end
    end

    function feasible_objective(tree::DRMSMIP.DR_Tree, feasible_prob, main_prob, alpha, beta, nodelist)
        objectives = Dict{Int,Float64}()
        for s in parallel.getpartition()
            m = feasible_prob[s]
            leaf = nodelist[nStage][s]
            hist = DD.get_history(tree, leaf)
            for stage in 1:nStage - 1
                node = hist[stage]
                for branch in Branches
                    for k in 1:MAX_LINE_ADDED
                        if alpha[node,branch,k] < 1e-9
                            set_to_0!(m,stage,branch,k)
                        else
                            set_to_1!(m,stage,branch,k)
                        end
                    end
                end
                for bus in Buses
                    set_beta!(m,stage,bus,beta[node,bus])
                end
            end
            JuMP.optimize!(m)
            if termination_status(m) != MOI.OPTIMAL
                open("modeltest.lp", "w") do f
                    print(f, m)
                end
                for node in 1:length(tree.nodes)
                    if tree.nodes[node].k < nStage
                        println(node)
                        for k in 1:MAX_LINE_ADDED
                            println(alpha[node,:,k])
                        end
                    end
                end
            end
            objectives[s] = JuMP.objective_value(m)
        end
        objvals_combined = parallel.combine_dict(objectives)

        if parallel.is_root()
            P = main_prob[:P]
            for s in 1:length(nodelist[nStage])
                id = nodelist[nStage][s]
                set_objective_coefficient(main_prob, P[id], objvals_combined[s])
            end
            JuMP.optimize!(main_prob)

            Pref = Dict()

            for id in 2:length(tree.nodes)
                Pref[id] = JuMP.value(P[id])
            end
            # println(objectives,Pref,objective_value(main_prob))
            return objective_value(main_prob)
        end
    end

    function heuristic_solution(tree::DRMSMIP.DR_Tree, alpha, beta)
        h_alpha = zeros(length(tree.nodes), nBranch, MAX_LINE_ADDED)
        for node in 1:length(tree.nodes)
            stage = DD.get_stage(tree,node)
            if stage < nStage
                for branch in Branches
                    for k in 1:MAX_LINE_ADDED
                        if stage >= 2 && h_alpha[tree.nodes[node].parent,branch,k] > 1-1e-9
                            h_alpha[node,branch,k] = 1
                        elseif k>=2 && h_alpha[node,branch,k-1] < 1e-9
                            h_alpha[node,branch,k] = 0
                        else
                            h_alpha[node,branch,k] = round(alpha[node,branch,k])
                        end
                    end
                end
            end
        end
        return h_alpha, beta
    end

    function branch_variable(tree::DRMSMIP.DR_Tree, LD::DRMSMIP.DRMS_LagrangeDual, alpha, nodelist, leafdict)
        if parallel.is_root()
            node, branch, k = choose_branch_variable(tree, alpha)
            parallel.bcast((node,branch,k))
        else
            node, branch, k = parallel.bcast(nothing)
        end
        leaves = DD.get_future(tree,node)
        stage = DD.get_stage(tree,node)
        next0 = hardcopy(LD,nodelist,leafdict)
        next1 = hardcopy(LD,nodelist,leafdict)
        for leaf in leaves
            block_id = leafdict[leaf]
            if block_id in parallel.getpartition()
                m0 = next0.block_model.model[block_id]
                set_to_0!(m0,stage,branch,k)

                m1 = next1.block_model.model[block_id]
                set_to_1!(m1,stage,branch,k)
            end
        end
        return next0, next1
    end

    function hardcopy(LD::DRMSMIP.DRMS_LagrangeDual, nodelist, leafdict)
        #next = DRMSMIP.DRMS_LagrangeDual(tree, BM.TrustRegionMethod)
        next = DRMSMIP.DRMS_LagrangeDual(tree, solMethod)
        models = Dict()
        for s in parallel.getpartition()
            id = nodelist[nStage][s]
            m = copy(LD.block_model.model[s])
            set_optimizer(m, Gurobi.Optimizer)
            JuMP.set_optimizer_attribute(m, "OutputFlag", 0)
            models[id] = m
        end

    for s in parallel.getpartition()
            id = nodelist[nStage][s]
            DD.add_block_model!(next, s, models[id])
        end

        coupling_variables = create_coupling_variables(tree, models, nodelist, leafdict)

        # Set nonanticipativity variables as an array of symbols.
        DD.set_coupling_variables!(next, coupling_variables)

        return next
    end

    function choose_branch_variable(tree::DRMSMIP.DR_Tree, alpha)
        for node in 1:length(tree.nodes)
            stage = DD.get_stage(tree,node)
            if stage < nStage
                for branch in Branches
                    for k in 1:MAX_LINE_ADDED
                        if alpha[node,branch,k] > 1e-9 && alpha[node,branch,k] < 1-1e-9
                            return node, branch, k
                        end
                    end
                end
            end
        end
    end

    function dual_decomp_results(tree::DRMSMIP.DR_Tree, feasible_prob, main_prob, solution, nodelist, time, z_hi, z_lo)
        objectives = Dict{Int,Float64}()
        if parallel.is_root()
            parallel.bcast(solution)
        else
            solution = parallel.bcast(nothing)
        end
        alpha_val = solution[1]
        beta_val = solution[2]

        for s in parallel.getpartition()
            m = feasible_prob[s]
            leaf = nodelist[nStage][s]
            hist = DD.get_history(tree, leaf)
            for stage in 1:nStage - 1
                node = hist[stage]
                for branch in Branches
                    for k in 1:MAX_LINE_ADDED
                        if alpha_val[node,branch,k] < 1e-9
                            set_to_0!(m,stage,branch,k)
                        else
                            set_to_1!(m,stage,branch,k)
                        end
                    end
                end
                for bus in Buses
                    set_beta!(m,stage,bus,beta_val[node,bus])
                end
            end
            JuMP.optimize!(m)
            objectives[s] = JuMP.objective_value(m)
        end
        objvals_combined = parallel.combine_dict(objectives)

        curtailment = Dict{Int,SparseVector{Float64}}()
        for s in parallel.getpartition()
            curtailment[s] = sparsevec(Dict{Int,Float64}(), nStage*nBus)
            m = feasible_prob[s]
            dc_ref = value.(m[:dc])
            for stage in 2:nStage
                for bus in Buses
                    curtailment[s][(stage-1)*nBus + bus] = dc_ref[stage,bus]
                end
            end
        end
        curtailment_combined = parallel.combine_dict(curtailment)

        if parallel.is_root()
            P = main_prob[:P]
            for s in 1:length(nodelist[nStage])
                id = nodelist[nStage][s]
                set_objective_coefficient(main_prob, P[id], objvals_combined[s])
            end
            JuMP.optimize!(main_prob)

            Pref = Dict()

            for id in 2:length(tree.nodes)
                Pref[id] = JuMP.value(P[id])
            end

            tot = 0

            open("./TEPresults/"*filename, "w") do io
                for s in 1:length(nodelist[nStage])
                    leaf = nodelist[nStage][s]

                    hist = DD.get_history(tree, leaf)
                    write(io, "scenario, ") + sum(write(io, string(hist[stage])*", ") for stage in 1:nStage) + write(io, "\n")
                    write(io, "P, , ") + sum(write(io, string(Pref[hist[stage]])*", ") for stage in 2:nStage) + write(io, "\n")
                    write(io, "lines 0, ") + sum(write(io, string(n0[branch])*", ") for branch in Branches) + write(io, "\n")
                    for stage in 1:nStage-1
                        lines = Int64[]
                        node = hist[stage]
                        for branch in Branches
                            push!(lines, sum(alpha_val[node,branch,k] for k in 1:MAX_LINE_ADDED))
                        end
                        write(io, "lines " * string(stage) * ", ") + sum(write(io, string(lines[branch])*", ") for branch in Branches) + write(io, "\n")
                    end
                    for stage in 2:nStage
                        write(io, "curtailments " * string(stage) * ", ") + sum(write(io, string(curtailment_combined[s][(stage-1)*nBus+bus])*", ") for bus in Buses) + write(io, "\n")
                    end
                    for stage in 1:nStage
                        node = hist[stage]
                        write(io, "storage " * string(stage) * ", ") + sum(write(io, string(beta_val[node,bus])*", ") for bus in Buses) + write(io, "\n")
                    end

                    write(io, "total, " * string(objvals_combined[s]) * "\n")
                    write(io, "\n")

                    id = nodelist[nStage][s]
                    tot += objectives[s] * Pref[id]
                end
                write(io, "feas objective, "*string(z_hi)* "\n")
                write(io, "dual objective, "*string(z_lo)* "\n")
                write(io,"time, "*string(time)*"\n")
            end
        end
    end

    # function dual_decomp_results(tree::DRMSMIP.DR_Tree, feasible_prob, main_prob, solution, nodelist, time)
    #     objectives = Dict()
    #     alpha_val = solution[1]
    #     beta_val = solution[2]
    #     for s in 1:length(feasible_prob)
    #         m = feasible_prob[s]
    #         leaf = nodelist[nStage][s]
    #         hist = DD.get_history(tree, leaf)
    #         for stage in 1:nStage - 1
    #             node = hist[stage]
    #             for branch in Branches
    #                 for k in 1:MAX_LINE_ADDED
    #                     if alpha_val[node,branch,k] < 1e-9
    #                         set_to_0!(m,stage,branch,k)
    #                     else
    #                         set_to_1!(m,stage,branch,k)
    #                     end
    #                 end
    #             end
    #             for bus in Buses
    #                 set_beta!(m,stage,bus,beta_val[node,bus])
    #             end
    #         end
    #         JuMP.optimize!(m)
    #         objectives[s] = JuMP.objective_value(m)
    #     end
    #     P = main_prob[:P]
    #     for s in 1:length(feasible_prob)
    #         id = nodelist[nStage][s]
    #         set_objective_coefficient(main_prob, P[id], objectives[s])
    #     end
    #     JuMP.optimize!(main_prob)

    #     Pref = Dict()

    #     for id in 2:length(tree.nodes)
    #         Pref[id] = JuMP.value(P[id])
    #     end

    #     tot = 0

    #     open("./TEPresults/dual_decomp_results.csv", "w") do io
    #         for s in 1:length(nodelist[nStage])
    #             m = feasible_prob[s]
    #             leaf = nodelist[nStage][s]

    #             alpha_ref = value.(m[:alpha])
    #             generation_ref = value.(m[:generation])
    #             dc_ref = value.(m[:dc])
    #             theta_ref = value.(m[:theta])
    #             flow_ref = value.(m[:flow])
    #             beta_ref = value.(m[:beta])

    #             hist = DD.get_history(tree, leaf)
    #             write(io, "scenario, ") + sum(write(io, string(hist[stage])*", ") for stage in 1:nStage) + write(io, "\n")
    #             write(io, "P, , ") + sum(write(io, string(Pref[hist[stage]])*", ") for stage in 2:nStage) + write(io, "\n")
    #             write(io, "lines 0, ") + sum(write(io, string(n0[branch])*", ") for branch in Branches) + write(io, "\n")
    #             for stage in 1:nStage-1
    #                 lines = Int64[]
    #                 for branch in Branches
    #                     push!(lines, sum(alpha_ref[stage,branch,k] for k in 1:MAX_LINE_ADDED))
    #                 end
    #                 write(io, "lines " * string(stage) * ", ") + sum(write(io, string(lines[branch])*", ") for branch in Branches) + write(io, "\n")
    #             end
    #             for stage in 2:nStage
    #                 write(io, "curtailments " * string(stage) * ", ") + sum(write(io, string(dc_ref[stage,bus])*", ") for bus in Buses) + write(io, "\n")
    #             end
    #             for stage in 1:nStage
    #                 write(io, "storage " * string(stage) * ", ") + sum(write(io, string(beta_ref[stage,bus])*", ") for bus in Buses) + write(io, "\n")
    #             end

    #             write(io, "total, " * string(objectives[s]) * "\n")
    #             write(io, "\n")

    #             id = nodelist[nStage][s]
    #             tot += objectives[s] * Pref[id]
    #         end
    #         write(io, "objective, "*string(tot)* "\n")
    #         write(io,"time, "*string(time)*"\n")
    #     end

    # end

    tree = create_tree()

    branch_and_bound(tree, 3600.0);
end

#run_main([3,2,10.0,"dual_decomp_results_.csv"])
run_main(ARGS)


text:
struct Inflows{P <: PlantCollection, T <: AbstractFloat}
    Q̃::Vector{T}

    function Inflows(plants::PlantCollection, Q̃::AbstractVector)
        nplants(plants) == length(Q̃) || error("Inconsistent number of power stations and inflows.")
        P = typeof(plants)
        T = eltype(Q̃)
        return new{P,T}(Q̃)
    end

    function Inflows(::Type{P}, Q̃::AbstractVector) where P <: PlantCollection
        nplants(P) == length(Q̃) || error("Inconsistent number of power stations and inflows.")
        T = eltype(Q̃)
        return new{P,T}(Q̃)
    end
end

function Inflows(plants::Vector{Plant}, Q̃::AbstractVector)
    return Inflows(PlantCollection(plants), Q̃)
end

function Base.zero(::Type{Inflows{P,T}}) where {P <: PlantCollection, T <: AbstractFloat}
    return Inflows(P,zeros(T,nplants(P)))
end

function Base.getindex(inflows::Inflows{P}, plant::Plant) where P <: PlantCollection
    idx = findfirst(p -> p == plant, plants(P))
    idx == nothing && throw(BoundsError(inflows, plant))
    return inflows.Q̃[idx]
end

plants(::Inflows{P}) where P <: PlantCollection = plants(P)
flows(inflows::Inflows) = inflows.Q̃
ninflows(::Inflows{P}) where P = nplants(P)

function Statistics.mean(inflow_collection::Vector{Inflows{P,T}}) where {P <: PlantCollection, T <: AbstractFloat}
    Q̃ = [mean([inflows.Q̃[p] for inflows in inflow_collection]) for p in 1:nplants(P)]
    return Inflows(P, Q̃)
end

function local_inflows(inflows::Inflows{P}, upstream_plants::Dict{Plant, Vector{Plant}}) where P <: PlantCollection
    local_inflows = [local_inflow(inflows, p, upstream_plants[p]) for p in plants(P)]
    return Inflows(P, local_inflows)
end

function local_inflow(inflows::Inflows, plant::Plant, upstream_plants::Vector{Plant})
    V = inflows[plant]
    if !isempty(upstream_plants)
        V -= sum(inflows[p] for p in upstream_plants)
    end
    return max(0, V)
end

struct InflowSequence{N, P <: PlantCollection, T <: AbstractFloat} <: AbstractVector{T}
    inflows::Vector{Inflows{P,T}}

    function InflowSequence(inflows::Vector{Inflows{P,T}}) where {P <: PlantCollection, T <: AbstractFloat}
        N = length(inflows)
        return new{N,P,T}(inflows)
    end

    function InflowSequence(plants::PlantCollection, inflows::AbstractMatrix)
        nplants(plants) == size(inflows,1) || error("Inconsistent number of power stations and inflows.")
        P = typeof(plants)
        N = size(inflows, 2)
        T = eltype(inflows)
        return new{N,P,T}([Inflows(plants, inflows[:,j]) for j in 1:size(inflows,2)])
    end

    function InflowSequence(::Type{P}, inflows::AbstractMatrix) where P <: PlantCollection
        nplants(P) == size(inflows,1) || error("Inconsistent number of power stations and inflows.")
        N = size(inflows, 2)
        T = eltype(inflows)
        return new{N,P,T}([Inflows(P, inflows[:,j]) for j in 1:size(inflows,2)])
    end
end

function InflowSequence(plants::Vector{Plant}, inflows::Matrix{T}) where T <: AbstractFloat
    return InflowSequence(PlantCollection(plants), inflows)
end

function Base.zero(::Type{InflowSequence{N,P,T}}) where {N, P <: PlantCollection, T <: AbstractFloat}
    return InflowSequence(P, zeros(T,nplants(P),N))
end
Base.iterate(inflowseq::InflowSequence) = iterate(inflowseq.inflows)
Base.length(inflowseq::InflowSequence{N}) where N = N
Base.size(inflowseq::InflowSequence) = size(inflowseq.inflows)
@inline function Base.getindex(inflowseq::InflowSequence, I...)
    @boundscheck checkbounds(inflowseq.inflows, I...)
    @inbounds return inflowseq.inflows[I...]
end
function Base.getindex(inflowseq::InflowSequence, horizon::Horizon)
    horizon <= HydroModels.horizon(inflowseq) || throw(BoundsError(inflowseq,horizon))
    return InflowSequence(inflowseq.inflows[1:ndays(horizon)])
end
@inline function Base.setindex!(inflowseq::InflowSequence, x, I...)
    @boundscheck checkbounds(inflowseq.inflows, I...)
    @inbounds inflowseq.inflows[I...] = x
end
Base.axes(inflowseq::InflowSequence) = axes(inflowseq.inflows)
Base.IndexStyle(::Type{<:InflowSequence}) = Base.IndexLinear()
horizon(inflowseq::InflowSequence{N}) where N = Days(N)

function Statistics.mean(sequences::Vector{InflowSequence{N,P,T}}) where {N, P <: PlantCollection, T <: AbstractFloat}
    inflow_sequence = [mean([sequence[i] for sequence in sequences]) for i in 1:N]
    return InflowSequence(inflow_sequence)
end

function local_inflow_sequence(sequence::InflowSequence{N,P}, upstream_plants::Dict{Plant, Vector{Plant}}) where {N, P <: PlantCollection}
    local_inflow_sequence = [local_inflows(sequence[d], upstream_plants) for d in 1:N]
    return InflowSequence(local_inflow_sequence)
end

function mean_flow(resolution::Resolution, Q̃::InflowSequence, t::Integer, p::Plant)
    if resolution.hours_in_period < 24
        hours_passed = t * resolution.hours_in_period
        d = div(hours_passed,24) + 1
        return Q̃[d][p]
    end
    days_in_period = div(resolution.hours_in_period, 24)
    idx = (t - 1) * days_in_period + 1
    return mean(Q̃[idx:idx+days_in_period-1])[p]
end


text:
#######################################################################
#    Script for testing sensitivities calculated using Sundials.jl    #
#                                                                     #
#     Sensitivity equations are solved with a finite difference       #
#         approximation to the Jacobian computed by Sundials.         #
#                               content                               #
#                                                                     #
#     The example is chosen to allow calculation of analytical        #
#         sensitivities for comparison. These are implemented         #
#                              as well.                               #
#                                                                     #
#     Christian Haargaard Olsen <christianhaargaard@gmail.com>        #
#                                                                     #
#######################################################################

using Sundials

# Uncomment this to use QT interface for PyPlot
using PyCall
@pyimport(matplotlib)
pygui(:qt)
using PyPlot

# First we create the datatype used for parsing parameter values to
# the right-hand-side function rhs.
type UserData
  parameters :: Array{Float64, 1}
end

# using the function f(t) = a*sin(b*t), f(0) = 0, we create an ODE by differentiating
# with respect to t, to get fdot (t) = a*b*cos(b*t). We will then use this ODE to test
# the sensitivity capabilities of sundials, as the correct sensitivites are given by
# f_a = sin(b*t) and f_b = a*t*cos(b*t).

function rhs(
  t         :: Float64,
  x         :: Sundials.N_Vector,
  xdot      :: Sundials.N_Vector,
  data
  )

  x         = Sundials.asarray(x)                 :: Array{Float64,1}
  xdot      = Sundials.asarray(xdot)              :: Array{Float64,1}
  user_data = unsafe_pointer_to_objref(data)      :: UserData

  a = user_data.parameters[1]
  b = user_data.parameters[2]

  x1 = x[1]
  xdot[1] = a*b*cos(b*t)

  return Int32(0)
end

function true_solution(
  t                                  :: Array{Float64,1},
  data                               :: UserData
  )

  a             = data.parameters[1] :: Float64
  b             = data.parameters[2] :: Float64
  result        = a*sin(b*t)         :: Array{Float64,1}

  return result
end

function true_sens(
  t                                  :: Array{Float64,1},
  data                               :: UserData
  )

  neq = 1
  NS = length(data.parameters)

  sens = zeros(length(t),neq,NS)

  a             = data.parameters[1] :: Float64
  b             = data.parameters[2] :: Float64
  sens[:,1,1] =  sin(b*t)            :: Array{Float64,1}
  sens[:,1,2] =  t.*a.*cos(b*t)      :: Array{Float64,1}

  return sens
end

function load_states()
  return ([0.0], UserData([1.0,1.0,1.0]))
end

function calcsens_finitedif(parameters :: Array{Float64,1}; step = sqrt(1e-6))
  (x, data)                 = load_states()               :: Tuple{Array{Float64,1}, UserData}
  data.parameters           = parameters                  :: Array{Float64,1}
  xdot                      = zeros(x)                    :: Array{Float64,1}
  t                         = collect(linspace(0,10,101)) :: Array{Float64,1}
  (reltol, abstol,sens_tol) = (1e-8, 1e-6, 1e-6)          :: Tuple{Float64,Float64, Float64}
  neq                       = length(x)
  NS                        = length(parameters)

  sols = zeros(length(t),neq,NS) # No need to copy initial condition, they are already zero
  sens = zeros(length(t),neq,NS) # No need to copy initial condition, they are already zero

  (t0, y0, x) = runmodel(parameters)
  @printf("calculating finite difference sensitivities for %i parameters.\n",NS)
  for i in 1:NS
    newparameters = zeros(parameters)
    newparameters = copy(parameters)
    newparameters[i] = (1.0+step)*newparameters[i]
    (t0, sols[:,:,i], x) = runmodel(newparameters)
    sens[:,:,i] = (sols[:,:,i]-y0)/(step*parameters[i])
    @printf("Done for parameter %i.\n",i)
    end

  return (sens)
end

function runsens(parameters :: Array{Float64,1})
  (x, data)                 = load_states()               :: Tuple{Array{Float64,1}, UserData}
  data.parameters           = parameters                  :: Array{Float64,1}
  xdot                      = zeros(x)                    :: Array{Float64,1}
  t                         = collect(linspace(0,10,101)) :: Array{Float64,1}
  (reltol, abstol,sens_tol) = (1e-8, 1e-6, 1e-6)          :: Tuple{Float64,Float64, Float64}
  neq                       = length(x)
  NS                        = length(parameters)

  #######################################################################
  #                         ODE solver settings                         #
  #######################################################################

  #cvode_mem = Sundials.CVodeCreate(Sundials.CV_BDF, Sundials.CV_FUNCTIONAL)
  cvode_mem = Sundials.CVodeCreate(Sundials.CV_BDF, Sundials.CV_NEWTON)
  flag = Sundials.CVodeInit(cvode_mem,rhs, t[1], x)
  flag = Sundials.CVodeSetUserData(cvode_mem, data)
  flag = Sundials.CVodeSStolerances(cvode_mem, reltol, abstol)
  flag = Sundials.CVDense(cvode_mem, neq)

  # List of parameters for which to calculate sensitivities
  plist = round(Int32,1:NS)
  println(plist)

  # scaling of parameters
  pbar = copy(data.parameters)

  # Tolerances for solving sensitivities for each parameter
  #sens_tol_vec = [abstol :: Float64 for i in 1:NS]

  # IDEA: Maybe these should be scaled with parameter value?
  sens_tol_vec = [abstol*data.parameters[i] :: Float64 for i in 1:NS]

  # Sundials want a (pointer to a) vector where each component corresponds to
  # a parameter, and where the element is a pointer to an array with a place for
  # each state variable.
  ySJ    = Array(Array{Float64,1}, NS);
  for i in 1:NS
    ySJ[i] = zeros(Float64,neq)*1.0;
  end
  ySjulia = [Sundials.nvector(ySJ[i]) for i in 1:NS]
  yS = pointer(ySjulia)

  #######################################################################
  #                         Sensitivity settigs                         #
  #######################################################################

  sensi_meth = Int32(Sundials.CV_SIMULTANEOUS) # CV_SIMULTANEOUS, CV_STAGGERED or CV_STAGGERED1
  flag = Sundials.CVodeSensInit1(cvode_mem, NS, sensi_meth, 0, yS);
  flag = Sundials.CVodeSensEEtolerances(cvode_mem);
  flag = Sundials.CVodeSetSensErrCon(cvode_mem, 0);
  #flag = Sundials.CVodeSensSStolerances(cvode_mem, reltol, sens_tol_vec);
  flag = Sundials.CVodeSetSensDQMethod(cvode_mem, Sundials.CV_CENTERED, 0.0);
  flag = Sundials.CVodeSetSensParams(cvode_mem, data.parameters, 0, 0);

  # Placeholder for solution and sensitivities
  solution = zeros(length(t), neq)
  solution[1,:] = copy(x)
  sens = zeros(length(t),neq,NS) # No need to copy initial condition, they are already zero

  tout = [t[1]] # Set first output time

  # Loop through all the output times
  for k in 2:length(t)
    # Extract the solution to x, and the sensitivities to yS
    flag = Sundials.CVode(cvode_mem, t[k], x, tout, Sundials.CV_NORMAL)
    flag = Sundials.CVodeGetSens(cvode_mem, tout, yS)

    #Save the results
    solution[k,:] = x
    for i in 1:NS
      sens[k,:,i] = Sundials.asarray(unsafe_load(yS,i))
    end
  end

  return (solution,sens)
end

#######################################################################
#                          Plot results                               #
#######################################################################

t = collect(linspace(0,10,101));
(x, data) = load_states()

(solution, sens) = runsens(data.parameters)
senstrue = true_sens(t,data)

figure(1)
plot(t,sens[:,1,1])
plot(t,senstrue[:,1,1])
show()

figure(2)
plot(t,sens[:,1,2])
plot(t,senstrue[:,1,2])
show()


text:
module SpectraDB

using ..Utils: build_interpolation, _readh5_file

export sl1_exp_spectrum,
	   sl1_ref_spectrum,
	   sl2_exp_spectrum,
	   sl2_ref_spectrum,
	   bragg_spectrum,
	   fp_spectrum,
	   hafnia_spectrum,
	   scandia_spectrum,
	   tantala_spectrum,
	   hafnia_ellips,
	   Info

function Info()
	tmp1 = "\n " *
	"\n " *
	"\n Available functions from SpectraDB module:" *
	"\n " *
	"\n     sl1_exp_spectrum(λ)" *
	"\n     sl1_ref_spectrum(λ)" *
	"\n     sl2_exp_spectrum(λ)" *
	"\n     sl2_ref_spectrum(λ)" *
	"\n     bragg_spectrum(λ)" *
	"\n     fp_spectrum(λ)" *
	"\n     hafnia_spectrum(λ)" *
	"\n     scandia_spectrum(λ)" *
	"\n     tantala_spectrum(λ)" *
	"\n "
	"\n     To use any of these functions type: SpectraDB.function(args)." *
	"\n "
	return println(tmp1)
end

"""

	Raw experimental reflectance spectrum of a porous silicon thin film on a crystalline
	silicon substrate. Interpolates for a valid input wavelength range.

		R = sl1_exp_spectrum(λ)

			λ: wavelength range (nm), ∈ [200, 1125]
			R: reflectance

"""
function sl1_exp_spectrum(λ::AbstractArray{T1,1}) where {T1<:Real}
	X = _readh5_file("SL1sample", :SP)
    itp = build_interpolation(hcat(vec(X["lambda"]), vec(X["reflectance"])))
    return itp.(λ)
end

"""

	Reference experimental reflectance spectrum of a crystalline silicon substrate.
	Interpolates for a valid input wavelength range.

		R = sl1_ref_spectrum(λ)

			λ: wavelength range (nm), ∈ [200, 1125]
			R: reflectance

"""
function sl1_ref_spectrum(λ::AbstractArray{T1,1}) where {T1<:Real}
	X = _readh5_file("SL1reference", :SP)
	itp = build_interpolation(hcat(vec(X["lambda"]), vec(X["reflectance"])))
    return itp.(λ)
end

"""

	Raw experimental reflectance spectrum of a porous silicon thin film on a crystalline
	silicon substrate. Interpolates for a valid input wavelength range.

		R = sl2_exp_spectrum(λ)

			λ: wavelength range (nm), ∈ [200, 1125]
			R: reflectance

"""
function sl2_exp_spectrum(λ::AbstractArray{T1,1}) where {T1<:Real}
	X = _readh5_file("SL2sample", :SP)
	itp = build_interpolation(hcat(vec(X["lambda"]), vec(X["reflectance"])))
    return itp.(λ)
end

"""

	Reference experimental reflectance spectrum of a crystalline silicon substrate.
	Interpolates for a valid input wavelength range.

		R = sl2_ref_spectrum(λ)

			λ: wavelength range (nm), ∈ [200, 1125]
			R: reflectance

"""
function sl2_ref_spectrum(λ)
	X = _readh5_file("SL2reference", :SP)
	itp = build_interpolation(hcat(vec(X["lambda"]), vec(X["reflectance"])))
    return itp.(λ)
end

"""

	Absolute experimental reflectance spectrum of a porous silicon thin film photonic
	crystal (DBR or Bragg stack) on a crystalline silicon substrate. Interpolates for
	a valid input wavelength range.

		R = bragg_spectrum(λ)

			λ: wavelength range (nm), ∈ [200, 1125]
			R: reflectance

"""
function bragg_spectrum(λ::AbstractArray{T1,1}) where {T1<:Real}
	X = _readh5_file("Bragg", :SP)
	itp = build_interpolation(hcat(vec(X["lambda"]), vec(X["reflectance"])))
    return itp.(λ)
end

"""

	Absolute experimental reflectance spectrum of a porous silicon thin film Fabry-Perot
	(MicroCavity stack) on a glass substrate. Interpolates for a valid input wavelength range.

		R = fp_spectrum(λ)

			λ: wavelength range (nm), ∈ [200, 1125]
			R: reflectance

"""
function fp_spectrum(λ::AbstractArray{T1,1}) where {T1<:Real}
	X = _readh5_file("MicroCavity", :SP)
	itp = build_interpolation(hcat(vec(X["lambda"]), vec(X["reflectance"])))
    return itp.(λ)
end

"""

	Absolute experimental transmitance spectrum of a hafnium oxide thin film on a fused
	silica UV substrate. Interpolates for a valid input wavelength range.

		T = hafnia_spectrum(λ)

			λ: wavelength range (nm), ∈ [190, 1100]
			T: transmitance

"""
function hafnia_spectrum(λ::AbstractArray{T1,1}) where {T1<:Real}
	X = _readh5_file("Hafnia", :SP)
	itp = build_interpolation(hcat(vec(X["lambda"]), vec(X["transmitance"])))
    return itp.(λ)./100.0
end

"""

	Absolute experimental transmitance spectrum of a scandium oxide thin film on a fused
	silica UV substrate. Interpolates for a valid input wavelength range.

		T = scandia_spectrum(λ)

			λ: wavelength range (nm), ∈ [190, 1100]
			T: transmitance

"""
function scandia_spectrum(λ::AbstractArray{T1,1}) where {T1<:Real}
	X = _readh5_file("Scandia", :SP)
	itp = build_interpolation(hcat(vec(X["lambda"]), vec(X["transmitance"])))
    return itp.(λ)./100.0
end

"""

	Absolute experimental ellipsomtry spectra of a tantalum oxide thin film on a fused
	silica UV substrate. Interpolates for a valid input wavelength range. Measured at 60
	degrees.

		Ψ, Δ = tantala_spectrum(λ)

			λ: wavelength range (nm), ∈ [200, 2100]
			Ψ: psi spectrum
			Δ: delta spectrum

"""
function tantala_spectrum(λ::AbstractArray{T1,1}) where {T1<:Real}
	X = _readh5_file("Tantala", :SP)
	itp = build_interpolation(hcat(vec(X["lambda"]), vec(X["psi"])))
	psi = itp.(λ)
	itp = build_interpolation(hcat(vec(X["lambda"]), vec(X["delta"])))
	delta = itp.(λ)
    return psi, delta
end

"""

	Absolute experimental ellipsomtry spectra of a tantalum oxide thin film on a fused
	silica UV substrate. Interpolates for a valid input wavelength range. Measured at 60
	degrees.

		Ψ, Δ = hafnia_ellips(λ)

			λ: wavelength range (nm), ∈ [200, 2100]
			Ψ: psi spectrum
			Δ: delta spectrum

"""
function hafnia_ellips(λ::AbstractArray{T1,1}) where {T1<:Real}
	X = _readh5_file("HafniaEllips", :SP)
	itp = build_interpolation(hcat(1240.0./vec(X["lambda"]), vec(X["psi"])))
	psi = itp.(λ)
	itp = build_interpolation(hcat(vec(1240.0./X["lambda"]), vec(X["delta"])))
	delta = itp.(λ)
    return psi, delta
end

end # module


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Rename(:col₁ => :newcol₁, :col₂ => :newcol₂, ..., :col₁ => :newcolₙ)

The transform that renames `col₁` to `newcol₁`, `col₂` to `newcol₂`, ...

# Examples

```julia
Rename(:a => :x, :c => :y)
Rename("a" => "x", "c" => "y")
```
"""
struct Rename <: Stateless
  names::Dict{Symbol,Symbol}
end

Rename(names::Pair) = _pairsyms(names) |> Dict |> Rename
Rename(names...) = _pairsyms.(names) |> Dict |> Rename

_pairsyms(x::Pair) = Symbol(first(x)) => Symbol(last(x))

function apply(transform::Rename, table)
  _rename(transform.names, table)
end

function revert(transform::Rename, table, cache)
  # reversing the key-value pairs of the Dict
  newnames = Dict(new => old for (old, new) in transform.names)
  _rename(newnames, table) |> first
end

function _rename(names, table)
  oldnames = Tables.columnnames(table)

  # check if requested renames exist in the table
  @assert keys(names) ⊆ oldnames "invalid column names"

  # use new names if necessary
  newnames = map(oldnames) do oldname
    oldname in keys(names) ? names[oldname] : oldname
  end

  # table with new tables
  cols = Tables.columns(table)
  vals = [Tables.getcolumn(cols, name) for name in oldnames]
  𝒯 = (; zip(newnames, vals)...) |> Tables.materializer(table)

  𝒯, nothing
end


text:
using Logging
using ProgressLogging
using TerminalLoggers
using TmuxDisplays
using UnicodePlots

d = tmuxdisplay(Text("This is the default tmux display"))
sleep(0.8)

smalldisplay = TmuxDisplays.split_window(size = 3, target = d)
display(smalldisplay, Text("This is a custom tmux display"))
sleep(0.8)

with_logger(TerminalLogger(IO(smalldisplay))) do
    @withprogress name = "Level 1" for i in 1:3
        plt = scatterplot(randn(10), randn(10), xlim = (-3, 3), ylim = (-3, 3))
        @withprogress name = "Level 2" for j in 1:30
            sleep(0.01)
            tmuxdisplay(scatterplot!(plt, randn(10), randn(10)))
            @logprogress j / 30
        end
        @logprogress i / 3
    end
end


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    mean(ensemble)

Mean of `ensemble`.
"""
function mean(ensemble::Ensemble)
  μs = (; (variable => mean(reals) for (variable, reals) in ensemble.reals)...)
  georef(μs, ensemble.domain)
end

"""
    var(ensemble)

Variance of `ensemble`.
"""
function var(ensemble::Ensemble)
  σs = (; (variable => var(reals) for (variable, reals) in ensemble.reals)...)
  georef(σs, ensemble.domain)
end

"""
    quantile(ensemble, p)

`p`-quantile of `ensemble`.
"""
function quantile(ensemble::Ensemble, p::Number)
  cols = []
  for (variable, reals) in ensemble.reals
    quantiles = map(1:nelements(ensemble.domain)) do location
      slice = getindex.(reals, location)
      quantile(slice, p)
    end
    push!(cols, variable => quantiles)
  end
  georef((; cols...), ensemble.domain)
end

quantile(ensemble::Ensemble, ps::AbstractVector) = [quantile(ensemble, p) for p in ps]


text:
export ODEVertex, StaticEdge, ODEEdge
export Symmetric, AntiSymmetric, Directed, Fiducial

"""
Abstract supertype for all vertex functions.
"""
abstract type VertexFunction end

abstract type Coupling end

"""
Abstract supertype for all edge functions.
"""
abstract type EdgeFunction{C<:Coupling} end

struct AntiSymmetric <: Coupling end
struct Symmetric <: Coupling end
struct Directed <: Coupling end
struct Fiducial <: Coupling end

coupling(::EdgeFunction{C}) where C = C()
coupling(::Type{<:EdgeFunction{C}}) where C = C()

Base.@kwdef struct ODEVertex{F} <: VertexFunction
    f::F
    dim::Int
    pdim::Int
end

struct StaticEdge{F,C} <: EdgeFunction{C}
    f::F # (e, v_s, v_t, p, F) -> nothing
    dim::Int # number of dimensions of e
    pdim::Int
end
function StaticEdge(; f, dim, pdim, coupling)
    # todo probably a method check?
    StaticEdge{typeof(f), typeof(coupling)}(f, dim, pdim)
end

struct ODEEdge{F,C} <: EdgeFunction{C}
    f::F # (de, e, v_s, v_t, p, t) -> nothing
    dim::Int # number of dimensions of e
    pdim::Int
end
function ODEEdge(; f, dim, pdim, coupling)
    # todo probably a method check?
    ODEEdge{typeof(f), typeof(coupling)}(f, dim, pdim)
end

dim(e::Union{EdgeFunction, VertexFunction}) = e.dim
pdim(e::Union{EdgeFunction, VertexFunction}) = e.pdim
accdim(e::EdgeFunction) = e.dim
accdim(e::EdgeFunction{Fiducial}) = Int(e.dim/2)


text:
module AFMEphys

include("fileManagement.jl")

export makeFileList

end


text:
module GeometricInterpretation

using GeometricFlux
using Flux
using CairoMakie
using GraphMakie

include("explainer.jl")
include("explainers/gnn_explainer.jl")

end


text:
import Adapt
import CUDA

Adapt.adapt_structure(to, data::IJKFVH{S, Nij, Nk}) where {S, Nij, Nk} =
    IJKFVH{S, Nij, Nk}(Adapt.adapt(to, parent(data)))

Adapt.adapt_structure(to, data::IJFH{S, Nij}) where {S, Nij} =
    IJFH{S, Nij}(Adapt.adapt(to, parent(data)))

Adapt.adapt_structure(to, data::IFH{S, Ni}) where {S, Ni} =
    IFH{S, Ni}(Adapt.adapt(to, parent(data)))

Adapt.adapt_structure(to, data::IJF{S, Nij}) where {S, Nij} =
    IJF{S, Nij}(Adapt.adapt(to, parent(data)))

Adapt.adapt_structure(to, data::IF{S, Ni}) where {S, Ni} =
    IF{S, Ni}(Adapt.adapt(to, parent(data)))

Adapt.adapt_structure(to, data::VF{S}) where {S} =
    VF{S}(Adapt.adapt(to, parent(data)))

parent_array_type(
    ::Type{A},
) where {A <: CUDA.GPUArrays.AbstractGPUArray{FT}} where {FT} = CUDA.CuArray{FT}

function knl_copyto!(dest, src)

    #=
    nij, nh = size(dest)

    thread_idx = CUDA.threadIdx().x
    block_idx = CUDA.blockIdx().x
    block_dim = CUDA.blockDim().x

    # mapping to global idx to make insensitive
    # to number of blocks / threads per device
    global_idx = thread_idx + (block_idx - 1) * block_dim

    nx, ny = nij, nij
    i = global_idx % nx == 0 ? nx : global_idx % nx
    j = cld(global_idx, nx)
    h = ((global_idx-1) % (nx*nx)) + 1
    =#

    i = CUDA.threadIdx().x
    j = CUDA.threadIdx().y

    h = CUDA.blockIdx().x

    p_dest = slab(dest, h)
    p_src = slab(src, h)

    @inbounds p_dest[i, j] = p_src[i, j]
    return nothing
end

function Base.copyto!(
    dest::IJFH{S, Nij},
    bc::Union{IJFH{S, Nij, A}, Base.Broadcast.Broadcasted{IJFHStyle{Nij, A}}},
) where {S, Nij, A <: CUDA.CuArray}
    _, _, _, _, Nh = size(bc)
    CUDA.@cuda threads = (Nij, Nij) blocks = (Nh,) knl_copyto!(dest, bc)
    return dest
end