text:
abstract type StochasticDiffEqAlgorithm <: AbstractSDEAlgorithm end
abstract type StochasticDiffEqAdaptiveAlgorithm <: StochasticDiffEqAlgorithm end
abstract type StochasticDiffEqCompositeAlgorithm <: StochasticDiffEqAlgorithm end

abstract type StochasticDiffEqRODEAlgorithm <: AbstractRODEAlgorithm end
abstract type StochasticDiffEqRODEAdaptiveAlgorithm <: StochasticDiffEqRODEAlgorithm end
abstract type StochasticDiffEqRODECompositeAlgorithm <: StochasticDiffEqRODEAlgorithm end

abstract type StochasticDiffEqNewtonAdaptiveAlgorithm{CS,AD,Controller} <: StochasticDiffEqAdaptiveAlgorithm end
abstract type StochasticDiffEqNewtonAlgorithm{CS,AD,Controller} <: StochasticDiffEqAlgorithm end

abstract type StochasticDiffEqJumpAlgorithm <: StochasticDiffEqAlgorithm end
abstract type StochasticDiffEqJumpAdaptiveAlgorithm <: StochasticDiffEqAlgorithm end
abstract type StochasticDiffEqJumpNewtonAdaptiveAlgorithm{CS,AD,Controller} <: StochasticDiffEqJumpAdaptiveAlgorithm end

abstract type StochasticDiffEqJumpDiffusionAlgorithm <: StochasticDiffEqAlgorithm end
abstract type StochasticDiffEqJumpDiffusionAdaptiveAlgorithm <: StochasticDiffEqAlgorithm end
abstract type StochasticDiffEqJumpNewtonDiffusionAdaptiveAlgorithm{CS,AD,Controller} <: StochasticDiffEqJumpDiffusionAdaptiveAlgorithm end

abstract type IteratedIntegralApprox end
struct IICommutative <:  IteratedIntegralApprox end
struct IIWiktorsson <:  IteratedIntegralApprox end

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

# Basics

struct EM{split} <: StochasticDiffEqAlgorithm end
EM(split=true) = EM{split}()

struct SplitEM <: StochasticDiffEqAlgorithm end
struct EulerHeun <: StochasticDiffEqAlgorithm end

struct LambaEM{split} <: StochasticDiffEqAdaptiveAlgorithm end
LambaEM(split=true) = LambaEM{split}()

struct LambaEulerHeun <: StochasticDiffEqAdaptiveAlgorithm end

"""
Kloeden, P.E., Platen, E., Numerical Solution of Stochastic Differential Equations.
Springer. Berlin Heidelberg (2011)
"""
struct SimplifiedEM <: StochasticDiffEqAlgorithm end

"""
Kloeden, P.E., Platen, E., Numerical Solution of Stochastic Differential Equations.
Springer. Berlin Heidelberg (2011)
"""
struct RKMil{interpretation} <: StochasticDiffEqAdaptiveAlgorithm end
RKMil(;interpretation=:Ito) = RKMil{interpretation}()

"""
Kloeden, P.E., Platen, E., Numerical Solution of Stochastic Differential Equations.
Springer. Berlin Heidelberg (2011)
"""
struct RKMilCommute{interpretation} <: StochasticDiffEqAdaptiveAlgorithm end
RKMilCommute(;interpretation=:Ito) = RKMilCommute{interpretation}()

"""
Kloeden, P.E., Platen, E., Numerical Solution of Stochastic Differential Equations.
Springer. Berlin Heidelberg (2011)
"""
struct RKMilGeneral{T<:IteratedIntegralApprox, TruncationType} <: StochasticDiffEqAdaptiveAlgorithm
  interpretation::Symbol
  ii_approx::T
  c::Int
  p::TruncationType
end
function RKMilGeneral(;interpretation=:Ito, ii_approx=IIWiktorsson(), c=1, p=nothing, dt=nothing)
  γ = 1//1
  p==true && (p = Int(floor(c*dt^(1//1-2//1*γ)) + 1))
  RKMilGeneral(interpretation, ii_approx, c, p)
end

struct WangLi3SMil_A <: StochasticDiffEqAlgorithm end
struct WangLi3SMil_B <: StochasticDiffEqAlgorithm end
struct WangLi3SMil_C <: StochasticDiffEqAlgorithm end
struct WangLi3SMil_D <: StochasticDiffEqAlgorithm end
struct WangLi3SMil_E <: StochasticDiffEqAlgorithm end
struct WangLi3SMil_F <: StochasticDiffEqAlgorithm end

#SROCK methods
struct SROCK1{interpretation,E} <: StochasticDiffEqAlgorithm
  eigen_est::E
end
SROCK1(;interpretation=:Ito,eigen_est=nothing) = SROCK1{interpretation,typeof(eigen_est)}(eigen_est)

# Weak Order 2
for Alg in [:SROCK2, :KomBurSROCK2, :SROCKC2]
  @eval begin
    struct $Alg{E} <: StochasticDiffEqAlgorithm
      eigen_est::E
    end
    $Alg(;eigen_est=nothing) = $Alg(eigen_est)
  end
end

# ROCK stabilization for EM
struct SROCKEM{E} <: StochasticDiffEqAlgorithm
  strong_order_1::Bool
  eigen_est::E
end
SROCKEM(;strong_order_1=true,eigen_est=nothing) = SROCKEM(strong_order_1,eigen_est)

struct SKSROCK{E} <: StochasticDiffEqAlgorithm
  post_processing::Bool
  eigen_est::E
end
SKSROCK(;post_processing=false,eigen_est=nothing) = SKSROCK(post_processing,eigen_est)

struct TangXiaoSROCK2{E} <: StochasticDiffEqAlgorithm
  version_num::Int
  eigen_est::E
end
TangXiaoSROCK2(;version_num=5,eigen_est=nothing) = TangXiaoSROCK2(version_num,eigen_est)
###############################################################################

# Predictor Corrector
struct PCEuler{T<:Real, F} <: StochasticDiffEqAlgorithm
  theta::T
  eta::T
  ggprime::F
end

"""
    PCEuler(ggprime; theta=1/2, eta=1/2)

Predictor Corrector Euler

# Arguments
- `ggprime::Function`:
  For scalar problems, `ggprime` ``= b\\partial_x(b)``
  For multi-dimensional problems
  `bbprime_k` ``= \\sum_{j=1...M, i=1...D} b^(j)_i \\partial_i b^(j)_k``
  where ``b^(j)`` correspond to the noise vector due to the j'th noise channel.
  If problem is in place - a in place ggprime should be supplied - and
  vice versa for not in place speicification of problem.
- `theta::Real`:
  Degree of implicitness in the drift term. Set to 0.5 by default.
- `eta::Real`:
  Degree of implicitness in the diffusion term. Set to 0.5 by default.

Reference: Stochastics and Dynamics, Vol. 8, No. 3 (2008) 561–581
Note that the original paper has a typo in the definition of ggprime...
"""
PCEuler(ggprime; theta=1/2, eta=1/2) = PCEuler(theta,eta,ggprime)

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

# Rossler

"""
Rößler A., Runge–Kutta Methods for the Strong Approximation of Solutions of
Stochastic Differential Equations, SIAM J. Numer. Anal., 48 (3), pp. 922–952.
DOI:10.1137/09076636X
"""
struct SRA{TabType} <: StochasticDiffEqAdaptiveAlgorithm
  tableau::TabType
end
SRA(;tableau=constructSRA1()) = SRA(tableau)

"""
Rößler A., Runge–Kutta Methods for the Strong Approximation of Solutions of
Stochastic Differential Equations, SIAM J. Numer. Anal., 48 (3), pp. 922–952.
DOI:10.1137/09076636X
"""
struct SRI{TabType} <: StochasticDiffEqAdaptiveAlgorithm
  tableau::TabType
  error_terms::Int
end
SRI(;tableau=constructSRIW1(),error_terms=4) = SRI(tableau,error_terms)

"""
Rößler A., Runge–Kutta Methods for the Strong Approximation of Solutions of
Stochastic Differential Equations, SIAM J. Numer. Anal., 48 (3), pp. 922–952.
DOI:10.1137/09076636X
"""
struct SRIW1 <: StochasticDiffEqAdaptiveAlgorithm end

"""
Rößler A., Runge–Kutta Methods for the Strong Approximation of Solutions of
Stochastic Differential Equations, SIAM J. Numer. Anal., 48 (3), pp. 922–952.
DOI:10.1137/09076636X
"""
struct SRIW2 <: StochasticDiffEqAdaptiveAlgorithm end
struct SOSRI <: StochasticDiffEqAdaptiveAlgorithm end
struct SOSRI2 <: StochasticDiffEqAdaptiveAlgorithm end

"""
Rößler A., Runge–Kutta Methods for the Strong Approximation of Solutions of
Stochastic Differential Equations, SIAM J. Numer. Anal., 48 (3), pp. 922–952.
DOI:10.1137/09076636X
"""
struct SRA1 <: StochasticDiffEqAdaptiveAlgorithm end

"""
Rößler A., Runge–Kutta Methods for the Strong Approximation of Solutions of
Stochastic Differential Equations, SIAM J. Numer. Anal., 48 (3), pp. 922–952.
DOI:10.1137/09076636X
"""
struct SRA2 <: StochasticDiffEqAdaptiveAlgorithm end

"""
Rößler A., Runge–Kutta Methods for the Strong Approximation of Solutions of
Stochastic Differential Equations, SIAM J. Numer. Anal., 48 (3), pp. 922–952.
DOI:10.1137/09076636X
"""
struct SRA3 <: StochasticDiffEqAdaptiveAlgorithm end
struct SOSRA <: StochasticDiffEqAdaptiveAlgorithm end
struct SOSRA2 <: StochasticDiffEqAdaptiveAlgorithm end

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

# Rossler second order for weak approx.

"""
Debrabant, K. and Rößler A., Families of efficient second order Runge–Kutta methods
for the weak approximation of Itô stochastic differential equations,
Applied Numerical Mathematics 59, pp. 582–594 (2009)
DOI:10.1016/j.apnum.2008.03.012
"""
struct DRI1 <: StochasticDiffEqAdaptiveAlgorithm end

"""
Debrabant, K. and Rößler A., Families of efficient second order Runge–Kutta methods
for the weak approximation of Itô stochastic differential equations,
Applied Numerical Mathematics 59, pp. 582–594 (2009)
DOI:10.1016/j.apnum.2008.03.012
"""
struct DRI1NM <: StochasticDiffEqAdaptiveAlgorithm end

"""
Rößler A., Second Order Runge–Kutta Methods for Itô Stochastic Differential Equations,
SIAM J. Numer. Anal., 47, pp. 1713-1738 (2009)
DOI:10.1137/060673308
"""
struct RI1 <: StochasticDiffEqAdaptiveAlgorithm end

"""
Rößler A., Second Order Runge–Kutta Methods for Itô Stochastic Differential Equations,
SIAM J. Numer. Anal., 47, pp. 1713-1738 (2009)
DOI:10.1137/060673308
"""
struct RI3 <: StochasticDiffEqAdaptiveAlgorithm end

"""
Rößler A., Second Order Runge–Kutta Methods for Itô Stochastic Differential Equations,
SIAM J. Numer. Anal., 47, pp. 1713-1738 (2009)
DOI:10.1137/060673308
"""
struct RI5 <: StochasticDiffEqAdaptiveAlgorithm end

"""
Rößler A., Second Order Runge–Kutta Methods for Itô Stochastic Differential Equations,
SIAM J. Numer. Anal., 47, pp. 1713-1738 (2009)
DOI:10.1137/060673308
"""
struct RI6 <: StochasticDiffEqAdaptiveAlgorithm end

"""
Debrabant, K. and Rößler A., Classification of Stochastic Runge–Kutta Methods for
the Weak Approximation of Stochastic Differential Equations,
Mathematics and Computers in Simulation 77, pp. 408-420 (2008)
DOI:10.1016/j.matcom.2007.04.016
"""
struct RDI1WM <: StochasticDiffEqAlgorithm end

"""
Debrabant, K. and Rößler A., Classification of Stochastic Runge–Kutta Methods for
the Weak Approximation of Stochastic Differential Equations,
Mathematics and Computers in Simulation 77, pp. 408-420 (2008)
DOI:10.1016/j.matcom.2007.04.016
"""
struct RDI2WM <: StochasticDiffEqAdaptiveAlgorithm end

"""
Debrabant, K. and Rößler A., Classification of Stochastic Runge–Kutta Methods for
the Weak Approximation of Stochastic Differential Equations,
Mathematics and Computers in Simulation 77, pp. 408-420 (2008)
DOI:10.1016/j.matcom.2007.04.016
"""
struct RDI3WM <: StochasticDiffEqAdaptiveAlgorithm end

"""
Debrabant, K. and Rößler A., Classification of Stochastic Runge–Kutta Methods for
the Weak Approximation of Stochastic Differential Equations,
Mathematics and Computers in Simulation 77, pp. 408-420 (2008)
DOI:10.1016/j.matcom.2007.04.016
"""
struct RDI4WM <: StochasticDiffEqAdaptiveAlgorithm end

# Stratonovich sense

"""
Rößler A., Second order Runge–Kutta methods for Stratonovich stochastic differential
equations, BIT Numerical Mathematics 47, pp. 657-680 (2007)
DOI:10.1007/s10543-007-0130-3
"""
struct RS1 <: StochasticDiffEqAlgorithm end

"""
Rößler A., Second order Runge–Kutta methods for Stratonovich stochastic differential
equations, BIT Numerical Mathematics 47, pp. 657-680 (2007)
DOI:10.1007/s10543-007-0130-3
"""
struct RS2 <: StochasticDiffEqAlgorithm end

"""
Kloeden, P.E., Platen, E., Numerical Solution of Stochastic Differential Equations.
Springer. Berlin Heidelberg (2011)
"""
struct PL1WM <: StochasticDiffEqAlgorithm end

"""
Kloeden, P.E., Platen, E., Numerical Solution of Stochastic Differential Equations.
Springer. Berlin Heidelberg (2011)
"""
struct PL1WMA <: StochasticDiffEqAlgorithm end

"""
Komori, Y., Weak second-order stochastic Runge–Kutta methods for non-commutative
stochastic differential equations, Journal of Computational and Applied
Mathematics 206, pp. 158 – 173 (2007)
DOI:10.1016/j.cam.2006.06.006
"""
struct NON <: StochasticDiffEqAlgorithm end

"""
Komori, Y., Weak order stochastic Runge–Kutta methods for commutative stochastic
differential equations, Journal of Computational and Applied Mathematics 203,
pp. 57 – 79 (2007)
DOI:10.1016/j.cam.2006.03.010
"""
struct COM <: StochasticDiffEqAlgorithm end

"""
Komori, Y., & Burrage, K. (2011). Supplement: Efficient weak second order stochastic
Runge–Kutta methods for non-commutative Stratonovich stochastic differential equations.
Journal of computational and applied mathematics, 235(17), pp. 5326 - 5329 (2011)
DOI:10.1016/j.cam.2011.04.021
"""
struct NON2 <: StochasticDiffEqAlgorithm end

"""
Tocino, A. and Vigo-Aguiar, J., Weak Second Order Conditions for Stochastic Runge-
Kutta Methods, SIAM Journal on Scientific Computing 24, pp. 507 - 523 (2002)
DOI:10.1137/S1064827501387814
"""
struct SIEA <: StochasticDiffEqAlgorithm end

"""
Tocino, A. and Vigo-Aguiar, J., Weak Second Order Conditions for Stochastic Runge-
Kutta Methods, SIAM Journal on Scientific Computing 24, pp. 507 - 523 (2002)
DOI:10.1137/S1064827501387814
"""
struct SMEA <: StochasticDiffEqAlgorithm end

"""
Tocino, A. and Vigo-Aguiar, J., Weak Second Order Conditions for Stochastic Runge-
Kutta Methods, SIAM Journal on Scientific Computing 24, pp. 507 - 523 (2002)
DOI:10.1137/S1064827501387814
"""
struct SIEB <: StochasticDiffEqAlgorithm end

"""
Tocino, A. and Vigo-Aguiar, J., Weak Second Order Conditions for Stochastic Runge-
Kutta Methods, SIAM Journal on Scientific Computing 24, pp. 507 - 523 (2002)
DOI:10.1137/S1064827501387814
"""
struct SMEB <: StochasticDiffEqAlgorithm end

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

# IIF

struct IIF1M{F} <: StochasticDiffEqAlgorithm
  nlsolve::F
end
IIF1M(;nlsolve=NLSOLVEJL_SETUP()) = IIF1M{typeof(nlsolve)}(nlsolve)

struct IIF2M{F} <: StochasticDiffEqAlgorithm
  nlsolve::F
end
IIF2M(;nlsolve=NLSOLVEJL_SETUP()) = IIF2M{typeof(nlsolve)}(nlsolve)

struct IIF1Mil{F} <: StochasticDiffEqAlgorithm
  nlsolve::F
end
IIF1Mil(;nlsolve=NLSOLVEJL_SETUP()) = IIF1Mil{typeof(nlsolve)}(nlsolve)

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

# SDIRK

struct ImplicitEM{CS,AD,F,F2,S,T2,Controller} <: StochasticDiffEqNewtonAdaptiveAlgorithm{CS,AD,Controller}
  linsolve::F
  nlsolve::F2
  diff_type::S
  theta::T2
  extrapolant::Symbol
  new_jac_conv_bound::T2
  symplectic::Bool
end
ImplicitEM(;chunk_size=0,autodiff=true,diff_type=Val{:central},
                          linsolve=DEFAULT_LINSOLVE,nlsolve=NLNewton(),
                          extrapolant=:constant,
                          theta = 1,symplectic=false,
                          new_jac_conv_bound = 1e-3,
                          controller = :Predictive) =
                          ImplicitEM{chunk_size,autodiff,
                          typeof(linsolve),typeof(nlsolve),typeof(diff_type),
                          typeof(new_jac_conv_bound),controller}(
                          linsolve,nlsolve,diff_type,
                          symplectic ? 1/2 : theta,
                          extrapolant,new_jac_conv_bound,symplectic)

STrapezoid(;kwargs...) = ImplicitEM(;theta=1/2,kwargs...)
SImplicitMidpoint(;kwargs...) = ImplicitEM(;theta=1/2,symplectic=true,kwargs...)

struct ImplicitEulerHeun{CS,AD,F,S,N,T2,Controller} <: StochasticDiffEqNewtonAdaptiveAlgorithm{CS,AD,Controller}
  linsolve::F
  diff_type::S
  nlsolve::N
  theta::T2
  extrapolant::Symbol
  new_jac_conv_bound::T2
  symplectic::Bool
end
ImplicitEulerHeun(;chunk_size=0,autodiff=true,diff_type=Val{:central},
                          linsolve=DEFAULT_LINSOLVE,nlsolve=NLNewton(),
                          extrapolant=:constant,
                          theta = 1,symplectic = false,
                          new_jac_conv_bound = 1e-3,
                          controller = :Predictive) =
                          ImplicitEulerHeun{chunk_size,autodiff,
                          typeof(linsolve),typeof(diff_type),
                          typeof(nlsolve),
                          typeof(new_jac_conv_bound),controller}(
                          linsolve,diff_type,nlsolve,
                          symplectic ? 1/2 : theta,
                          extrapolant,
                          new_jac_conv_bound,symplectic)

struct ImplicitRKMil{CS,AD,F,S,N,T2,Controller,interpretation} <: StochasticDiffEqNewtonAdaptiveAlgorithm{CS,AD,Controller}
  linsolve::F
  diff_type::S
  nlsolve::N
  theta::T2
  extrapolant::Symbol
  new_jac_conv_bound::T2
  symplectic::Bool
end
ImplicitRKMil(;chunk_size=0,autodiff=true,diff_type=Val{:central},
                          linsolve=DEFAULT_LINSOLVE,nlsolve=NLNewton(),
                          extrapolant=:constant,
                          theta = 1,symplectic = false,
                          new_jac_conv_bound = 1e-3,
                          controller = :Predictive,interpretation=:Ito) =
                          ImplicitRKMil{chunk_size,autodiff,
                          typeof(linsolve),typeof(diff_type),
                          typeof(nlsolve),typeof(new_jac_conv_bound),
                          controller,interpretation}(
                          linsolve,diff_type,nlsolve,
                          symplectic ? 1/2 : theta,
                          extrapolant,
                          new_jac_conv_bound,symplectic)

struct ISSEM{CS,AD,F,S,N,T2,Controller} <: StochasticDiffEqNewtonAdaptiveAlgorithm{CS,AD,Controller}
  linsolve::F
  diff_type::S
  nlsolve::N
  theta::T2
  extrapolant::Symbol
  new_jac_conv_bound::T2
  symplectic::Bool
end
ISSEM(;chunk_size=0,autodiff=true,diff_type=Val{:central},
                       linsolve=DEFAULT_LINSOLVE,nlsolve=NLNewton(),
                       extrapolant=:constant,
                       theta = 1,symplectic=false,
                       new_jac_conv_bound = 1e-3,
                       controller = :Predictive) =
                       ISSEM{chunk_size,autodiff,
                       typeof(linsolve),typeof(diff_type),
                       typeof(nlsolve),
                       typeof(new_jac_conv_bound),controller}(
                       linsolve,diff_type,nlsolve,
                       symplectic ? 1/2 : theta,
                       extrapolant,
                       new_jac_conv_bound,symplectic)

struct ISSEulerHeun{CS,AD,F,S,N,T2,Controller} <: StochasticDiffEqNewtonAdaptiveAlgorithm{CS,AD,Controller}
 linsolve::F
 diff_type::S
 nlsolve::N
 theta::T2
 extrapolant::Symbol
 new_jac_conv_bound::T2
 symplectic::Bool
end
ISSEulerHeun(;chunk_size=0,autodiff=true,diff_type=Val{:central},
                      linsolve=DEFAULT_LINSOLVE,nlsolve=NLNewton(),
                      extrapolant=:constant,
                      theta = 1,symplectic=false,
                      new_jac_conv_bound = 1e-3,
                      controller = :Predictive) =
                      ISSEulerHeun{chunk_size,autodiff,
                      typeof(linsolve),typeof(diff_type),
                      typeof(nlsolve),typeof(new_jac_conv_bound),controller}(
                      linsolve,diff_type,nlsolve,
                      symplectic ? 1/2 : theta,
                      extrapolant,
                      new_jac_conv_bound,symplectic)

struct SKenCarp{CS,AD,F,FDT,N,T2,Controller} <: StochasticDiffEqNewtonAdaptiveAlgorithm{CS,AD,Controller}
  linsolve::F
  diff_type::FDT
  nlsolve::N
  smooth_est::Bool
  extrapolant::Symbol
  new_jac_conv_bound::T2
  ode_error_est::Bool
end

SKenCarp(;chunk_size=0,autodiff=true,diff_type=Val{:central},
                   linsolve=DEFAULT_LINSOLVE,nlsolve=NLNewton(),
                   smooth_est=true,extrapolant=:min_correct,
                   new_jac_conv_bound = 1e-3,controller = :Predictive,
                   ode_error_est = true) =
 SKenCarp{chunk_size,autodiff,typeof(linsolve),typeof(diff_type),
        typeof(nlsolve),typeof(new_jac_conv_bound),controller}(
        linsolve,diff_type,nlsolve,smooth_est,extrapolant,new_jac_conv_bound,
        ode_error_est)

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

# Jumps

struct TauLeaping <: StochasticDiffEqJumpAdaptiveAlgorithm end
struct CaoTauLeaping <: StochasticDiffEqJumpAdaptiveAlgorithm end

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

# Etc.

struct StochasticCompositeAlgorithm{T,F} <: StochasticDiffEqCompositeAlgorithm
  algs::T
  choice_function::F
end

struct RandomEM <: StochasticDiffEqRODEAlgorithm end

const SplitSDEAlgorithms = Union{IIF1M,IIF2M,IIF1Mil,SKenCarp,SplitEM}

@doc raw"""
Leimkuhler B., Matthews C., Robust and efficient configurational molecular sampling via
Langevin dynamics, J. Chem. Phys. 138, 174102 (2013)
DOI:10.1063/1.4802990

```math
du = vdt \\
dv = f(v,u) dt - \gamma v dt + g(u) \sqrt{2\gamma} dW
```
"""
struct BAOAB{T} <: StochasticDiffEqAlgorithm
  gamma::T
end
BAOAB(;gamma=1.0) = BAOAB(gamma)


text:
function Base.show(f::IO, m::MIME"text/html", t::D3Tree)
    tree_json = json(t)
    root_id = 1
    css = read(joinpath(dirname(@__FILE__()), "..", "css", "tree_vis.css"), String)
    js = read(joinpath(dirname(@__FILE__()), "..", "js", "tree_vis.js"), String)
    div = "treevis$(randstring())"

    html_string = """
        <!DOCTYPE html>

            <meta charset="UTF-8">
            <title>$(t.title)</title>

        Attempting to display the tree. If the tree is large, this may take some time.

        Note: D3Trees.jl requires an internet connection. If no tree appears, please check your connection. To help fix this, please see this issue. You may also diagnose errors with the javascript console (Ctrl-Shift-J in chrome).

    """

    println(f,html_string)
end

# fallback when only the repl is available
function Base.show(io::IO, m::MIME"text/plain", t::D3Tree)
    show(io, m, D3TreeView(D3TreeNode(t, 1), get(t.options, :init_expanded, false) ? typemax(Int) : 3))
end
Base.show(io::IO, m::MIME"text/plain", v::D3TreeView) = shownode(io, v.root, v.depth, "", "")

ext: jl


text:
## Expression parsing ########################################################

# Very simple, overly permissive S-expression parsing, without involving the
# parser framework. Does not permit empty expresions. If rpn is true, the head
# is assumed to be last rather than first in each expression.

function parse(::Type{Ex}, str::AbstractString; rpn=false)

    token = r"\( | \) | (\"([^\"\\]|\\[\"\\])*\") | [^\s()]+"x

    stack = [[], []]

    function close()
        v = pop!(stack)
        if rpn
            head = last(v)
            args = length(v) > 1 ? v[1:end-1] : []
        else
            head = first(v)
            args = length(v) > 1 ? v[2:end] : []
        end
        push!(last(stack), Call(head, args))
    end

    for m in eachmatch(token, str)
        t = m.match
        if t == "("
            push!(stack, [])
        elseif t == ")"
            close()
        elseif m[1] ≢ nil
            push!(last(stack), unescape_string(t[2:end-1]))
        else
            n = tryparse(Float64, t)
            push!(last(stack), n ≡ nil ? Symbol(t) : n)
        end
    end

    length(last(stack)) > 1 && close()

    return last(stack)[1]

end

## Expression pretty-printing ################################################

# Display rules loosely based on
# https://github.com/TotalVerb/SExpressions.jl/blob/master/src/Lists/show.jl

ppfmt(x)         = string(x)
ppfmt(x::Nil)    = "nil"
ppfmt(x::Number) = string(isinteger(x) ? Int(x) : x)
ppfmt(x::Bool)   = string(x) # To override Number
ppfmt(x::String) = repr(x)
ppfmt(x::Ex)     = string("(", join(ppfmt.([head(x), args(x)...]), " "), ")")

function pprint(io::IO, x; hang=false, indent=0)
    str = ppfmt(x)
    hang || print(io, " "^indent)
    print(io, str)
end

function pprint(io::IO, x::Ex; hang=false, indent=0)

    hang || print(io, " "^indent)

    ## Short enough to fit on the current line?

    str     = ppfmt(x)
    width   = textwidth(str)
    _, cols = displaysize(io)

    if indent + width ≤ cols && width ≤ cols/2  # Heuristic
        print(io, str)
        return
    end

    ## Otherwise, we must break it up:

    print(io, "(")

    headstr = ppfmt(head(x))                     # Never break the head
    print(io, headstr)

    headw = textwidth(headstr)
    sameln = headw ≤ 3                          # Heuristic
    extra = sameln ? headw + 2 : 3              # Extra indentation

    for arg in args(x)
        print(io, sameln ? " " : "\n")
        pprint(io, arg, hang=sameln, indent=indent + extra)
        sameln = false
    end

    print(io, ")")

end


text:
export freqCont, freqContExtendedSources, freqContExtendedSourcesSS, freqCont;
using JLD
using Multigrid.ParallelJuliaSolver
using Statistics

export FreqContParam, getFreqContParams
"""
	FreqContParam

	parameters for the frequency continuation method

	Input:
		mc    		- current model
		itersNum	- number of GN iters for each frequency continuation step
		originalSources	- original sources before sources extension
		nrcv		- number of receivers
		pInv		- Inverse param
		pMis 		- misfit params (remote)
		windowSize  - How many frequencies to treat at once at the most.
		resultsFilename - a filename for saving the intermediate results according to the GN and continuation (FC) iterations (done in dumpFun())
		dumpFun     - a function for plotting, saving and doing all the things with the intermidiate results.
		Iact		- active set locations
		mback		- active set model values
		Z1			- first part of the extended sources initial guess
		simSrcDim	- dimensions of the simultaneous sources
		alpha1		- Z1 penalty coeficient
		alpha2Orig	- Z2 penalty coeficient
		stepReg		- regularization coeficient on the step size
		mode        - either "1stInit" or anything else.
					  1stInit will use the first group of misfits as an initialization and will not include it together with the next group of misfits.
		startFrom   - a start index for the continuation. Usefull when our run broke in the middle of the iterations.
		endAt		- index of the final frequency for the continuation.Usefull when we want to use only low frequencies.
		cycle       - just for identifying different runs when saving the results.
		method      - "projGN" or "barrierGN"
		FWImethod	- "FWI" or "FWI_ES"
		updateMref	- wheter to update the mref after each GN iteration, relevant for the ES case.

"""
mutable struct FreqContParam
	mc						:: Vector{Float64}
	itersNum				:: Int64
	originalSources			:: SparseMatrixCSC
	nrcv					:: Int64
	pInv					:: InverseParam
	pMis					:: Array{RemoteChannel}
	windowSize				:: Int64
	resultsFilename			:: String
	dumpFun					:: Function
	Iact					:: SparseMatrixCSC
	mback					:: Vector{Float64}
	Z1						:: Union{Int64, Array{ComplexF64, 2}}
	simSrcDim				:: Int64
	alpha1					:: Float64
	alpha2Orig				:: Float64
	stepReg					:: Float64
	mode					:: String
	startFrom				:: Int64
	endAt					:: Int64
	cycle					:: Int64
	method					:: String
	FWImethod				:: String
	updateMref				:: Bool
end

function getFreqContParams(mc,itersNum::Int64,
		originalSources::SparseMatrixCSC,nrcv, pInv::InverseParam,
		pMis::Array{RemoteChannel}, windowSize::Int64, resultsFilename::String,
		dumpFun::Function,Iact,mback;Z1=0,simSrcDim=1,
		alpha1=0.0, alpha2Orig=0.0, stepReg=0.0, mode::String="", startFrom::Int64 = 1,
		endAt=length(pMis),cycle::Int64=0, method::String="projGN",
		FWImethod::String="FWI",updateMref=false)
		return FreqContParam(mc, itersNum, originalSources,
				nrcv, pInv, pMis, windowSize, resultsFilename, dumpFun,
				Iact, mback, Z1, simSrcDim, alpha1, alpha2Orig, stepReg,
				mode, startFrom, endAt, cycle, method, FWImethod, updateMref)
end

function calculateZ2(m,misfitCalc::Function, p::Integer, nsrc::Integer, nfreq::Integer,
	nrec::Integer, nwork::Integer,
	numOfCurrentProblems::Integer, mergedWd::Array, mergedRc::Array, pMis::Array,
	pMisCurrent::Array{MisfitParam}, Z1::Matrix, alpha::Float64)

	rhs = zeros(ComplexF64, (p, nsrc));
	lhs = zeros(ComplexF64, (p,p));

	HinvPsZ1 = computeHinvTRecX(pMis, Z1, m, 0);

	function calculateLhs(meanWd,Z1,HinvP)
		return meanWd .* HinvP' * HinvP;
	end

	function calculateRhs(meanWd,HinvP,mergedRc)
		return meanWd .* HinvP' * mergedRc;
	end

	lhsides = Array{Array{ComplexF64}}(undef,nfreq)
	rhsides = Array{Array{ComplexF64}}(undef,nfreq)

	@sync begin
	 		for k=1:nfreq
		@async begin
			meanWd = mean(mergedWd[k])
			meanWdC = (real(meanWd).^2)
				lhsides[k] = remotecall_fetch(calculateLhs, k % nworkers() + 1,meanWdC,Z1,HinvPsZ1[k]);
				rhsides[k] = remotecall_fetch(calculateRhs, k % nworkers() + 1,meanWdC,HinvPsZ1[k], mergedRc[k]);
	 		end
	 	end
	end

	lhs = sum(lhsides)
	rhs = sum(rhsides)
	lhs += alpha * I;

	return lhs\rhs;
end

function MyPCG(A::Function,b::Array,x::Array,M::Array,numiter)
r = b-A(x);
z = M.*r;
p = copy(z);
norms = [norm(r)];

# use 20 iters if starting residual is high
if norms[1] > 100
	numiter=20;
end
for k=1:numiter
	Ap = A(p);
	alpha = real(dot(z,r)/dot(p,Ap));
	x .+= alpha*p;
	beta = real(dot(z,r));
	r .-= alpha*Ap
	z = M.*r
	nn = norm(r);
	if (nn < 0.1 * norms[1] || nn < 1e-2)
		break
	end
	norms = [norms; nn];
	beta = real(dot(z,r)) / beta;
	p = z + beta*p;
end
return x,norms
end

function misfitCalc(m, Z1,Z2,mergedWd,mergedRc,nfreq,alpha1,alpha2,pMis)
	## HERE WE  NEED TO MAKE SURE THAT Wd is equal in its real and imaginary parts.
	sum = 0.0;
	HinvPs = computeHinvTRecX(pMis, Z1, m, 0);
	for i = 1:nfreq
		res, = SSDFun(HinvPs[i] * Z2,mergedRc[i],mean(mergedWd[i]) .* ones(size(mergedRc[i])));
		sum += res;
	end
	return sum;
end

function objectiveCalc(Z1,Z2,misfit,alpha1,alpha2)
	return misfit + 0.5*alpha1 * norm(Z1)^2 + 0.5*alpha2 * norm(Z2)^2;
end

function MultOp(HPinv, R, Z2)
	return (HPinv' * R) * Z2
end

function MultOpT(HPinv, R, Z2)
	return HPinv * (R * Z2')
end

function MultAll(m, avgWds, pMis, R, Z1, Z2, alpha, stepReg)
	sum = zeros(ComplexF64, size(R))
	inner = computeHinvTRecX(pMis, R,m, 0);

	function calculateInner(avgWd, inner, Z2)
		return avgWd^2 .* inner * (Z2 * Z2');
	end
	@sync begin
	 		for k=1:length(avgWds)
		@async begin
	 			inner[k] = remotecall_fetch(calculateInner,
	 				k % nworkers() + 1,avgWds[k],inner[k],Z2);
	 		end
	 	end
	end

	outer = computeHinvTRecXarr(pMis, inner,m, 1);
	for i = 1:length(avgWds)
		sum += outer[i]
	end
	eps = 1e-2
	return sum + (alpha./(abs.(Z1) .+ eps * maximum(abs.(Z1)))).*R + stepReg*R;
end

function calculateZ1(m, misfitCalc::Function, nfreq::Integer, mergedWd::Array, mergedRc::Array, pMis::Array, Z1::Matrix,Z2::Matrix, alpha1::Float64,stepReg::Float64)
	rhsElement = computeHinvTRecXarr(pMis, mergedRc, m, 1);
	rhsides = Array{Array{ComplexF64}}(undef,nfreq)

	function calculateRhs(meanWd,Z2,rhsElement)
		return meanWd.*  (rhsElement * Z2');
	end

	@sync begin
	 		for k=1:nfreq
		@async begin
			meanWd = mean(mergedWd[k])
			meanWdC = (real(meanWd).^2)
			rhsides[k] = remotecall_fetch(calculateRhs, k % nworkers() + 1,meanWdC,Z2, rhsElement[k]);
	 		end
	 	end
	end

	rhs = sum(rhsides) + (stepReg) .* Z1

	OP = x-> MultAll(m, mean.(real.(mergedWd)), pMis, x, Z1, Z2, alpha1, stepReg);
	eps = 10.0
	M = (abs.(Z1) .+ eps * maximum(abs.(Z1)));
	Z1, = MyPCG(OP,rhs,Z1,M,5);
	return Z1;
end

function standardGNrun(method, mc, simSrcDim, originalSources, pInv, pMisTemp,
		HIS, resultsFilename, cycle, freqIdx, dumpFun)
pMisTempFetched = map(fetch, pMisTemp);
dobs = map(x -> x.dobs,pMisTempFetched)
wd = map(x-> x.Wd,pMisTempFetched)
nsrc = size(originalSources,2);

if simSrcDim==1
	TEmat = Matrix(1.0I,nsrc,nsrc)
else
	TEmat = rand([-1,1],(nsrc,simSrcDim));
end

if simSrcDim > 1
	reducedSources = originalSources * TEmat;
	reducedDobs = map(x-> x*TEmat, dobs);
	sizeWD = size(reducedDobs[1])
	reducedWd = map(x-> mean(x)/sqrt(simSrcDim) * ones(sizeWD), wd);

	pMisTemp = setSourcesSame(pMisTemp,reducedSources);
	pMisTemp = setDobs(pMisTemp,reducedDobs);
	pMisTemp = setWd(pMisTemp,reducedWd);
end

if resultsFilename == ""
	filename = "";
	hisMatFileName = "";
else
	Temp = splitext(resultsFilename);
	if cycle==0
		filename = string(Temp[1],"_FC",freqIdx,"_GN",Temp[2]);
		hisMatFileName  = string(Temp[1],"_FC",freqIdx);
	else
		filename = string(Temp[1],"_Cyc",cycle,"_FC",freqIdx,"_GN",Temp[2]);
		hisMatFileName  =  string(Temp[1],"_Cyc",cycle,"_FC",freqIdx);
	end
end

# Here we set a dump function for GN for this iteracion of FC
function dumpGN(mc,Dc,iter,pInv,PF)
	dumpFun(mc,Dc,iter,pInv,PF,filename);
end

if method == "projGN"
	mc,Dc,flag,His = projGNCG(mc,pInv,pMisTemp,dumpResults = dumpGN);
elseif method == "barrierGN"
	mc,Dc,flag,His = barrierGNCG(mc,pInv,pMisTemp,rho=1.0,dumpResults = dumpGN);
end

if simSrcDim > 1
	# set full data back
	pMisTemp = setSourcesSame(pMisTemp,originalSources);
	pMisTemp = setDobs(pMisTemp,dobs)
	pMisTemp = setWd(pMisTemp,wd)
end

if hisMatFileName != ""
	file = matopen(string(hisMatFileName,"_HisGN.mat"), "w");
	### PUT THE CONTENT OF "HIS" INTO THE MAT FILE ###
	His.Dc = [];
	write(file,"His",His);
	close(file);
end
His.Dc = []
push!(HIS,His)

clear!(pMisTemp);

return mc,Dc,flag,HIS
end

function ExtendedSourcesGNRun(method, mc, simSrcDim, nrcv, itersNum, originalSources,
		Z1, alpha1, alpha2, stepReg, pInv, pMisTemp, HIS, resultsFilename,
		cycle, freqIdx, dumpFun)
numOfCurrentProblems = length(pMisTemp);
N_nodes = prod(pInv.MInv.n .+ 1);
nsrc = size(originalSources, 2);
p = size(Z1,2);
nwork = nworkers()
nsrc = size(originalSources, 2);
nfreq = length(pMisTemp);
wheres = map(x -> x.where, pMisTemp);
pMisTempFetched = map(fetch, pMisTemp);
mergedDobs = map(x -> x.dobs,pMisTempFetched)
mergedWd = map(x-> x.Wd,pMisTempFetched)

##########################
# Initialize outputs	 #
##########################
Dc = 0;
flag = -1;
HIS = [];
Z2 = 0;
##########################

FafterGN = 0.0;
F_zero_prev = 0.0;
mc_prev = convert(Array{Float16},mc);
for j = 1:itersNum
	println("============================== New ALM Iter ======================================");
	flush(Base.stdout)

	if simSrcDim==1
		TEmat = Matrix(1.0I,nsrc,nsrc)
	else
		TEmat = rand([-1,1],(nsrc,simSrcDim));
	end

	reducedSources = originalSources * TEmat;
	reducedDobs = map(x-> x*TEmat, mergedDobs);
	sizeWD = size(reducedDobs[1])
	reducedWd = map(x-> mean(x)/sqrt(simSrcDim) * ones(sizeWD), mergedWd);

	pMisTemp = setSourcesSame(pMisTemp,reducedSources);
	pMisTemp = setDobs(pMisTemp,reducedDobs);
	pMisTemp = setWd(pMisTemp,reducedWd);

	t1 = time_ns();
	Dc,F_zero, = computeMisfit(mc,pMisTemp,false);
	e1 = time_ns();
	println("Computed Misfit with orig sources : ",F_zero, " [Time: ",(e1 - t1)/1.0e9," sec]");

	mc_prev = convert(Array{Float16},mc);
	pMisHps = pMisTemp

	# Rc = Dc-Dobs, where Dc is the clean (wrt Z1,Z2) simulated data
	mergedRc = Array{Array{ComplexF64}}(undef,nfreq);
	for f = 1:nfreq
		mergedRc[f] = reducedDobs[f] .- fetch(Dc[f]);
	end

	println("Misfit with Zero Z2 (our computation): ", misfitCalc(mc,zeros(ComplexF64, (N_nodes, p)),zeros(ComplexF64, (p, size(TEmat,2))),mergedWd ./ sqrt(simSrcDim),mergedRc,nfreq,alpha1,alpha2, pMisHps))
	mergedRcReduced = mergedRc;

	pMisTempFetched = map(fetch, pMisTemp)
	t1 = time_ns();
	srcNum = simSrcDim == 1 ? nsrc : simSrcDim;
	Z2 = calculateZ2(mc,misfitCalc, p, srcNum, nfreq, nrcv,nwork, numOfCurrentProblems, mergedWd ./ sqrt(simSrcDim), mergedRcReduced, pMisHps, pMisTempFetched, Z1, alpha2);
	e1 = time_ns();

	print("After First Z2 update: ");
	mis = misfitCalc(mc,Z1,Z2,mergedWd ./ sqrt(simSrcDim),mergedRcReduced,nfreq,alpha1,alpha2, pMisHps);
	obj = objectiveCalc(Z1,Z2,mis,alpha1,alpha2);
	initialMis = mis
	initialObj = obj

	println("mis: ",mis,", obj: ",obj,", norm Z2 = ", norm(Z2)^2," norm Z1: ", norm(Z1)^2, ", [Time: ",(e1 - t1)/1.0e9," sec]")

	innerIters = j == 1 ? 5 : 1; # do five in first outer iteration only for better performance
	for iters = 1:innerIters
		###################################################
		#### COMPUTING Z1:
		###################################################
		t1 = time_ns();
		Z1 = calculateZ1(mc,misfitCalc, nfreq, mergedWd ./ sqrt(simSrcDim), mergedRcReduced, pMisHps, Z1, Z2, alpha1, stepReg);
		e1 = time_ns();
		mis = misfitCalc(mc,Z1,Z2,mergedWd ./ sqrt(simSrcDim) ,mergedRcReduced,nfreq,alpha1,alpha2, pMisHps);
		obj = objectiveCalc(Z1,Z2,mis,alpha1,alpha2);
		println("After Z1: mis: ",mis,", obj: ",obj,", norm Z2 = ", norm(Z2)^2," norm Z1: ", norm(Z1)^2, ", [Time: ",(e1 - t1)/1.0e9," sec]")

		###################################################
		#### COMPUTING Z2:
		###################################################
		t1 = time_ns();
		Z2 = calculateZ2(mc,misfitCalc, p, srcNum, nfreq, nrcv,nwork, numOfCurrentProblems, mergedWd ./ sqrt(simSrcDim), mergedRcReduced, pMisHps, pMisTempFetched, Z1, alpha2);
		e1 = time_ns();
		mis = misfitCalc(mc,Z1,Z2,mergedWd ./ sqrt(simSrcDim) ,mergedRcReduced,nfreq,alpha1,alpha2, pMisHps);
		obj = objectiveCalc(Z1,Z2,mis,alpha1,alpha2);
		println("After Z2: mis: ",mis,", obj: ",obj,", norm Z2 = ", norm(Z2)^2," norm Z1: ", norm(Z1)^2, ", [Time: ",(e1 - t1)/1.0e9," sec]")

		# allow extended source to reduce misfit by up to 20%
		if mis < 0.8 * initialMis
			break
		end
	end

	if mis / F_zero > 0.95
			alpha1 = alpha1 / 1.5;
			alpha2 = alpha2 / 1.5;
			println("Ratio mis/F_zero is: ",mis/F_zero,", hence decreasing alphas by 1.5: ",alpha1,",",alpha2);
	end

	if mis / F_zero < 0.3
			alpha1 = alpha1*1.5;
			alpha2 = alpha2*1.5;
			println("Ratio mis/F_zero is: ",mis/F_zero,", hence increasing alpha1 by 1.5: ",alpha1,",",alpha2);
	end

	newSources = originalSources * TEmat + Z1 * Z2;

	pMisTemp = setSourcesSame(pMisTemp,newSources);
	pMisTemp = setDobs(pMisTemp,reducedDobs)
	pMisTemp = setWd(pMisTemp,reducedWd)

	if resultsFilename == ""
		filename = "";
		hisMatFileName = "";
	else
		Temp = splitext(resultsFilename);
		filename = string(Temp[1],"_Cyc",cycle,"_FC",freqIdx,"_",j,"_GN",Temp[2]);
		hisMatFileName  =  string(Temp[1],"_Cyc",cycle,"_FC",freqIdx);
	end

	# Here we set a dump function for GN for this iteracion of FC
	function dumpGN(mc,Dc,iter,pInv,PF)
		dumpFun(mc,Dc,iter,pInv,PF,filename);
	end

	pMisTE = pMisTemp;

	flush(Base.stdout)
	t1 = time_ns();
	if method == "projGN"
		mc,Dc,flag,His = projGNCG(mc,pInv,pMisTE,dumpResults = dumpGN);
	elseif method == "barrierGN"
		mc,Dc,flag,His = barrierGNCG(mc,pInv,pMisTE,rho=1.0,dumpResults = dumpGN);
	end

	e1 = time_ns();
	print("runtime of GN:"); println((e1 - t1)/1.0e9);

	FafterGN = His.F[end];
	println("Computed Misfit with new sources after GN : ",FafterGN);

	if hisMatFileName != ""
		file = matopen(string(hisMatFileName,"_HisGN.mat"), "w");
		### PUT THE CONTENT OF "HIS" INTO THE MAT FILE ###
		His.Dc = [];
		write(file,"His",His);
		close(file);
	end
	His.Dc = []
	push!(HIS,His)
end

pMisTemp = setSourcesSame(pMisTemp,originalSources);
pMisTemp = setDobs(pMisTemp,mergedDobs)
pMisTemp = setWd(pMisTemp,mergedWd)

return mc,Z1,Z2,alpha1,alpha2*simSrcDim,Dc,flag,HIS
end

"""
	function freqCont

	Frequency continuation procedure for running FWI.
	This function runs GaussNewton on misfit functions defined by the frequencies.
"""
function freqCont(freqContParam::FreqContParam)
mc = freqContParam.mc
itersNum = freqContParam.itersNum
originalSources = freqContParam.originalSources
nrcv = freqContParam.nrcv
pInv = freqContParam.pInv
pMis = freqContParam.pMis
windowSize = freqContParam.windowSize
resultsFilename = freqContParam.resultsFilename
dumpFun = freqContParam.dumpFun
Iact = freqContParam.Iact
mback = freqContParam.mback
Z1 = freqContParam.Z1
simSrcDim = freqContParam.simSrcDim
alpha1 = freqContParam.alpha1
alpha2Orig = freqContParam.alpha2Orig
stepReg = freqContParam.stepReg
mode = freqContParam.mode
startFrom = freqContParam.startFrom
endAt = freqContParam.endAt
cycle = freqContParam.cycle
method = freqContParam.method
FWImethod = freqContParam.FWImethod
updateMref = freqContParam.updateMref

Z2 = 0;
Dc = 0;
flag = -1;
HIS = [];
alpha2 = alpha2Orig / simSrcDim;

println("~~~~~~~~~~~~~~~ SEG1  FreqCont: Regs are: ",alpha1,",",alpha2,",",stepReg);
regfun = pInv.regularizer
originalItersNum = itersNum
for freqIdx = startFrom:endAt
	if updateMref
		pInv.mref = copy(mc[:]);
	end
	if mode=="1stInit"
		reqIdx1 = freqIdx;
		if freqIdx > 1
			reqIdx1 = max(2,freqIdx-windowSize+1);
		end
		reqIdx2 = freqIdx;
	else
		reqIdx1 = freqIdx;
		if freqIdx > 1
			reqIdx1 = max(1,freqIdx-windowSize+1);
		end
		reqIdx2 = freqIdx;
	end

	currentProblems = reqIdx1:reqIdx2;
	pMisTemp = pMis[currentProblems];
	println("\n======= New Continuation Stage: selecting continuation batches: ",reqIdx1," to ",reqIdx2,"=======\n");
	if FWImethod == "FWI"
		pInv.maxIter = itersNum;
		mc,Dc,flag,HIS = standardGNrun(method, mc, simSrcDim, originalSources,
		 		pInv, pMisTemp, HIS, resultsFilename, cycle, freqIdx, dumpFun)
	elseif FWImethod == "FWI_ES"
		pInv.maxIter = 1;
		mc,Z1,Z2,alpha1,alpha2,Dc,flag,HIS = ExtendedSourcesGNRun(
				method, mc, simSrcDim, nrcv, itersNum, originalSources,
				Z1,alpha1, alpha2, stepReg, pInv, pMisTemp, HIS,
				resultsFilename, cycle, freqIdx, dumpFun)
	end

end

pInv.regularizer = regfun;
return mc,Z1,Z2,alpha1,alpha2*simSrcDim,Dc,flag,HIS;
end


text:
# it's too risky
using Iterators
import Iterators.filter

# for number of groups << 2^16
function sumby_fewgroups{T,S}(by::AbstractArray{T,1}, val::AbstractArray{S,1}; powerreduction = 48, seed = UInt(0))
  by_index = (hash.(by, seed) .>>> powerreduction) .+ UInt(1)
  res = zeros(S, Int64(2)^(64-powerreduction))
  # resize the Dict to a larger size
  for (byi, vali) in zip(by_index, val)
      @inbounds res[byi] += vali
  end
  collect(Iterators.filter(x->x[2] != 0, zip(by, res)))
end

@time sumby_fewgroups(id4, v1) # 3.4

@time sumby_fewgroups(id4, v1, seed = UInt(3)) # 2.7
@time sumby_fewgroups(id4, v1) # 3.4
@time sumby_fewgroups(id4, v1) # 2.7
using BenchmarkTools
@benchmark sumby_fewgroups(id4, v1) # 2.7

# complete failure for id6
@time sumby_fewgroups(id6, v1, 32)
@time sumby_fewgroups(id6, v1, 37, UInt(1))

using StatsBase
@time countmap(id4, weights(v1))
@time countmap(id6, weights(v1))
@time sumby(id6, v1)
addprocs()
@everywhere using FastGroupBy
@time psumby(id6, v1)

unique(sample(id6, 1_000_000, replace=false))
unique(sample(id6, 2_000_000, replace=false))

unique(rand(id6,1_000_000))
unique(rand(id6,2_000_000))
unique(rand(id6,3_000_000))
unique(rand(id6,4_000_000))
unique(rand(id6,5_000_000))
unique(rand(id6,6_000_000))

unique(rand(id4,1_000_000))
unique(rand(id4,2_000_000))
unique(rand(id4,3_000_000))
unique(rand(id4,4_000_000))
unique(rand(id4,5_000_000))
unique(rand(id4,6_000_000))


text:
# The following is largely inspired from JuMP/test/JuMPExtension.jl

@enum Level BOTH LOWER UPPER

abstract type AbstractBilevelModel <: JuMP.AbstractModel end

Base.broadcastable(model::AbstractBilevelModel) = Ref(model)

mutable struct BilevelModel <: AbstractBilevelModel
    # Structured data
    upper::JuMP.AbstractModel
    lower::JuMP.AbstractModel

    # Model data
    nextvaridx::Int                                 # Next variable index is nextvaridx+1
    variables::Dict{Int, JuMP.AbstractVariable}     # Map varidx -> variable
    varnames::Dict{Int, String}                     # Map varidx -> name
    var_level::Dict{Int, Level}
    var_upper::Dict{Int, JuMP.AbstractVariableRef}
    var_lower::Dict{Int, JuMP.AbstractVariableRef}

    # upper level decisions that are "paramters" of the second level
    upper_to_lower_link::Dict{JuMP.AbstractVariableRef, JuMP.AbstractVariableRef}
    # lower level decisions that are input to upper
    lower_to_upper_link::Dict{JuMP.AbstractVariableRef, JuMP.AbstractVariableRef}
    # joint link
    link::Dict{JuMP.AbstractVariableRef, JuMP.AbstractVariableRef}

    nextconidx::Int                                 # Next constraint index is nextconidx+1
    constraints::Dict{Int, JuMP.AbstractConstraint} # Map conidx -> variable
    connames::Dict{Int, String}                     # Map varidx -> name
    ctr_level::Dict{Int, Level}
    ctr_upper::Dict{Int, JuMP.ConstraintRef}
    ctr_lower::Dict{Int, JuMP.ConstraintRef}

    upper_objective_sense::MOI.OptimizationSense
    upper_objective_function::JuMP.AbstractJuMPScalar

    lower_objective_sense::MOI.OptimizationSense
    lower_objective_function::JuMP.AbstractJuMPScalar

    # solution data
    solver#::MOI.ModelLike
    upper_to_sblm
    lower_to_sblm
    sblm_to_solver

    objdict::Dict{Symbol, Any}                      # Same that JuMP.Model's field `objdict`

    function BilevelModel()

        model = new(
            JuMP.Model(),
            JuMP.Model(),

            # var
            0, Dict{Int, JuMP.AbstractVariable}(),   Dict{Int, String}(),    # Model Variables
            Dict{Int, Level}(), Dict{Int, JuMP.AbstractVariable}(), Dict{Int, JuMP.AbstractVariable}(),
            # links
            Dict{Int, JuMP.AbstractVariable}(), Dict{Int, JuMP.AbstractVariable}(),
            Dict{Int, JuMP.AbstractVariable}(),
            #ctr
            0, Dict{Int, JuMP.AbstractConstraint}(), Dict{Int, String}(),    # Model Constraints
            Dict{Int, Level}(), Dict{Int, JuMP.AbstractConstraint}(), Dict{Int, JuMP.AbstractConstraint}(),
            #obj
            MOI.FEASIBILITY_SENSE,
            zero(JuMP.GenericAffExpr{Float64, BilevelVariableRef}), # Model objective
            MOI.FEASIBILITY_SENSE,
            zero(JuMP.GenericAffExpr{Float64, BilevelVariableRef}), # Model objective

            nothing,
            nothing,
            nothing,
            nothing,
            Dict{Symbol, Any}(),
            )

        return model
    end
end

abstract type InnerBilevelModel <: AbstractBilevelModel end
struct UpperModel <: InnerBilevelModel
    m::BilevelModel
end
Upper(m::BilevelModel) = UpperModel(m)
UpperToLower(m::BilevelModel) = UpperModel(m)
struct LowerModel <: InnerBilevelModel
    m::BilevelModel
end
Lower(m::BilevelModel) = LowerModel(m)
LowerToUpper(m::BilevelModel) = LowerModel(m)
bilevel_model(m::InnerBilevelModel) = m.m
mylevel_model(m::UpperModel) = bilevel_model(m).upper
mylevel_model(m::LowerModel) = bilevel_model(m).lower
level(m::LowerModel) = LOWER
level(m::UpperModel) = UPPER
mylevel_ctr_list(m::LowerModel) = bilevel_model(m).ctr_lower
mylevel_ctr_list(m::UpperModel) = bilevel_model(m).ctr_upper
mylevel_var_list(m::LowerModel) = bilevel_model(m).var_lower
mylevel_var_list(m::UpperModel) = bilevel_model(m).var_upper

# obj

mylevel_obj_sense(m::LowerModel) = bilevel_model(m).lower_objective_sense
mylevel_obj_function(m::LowerModel) = bilevel_model(m).lower_objective_function
mylevel_obj_sense(m::UpperModel) = bilevel_model(m).upper_objective_sense
mylevel_obj_function(m::UpperModel) = bilevel_model(m).upper_objective_function

set_mylevel_obj_sense(m::LowerModel, val) = bilevel_model(m).lower_objective_sense = val
set_mylevel_obj_function(m::LowerModel, val) = bilevel_model(m).lower_objective_function = val
set_mylevel_obj_sense(m::UpperModel, val) = bilevel_model(m).upper_objective_sense = val
set_mylevel_obj_function(m::UpperModel, val) = bilevel_model(m).upper_objective_function = val

function set_link!(m::UpperModel, upper::JuMP.AbstractVariableRef, lower::JuMP.AbstractVariableRef)
    bilevel_model(m).upper_to_lower_link[upper] = lower
    bilevel_model(m).link[upper] = lower
    nothing
end
function set_link!(m::LowerModel, upper::JuMP.AbstractVariableRef, lower::JuMP.AbstractVariableRef)
    bilevel_model(m).lower_to_upper_link[lower] = upper
    bilevel_model(m).link[upper] = lower
    nothing
end

# Models to deal with variables that are not exchanged between models
abstract type SingleBilevelModel <: AbstractBilevelModel end
struct UpperOnlyModel <: SingleBilevelModel
    m::BilevelModel
end
UpperOnly(m::BilevelModel) = UpperOnlyModel(m)
struct LowerOnlyModel <: SingleBilevelModel
    m::BilevelModel
end
LowerOnly(m::BilevelModel) = LowerOnlyModel(m)

bilevel_model(m::SingleBilevelModel) = m.m
mylevel_model(m::UpperOnlyModel) = bilevel_model(m).upper
mylevel_model(m::LowerOnlyModel) = bilevel_model(m).lower
level(m::LowerOnlyModel) = LOWER
level(m::UpperOnlyModel) = UPPER
mylevel_var_list(m::LowerOnlyModel) = bilevel_model(m).var_lower
mylevel_var_list(m::UpperOnlyModel) = bilevel_model(m).var_upper

#### Model ####

# Variables
struct BilevelVariableRef <: JuMP.AbstractVariableRef
    model::BilevelModel # `model` owning the variable
    idx::Int       # Index in `model.variables`
    level::Level
end
mylevel(v::BilevelVariableRef) = v.level
in_level(v::BilevelVariableRef, level::Level) = (v.level === BOTH || v.level === level)
function solver_ref(v::BilevelVariableRef)
    m = v.model
    if mylevel(v) == LOWER
        return m.sblm_to_solver[
            m.lower_to_sblm[JuMP.index(m.var_lower[v.idx])]]
    else
        return m.sblm_to_solver[
            m.upper_to_sblm[JuMP.index(m.var_upper[v.idx])]]
    end
end
Base.broadcastable(v::BilevelVariableRef) = Ref(v)
Base.copy(v::BilevelVariableRef) = v
# TODO
# Base.copy(v::BilevelVariableRef, new_model::BilevelModel) = BilevelVariableRef(new_model, v.idx, v.level) # also requires a map
Base.:(==)(v::BilevelVariableRef, w::BilevelVariableRef) =
    v.model === w.model && v.idx == w.idx && v.level == w.level
JuMP.owner_model(v::BilevelVariableRef) = v.model
JuMP.isequal_canonical(v::BilevelVariableRef, w::BilevelVariableRef) = v == w
JuMP.variable_type(::AbstractBilevelModel) = BilevelVariableRef
# add in BOTH levels
function JuMP.add_variable(inner::InnerBilevelModel, v::JuMP.AbstractVariable, name::String="")
    m = bilevel_model(inner)
    m.nextvaridx += 1
    vref = BilevelVariableRef(m, m.nextvaridx, BOTH)
    v_upper = JuMP.add_variable(m.upper, v, name)
    m.var_upper[vref.idx] = v_upper
    v_lower = JuMP.add_variable(m.lower, v, name)
    m.var_lower[vref.idx] = v_lower
    m.var_level[vref.idx] = BOTH
    set_link!(inner, v_upper, v_lower)
    m.variables[vref.idx] = v
    JuMP.set_name(vref, name)
    vref
end
function JuMP.add_variable(single::SingleBilevelModel, v::JuMP.AbstractVariable, name::String="")
    m = bilevel_model(single)
    m.nextvaridx += 1
    vref = BilevelVariableRef(m, m.nextvaridx, level(single))
    v_level = JuMP.add_variable(mylevel_model(single), v, name)
    mylevel_var_list(single)[vref.idx] = v_level
    m.var_level[vref.idx] = level(single)
    m.variables[vref.idx] = v
    JuMP.set_name(vref, name)
    vref
end
function MOI.delete!(m::AbstractBilevelModel, vref::BilevelVariableRef)
    error("No deletion on bilevel models")
    delete!(m.variables, vref.idx)
    delete!(m.varnames, vref.idx)
end
MOI.is_valid(m::BilevelModel, vref::BilevelVariableRef) = vref.idx in keys(m.variables)
JuMP.num_variables(m::BilevelModel) = length(m.variables)
JuMP.num_variables(m::InnerBilevelModel) = JuMP.num_variables(bilevel_model(m))

# -------------------------
# begin
# Unchanged from StructJuMP
# -------------------------

# Internal function
variable_info(vref::BilevelVariableRef) = vref.model.variables[vref.idx].info
function update_variable_info(vref::BilevelVariableRef, info::JuMP.VariableInfo)
    vref.model.variables[vref.idx] = JuMP.ScalarVariable(info)
end

JuMP.has_lower_bound(vref::BilevelVariableRef) = variable_info(vref).has_lb
function JuMP.lower_bound(vref::BilevelVariableRef)
    @assert !JuMP.is_fixed(vref)
    variable_info(vref).lower_bound
end
function JuMP.set_lower_bound(vref::BilevelVariableRef, lower)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(true, lower,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.delete_lower_bound(vref::BilevelVariableRef)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(false, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
JuMP.has_upper_bound(vref::BilevelVariableRef) = variable_info(vref).has_ub
function JuMP.upper_bound(vref::BilevelVariableRef)
    @assert !JuMP.is_fixed(vref)
    variable_info(vref).upper_bound
end
function JuMP.set_upper_bound(vref::BilevelVariableRef, upper)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           true, upper,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.delete_upper_bound(vref::BilevelVariableRef)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           false, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
JuMP.is_fixed(vref::BilevelVariableRef) = variable_info(vref).has_fix
JuMP.fix_value(vref::BilevelVariableRef) = variable_info(vref).fixed_value
function JuMP.fix(vref::BilevelVariableRef, value)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           true, value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
function JuMP.unfix(vref::BilevelVariableRef)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           false, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, info.integer))
end
JuMP.start_value(vref::BilevelVariableRef) = variable_info(vref).start
function JuMP.set_start_value(vref::BilevelVariableRef, start)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           true, start,
                                           info.binary, info.integer))
end
JuMP.is_binary(vref::BilevelVariableRef) = variable_info(vref).binary
function JuMP.set_binary(vref::BilevelVariableRef)
    @assert !JuMP.is_integer(vref)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           true, info.integer))
end
function JuMP.unset_binary(vref::BilevelVariableRef)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           false, info.integer))
end
JuMP.is_integer(vref::BilevelVariableRef) = variable_info(vref).integer
function JuMP.set_integer(vref::BilevelVariableRef)
    @assert !JuMP.is_binary(vref)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, true))
end
function JuMP.unset_integer(vref::BilevelVariableRef)
    info = variable_info(vref)
    update_variable_info(vref,
                         JuMP.VariableInfo(info.has_lb, info.lower_bound,
                                           info.has_ub, info.upper_bound,
                                           info.has_fix, info.fixed_value,
                                           info.has_start, info.start,
                                           info.binary, false))
end

# -------------------------
# end
# Unchanged from StructJuMP
# -------------------------

# Constraints
struct BilevelConstraintRef
    model::BilevelModel # `model` owning the constraint
    idx::Int       # Index in `model.constraints`
    # level
end
JuMP.constraint_type(::AbstractBilevelModel) = BilevelConstraintRef
function JuMP.add_constraint(m::BilevelModel, c::JuMP.AbstractConstraint, name::String="")
    error(
        "Can't add constraint directly to the bilevel model `m`, "*
        "attach the constraint to the upper or lower model "*
        "with @constraint(Upper(m), ...) or @constraint(Lower(m), ...)")
end
# function constraint_object(con_ref::ConstraintRef{Model, _MOICON{FuncType, SetType}}) where
#     {FuncType <: MOI.AbstractScalarFunction, SetType <: MOI.AbstractScalarSet}
#     model = con_ref.model
#     f = MOI.get(model, MOI.ConstraintFunction(), con_ref)::FuncType
#     s = MOI.get(model, MOI.ConstraintSet(), con_ref)::SetType
#     return ScalarConstraint(jump_function(model, f), s)
# end
JuMP.add_constraint(m::UpperModel, c::JuMP.VectorConstraint, name::String="") =
error("no vec ctr")
function JuMP.add_constraint(m::InnerBilevelModel, c::JuMP.ScalarConstraint{F,S}, name::String="") where {F,S}
    blm = bilevel_model(m)
    blm.nextconidx += 1
    cref = BilevelConstraintRef(blm, blm.nextconidx)
    func = JuMP.jump_function(c)
    level_func = replace_variables(func, bilevel_model(m), mylevel_model(m), mylevel_var_list(m), level(m))
    level_c = JuMP.build_constraint(error, level_func, c.set)
    level_cref = JuMP.add_constraint(mylevel_model(m), level_c, name)
    blm.ctr_level[cref.idx] = level(m)
    mylevel_ctr_list(m)[cref.idx] = level_cref
    blm.constraints[cref.idx] = c
    JuMP.set_name(cref, name)
    cref
end
function MOI.delete!(m::AbstractBilevelModel, cref::BilevelConstraintRef)
    error("can't delete")
    delete!(m.constraints, cref.idx)
    delete!(m.connames, cref.idx)
end
MOI.is_valid(m::BilevelModel, cref::BilevelConstraintRef) = cref.idx in keys(m.constraints)
MOI.is_valid(m::InnerBilevelModel, cref::BilevelConstraintRef) =
    MOI.is_valid(bilevel_model(m), cref) && bilevel_model(m).ctr_level[cref.idx] == level(m)
function JuMP.constraint_object(cref::BilevelConstraintRef, F::Type, S::Type)
    c = cref.model.constraints[cref.idx]
    # `TypeError` should be thrown is `F` and `S` are not correct
    # This is needed for the tests in `constraints.jl`
    c.func::F
    c.set::S
    c
end

# Objective
function JuMP.set_objective(m::InnerBilevelModel, sense::MOI.OptimizationSense,
                            f::JuMP.AbstractJuMPScalar)
    set_mylevel_obj_sense(m, sense)
    set_mylevel_obj_function(m, f)
    level_f = replace_variables(f, bilevel_model(m), mylevel_model(m), mylevel_var_list(m), level(m))
    JuMP.set_objective(mylevel_model(m), sense, level_f)
end
JuMP.objective_sense(m::InnerBilevelModel) = mylevel_obj_sense(m)
JuMP.objective_function_type(m::InnerBilevelModel) = typeof(mylevel_obj_function(m))
JuMP.objective_function(m::InnerBilevelModel) = mylevel_obj_function(m)
function JuMP.objective_function(m::InnerBilevelModel, FT::Type)
    mylevel_obj_function(m) isa FT || error("The objective function is not of type $FT")
    mylevel_obj_function(m)
end

# todo remove
JuMP.objective_sense(m::AbstractBilevelModel) = MOI.FEASIBILITY_SENSE
# end todo remove
JuMP.num_variables(m::AbstractBilevelModel) = JuMP.num_variables(bilevel_model(m))
JuMP.show_constraints_summary(::Any, ::AbstractBilevelModel) = "no summary"
JuMP.show_backend_summary(::Any, ::AbstractBilevelModel) = "no summary"
JuMP.object_dictionary(m::BilevelModel) = m.objdict
JuMP.object_dictionary(m::AbstractBilevelModel) = JuMP.object_dictionary(bilevel_model(m))
JuMP.show_objective_function_summary(::IO, ::AbstractBilevelModel) = "no summary"
# TODO
# function JuMP.constraints_string(print_mode, model::MyModel)

bileve_obj_error() = error("There is no objective for BilevelModel use Upper(.) and Lower(.)")

function JuMP.set_objective(m::BilevelModel, sense::MOI.OptimizationSense,
    f::JuMP.AbstractJuMPScalar)
    bileve_obj_error()
end
JuMP.objective_sense(m::BilevelModel) = JuMP.objective_sense(m.upper)#bileve_obj_error()
JuMP.objective_function_type(model::BilevelModel) = bileve_obj_error()
JuMP.objective_function(model::BilevelModel) = bileve_obj_error()
function JuMP.objective_function(model::BilevelModel, FT::Type)
    bileve_obj_error()
end

# Names
JuMP.name(vref::BilevelVariableRef) = vref.model.varnames[vref.idx]
function JuMP.set_name(vref::BilevelVariableRef, name::String)
    vref.model.varnames[vref.idx] = name
end
JuMP.name(cref::BilevelConstraintRef) = cref.model.connames[cref.idx]
function JuMP.set_name(cref::BilevelConstraintRef, name::String)
    cref.model.connames[cref.idx] = name
end
# TODO
# function JuMP.variable_by_name(model::MyModel, name::String)
# see jump extensions
# function JuMP.constraint_by_name(model::MyModel, name::String)

# replace variables
function replace_variables(var::BilevelVariableRef,
    model::BilevelModel,
    inner::JuMP.AbstractModel,
    variable_map::Dict{Int, V},
    level::Level) where {V<:JuMP.AbstractVariableRef}
    if var.model === model && in_level(var, level)
        return variable_map[var.idx]
    elseif var.model === model
        error("Variable $(var) belonging Only to $(var.level) level, was added in the $(level) level.")
    else
        error("A BilevelModel cannot have expression using variables of a BilevelModel different from itself")
    end
end
function replace_variables(aff::JuMP.GenericAffExpr{C, BilevelVariableRef},
    model::BilevelModel,
    inner::JuMP.AbstractModel,
    variable_map::Dict{Int, V},
    level::Level) where {C,V<:JuMP.AbstractVariableRef}
    result = JuMP.GenericAffExpr{C, JuMP.VariableRef}(0.0)#zero(aff)
    result.constant = aff.constant
    for (coef, var) in JuMP.linear_terms(aff)
        JuMP.add_to_expression!(result,
        coef,
        replace_variables(var, model, model, variable_map, level))
    end
    return result
end
# TODO allow quadratic obj
function replace_variables(quad::JuMP.GenericQuadExpr{C, BilevelVariableRef},
    model::BilevelModel,
    inner::JuMP.AbstractModel,
    variable_map::Dict{Int, V},
    level::Level) where {C,V<:JuMP.AbstractVariableRef}
    aff = replace_variables(quad.aff, model, model, variable_map, level)
    quadv = JuMP.GenericQuadExpr{C, JuMP.VariableRef}(aff)
    for (coef, var1, var2) in JuMP.quad_terms(quad)
        JuMP.add_to_expression!(quadv,
        coef,
        replace_variables(var1, model, model, variable_map, level),
        replace_variables(var1, model, model, variable_map, level))
    end
    return quadv
end
function replace_variables(quad::C,
    model::BilevelModel,
    inner::JuMP.AbstractModel,
    variable_map::Dict{Int, V},
    level::Level) where {C,V<:JuMP.AbstractVariableRef}
    error("A BilevelModel cannot have $(C) function")
end
replace_variables(funcs::Vector, args...) = map(f -> replace_variables(f, args...), funcs)

# using MathOptFormat
# function print_lp(m, name)
#     lp_model = MathOptFormat.MOF.Model()
#     MOI.copy_to(lp_model, m)
#     MOI.write_to_file(lp_model, name)
# end

JuMP.optimize!(::T) where {T<:AbstractBilevelModel} =
    error("cant solve a model of type: $T ")
function JuMP.optimize!(model::BilevelModel, optimizer, mode::BilevelSolverMode{T} = SOS1Mode()) where T

    if !MOI.is_empty(optimizer)
        error("An empty optimizer must be provided")
    end

    upper = JuMP.backend(model.upper)
    lower = JuMP.backend(model.lower)

    # print_lp(upper, "upper.lp")
    # print_lp(lower, "lower.lp")

    moi_upper = JuMP.index.(
        collect(values(model.upper_to_lower_link)))
    moi_link = JuMP.index(model.link)

    single_blm, upper_to_sblm, lower_to_sblm = build_bilevel(
        upper, lower, moi_link, moi_upper, mode)

    solver = optimizer#MOI.Bridges.full_bridge_optimizer(optimizer, Float64)
    sblm_to_solver = MOI.copy_to(solver, single_blm)

    MOI.optimize!(solver)

    # map from bridged to single_blm
    # map from single_blm to upper & lowel
    # map from upper & lower to model
    model.solver  = solver#::MOI.ModelLike
    model.upper_to_sblm = upper_to_sblm
    model.lower_to_sblm = lower_to_sblm
    model.sblm_to_solver = sblm_to_solver

    # feasible = JuMP.primal_status(single_blm)# == MOI.FEASIBLE_POINT
    # termination = JuMP.termination_status(single_blm)# == MOI.OPTIMAL
    # objective_value = NaN
    # if feasible == MOI.FEASIBLE_POINT
    #     objective_value = JuMP.objective_value(single_blm)
    # end
    # variable_value = Dict{JuMP.VariableRef, Float64}()
    return nothing
end

function JuMP.index(d::Dict)
    ret = Dict{VI,VI}()
    # sizehint!(ret, length(d))
    for (k,v) in d
        ret[JuMP.index(k)] = JuMP.index(v)
    end
    return ret
end

function JuMP.value(v::BilevelVariableRef)::Float64
    m = owner_model(v)
    solver = m.solver
    ref = solver_ref(v)
    return MOI.get(solver, MOI.VariablePrimal(), ref)
end
function JuMP.value(v::BilevelConstraintRef)::Float64
    error("value of BilevelConstraintRef not enabled")
end
function JuMP.dual(v::BilevelConstraintRef)::Float64
    error("dual of BilevelConstraintRef not enabled")
end
function JuMP.primal_status(model::BilevelModel)
    return MOI.get(model.solver, MOI.PrimalStatus())
end
JuMP.dual_status(model::BilevelModel) = error("dual status cant be queried for BilevelModel")
function JuMP.termination_status(model::BilevelModel)
    return MOI.get(model.solver, MOI.TerminationStatus())
end
function JuMP.objective_value(model::BilevelModel)
    return MOI.get(model.solver, MOI.ObjectiveValue())
end


text:
##### Looking at the biomass dynamics of combinations of DR params and dose that result in unstable dynamics

## Initial conditions
Z = 1
I_K = 3
T = 293.15

# time span for each simulation
tspan = (0.0, 315360000000.0)

# define extinction threshold
threshold = 10e-12

# DR params
sl = 0.0
asym  = 0.0
xmid = 50.0

dose = 98

# generate parameters specific to this combo of target rate, DR slope and dose
p = ModelParameters_CG(param, T, I_K, Z, dose, "x", slope = sl, lower = asym, xmid = xmid)

# Calculate starting biomasses
u0 = [p.K[1]/2, p.K[1]/2, p.K[1]/2]

# Define the problem
prob = ODEProblem(BEFW, u0, tspan, p)

# Solve the problem
sol = solve(prob, maxiters = 1e7, saveat = t_keep)

plot(sol, title = "sl = $sl, as = $asym, xmid = $xmid, dose = $dose")
plot(sol[1:500],title = "sl = $sl, as = $asym, xmid = $xmid, dose = $dose")

u0

######### Save at interval
### Using Eva's code to make an interval to collect biomasses
interval_tkeep = 5
t_keep = collect(0:interval_tkeep:tspan[2])

tspan[2]

max_stars_repo_name: hanamayall/simple-feeding-DR

max_issues_repo_name: hanamayall/simple-feeding-DR

max_forks_repo_name: hanamayall/simple-feeding-DR


text:
using Forge.DSL
using Base.Test

@test isa(create_window, Function)
@test isa(isclosed, Function)
@test isa(!isclosed, Function)


text:
# computation using analytical expression

# define cache
struct PBM_cache{cnumtype, cdentype, cprobtype, cpredxtype, crestype}
  cnum::cnumtype
  cden::cdentype
  cprob::cprobtype
  cpredx::cpredxtype
  cres::crestype
end

# get optimal predictions for specific sample
function get_pred_opt_x_PBM!(x, distribution, p_range,PBMc)
  @unpack cnum, cden, cprob, cpredx, cres = PBMc

  cnum[1] = zero(eltype(cnum[1]))
  cden[1] = zero(eltype(cden[1]))

  # compute numerator and denominator
  for p in p_range
    cprob[1] = distribution(x, p)
    cnum[1] += cprob[1]*p
    cden[1] += cprob[1]
  end

  if cden[1] == zero(eltype(cden[1]))
    return zero(eltype(cden[1]))
  else
    return cnum[1]/cden[1]
  end
end

# get mean optimal predictions for fixed value of tuning parameter
function get_pred_opt_p_PBM!(samples, distribution, p_range, p_tar,PBMc)
  @unpack cnum, cden, cprob, cpredx, cres = PBMc

  cres[1] = zero(eltype(cres[1]))

  for x in samples
    cpredx[1] = get_pred_opt_x_PBM!(x, distribution, p_range,PBMc)
    cprob[1] = distribution(x, p_tar)

    # contribution to optimal prediction
    cres[1] += cprob[1]*cpredx[1]

    # contribution to optimal loss
    cres[2] += cprob[1]*(p_tar-cpredx[1])^2
  end
  return nothing
end

# compute optimal predictions and indicators, as well as optimal loss of PBM
function get_indicators_PBM_analytical(samples, distribution, p_range, dp)

  caches=[PBM_cache([zero(eltype(dp))], [zero(eltype(dp))], [zero(eltype(dp))],[zero(eltype(dp))], [zero(eltype(dp)), zero(eltype(dp))]) for i in 1:Threads.nthreads()]
  pred_PBM_opt = zeros(eltype(dp), length(p_range))

  # start parallel computation for sampled values of tuning parameter
  Threads.@threads for indxp in collect(1:length(p_range))
    get_pred_opt_p_PBM!(samples, distribution, p_range, p_range[indxp], caches[Threads.threadid()])
    pred_PBM_opt[indxp] = caches[Threads.threadid()].cres[1]
  end

  # sum up all individual loss contributions
  cres2 = zero(eltype(dp))
  for i in 1:Threads.nthreads()
    cres2 += caches[i].cres[2]
  end

  # compute indicator using symmetric difference quotient
  return pred_PBM_opt, map(x -> x-1,(circshift(pred_PBM_opt, -1).-circshift(pred_PBM_opt, 1))./(2*dp))[2:end-1], cres2/length(p_range)
end

# computation using neural networks

# loss function where each prediction is weighted by probability associated with input
function main_loss_PBM_weighted(NN, pnn, data, p_range, inputs)
  indices = convert.(Int, data[1, :])
  input = inputs[:, indices]

  pred = NN(pnn)(input)
  labels = reshape([p_range[Int(i)] for i in data[3, :]], 1, length(data[3, :]))
  freq = reshape(data[2, :], 1, length(data[2, :]))

  return sum(freq.*((pred.-labels).^2))
end

# loss function without weighting
function main_loss_PBM_stochastic(NN, pnn, dataset, p_range, inputs)
  indices = convert.(Int, dataset[1, :])
  input = inputs[:, indices]

  pred = NN(pnn)(input)
  labels = reshape([p_range[Int(i)] for i in dataset[3, :]], 1, length(dataset[3, :]))

  return sum((pred.-labels).^2)/length(dataset[3, :])
end

# train NN in PBM based on weighted loss function
function train_PBM_weighted(NN, pnn, dataset, epochs, p_range, dp, opt, batchsize, n_batches, inputs; verbose=false, saveat=epochs)
  NN_logger = [zeros(eltype(p_range[1]), length(pnn))]
  pred_logger = [(zeros(eltype(p_range[1]), length(p_range)), zeros(eltype(p_range[1]), length(p_range)-2), zeros(eltype(p_range[1]), 1))]

  pnn_best = zeros(eltype(pnn[1]), length(pnn))
  best_loss = zero(eltype(pnn[1]))

  losses = zeros(eltype(p_range[1]), epochs)
  indices = collect(1:size(dataset)[2])
  for epoch in 1:epochs
    Random.shuffle!(indices)

    # compute gradient batch-wise
    grad = zeros(eltype(p_range[1]), length(pnn))
    for batch in 1:n_batches
      randint = get_batches(indices, batchsize, batch)
      data = reshape(dataset[:, randint], 3, length(randint))
      val, back = Flux.Zygote.pullback(p -> main_loss_PBM_weighted(NN, p, data, p_range, inputs), pnn)
      grad .+= back(one(val))[1]
      losses[epoch] += val
    end

    # update NN parameters based on overall gradient
    Flux.Optimise.update!(opt, pnn, grad./length(p_range))
    losses[epoch] =losses[epoch]/length(p_range)

    # keep track of best performing NN
    if epoch == 1
      best_loss = losses[epoch]
    elseif losses[epoch] < best_loss
      println("better")
      println(best_loss)
      println(losses[epoch])
      pnn_best = deepcopy(pnn)
      best_loss = losses[epoch]
    end

    # save at regular intervals
    if epoch % saveat == 0
      push!(NN_logger, pnn_best)
      push!(pred_logger, predict_PBM(dataset, p_range, pnn_best, dp, NN, batchsize, n_batches, inputs, calc_loss=true))
    end

    verbose && println("epoch: $epoch / $(epochs)")
    verbose && println("loss: "*string(losses[epoch]))
  end
  return  losses, NN_logger, pred_logger
end

# train NN in PBM based on ``stochastic'' loss function
function train_PBM_stochastic(NN,pnn, dataset, epochs, p_range, dp, opt, batchsize, batchsize_stochastic, n_batches, n_batches_train_stochastic, inputs; verbose=false, saveat=epochs)
  NN_logger = [zeros(eltype(p_range[1]), length(pnn))]
  pred_logger = [(zeros(eltype(p_range[1]), length(p_range)), zeros(eltype(p_range[1]), length(p_range)-2), zeros(eltype(p_range[1]), 1))]

  pnn_best = zeros(eltype(pnn[1]), length(pnn))
  best_loss = zero(eltype(pnn[1]))

  losses = zeros(eltype(p_range[1]), epochs)
  for epoch in 1:epochs

    # compute loss and gradient batchwise
    for batch in 1:n_batches_train_stochastic
      randint = sample(1:length(dataset[1, :]), Weights(dataset[2, :]), batchsize_stochastic)
      data = reshape(dataset[:, randint], 3, length(randint))
      val, back = Flux.Zygote.pullback(p -> main_loss_PBM_stochastic(NN, p, data, p_range, inputs), pnn)
      grad = back(one(val))[1]

      # update NN parameters based on batch gradient
      Flux.Optimise.update!(opt, pnn, grad)
      losses[epoch] += val
    end
    losses[epoch] =losses[epoch]/n_batches_train_stochastic

    # keep track of best performing NN
    if epoch == 1
      best_loss = losses[epoch]
    elseif losses[epoch] < best_loss
      println("better")
      println(best_loss)
      println(losses[epoch])
      pnn_best = deepcopy(pnn)
      best_loss = losses[epoch]
    end

    # save at regular intervals
    if epoch % saveat == 0
      push!(NN_logger, pnn_best)
      push!(pred_logger, predict_PBM(dataset, p_range, pnn_best, dp, NN, batchsize, n_batches, inputs, calc_loss=true))
    end

    verbose && println("epoch: $epoch / $(epochs)")
    verbose && println("loss: "*string(losses[epoch]))
  end
  return  losses, NN_logger, pred_logger
end

# compute predictions, indicators, and loss value based on current NN
function predict_PBM(dataset, p_range, pnn, dp, NN, batchsize, n_batches, inputs; calc_loss=false, loss=zero(eltype(p_range[1])))

  # compute predictions and loss batch-wise
  predictions = zeros(eltype(p_range[1]), length(p_range))
  indices = collect(1:size(dataset)[2])
  for batch in 1:n_batches
    rand_int = get_batches(indices, batchsize ,batch)
    data = reshape(dataset[:, rand_int], 3, length(rand_int))

    indicess = convert.(Int,data[1, :])
    input = inputs[:, indicess]

    pred = NN(pnn)(input)
    for indxp in 1:length(rand_int)
      predictions[Int(data[3, indxp])] += data[2, indxp]*pred[1, indxp]
    end

    if calc_loss
      loss += main_loss_PBM_weighted(NN, pnn, data, p_range, inputs)
    end
  end

  if calc_loss
    loss = loss/length(p_range)
  end

  # compute indicator using symmetric difference quotient
  return predictions,map(x -> x-1,(circshift(predictions, -1).-circshift(predictions, 1))./(2*dp))[2:end-1], [loss]
end

# apply PBM using neural networks
function get_indicators_PBM_numerical(pnn, NN, dataset, epochs, p_range, dp, opt, inputs; verbose=false, trained=false, saveat=epochs, batchsize=length(dataset[1,:]), stochastic=false, n_batches_train_stochastic=10, batchsize_stochastic=length(dataset[1,:]), dataset_train=dataset)

  n_batches = ceil(eltype(batchsize), size(dataset)[2]/batchsize)

  if !trained
    if stochastic
      losses, NN_logger, pred_logger = train_PBM_stochastic(NN, pnn, dataset, epochs, p_range, dp, opt, batchsize, batchsize_stochastic, n_batches, n_batches_train_stochastic, inputs, verbose=verbose, saveat=saveat)
    else
      losses, NN_logger, pred_logger = train_PBM_weighted(NN, pnn, dataset, epochs, p_range, dp, opt, batchsize, n_batches, inputs, verbose=verbose, saveat=saveat)
    end

    pred_logger = pred_logger[2:end]
    NN_logger = NN_logger[2:end]

  else
    predictions, indicator, loss = predict_PBM(dataset, p_range, pnn, dp, NN, batchsize, n_batches, inputs, calc_loss=true)
    losses = zeros(eltype(dp), epochs)
    NN_logger = [pnn]
    pred_logger = [(predictions, indicator, loss)]
  end

  return pred_logger, losses, NN_logger
end


text:
function CoordinateDescent3{F<:AbstractFloat}(X::Array{F}, Y::Array{F}, LASSOpath::Array{F}, λ_max::F, w::Array{F}, σ::Array{F})
    ## Number of available lambdas: Granularity of the regulation parameter
    MaxIterλ = length(LASSOpath)

    n, p = size(X)

    ## Preallocating vector β_lasso for speed
    β_lasso = zeros(p, MaxIterλ)

    ## Initializing parameters
    β_tilde = zeros(p,1)

    ## Active set
    A= Int[]

    innerX = X' * X

    ## Calcula os produtos internos das colunas de X com Y
    innerY = sum(X.*repmat(Y,1,p),1)

    bestB = 1
    BICold = 100000000
    teste = 0

    β_ols = zeros(p,1)

    ## First for runs through the LASSO path
    for cont = 1:MaxIterλ
        ## Computing next lambda
        λ = LASSOpath[cont] * λ_max
        lastchange = 0
        it_while = 1

        while (teste != 2)
            ##  Reference active set
            A_anterior = A

            ##  Cycle
            if it_while == 1
                j_range = collect(1:p)'
            else
                #j_range = A_actual
                j_range = collect(1:p)'
            end

            for j = j_range

                aux, λEff, β_ols = CoordinateDescentCore(A, innerX, j, β_tilde, β_ols, n, innerY, λ, σ)

                #updateCandidates{T<:Signed, F<:AbstractFloat}(λEff::F, β_ols::Array{F},
                #                 j::T, A_actual::Array{T}, lastchange::T, β_tilde::Array{F}; w::Array{F} = [])

                A, β_tilde, lastchange, breakThis = updateCandidates(λEff, β_ols, j, A, lastchange, β_tilde, w = w)

                if breakThis == 1
                    break
                end

            end
            ##  Check if the active set has been altered in the last cycle
            teste = convCheck(A_anterior, A)
            it_while = it_while + 1

        end
        teste = 0

        BICnew = n * ( log(var( Y - X*β_tilde )) ) + sum( abs(β_tilde./σ) .> 1e-3 ) * log(n)
        if BICnew < BICold
            BICold = BICnew
            bestB = cont
        end

        β_lasso[:,cont] = β_tilde
    end

    return β_lasso, bestB
end


text:
# ---
# title: 1031. Maximum Sum of Two Non-Overlapping Subarrays
# id: problem1031
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array
# link: <https://leetcode.com/problems/maximum-sum-of-two-non-overlapping-subarrays/description/>
# hidden: true
# ---
#
# Given an array `A` of non-negative integers, return the maximum sum of
# elements in two non-overlapping (contiguous) subarrays, which have lengths `L`
# and `M`.  (For clarification, the `L`-length subarray could occur before or
# after the `M`-length subarray.)
#
# Formally, return the largest `V` for which `V = (A[i] + A[i+1] + ... +
# A[i+L-1]) + (A[j] + A[j+1] + ... + A[j+M-1])` and either:
#
#   * `0 <= i < i + L - 1 < j < j + M - 1 < A.length`, **or**
#   * `0 <= j < j + M - 1 < i < i + L - 1 < A.length`.
#
#
#
# **Example 1:**
#
#
#
#     Input: A = [0,6,5,2,2,5,1,9,4], L = 1, M = 2
#     Output: 20
#     Explanation: One choice of subarrays is [9] with length 1, and [6,5] with length 2.
#
#
# **Example 2:**
#
#
#
#     Input: A = [3,8,1,3,2,1,8,9,0], L = 3, M = 2
#     Output: 29
#     Explanation: One choice of subarrays is [3,8,1] with length 3, and [8,9] with length 2.
#
#
# **Example 3:**
#
#
#
#     Input: A = [2,1,5,6,0,9,5,0,3,8], L = 4, M = 3
#     Output: 31
#     Explanation: One choice of subarrays is [5,6,0,9] with length 4, and [3,8] with length 3.
#
#
#
#
# **Note:**
#
#   1. `L >= 1`
#   2. `M >= 1`
#   3. `L + M <= A.length <= 1000`
#   4. `0 <= A[i] <= 1000`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
# Return the binary header from a SEGY file
export read_fileheader

"""
# Info

Use: fileheader = read_fileheader(s::IO; bigendian::Bool = true)

Returns a binary file header formed using bytes 3200-3600 from the stream 's'.
"""
function read_fileheader(s::IO; bigendian::Bool = true)

    # Return to start of stream
    seekstart(s)

    # Read text header
    th = String(read(s, 3200))

    # Initialize binary file header
    bfh = BinaryFileHeader()
    fh_b2s = fh_byte2sample()

    if bigendian
        # Read all header values and put in fileheader
        for k in keys(fh_b2s)

           # Seek to this header value location
           seek(s, fh_b2s[k])

           # Read into file header
           sym = Symbol(k)
           setfield!(bfh, sym, bswap(read(s, typeof(getfield(bfh, sym)))))
        end
    else
        # Read all header values and put in fileheader
        for k in keys(fh_b2s)

           # Seek to this header value location
           seek(s, fh_b2s[k])

           # Read into file header
           sym = Symbol(k)
           setfield!(bfh, sym, read(s, typeof(getfield(bfh, sym))))
        end
    end

    # Move to end of file header
    seek(s, 3600)
    return FileHeader(th, bfh)
end

"""
Use: fileheader = read_fileheader(s::IO, keys::Array{String,1}; bigendian::Bool = true)

Return a fileheader from stream 's' with the fields defined in 'keys'.

# Examples

Read the entire file header.

    julia> s = open("data/testdata.segy")
    IO(<file data/testdata.segy>)

    julia> fh = SegyIO.read_fileheader(s)
    SegyIO.BinaryFileHeader(9999, 9999, 1, 400, 0, 4000, 4000, 560, 560, 1, -13922, 4, 1, 0,
    0, 0, 0, 0, 0, 0, 0, 2, 1, 4, 2, 0, 0, 0, 0, 0, Dict("expf"=>3226,"sfe"=>3234,
    "rgc"=>3250,"jobid"=>3200,"dt"=>3216,"nsfr"=>3222,"slen"=>3236,"vpol"=>3258,"renum"=>3208,
    "dsf"=>3224…))

Read only the sample interval and number of traces from the file header.

    julia> s = open("data/testdata.segy")
    IO(<file data/testdata.segy>)

    julia> fh = SegyIO.read_fileheader(s, ["dt"; "ns"])
    SegyIO.BinaryFileHeader(0, 0, 0, 0, 0, 4000, 0, 560, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Dict("expf"=>3226,"sfe"=>3234,"rgc"=>3250,"jobid"=>3200,
    "dt"=>3216,"nsfr"=>3222,"slen"=>3236,"vpol"=>3258,"renum"=>3208,"dsf"=>3224…))
"""
function read_fileheader(s::IO, keys::Array{String,1}; bigendian::Bool = true)

    # Return to start of stream
    seekstart(s)

    # Read text header
    th = String(read(s, 3200))

    # Initialize binary file header
    bfh = BinaryFileHeader()
    fh_b2s = fh_byte2sample()

    if bigendian
        # Read all header values and put in fileheader
        for k in keys

           # Seek to this header value location
           seek(s, fh_b2s[k])

           # Read into file header
           sym = Symbol(k)
           setfield!(bfh, sym, bswap(read(s, typeof(getfield(bfh, sym)))))
        end
    else
        # Read all header values and put in fileheader
        for k in keys

           # Seek to this header value location
           seek(s, fh_b2s[k])

           # Read into file header
           sym = Symbol(k)
           setfield!(bfh, sym, read(s, typeof(getfield(bfh, sym))))
        end
    end

    # Move to end of file header
    seek(s, 3600)

    return FileHeader(th, bfh)
end


text:
# by lazarusA # HIDE
# using GLMakie # HIDE
using GLMakie
GLMakie.activate!() # HIDE
let
    x = y =  LinRange(-2, 2, 51)
    z = (-x .* exp.(-x .^ 2 .- (y') .^ 2)) .* 4
    zmin, zmax = minimum(z), maximum(z)
    cmap = :viridis
    fig = Figure(resolution = (900,900))
    ax = Axis3(fig, aspect = :data, perspectiveness = 0.5, elevation = π/9,
        xzpanelcolor= (:black, 0.75), yzpanelcolor= (:black,0.75),
        zgridcolor = :grey, ygridcolor = :grey,xgridcolor = :grey)
    surface!(ax, x, y, z, colormap = cmap, colorrange = (zmin, zmax))
    xm, ym, zm = minimum(ax.finallimits[])
    contour!(ax, x, y, z, levels = 20, colormap = cmap, linewidth = 2,
        colorrange=(zmin, zmax), transformation = (:xy, zm))
    wireframe!(ax, x, y, z, overdraw = true, transparency = true,
        color = (:black, 0.1))
    fig[1,1] = ax
    fig
    save(joinpath(@__DIR__, "output", "surfWireContour.png"), fig, px_per_unit = 2.0) # HIDE
end
using Pkg # HIDE
Pkg.status(["GLMakie"]) # HIDE

max_stars_repo_name: MisterBiggs/BeautifulMakie

max_issues_repo_name: MisterBiggs/BeautifulMakie

max_forks_repo_name: MisterBiggs/BeautifulMakie


text:
include("utils.jl")

mutable struct CST{T,N,N1,N2} <: AbstractAirfoil{T}
    class::CSTClassFunction{N1,N2}
    shape::CSTShapeFunction{N,T}
    au::Vector{T}
    al::Vector{T}
    a0::T
    yu::T
    yl::T
end

function CST(::UndefInitializer, ::Type{T}, N::Int, N1=0.5, N2=1) where {T <: AbstractFloat}
    N1_ = convert(T, N1)
    N2_ = convert(T, N2)
    class = CSTClassFunction(N1_, N2_)
    shape = CSTShapeFunction(N, T)
    CST{T,N,N1_,N2_}(class, shape, Vector{T}(undef, N), Vector{T}(undef, N), 0, 0, 0)
end

function CST(a::AbstractVector{T}, N1=0.5, N2=1, T_TE=0, Y_TE=0) where {T <: AbstractFloat}
    N::Int, rem = divrem(length(a), 2)
    rem == 1 || error("CST params must be of length 2*N+1")
    N1_ = convert(T, N1)
    N2_ = convert(T, N2)
    class = CSTClassFunction(N1_, N2_)
    shape = CSTShapeFunction(N, T)
    yu = convert(T, Y_TE + T_TE / 2)
    yl = convert(T, Y_TE - T_TE / 2)
    a0 = a[1]
    au = a[2:N + 1]
    al = a[N + 2:2N + 1]
    CST{T,N,N1_,N2_}(class, shape, au, al, a0, yu, yl)
end

function set_param!(cst::CST{T,N,N1,N2}, a::AbstractVector{T}) where {T,N,N1,N2}
    length(a) == 2N + 1 || error("CST params must be of length 2*N+1")
    cst.a0 = a[1]
    cst.au = a[2:N + 1]
    cst.al = a[N + 2:2N + 1]
end

get_param(cst::CST{T}) where {T} = vcat(cst.a0, cst.au, cst.al)

function modify_te!(cst::CST{T}; T_TE=nothing, Y_TE=nothing) where T
    if T_TE === nothing
        T_TE = cst.yu - cst.yl
    end
    if Y_TE === nothing
        Y_TE = (cst.yu + cst.yl) / 2
    end
    cst.yu = convert(T, Y_TE + T_TE / 2)
    cst.yl = convert(T, Y_TE - T_TE / 2)
end

@inline shape_f(cst::CST, i, x) = shape_f(cst.shape, i, x)
@inline shape_d(cst::CST, i, x) = shape_d(cst.shape, i, x)
@inline class_f(cst::CST, x) = class_f(cst.class, x)
@inline class_d(cst::CST, x) = class_d(cst.class, x)

function Base.show(io::IO, cst::CST{T,N,N1,N2}) where {T,N,N1,N2}
    println(io, "CST{$T,$N,$N1,$N2}")
    println(io, "  Upper surface:", vcat(cst.a0, cst.au))
    println(io, "  Lower surface:", vcat(-cst.a0, cst.al))
    println(io, "  Trailing edge: yu=", cst.yu, " yl=", cst.yl)
end

"""
    fit(:CST, xu, yu, xl, yl; type=Float64, N=9, N1=0.5, N2=1.0)
"""
function fit(::Val{:CST}, xx, yy; T=Float64, N=9, N1=0.5, N2=1.0, kw...)
    if (!isempty(kw))
        @warn "Unused keyword arguments:" kw...
    end
    N1_ = convert(T, N1)
    N2_ = convert(T, N2)
    xu, yu, xl, yl = cst_preprocess(convert.(T, xx), convert.(T, yy))
    @assert issorted(xu) && issorted(xl)
    @assert xu[1] == xl[1] == yu[1] == yl[1] == 0
    @assert (xu[end] + xl[end])/2 ≈ 1
    class = CSTClassFunction(N1_, N2_)
    shape = CSTShapeFunction(N, T)
    au, al, a0, yu1, yl1 = cst_fit(class, shape, xu, yu, xl, yl)
    CST{T,N,N1,N2}(class, shape, au, al, a0, yu1, yl1)
end

"""
    fit!(cst, xu, yu, xl, yl)

"""
function fit!(cst::CST{T,N,N1,N2}, xx, yy) where {T,N,N1,N2}
    xu, yu, xl, yl = cst_preprocess(convert.(T, xx), convert.(T, yy))
    @assert issorted(xu) && issorted(xl)
    @assert xu[1] == xl[1] == yu[1] == yl[1] == 0
    @assert (xu[end] + xl[end])/2 ≈ 1
    cst.au, cst.al, cst.a0, cst.yu, cst.yl = cst_fit(cst.class, cst.shape, xu, yu, xl, yl)
    cst
end

function y_upper(cst::CST{T,N}, x::T)::T where {T <: AbstractFloat,N}
    y = cst.a0 * shape_f(cst, 0, x)
    for i = 1:N
        y += cst.au[i] * shape_f(cst, i, x)
    end
    y * class_f(cst, x) + cst.yu * x
end

function y_lower(cst::CST{T,N}, x::T)::T where {T <: AbstractFloat,N}
    y = -cst.a0 * shape_f(cst, 0, x)
    for i = 1:N
        y += cst.al[i] * shape_f(cst, i, x)
    end
    y * class_f(cst, x) + cst.yl * x
end

function dy_upper(cst::CST{T,N}, x::T)::T where {T <: AbstractFloat,N}
    dy = cst.a0 * shape_d(cst, 0, x)
    y = cst.a0 * shape_f(cst, 0, x)
    for i = 1:N
        dy += cst.au[i] * shape_d(cst, i, x)
        y += cst.au[i] * shape_f(cst, i, x)
    end
    y * class_d(cst, x) + dy * class_f(cst, x) + cst.yu
end

function dy_lower(cst::CST{T,N}, x::T)::T where {T <: AbstractFloat,N}
    dy = -cst.a0 * shape_d(cst, 0, x)
    y = -cst.a0 * shape_f(cst, 0, x)
    for i = 1:N
        dy += cst.al[i] * shape_d(cst, i, x)
        y += cst.al[i] * shape_f(cst, i, x)
    end
    y * class_d(cst, x) + dy * class_f(cst, x) + cst.yl
end

function fy_upper(cst::CST{T,N}, x::T)::Tuple{T,T} where {T <: AbstractFloat,N}
    dy = cst.a0 * shape_d(cst, 0, x)
    y = cst.a0 * shape_f(cst, 0, x)
    for i = 1:N
        dy += cst.au[i] * shape_d(cst, i, x)
        y += cst.au[i] * shape_f(cst, i, x)
    end
    C = class_f(cst, x)
    dC = class_d(cst, x)
    y * C + cst.yu * x, y * dC + dy * C + cst.yu
end

function fy_lower(cst::CST{T,N}, x::T)::Tuple{T,T} where {T <: AbstractFloat,N}
    dy = -cst.a0 * shape_d(cst, 0, x)
    y = -cst.a0 * shape_f(cst, 0, x)
    for i = 1:N
        dy += cst.al[i] * shape_d(cst, i, x)
        y += cst.al[i] * shape_f(cst, i, x)
    end
    C = class_f(cst, x)
    dC = class_d(cst, x)
    y * C + cst.yl * x, y * dC + dy * C + cst.yl
end


text:
##
## COMPUTE FUNDFAIR16 SCGHGS
##

include(joinpath(@__DIR__, "..", "src", "helper.jl"))
include(joinpath(@__DIR__, "marginaldamages_MimiFAIR162.jl"))
include(joinpath(@__DIR__, "../src/MimiFUND-MimiFAIR162-Flat.jl"))

scenarios = ["ssp245", "ssp370"]
years =[2020, 2030, 2040]
rates = [0.02, 0.03]
gases = [:CO2, :CH4, :N2O]

scghg_dict = Dict()

for scen in scenarios
    m = get_fundfair162(ar6_scenario = scen)
    run(m)
    for gas in gases
        for year in years
            for rate in rates
                mm = create_marginal_fundfair162_model(m; year = year, gas = gas)
                scghg_dict[gas, year, rate, scen] = compute_fundfair162_sc(mm, year = year, rate = rate)
            end
        end
    end
end

scghg_dict

# ## save
# using CSV # sorry david
# path = "C:/Users/TTAN/Environmental Protection Agency (EPA)/NCEE Social Cost of Carbon - General/models/Notes/Code/output/fair1.6"
# CSV.write(joinpath(path,"fund39_fair16_scghgs.csv"), scghg_dict)


text:
module ExtendedExtremes

using CSV, Distributions, Extremes, Optim
using Gadfly
using TimeSeries
import Distributions: cdf, insupport, logcdf, logpdf, params, partype, pdf, quantile, sampler
import Extremes: qqplot, probplot, histplot, returnlevelplot, diagnosticplots
import Random: AbstractRNG
import Base: maximum, minimum, rand

include("parameterestimation.jl")
include("plots.jl")
include("utils.jl")

# distributions
include("distributions/naveau2016_type1.jl");
include("distributions/naveau2016_type1_nonstat.jl");
include("distributions/naveau2016_type2.jl");
include("distributions/naveau2016_type3.jl");
include("distributions/naveau2016_type4.jl");
include("distributions/EGPnormal.jl");
include("distributions/kallache2011.jl");

# bias correction
include("biascorrect.jl");
include("valid_indices.jl")

export
    # distribution types
    EGPpower,
    EGPpowermix,
    EGPbeta,
    EGPbetapower,
    EGPnormal,

    # methods
    cdf,         # cumulative distribution function
    insupport,   # predicate, is x in the support of the distribution?
    logpdf,      # log probability density
    logcdf,      # cdf returning log-probability
    maximum,
    minimum,
    params,      # get the tuple of parameters
    partype,
    pdf,         # probability density function
    quantile,    # inverse of cdf (defined for p in (0,1))
    rand,
    sampler,      # create a Sampler object for efficient samples

    # distribution fitting
    EGPpowerfit,
    EGPpowermixfit,
    EGPbetafit,
    EGPbetapowerfit,
    EGPnormalfit,

    # validation
    BIC,
    qqbuild,
    qqplot,
    ppbuild,
    probplot,
    histplot,
    returnlevelplot,
    diagnosticplots,
    Guide,

    # bias correction
    T_transform,
    bias_correct,
    CDFT,

    # validation indices
    index_skew,
    index_Rx,
    index_R10,
    index_pXwet,
    index_p98wet,
    index_pXwetamount,
    index_p98wetamount,
    index_RVxmax,
    index_RV20max,
    index_SDII,
    index_p90wet

end # module


text:
abstract type AbstractDiffSys end

abstract type SpecialDiffSys end

function Base.show(io::IO, ::MIME"text/plain", sys::SpecialDiffSys)
    println(io, typeof(sys))
    for fname in fieldnames(typeof(sys))
        fval = getfield(sys, fname)
        println(io, "    $fname = $fval")
    end
end

mutable struct DiffSys <: AbstractDiffSys
    n_dim::Int64
    dt::Float64
    t_max::Float64
    func::Function
    noise::AbstractArray
    name::String
end

function Base.show(io::IO, ::MIME"text/plain", sys::DiffSys)
    println(io, typeof(sys))
    println(io, "    Differential equation name: ", sys.name)
    println(io, "    Differential equation dimension: ", sys.n_dim)
    println(io, "    time step: ", sys.dt)
    println(io, "    Calculation time: ", sys.t_max)
    println(io, "    noise: ", sys.noise)
end

# DiffSys(sys::SpecialDiffSys; dt=0.01, t_max=200) = error("unimplemented")
function DiffSys(sys::SpecialDiffSys; dt=0.01, t_max=200)
    n_dim = dimension(sys)
    func = x -> vec_field(sys, x)
    noise = sys.noise
    nm = name(sys)
    return DiffSys(n_dim, dt, t_max, func, noise, nm)
end
DiffSys(sys::DiffSys) = sys

function EulerMaruyama(func::Function, noise::AbstractArray, dt::Float64, x0::AbstractArray)
    return x0 + dt*func(x0) + √(dt)*noise.*randn(length(noise))
end

function RungeKutta(func::Function, dt::Float64, x0::AbstractArray)
    k1 = func(x0)
    k2 = func(x0+0.5dt*k1)
    k3 = func(x0+0.5dt*k2)
    k4 = func(x0+dt*k3)
    dx = dt*(k1+2k2+2k3+k4)/6
    return x0+dx
end

EulerMaruyama(sys::DiffSys, x0::AbstractArray) = EulerMaruyama(sys.func, sys.noise, sys.dt, x0)

function integrate(sys::DiffSys, x0::AbstractArray)
    ts = [0:sys.dt:sys.t_max;]
    n_step = length(ts)
    n_dim = sys.n_dim
    orbits = zeros(n_step, n_dim)
    x = copy(x0)
    orbits[1, :] = x
    for i in 2:n_step
        x = EulerMaruyama(sys, x)
        orbits[i, :] = x
    end
    return ts, orbits
end

integrate(sys::SpecialDiffSys, x0::AbstractArray; dt=0.01, t_max=200) = integrate(DiffSys(sys, dt=dt, t_max=t_max), x0)

mutable struct CoupledDiffSys <: AbstractDiffSys
    n_body::Int64
    dt::Float64
    t_max::Float64
    SelfDiffSyss::Array{DiffSys}
    OriginalSpecialDiffSyss::Array{SpecialDiffSys}
    CouplingFuncDict::Dict
end

function Base.show(io::IO, ::MIME"text/plain", sys::CoupledDiffSys)
    println(io, typeof(sys))
    println(io, "    Differential equation number of body: ", sys.n_body)
    total_dim = sum([s.n_dim for s in sys.SelfDiffSyss])
    println(io, "    Differential equation dimension: ", total_dim)
    println(io, "    time step: ", sys.dt)
    println(io, "    Calculation time: ", sys.t_max)
    println(io, "    Name of each systems")
    for i in 1:sys.n_body
        println(io, "        Number $i: $(sys.SelfDiffSyss[i].name)")
    end
    couplings = keys(sys.CouplingFuncDict)
    println(io, "    Coupled through")
    for (i,j) in couplings
        println(io, "        $i -> $j")
    end
end

function CoupledDiffSys(OriginalSpecialDiffSyss::Array{SpecialDiffSys}, CouplingFuncDict::Dict; dt=0.01, t_max=200)
    n_body = length(OriginalSpecialDiffSyss)
    SelfDiffSyss = [DiffSys(sys) for sys in OriginalSpecialDiffSyss]
    return CoupledDiffSys(n_body, dt, t_max, SelfDiffSyss, OriginalSpecialDiffSyss, CouplingFuncDict)
end

function ith_coupling_idxs(CouplingFuncDict::Dict, i::Int64)
    idxs = [couplings[2] for couplings in collect(keys(CouplingFuncDict)) if couplings[1]==i]
    return idxs
end

function create_vec(sys::CoupledDiffSys)
    dims = [s.n_dim for s in sys.SelfDiffSyss]
    cumdims = cumsum(dims)
    self_vecs = [s.func for s in sys.SelfDiffSyss]
    func = x -> begin
        xs = [x[1:cumdims[1]]]
        for i in 2:sys.n_body
            push!(xs, x[cumdims[i-1]+1:cumdims[i]])
        end
        vs = []
        for i in 1:sys.n_body
            v = self_vecs[i](xs[i])
            coupling_idxs = ith_coupling_idxs(sys.CouplingFuncDict, i)
            for j in coupling_idxs
                v += sys.CouplingFuncDict[(i,j)](xs[i], xs[j])
            end
            push!(vs, v)
        end
        return vcat(vs...)
    end
    return func
end

function DiffSys(sys::CoupledDiffSys; dt=0.01, t_max=200)
    dims = [s.n_dim for s in sys.SelfDiffSyss]
    n_dim = sum(dims)
    func = create_vec(sys)
    noises = vcat([s.noise for s in sys.SelfDiffSyss]...)
    names = ""
    return DiffSys(n_dim, dt, t_max, func, noises, names)
end

integrate(sys::CoupledDiffSys, x0::AbstractArray; dt=0.01, t_max=200) = integrate(DiffSys(sys, dt=dt, t_max=t_max), x0)