text:
module LazyLinSpace

    export linspace
    import Base: map, float, big, show, isempty, step, length, first, last, start, next, done,
        getindex, unsafe_getindex, promote_rule, convert, reverse, eltype
    import Base.Operators: +, -, .+, .-, *, .*, ==, ./

    import ..AbstractFloat
    import ..@compat

    ## linspace and logspace

    immutable LinSpace{T<:AbstractFloat} <: Range{T}
        start::T
        stop::T
        len::T
        divisor::T
    end

    # float rationalization helper
    function rat(x)
        y = x
        a = d = 1
        b = c = 0
        m = maxintfloat(Float32)
        while abs(y) <= m
            f = trunc(Int,y)
            y -= f
            a, c = f*a + c, a
            b, d = f*b + d, b
            max(abs(a),abs(b)) <= convert(Int,m) || return c, d
            oftype(x,a)/oftype(x,b) == x && break
            y = inv(y)
        end
        return a, b
    end

    function linspace{T<:AbstractFloat}(start::T, stop::T, len::T)
        len == round(len) || throw(InexactError())
        0 <= len || error("linspace($start, $stop, $len): negative length")
        if len == 0
            n = convert(T, 2)
            if isinf(n*start) || isinf(n*stop)
                start /= n; stop /= n; n = one(T)
            end
            return LinSpace(-start, -stop, -one(T), n)
        end
        if len == 1
            start == stop || error("linspace($start, $stop, $len): endpoints differ")
            return LinSpace(-start, -start, zero(T), one(T))
        end
        n = convert(T, len - 1)
        len - n == 1 || error("linspace($start, $stop, $len): too long for $T")
        a0, b = rat(start)
        a = convert(T,a0)
        if a/convert(T,b) == start
            c0, d = rat(stop)
            c = convert(T,c0)
            if c/convert(T,d) == stop
                e = lcm(b,d)
                a *= div(e,b)
                c *= div(e,d)
                s = convert(T,n*e)
                if isinf(a*n) || isinf(c*n)
                    s, p = frexp(s)
                    p2 = oftype(s,2)^p
                    a /= p2; c /= p2
                end
                if a*n/s == start && c*n/s == stop
                    return LinSpace(a, c, len, s)
                end
            end
        end
        a, c, s = start, stop, n
        if isinf(a*n) || isinf(c*n)
            s, p = frexp(s)
            p2 = oftype(s,2)^p
            a /= p2; c /= p2
        end
        if a*n/s == start && c*n/s == stop
            return LinSpace(a, c, len, s)
        end
        return LinSpace(start, stop, len, n)
    end
    function linspace{T<:AbstractFloat}(start::T, stop::T, len::Real)
        T_len = convert(T, len)
        T_len == len || throw(InexactError())
        linspace(start, stop, T_len)
    end
    linspace(start::Real, stop::Real, len::Real=50) =
        linspace(promote(AbstractFloat(start), AbstractFloat(stop))..., len)

    function show(io::IO, r::LinSpace)
        print(io, "linspace(")
        show(io, first(r))
        print(io, ',')
        show(io, last(r))
        print(io, ',')
        show(io, length(r))
        print(io, ')')
    end

    logspace(start::Real, stop::Real, n::Integer=50) = 10.^linspace(start, stop, n)

    eltype{T}(r::LinSpace{T}) = T
    isempty(r::LinSpace) = length(r) == 0
    step{T}(r::LinSpace{T}) = ifelse(r.len <= 0, convert(T,NaN), (r.stop-r.start)/r.divisor)
    length(r::LinSpace) = convert(Integer,r.len + signbit(r.len - 1))
    first{T}(r::LinSpace{T}) = convert(T, (r.len-1)*r.start/r.divisor)
    last{T}(r::LinSpace{T}) = convert(T, (r.len-1)*r.stop/r.divisor)
    start(r::LinSpace) = 1
    done(r::LinSpace, i::Int) = length(r) < i
    next{T}(r::LinSpace{T}, i::Int) =
        (convert(T, ((r.len-i)*r.start + (i-1)*r.stop)/r.divisor), i+1)

    getindex{T}(r::LinSpace{T}, i::Int) = (checkbounds(r, i); unsafe_getindex(r, i))
    getindex{T}(r::LinSpace{T}, i::Integer) = (checkbounds(r, i); unsafe_getindex(r, convert(Int,i)))
    unsafe_getindex{T}(r::LinSpace{T}, i::Int) = convert(T, ((r.len-i)*r.start + (i-1)*r.stop)/r.divisor)

    getindex(r::LinSpace, s::OrdinalRange) = (checkbounds(r, s); unsafe_getindex(r, s))
    function unsafe_getindex{T}(r::LinSpace{T}, s::OrdinalRange)
        sl::T = length(s)
        ifirst = first(s)
        ilast = last(s)
        vfirst::T = ((r.len - ifirst) * r.start + (ifirst - 1) * r.stop) / r.divisor
        vlast::T = ((r.len - ilast) * r.start + (ilast - 1) * r.stop) / r.divisor
        return linspace(vfirst, vlast, sl)
    end

    function map{T<:AbstractFloat}(::Type{T}, r::LinSpace)
        new_len = T(r.len)
        new_len == r.len || error("$r: too long for $T")
        LinSpace(T(r.start), T(r.stop), new_len, T(r.divisor))
    end

    for fn in (:float,:big)
        @eval function $fn(r::LinSpace)
            new_len = $fn(r.len)
            new_len == r.len || error(string(r, ": too long for ", $fn))
            LinSpace($fn(r.start), $fn(r.stop), new_len, $fn(r.divisor))
        end
    end

    for fn in (:+, :-)
        @eval function $fn{T<:AbstractFloat}(r1::LinSpace{T}, r2::LinSpace{T})
                len = r1.len
                (len == r2.len ||
                 throw(DimensionMismatch("argument dimensions must match")))
                divisor1, divisor2 = r1.divisor, r2.divisor
                if divisor1 == divisor2
                    LinSpace{T}($fn(r1.start, r2.start), $fn(r1.stop, r2.stop),
                                len, divisor1)
                else
                    linspace(convert(T, $fn(first(r1), first(r2))),
                             convert(T, $fn(last(r1), last(r2))), len)
                end
            end

        @eval $fn(r1::(@compat Union{FloatRange, OrdinalRange, LinSpace}),
               r2::(@compat Union{FloatRange, OrdinalRange, LinSpace})) =
                   $fn(promote(r1, r2)...)
    end

    (==){T<:LinSpace}(r::T, s::T) = (first(r) == first(s)) & (length(r) == length(s)) & (last(r) == last(s))
    -(r::LinSpace)     = LinSpace(-r.start, -r.stop, r.len, r.divisor)

    function .+{T}(x::Real, r::LinSpace{T})
        x2 = x * r.divisor / (r.len - 1)
        LinSpace(x2 + r.start, x2 + r.stop, r.len, r.divisor)
    end

    function .-(x::Real, r::LinSpace)
        x2 = x * r.divisor / (r.len - 1)
        LinSpace(x2 - r.start, x2 - r.stop, r.len, r.divisor)
    end

    function .-(r::LinSpace, x::Real)
        x2 = x * r.divisor / (r.len - 1)
        LinSpace(r.start - x2, r.stop - x2, r.len, r.divisor)
    end

    .*(x::Real, r::LinSpace)     = LinSpace(x * r.start, x * r.stop, r.len, r.divisor)
    .*(r::LinSpace, x::Real)     = x .* r
    ./(r::LinSpace, x::Real)     = LinSpace(r.start / x, r.stop / x, r.len, r.divisor)

    promote_rule{T1,T2}(::Type{LinSpace{T1}},::Type{LinSpace{T2}}) =
        LinSpace{promote_type(T1,T2)}
    convert{T}(::Type{LinSpace{T}}, r::LinSpace{T}) = r
    convert{T}(::Type{LinSpace{T}}, r::LinSpace) =
        LinSpace{T}(r.start, r.stop, r.len, r.divisor)

    promote_rule{F,OR<:OrdinalRange}(::Type{LinSpace{F}}, ::Type{OR}) =
        LinSpace{promote_type(F,eltype(OR))}
    convert{T}(::Type{LinSpace{T}}, r::OrdinalRange) =
        linspace(convert(T, first(r)), convert(T, last(r)), convert(T, length(r)))
    convert{T}(::Type{LinSpace}, r::OrdinalRange{T}) =
        convert(LinSpace{typeof(float(first(r)))}, r)

    # Promote FloatRange to LinSpace
    promote_rule{F,OR<:FloatRange}(::Type{LinSpace{F}}, ::Type{OR}) =
        LinSpace{promote_type(F,eltype(OR))}
    convert{T}(::Type{LinSpace{T}}, r::FloatRange) =
        linspace(convert(T, first(r)), convert(T, last(r)), convert(T, length(r)))
    convert{T}(::Type{LinSpace}, r::FloatRange{T}) =
        convert(LinSpace{T}, r)

    reverse(r::LinSpace)     = LinSpace(r.stop, r.start, r.len, r.divisor)

end


text:
# elastic-net regularization

export ElasticNet

"""
**Elastic-net regularization**

    ElasticNet(μ=1.0, λ=1.0)

Returns the function
```math
f(x) = μ\\|x\\|_1 + (λ/2)\\|x\\|^2,
```
for nonnegative parameters `μ` and `λ`.
"""
struct ElasticNet{R <: Real} <: ProximableFunction
    mu::R
    lambda::R
    function ElasticNet{R}(mu::R, lambda::R) where {R <: Real}
        if lambda < 0 || mu < 0
            error("parameters `μ` and `λ` must be nonnegative")
        else
            new(mu, lambda)
        end
    end
end

is_separable(f::ElasticNet) = true
is_prox_accurate(f::ElasticNet) = true
is_convex(f::ElasticNet) = true

ElasticNet(mu::R=1.0, lambda::R=1.0) where {R <: Real} = ElasticNet{R}(mu, lambda)

function (f::ElasticNet)(x::AbstractArray{R}) where R <: RealOrComplex
    return f.mu*norm(x,1) + (f.lambda/2)*norm(x,2)^2
end

function prox!(y::AbstractArray{R}, f::ElasticNet{R}, x::AbstractArray{R}, gamma::R=R(1)) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    gm = gamma*f.mu
    gl = gamma*f.lambda
    for i in eachindex(x)
        y[i] = (x[i] + (x[i] <= -gm ? gm : (x[i] >= gm ? -gm : -x[i])))/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function prox!(y::AbstractArray{R}, f::ElasticNet{R}, x::AbstractArray{R}, gamma::AbstractArray{R}) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    for i in eachindex(x)
        gm = gamma[i]*f.mu
        gl = gamma[i]*f.lambda
        y[i] = (x[i] + (x[i] <= -gm ? gm : (x[i] >= gm ? -gm : -x[i])))/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function prox!(y::AbstractArray{Complex{R}}, f::ElasticNet{R}, x::AbstractArray{Complex{R}}, gamma::R=R(1)) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    gm = gamma*f.mu
    gl = gamma*f.lambda
    for i in eachindex(x)
        y[i] = sign(x[i])*max(0, abs(x[i]) - gm)/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function prox!(y::AbstractArray{Complex{R}}, f::ElasticNet{R}, x::AbstractArray{Complex{R}}, gamma::AbstractArray{R}) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    for i in eachindex(x)
        gm = gamma[i]*f.mu
        gl = gamma[i]*f.lambda
        y[i] = sign(x[i])*max(0, abs(x[i]) - gm)/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function gradient!(y::AbstractArray{T}, f::ElasticNet{R}, x::AbstractArray{T}) where {T <: RealOrComplex, R <: Real}
    # Gradient of 1 norm
    y .= f.mu.*sign.(x)
    # Gradient of 2 norm
    y .+= f.lambda.*x
    return f.mu*norm(x,1) + (f.lambda/2)*norm(x,2)^2
end

fun_name(f::ElasticNet) = "elastic-net regularization"
fun_dom(f::ElasticNet) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::ElasticNet) = "x ↦ μ||x||_1 + (λ/2)||x||²"
fun_params(f::ElasticNet) = "μ = $(f.mu), λ = $(f.lambda)"

function prox_naive(f::ElasticNet, x::AbstractArray{R}, gamma::Real=1.0) where R <: RealOrComplex
    uz = max.(0, abs.(x) .- gamma*f.mu)/(1 + f.lambda * gamma)
    return sign.(x) .* uz, f.mu * norm(uz,1) + (f.lambda/2) * norm(uz)^2
end

function prox_naive(f::ElasticNet, x::AbstractArray{R}, gamma::AbstractArray) where R <: RealOrComplex
    uz = max.(0, abs.(x) .- gamma.*f.mu)./(1 .+ f.lambda .* gamma)
    return sign.(x) .* uz, f.mu * norm(uz,1) + (f.lambda/2) * norm(uz)^2
end

ext: jl


text:
using YaoBase, YaoBlocks, YaoArrayRegister
using BitBasis: BitStr
# NOTE: this file should only exists in v0.4.x
#       since in v0.4.x the conponent packages
#       are all new packages and should not have
#       deprecation warning itself.

# Base
@deprecate addbit!(r::AbstractRegister, n::Int) addbits!(r, n)
@deprecate addbit!(n::Int) addbits!(n)
@deprecate reset!(r::AbstractRegister; val::Integer=0)  collapseto!(r, val)
@deprecate measure_reset!(r::AbstractRegister; val::Int=0) measure_collapseto!(r; config=val)
@deprecate measure_reset!(r::AbstractRegister, locs; val::Int=0) measure_collapseto!(r, locs; config=val)

# ArrayReg
@deprecate register(raw::AbstractMatrix; B=size(raw,2)) ArrayReg{B}(raw)
@deprecate register(raw::AbstractVector) ArrayReg(raw)
@deprecate register(bits::BitStr, nbatch::Int=1) ArrayReg{nbatch}(bits)
@deprecate register(::Type{T}, bits::BitStr, nbatch::Int) where T ArrayReg{nbatch}(T, bits)

@deprecate usedbits(block::AbstractBlock{N}) where N occupied_locs(block)

################ Compatibility Code ###################
export DefaultRegister, MatrixBlock, Sequential, ReflectBlock, GeneralMatrixGate, AddBlock, ⊗
const DefaultRegister = ArrayReg
const MatrixBlock = AbstractBlock
const Sequential = ChainBlock
const ReflectBlock = ReflectGate
const GeneralMatrixGate = GeneralMatrixBlock
const AddBlock = YaoBlocks.Sum

@deprecate sequence(args...) chain(args...)
@deprecate matrixgate(args...) matblock(args...)

import Base: collect, parent
@deprecate collect(t::Type, x::AbstractBlock) collect_blocks(t, x)

# joining two registers
⊗(reg::AbstractRegister, reg2::AbstractRegister) = join(reg, reg2)
⊗(A::AbstractArray, B::AbstractArray) = kron(A, B)

# Originally addrs means return the addression stored in a composite block
# this is not necessary since not all composite block store the address, e.g
# roller, etc.
# @deprecate addrs(block::AbstractBlock) occupied_locs(block)

# NOTE: block is frequently used as variable name, to make sure there is
#       no conflicts, this is commented.
@deprecate block(x::AbstractContainer) content(x)
@deprecate parent(x::AbstractContainer) content(x)
@deprecate chblock(x::AbstractContainer, blk::AbstractBlock) chsubblocks(x, blk)

@deprecate timeevolve(block::MatrixBlock, t::Number; tol::Real=1e-7) time_evolve(block, t; tol=tol)
@deprecate timeevolve(t::Number; tol::Real=1e-7) time_evolve(t; tol=tol)

import Base: collect, *
@deprecate collect(x::AbstractBlock, ::Type{T}) where {T <: AbstractBlock} collect_blocks(T, x)

function Base.:*(op::AbstractMatrix, r::ArrayRegOrAdjointArrayReg)
    Base.depwarn("`matrix * register` is deprecated", :*)
    return op * state(r)
end


text:
# Autogenerated wrapper script for LMDB_jll for x86_64-w64-mingw32
export liblmdb, mdb_copy, mdb_dump, mdb_load, mdb_stat

JLLWrappers.@generate_wrapper_header("LMDB")
JLLWrappers.@declare_library_product(liblmdb, "liblmdb.dll")
JLLWrappers.@declare_executable_product(mdb_copy)
JLLWrappers.@declare_executable_product(mdb_dump)
JLLWrappers.@declare_executable_product(mdb_load)
JLLWrappers.@declare_executable_product(mdb_stat)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        liblmdb,
        "bin\\liblmdb.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_executable_product(
        mdb_copy,
        "bin\\mdb_copy.exe",
    )

    JLLWrappers.@init_executable_product(
        mdb_dump,
        "bin\\mdb_dump.exe",
    )

    JLLWrappers.@init_executable_product(
        mdb_load,
        "bin\\mdb_load.exe",
    )

    JLLWrappers.@init_executable_product(
        mdb_stat,
        "bin\\mdb_stat.exe",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
"""
Type: NLPStopping (specialization of GenericStopping)
Methods: start!, stop!, update_and_start!, update_and_stop!, fill_in!, reinit!, status

Stopping structure for non-linear programming problems using NLPModels.
    Input :
       - pb         : an AbstractNLPModel
       - optimality_check : a stopping criterion through an admissibility function
       - state      : The information relative to the problem, see GenericState
       - max_cntrs  : Dict contains the max number of evaluations
       - (opt) meta : Metadata relative to stopping criterion.
       - (opt) main_stp : Stopping of the main loop in case we consider a Stopping
                          of a subproblem.
                          If not a subproblem, then nothing.

 Note:
 * optimality_check : takes two inputs (AbstractNLPModel, NLPAtX)
 and returns a Float64 to be compared at 0.
 * designed for NLPAtX State. Constructor checks that the State has the
 required entries.

 Warning:
 * optimality_check does not necessarily fill in the State.
 """
mutable struct NLPStopping <: AbstractStopping

    # problem
    pb :: AbstractNLPModel

    # stopping criterion
    optimality_check :: Function

    # Common parameters
    meta      :: AbstractStoppingMeta
    # Parameters specific to the NLPModels
    max_cntrs :: Dict #contains the max number of evaluations

    # current state of the problem
    current_state :: AbstractState

    # Stopping of the main problem, or nothing
    main_stp :: Union{AbstractStopping, Nothing}

    function NLPStopping(pb             :: AbstractNLPModel,
                         admissible     :: Function,
                         current_state  :: AbstractState;
                         meta           :: AbstractStoppingMeta = StoppingMeta(),
                         max_cntrs      :: Dict = _init_max_counters(),
                         main_stp       :: Union{AbstractStopping, Nothing} = nothing,
                         kwargs...)

        if !(isempty(kwargs))
           meta = StoppingMeta(;kwargs...)
        end

        #current_state is an AbstractState with requirements
        try
            current_state.evals
            current_state.fx, current_state.gx, current_state.Hx
            #if there are bounds:
            current_state.mu
            if pb.meta.ncon > 0 #if there are constraints
               current_state.Jx, current_state.cx, current_state.lambda
            end
        catch
            throw("error: missing entries in the given current_state")
        end

        return new(pb, admissible, meta, max_cntrs, current_state, main_stp)
    end

end

"""
NLPStopping(pb): additional default constructor
The function creates a Stopping where the State is by default and the
optimality function is the function KKT().

key arguments are forwarded to the classical constructor.
"""
function NLPStopping(pb :: AbstractNLPModel; kwargs...)
 #Create a default NLPAtX
 nlp_at_x = NLPAtX(pb.meta.x0)
 admissible = KKT

 return NLPStopping(pb, admissible, nlp_at_x; kwargs...)
end

"""
_init_max_counters(): initialize the maximum number of evaluations on each of
                        the functions present in the Counters (NLPModels).
"""
function _init_max_counters(; obj    :: Int64 = 20000,
                              grad   :: Int64 = 20000,
                              cons   :: Int64 = 20000,
                              jcon   :: Int64 = 20000,
                              jgrad  :: Int64 = 20000,
                              jac    :: Int64 = 20000,
                              jprod  :: Int64 = 20000,
                              jtprod :: Int64 = 20000,
                              hess   :: Int64 = 20000,
                              hprod  :: Int64 = 20000,
                              jhprod :: Int64 = 20000,
                              sum    :: Int64 = 20000*11)

  cntrs = Dict([(:neval_obj,       obj), (:neval_grad,   grad),
                (:neval_cons,     cons), (:neval_jcon,   jcon),
                (:neval_jgrad,   jgrad), (:neval_jac,    jac),
                (:neval_jprod,   jprod), (:neval_jtprod, jtprod),
                (:neval_hess,     hess), (:neval_hprod,  hprod),
                (:neval_jhprod, jhprod), (:neval_sum,    sum)])

 return cntrs
end

"""
_init_max_counters_NLS(): initialize the maximum number of evaluations on each of
                          the functions present in the NLSCounters (NLPModels).
https://github.com/JuliaSmoothOptimizers/NLPModels.jl/blob/master/src/NLSModels.jl
"""
function _init_max_counters_NLS(; residual        :: Int = 20000,
                                  jac_residual    :: Int = 20000,
                                  jprod_residual  :: Int = 20000,
                                  jtprod_residual :: Int = 20000,
                                  hess_residual   :: Int = 20000,
                                  jhess_residual  :: Int = 20000,
                                  hprod_residual :: Int = 20000,
                                  kwargs...)

  cntrs_nlp = _init_max_counters(;kwargs...)
  cntrs = Dict([(:neval_residual, residual),
                (:neval_jac_residual, jac_residual),
                (:neval_jprod_residual, jprod_residual),
                (:neval_jtprod_residual, jtprod_residual),
                (:neval_hess_residual, hess_residual),
                (:neval_jhess_residual, jhess_residual),
                (:neval_hprod_residual, hprod_residual)])

 return merge(cntrs_nlp, cntrs)
end

"""
fill_in!: a function that fill in the required values in the State
"""
function fill_in!(stp  :: NLPStopping,
                  x    :: Iterate;
                  fx   :: Iterate     = nothing,
                  gx   :: Iterate     = nothing,
                  Hx   :: Iterate     = nothing,
                  cx   :: Iterate     = nothing,
                  Jx   :: Iterate     = nothing,
                  lambda :: Iterate   = nothing,
                  mu     :: Iterate   = nothing,
                  matrix_info :: Bool = true,
                  kwargs...)

 gfx = fx == nothing  ? obj(stp.pb, x)   : fx
 ggx = gx == nothing  ? grad(stp.pb, x)  : gx

 if Hx == nothing && matrix_info
   gHx = hess(stp.pb, x)
 else
   gHx = Hx
 end

 if stp.pb.meta.ncon > 0
     gJx = Jx == nothing ? jac(stp.pb, x)  : Jx
     gcx = cx == nothing ? cons(stp.pb, x) : cx
 else
     gJx = stp.current_state.Jx
     gcx = stp.current_state.cx
 end

 #update the Lagrange multiplier if one of the 2 is asked
 if (stp.pb.meta.ncon > 0 || has_bounds(stp.pb)) && (lambda == nothing || mu == nothing)
  lb, lc = _compute_mutliplier(stp.pb, x, ggx, gcx, gJx; kwargs...)
 elseif  stp.pb.meta.ncon == 0 && !has_bounds(stp.pb) && lambda == nothing
  lb, lc = mu, stp.current_state.lambda
 else
  lb, lc = mu, lambda
 end

 return update!(stp.current_state, x=x, fx = gfx,    gx = ggx, Hx = gHx,
                                        cx = gcx,    Jx = gJx, mu = lb,
                                        lambda = lc)
end

"""
_resources_check!: check if the optimization algorithm has exhausted the resources.
                   This is the NLP specialized version that takes into account
                   the evaluation of the functions following the sum_counters
                   structure from NLPModels.

Note:
* function uses counters in stp.pb, and update the counters in the state.
* function is compatible with Counters, NLSCounters, and any type whose entries
match the entries in stp.max_cntrs.
* all the problems have an entry "pb.counters" and a function "sum_counters(pb)"
"""
function _resources_check!(stp    :: NLPStopping,
                           x      :: Iterate)

  cntrs = stp.pb.counters
  update!(stp.current_state, evals = cntrs)

  max_cntrs = stp.max_cntrs

  # check all the entries in the counter
  max_f = false
  if typeof(stp.pb.counters) == Counters
   for f in fieldnames(Counters)
       max_f = max_f || (getfield(cntrs, f) > max_cntrs[f])
   end
  elseif typeof(stp.pb.counters) == NLSCounters
    for f in fieldnames(NLSCounters)
     max_f = f != :counters ? (max_f || (getfield(cntrs, f) > max_cntrs[f])) : max_f
    end
    for f in fieldnames(Counters)
        max_f = max_f || (getfield(cntrs.counters, f) > max_cntrs[f])
    end
  else #Unknown counters type
   for f in fieldnames(typeof(stp.pb.counters))
    max_f = max_f || (getfield(cntrs, f) > max_cntrs[f])
   end
  end

 # Maximum number of function and derivative(s) computation
 max_evals = sum_counters(stp.pb) > max_cntrs[:neval_sum]

 # global user limit diagnostic
 stp.meta.resources = max_evals || max_f

 return stp
end

"""
_unbounded_problem_check!: This is the NLP specialized version that takes into account
                   that the problem might be unbounded if the objective or the
                   constraint function are unbounded.

Note: * evaluate the objective function if state.fx is void.
      * evaluate the constraint function if state.cx is void.
"""
function _unbounded_problem_check!(stp  :: NLPStopping,
                                   x    :: Iterate)

 if stp.current_state.fx == nothing
	 stp.current_state.fx = obj(stp.pb, x)
 end
 f_too_large = norm(stp.current_state.fx) >= stp.meta.unbounded_threshold

 c_too_large = false
 if stp.pb.meta.ncon != 0 #if the problems has constraints, check |c(x)|
  if stp.current_state.cx == nothing
   stp.current_state.cx = cons(stp.pb, x)
  end
  c_too_large = norm(stp.current_state.cx) >= abs(stp.meta.unbounded_threshold)
 end

 stp.meta.unbounded_pb = f_too_large || c_too_large

 return stp
end

"""
_optimality_check: compute the optimality score.

This is the NLP specialized version that takes into account the structure of the
NLPStopping where the optimality_check function is an input.
"""
function _optimality_check(stp  :: NLPStopping; kwargs...)

 optimality = stp.optimality_check(stp.pb, stp.current_state; kwargs...)
 stp.current_state.current_score = optimality

 return optimality
end

################################################################################
# Nonlinear problems admissibility functions
# Available: unconstrained_check(...), optim_check_bounded(...), KKT
################################################################################
include("nlp_admissible_functions.jl")

################################################################################
# Functions computing Lagrange multipliers of a nonlinear problem
# Available: _compute_mutliplier(...)
################################################################################
include("nlp_compute_multiplier.jl")


text:
struct Equirectangular
end

struct Mercator
    central_meridian_deg::Float64
end

struct Lambert
    central_meridian_deg::Float64
end

struct Gnomonic
    center_lon_deg::Float64
    center_lat_deg::Float64
end

"""
projection(lon_deg::Float64, lat_deg::Float64,
    projparam::Equirectangular=Equirectangular())
Equirectangular projection
"""
function projection(lon_deg::Float64, lat_deg::Float64,
    projparam::Equirectangular=Equirectangular())
    [lon_deg lat_deg]
end

"""
projection_inv(x::Float64, y::Float64,
    projparam::Equirectangular=Equirectangular())
Equirectangular projection (inverse)
"""
function projection_inv(x::Float64, y::Float64,
    projparam::Equirectangular=Equirectangular())
    [x y]
end

"""
projection(lon_deg::Float64, lat_deg::Float64, projparam::Mercator)
Mercator projection

Source: https://mathworld.wolfram.com/MercatorProjection.html
"""
function projection(lon_deg::Float64, lat_deg::Float64, projparam::Mercator)
    [(lon_deg-projparam.central_meridian_deg) rad2deg(log(tan(π/4.0 + deg2rad(lat_deg) / 2.0)))]
end

"""
projection_inv(x::Float64, y::Float64, projparam::Mercator)
Mercator projection (inverse)

Source: https://mathworld.wolfram.com/MercatorProjection.html
"""
function projection_inv(x::Float64, y::Float64, projparam::Mercator)
    [(x + projparam.central_meridian_deg) rad2deg(atan(sinh(deg2rad(y))))]
end

"""
projection(lon_deg::Float64, lat_deg::Float64, projparam::Lambert)
Lambert cylindrical equal-area

Source: https://mathworld.wolfram.com/CylindricalEqual-AreaProjection.html
"""
function projection(lon_deg::Float64, lat_deg::Float64, projparam::Lambert)
    [(lon_deg-projparam.central_meridian_deg) rad2deg(sin(deg2rad(lat_deg)))]
end

"""
projection_inv(x::Float64, y::Float64, projparam::Lambert)
Lambert cylindrical equal-area (inverse)

Source: https://mathworld.wolfram.com/CylindricalEqual-AreaProjection.html
"""
function projection_inv(x::Float64, y::Float64, projparam::Lambert)
    [(x + projparam.central_meridian_deg) rad2deg(asin(deg2rad(y)))]
end

"""
projection(lat_deg::Float64, lon_deg::Float64, projparam::Gnomonic)

Source: https://mathworld.wolfram.com/GnomonicProjection.html
"""
function projection(lon_deg::Float64, lat_deg::Float64, projparam::Gnomonic)

    limited_point = limitdistance(projparam.center_lon_deg,
    projparam.center_lat_deg, lon_deg, lat_deg)

    ϕ₀ = deg2rad(projparam.center_lat_deg)
    λ₀ = deg2rad(projparam.center_lon_deg)
    ϕ = deg2rad(limited_point[2])
    λ = deg2rad(limited_point[1])

    cosc = sin(ϕ₀)sin(ϕ) + cos(ϕ₀)cos(ϕ)cos(λ - λ₀)
    x = cos(ϕ)sin(λ - λ₀) / cosc
    y = (cos(ϕ₀)sin(ϕ) - sin(ϕ₀)cos(ϕ)cos(λ - λ₀)) / cosc
    return [rad2deg(x) rad2deg(y)]
end

"""
projection_inv(x::Float64, y::Float64, projparam::Gnomonic) inverse

Source: https://mathworld.wolfram.com/GnomonicProjection.html
"""
function projection_inv(x::Float64, y::Float64, projparam::Gnomonic)

    # limited_point = limitdistance(projparam.center_lon_deg,
    # projparam.center_lat_deg, lon_deg, lat_deg)

    ϕ₀ = deg2rad(projparam.center_lat_deg)
    λ₀ = deg2rad(projparam.center_lon_deg)

    xr = deg2rad(x)
    yr = deg2rad(y)

    ρ = √(xr^2 + yr^2)
    if ρ != 0.0
        c = atan(ρ)
        ϕ = asin(cos(c)sin(ϕ₀) + yr*sin(c)cos(ϕ₀)/ρ)
        λ = λ₀ + atan(xr*sin(c),ρ*cos(ϕ₀)cos(c)-yr*sin(ϕ₀)sin(c))
    else  # ρ == 0.0
        ϕ = ϕ₀
        λ = λ₀
    end
    return [rad2deg(λ) rad2deg(ϕ)]
end

function limitdistance(centre_point_lon_deg::Float64,
    centre_point_lat_deg::Float64, lon_deg::Float64, lat_deg::Float64)
    center_point = Navigation.Point_deg(centre_point_lat_deg,
    centre_point_lon_deg)
    point = Navigation.Point_deg(lat_deg, lon_deg)
    angular_distance = Navigation.distance(center_point, point, 1.0)
    max_angular_distance = π * 0.4
    if angular_distance > max_angular_distance
        fraction = max_angular_distance / angular_distance
        new_point = Navigation.intermediate_point(center_point, point, fraction)
        return [new_point.λ new_point.ϕ]
    else
        return [lon_deg lat_deg]
    end
end


text:
@testset "LBFGS" begin
    model = local_level(rand(1000))

    # Test error
    opt_method = LBFGS(model, 10)
    @test opt_method.n_seeds == 10
    @test length(opt_method.seeds) == 10

    opt_method = LBFGS(model, [[1.0; 2.0]])
    @test opt_method.n_seeds == 1
    @test opt_method.seeds == [[1.0; 2.0]]

    err = ErrorException("Seed 1 has 1 elements and the model has 2 unknowns.")
    @test_throws err LBFGS(model, [[1.0]])
end

@testset "BFGS" begin
    model = local_level(rand(1000))

    # Test error
    opt_method = BFGS(model, 10)
    @test opt_method.n_seeds == 10
    @test length(opt_method.seeds) == 10

    opt_method = BFGS(model, [[1.0; 2.0]])
    @test opt_method.n_seeds == 1
    @test opt_method.seeds == [[1.0; 2.0]]

    err = ErrorException("Seed 1 has 1 elements and the model has 2 unknowns.")
    @test_throws err BFGS(model, [[1.0]])
end


text:
export alphabet

"""
	abstract type AbstractStrain end

The only requirement for concrete implementations of `AbstractStrain` is the existence of two fields: `seq` and `data::Dict`. `data` is often expected to contain a key `"strain"` referring to the name of the strain by default, but I should try not to enforce it.
"""
abstract type AbstractStrain end

getindex(s::AbstractStrain, key::String) = get(s.data, key) do
	haskey(s.data, Symbol(key)) ? s.data[Symbol(key)] : @error "Keys $key or $(Symbol(key)) not found."
end
getindex(s::AbstractStrain, key::Symbol) = get(s.data, key) do
	haskey(s.data, String(key)) ? s.data[String(key)] : @error "Keys $key or $(String(key)) not found."
end
"""
	isempty(st::AbstractStrain)
"""
function isempty(st::AbstractStrain)
	return isempty(st.seq) && isempty(st.data)
end

"""
	mutable struct Strain{A} <: AbstractStrain

Fields:
- `seq::BioSequence{A}`
- `data::Dict`
"""
mutable struct Strain{A} <: AbstractStrain
	seq::LongSequence{A}
	data::Dict
end
function Strain(seqtype)
	if seqtype == :aa
		return Strain(LongAminoAcidSeq(), Dict())
	elseif seqtype == :dna
		Strain(LongDNASeq(), Dict())
	elseif seqtype == :rna
		Strain(LongRNASeq(), Dict())
	else
		unknown_seqtype()
	end
end
"""
	Strain(seq, dat, seqtype)

Construct `Strain` object from `seq` and `dat`, when `seq` needs to be converted to a `BioSequence` object first. Values of `seqtype`: `aa`, `dna` or `rna`.
"""
function Strain(seq, dat, seqtype ; verbose=false)
	if seqtype == :aa
		try
			return Strain(LongAminoAcidSeq(seq), dat)
		catch
			cannot_read(seq, dat, seqtype)
			return Strain(seqtype)
		end
	elseif seqtype == :dna
		try
			return Strain(LongDNASeq(seq), dat)
		catch
			cannot_read(seq, dat, seqtype)
			return Strain(seqtype)
		end
	elseif seqtype == :rna
		return Strain(LongRNASeq(seq), dat)
	else
		unknown_seqtype()
	end
end
function cannot_read(seq, dat, seqtype)
	@warn "Could not read the following sequence:"
	println("$seq")
	println("Associated data:\n$dat")
	println("Called seq. type: $seqtype")
	println()
	return nothing
end

"""
	mutable struct ArtificialStrain{T} <: AbstractStrain
		seq::Array{T,1}
		data::Dict
	end
"""
mutable struct ArtificialStrain{T} <: AbstractStrain
	seq::Array{T,1}
	data::Dict
end
function ArtificialStrain(T::DataType)
	return ArtificialStrain(T[], Dict())
end
"""
	ArtificialStrain(seq, dat, T::DataType = Int64)
"""
function ArtificialStrain(seq, dat, T::DataType = Int64)
	return ArtificialStrain(convert(Array{T,1}, seq), dat)
end

# abstract type SequenceProfile end
# For now I don't see an efficient application of profile outside of this.
# For DCA like stuff, it's a float vector, and it's highly unpractical to have something both general and efficient here.
# So no abstract type

# One way would be to have `data` be an array of subtype `colProfile`.
# `colProfile` could then be designed independently
# As long as the getindex function is overloaded properly this could be practical.
"""
	mutable struct SiteFrequency{A}
		i::Int64
		M::Int64
		freq::Dict{A,Float64}
	end

## Note
- From outside, the type should behave like a `Dict{A,Float64}`. That is, `keys(F::SiteFrequency)` should return the keys of `F.freq`.
- TRY: Calling `f[a]` returns `0.` as a default value if `a` does not appear in `f.freq`.
"""
mutable struct SiteFrequency{A}
	i::Int64
	M::Int64
	freq::Dict{A,Float64}
end
"""
	SiteFrequency(T::DataType, i::Int64=0, M::Int64=0)

Constructor for empty structure.
"""
function SiteFrequency(T::DataType, i::Int64=0, M::Int64=0)
	return SiteFrequency(i, M, Dict{T,Float64}())
end
setindex!(F::SiteFrequency{A}, f::Float64, a::A) where A = setindex!(F.freq, f, a)
get(F::SiteFrequency{A}, a::A, def) where A = get(F.freq, a, def)
getindex(F::SiteFrequency{A}, a::A) where A = get(F, a, 0.)
count(F::SiteFrequency{A}, a::A) where A = (F.M == 0) ? error("SiteFrequency: Unknown sequence number") : round(Int64, F.freq[a] * F.M)
counts(F::SiteFrequency) = (F.M == 0) ? error("SiteFrequency: Unknown sequence number") : Dict(k=>v*F.M for (k,v) in F.freq)
keys(F::SiteFrequency) = keys(F.freq)
alphabet(F::SiteFrequency) = keys(F)
values(F::SiteFrequency) = values(F.freq)
isempty(F::SiteFrequency) = (F.M == 0) || isempty(F.freq)

"""
	mutable struct Profile{A}
		data::Array{Dict{A, Float64},1}
		M::Int64
	end
Access `data` by indexing: `p[i,a]` --> `p.data[i][a]`, or `p[i][a]` --> `p.data[i][a]`. This does not return a default value. For a default value, try `get(p, i, a, def)` or `get(p[i], a, def)` (the former might be slower).
"""
mutable struct Profile{A}
	data::Array{SiteFrequency{A},1}
	M::Int64 # Number of sequences the profile is based on
end
function Profile(T::DataType, L::Int64, M::Int64=0)
	return Profile([SiteFrequency(T,i,M) for i in 1:L], M)
end
getindex(P::Profile, i::Int64, a) = P.data[i][a]
getindex(P::Profile, i::Int64) = P.data[i]
get(P::Profile, i::Int64, default) = get(P.data, i, default)
get(P::Profile, i::Int64, a, default) = get(get(P.data, i, Dict()), a, default)
length(P::Profile) = length(P.data)
iterate(P::Profile, n=1) = iterate(P.data, n)

"""
	Profile(S::Array{<:BioSequence{A}}; ambiguous=false) where A
	Profile(S::Array{<:Strain{A}}; ambiguous=false) where A
"""
function Profile(S::Array{<:BioSequence{A}}; ambiguous=false) where A
	if isempty(S)
		@error "Cannot build a profile  from empty alignment"
	else
		prof = Profile(eltype(A), length(S[1]), length(S))
		for (m,s) in enumerate(S)
			for (i,a) in enumerate(s)
				if ambiguous || !isambiguous(a)
					prof.data[i][a] = get(prof[i], a, 0.) + 1.
				else
					prof.data[i].M -= 1
				end
			end
		end
		# Normalizing
		for i in 1:length(prof)
			for a in alphabet(prof.data[i])
				prof.data[i][a] /= prof.data[i].M
			end
		end
	end
	return prof
end
function Profile(S::Array{<:Strain{A}}; ambiguous = false) where A
	M = length(S)
	if isempty(S)
		@error "Cannot build a profile from empty alignment"
	else
		prof = Profile(eltype(A), length(S[1].seq), length(S))
		for (m,s) in enumerate(S)
			for (i,a) in enumerate(s.seq)
				if ambiguous || !isambiguous(a)
					prof.data[i][a] = get(prof[i], a, 0.) + 1.
				else
					prof.data[i].M -= 1
				end
			end
		end
		# Normalizing
		for i in 1:length(prof)
			for a in alphabet(prof.data[i])
				prof.data[i][a] /= prof.data[i].M
			end
		end
	end
	return prof
end

max_stars_repo_name: PierreBarrat/FluPredictibility

max_issues_repo_name: PierreBarrat/FluPredictibility

max_forks_repo_name: PierreBarrat/FluPredictibility


text:
function oz_solve(st::Structure, γ, br, params::Parameters)
    ur = apply_potential(st.pot, st.r) ./ params.ktemp
    broadcast!(x -> isnan(x) ? 0.0 : x, ur, ur)
    cr = closure_relation(γ, br, ur, st.r)
    broadcast!(x -> isnan(x) ? 0.0 : x, cr, cr)

    fits = continuous!(cr, st.r, params.diam)

    (r, p) = make_fft_plan(cr)
    ck = fft_oz(cr, r, params.rmax, params.mr, p)
    remove_continuous!(ck, st.q, params.diam, fits)

    δ = @. 1.0 - params.δρ * ck
    γk = @. params.δρ * ck^2 / δ

    (r, p) = make_fft_plan(γk)
    new_gamma = ifft_oz(γk, r, params.rmax, params.mr, p)

    cr = closure_relation(new_gamma, br, ur, st.r)
    broadcast!(x -> isnan(x) ? 0.0 : x, cr, cr)

    return cr, new_gamma
end

function build_solve_pairwise(x::AbstractMatrix, y::AbstractVector, r)
    pairwise_product_matrix = pairwise_inner_product(x, r)

    n = size(pairwise_product_matrix, 1)
    difference_product = zeros(eltype(x), n)

    for i in 1:n
        temp_product = y .* x[i, :]
        difference_product[i] = inner_product(temp_product, r)
    end

    # Solve the linear system
    return pairwise_product_matrix \ difference_product
end

function create_solution(coef::AbstractVector, gmat::AbstractMatrix)
    n = size(gmat, 1)
    result = (1.0 - sum(coef)) .* gmat[end, :]

    for i in 1:(n - 1)
        @. result += coef[i] * gmat[n - i, :]
    end

    return result
end

function check_precision(t::Interaction, y::AbstractVector, r)
    cr, g = oz_solve(t.structure, y, t.bridge, t.params)
    diff_vector = g .- y
    prec = inner_product(diff_vector .^ 2, r)

    return (√prec, diff_vector, g, cr)
end

function differences_coefs!(cmat, dmat, y)
    for i in 1:(size(cmat, 1))
        cmat[i, :] = y .- dmat[end - i, :]
    end
end

function solve_to_precision(t::Interaction, cmat, dmat, gmat; prec=1e-4)
    new_prec = Inf
    view_gmat = gmat[2:end, :]
    last_diff = view_gmat[end, :] .- view_gmat[end - 1, :]
    differences_coefs!(cmat, dmat, last_diff)
    cr = zeros(eltype(gmat), size(gmat[1, :]))

    while (new_prec - prec) > 0
        coefs = build_solve_pairwise(cmat, last_diff, t.structure.r)
        fnew = create_solution(coefs, view_gmat)
        new_prec, diff_vector, g, cr = check_precision(t, fnew, t.structure.r)
        dmat[end, :] = diff_vector
        view_gmat[end, :] = g
        view_gmat = circshift(view_gmat, (-1, 0))
        dmat = circshift(dmat, (-1, 0))
        cr, view_gmat[end, :] = oz_solve(
            t.structure, view_gmat[end - 1, :], t.bridge, t.params
        )
        last_diff = view_gmat[end, :] .- view_gmat[end - 1, :]
        differences_coefs!(cmat, dmat, last_diff)
    end

    gmat[2:end, :] = view_gmat

    return cr
end

function ng_method_first!(t::Interaction, γ_matrix)
    for i in 2:size(γ_matrix, 1)
        _, γ_matrix[i, :] = oz_solve(t.structure, γ_matrix[i - 1, :], t.bridge, t.params)
    end

    γ_matrix[1, :] = γ_matrix[end, :]

    return nothing
end

function ng_method!(t::Interaction, γ_matrix)
    for i in 2:size(γ_matrix, 1)
        _, γ_matrix[i, :] = oz_solve(t.structure, γ_matrix[i - 1, :], t.bridge, t.params)
    end
    dif_matrix = diff(γ_matrix; dims=1)

    (n, m) = size(dif_matrix)
    coeff_matrix = zeros(n - 1, m)

    cr = solve_to_precision(t, coeff_matrix, dif_matrix, γ_matrix)

    return cr
end

function solve(t::Interaction; npoints=5)
    γ_matrix = zeros(npoints + 2, t.params.mr) # +2 for initial and one last computation
    temp_gamma = zeros(t.params.mr)

    δρ = t.params.ρ / t.params.nρ
    rho_range = range(δρ, t.params.nρ * δρ; length=t.params.nρ)

    # First iteration
    t.params.δρ = rho_range[1]
    ng_method_first!(t, γ_matrix)
    temp_gamma = γ_matrix[end, :]

    r = Result()

    for j in rho_range[2:end]
        t.params.δρ = j

        cr = ng_method!(t, γ_matrix)
        r = Result(cr, γ_matrix[end, :], t.structure, t.params)

        extrapolated = extrapolation(temp_gamma, γ_matrix[end, :], t.params.ρ, j)
        γ_matrix[1, :] = extrapolated
        temp_gamma = γ_matrix[end, :]
    end

    return r
end


text:
#######################################################################
# Analysing expressions for lower bounds
###############################
#
# The goal is to find all occurrences of lower bounds
#   on type variables in expresions.
#
# This can happen in 3 cases:
#
#   1) where T >: Int
#   2) where Int <: T <: Number
#   3) Foo{>:Int} (shorthand for Foo{T} where T>:Int)
#######################################################################

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# Parsing and detecting lower bounds
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

# Any → LowerBound|nothing
# If `expr` is an expression that directly contains a lower bound,
# returns it; otherwise, returns `nothing`
# NOTE: Matches three patterns: `>:LB`, `T>:LB`, and `LB<:T<:UB`
#   Although syntactically it's possible to write `UB>:T` or `UB>:T>:LB`,
#   Julia doesn't accept such a syntax
extractLowerBound(expr) = begin
    hasLB = @capture(expr, >: LB_) # shorthand
    hasLB && return LB # if the first pattern worked, we are done
    hasLB = @capture(expr, T_ >: LB_)
    hasLB && return LB # if the second pattern worked, also done
    hasLB = @capture(expr, LB_ <: T_ <: UB_)
    hasLB && return LB
    nothing # if neither pattern worked out
end

# Computes frequencies of lower bounds in `expr`
extractLowerBounds(expr) :: LBValsFreq = begin
    bounds = LBValsFreq()
    # we need to capture fresh state, so the local function
    # closes over `bounds`
    recordLBs(e) = begin
        lb = extractLowerBound(e)
        # if there is bound, save it
        lb === nothing || push!(bounds, lb)
        # return the same expr for `prewalk` to work
        e
    end
    MacroTools.prewalk(recordLBs, expr)
    bounds
end

max_stars_repo_name: prl-julia/julia-sub

max_issues_repo_name: prl-julia/julia-sub

max_forks_repo_name: prl-julia/julia-sub


text:
################################################################################
#
#	CONCRETE TYPE
#
#   Bandstructure{path} <: AbstractBandstructure{path}
#   --> path is the path in k-space
#
#   FILE CONTAINS
#       - concrete struct definition
#        - interface implementation
#
################################################################################

mutable struct Bandstructure{P,H} <: AbstractBandstructure{P,H}

    # the path along which the band structure is calcualted
    path :: P
    # the hamiltonian
    h :: H

    # bands for each segment
    # bands[i] gives all bands of segment i
    # bands[i][j] gives all energy values for band j of segment i
    # bands[i][j][k] gives the energy value at kpoint index k of band j in segment i
    bands :: Array{Array{Array{Float64, 1}, 1}, 1}

end

# export the type
export Bandstructure

################################################################################
#
#	INTERFACING / ACCESSING BANDSTRUCTURES
#	(functions have to be overwritten by concrete types)
#
################################################################################

# obtaining the hamiltonian
function hamiltonian(
            bs :: Bandstructure{P,H}
        ) :: H where {RP, P<:AbstractReciprocalPath{RP}, L,UC,HB,H<:AbstractHamiltonian{L,UC,HB}}

    # return the field
    return bs.h
end

# export the function
export hamiltonian

# obtaining the path
function path(
            bs :: Bandstructure{P,H}
        ) :: P where {RP, P<:AbstractReciprocalPath{RP}, L,UC,HB,H<:AbstractHamiltonian{L,UC,HB}}

    # return the field
    return bs.path
end

# export the function
export path

# obtaining the energy values
function energies(
            bs :: Bandstructure{P,H}
        ) :: Array{Array{Array{Float64, 1}, 1}, 1} where {RP, P<:AbstractReciprocalPath{RP}, L,UC,HB,H<:AbstractHamiltonian{L,UC,HB}}

    # return the energy array
    return bs.bands
end

function saveBandstructure(
        bs :: BS,
        fn :: AbstractString,
        group :: AbstractString = "bandstructure"
        ;
        append :: Bool = false
    ) where {RP, P<:AbstractReciprocalPath{RP}, L,UC,HB,H<:AbstractHamiltonian{L,UC,HB}, BS<:Bandstructure{P,H}}

    # determine the mode based on if one wants to append stuff
    if append
        mode = "r+"
    else
        mode = "w"
    end

    # group for hamiltonian
    group_h = group*"/hamiltonian"
    # group for path
    group_p = group*"/path"
    # group for energy bands
    group_e = group*"/bands"

    # open the file in mode
    h5open(fn, mode) do file
        # create the group in which the bonds are saved
        group_bs = g_create(file, group)
        # write number of segments
        attrs(group_bs)["segments"]    = length(bs.bands)
        # save the groups
        attrs(group_bs)["hamiltonian"] = group_h
        attrs(group_bs)["path"]        = group_p
        attrs(group_bs)["bands"]       = group_e
    end

    # save hamiltonian
    saveHamiltonian(hamiltonian(bs), fn, group_h, append=true)
    # save reciprocal path
    saveReciprocalPath(path(bs), fn, group_p, append=true)

    # save the individual segments within the bands group
    h5open(fn, "r+") do file
        # create the group in which the bonds are saved
        group_bands = g_create(file, group_e)
        # save the individual segments
        for si in 1:length(bs.bands)
            s = bs.bands[si]
            # reformat segment and save as matrix
            #segmat = [s[j][i] for i in 1:length(s[1]) for j in 1:length(s)]
            segmat = zeros(length(s[1]), length(s))
            for i in 1:length(s[1])
                for j in 1:length(s)
                    segmat[i,j] = s[j][i]
                end
            end
            # save
            group_bands["segment_$(si)"] = segmat
        end
    end

    # return nothing
    return nothing
end

function loadBandstructure(
        ::Type{BS},
        fn :: AbstractString,
        group :: AbstractString = "bandstructure"
    ) where {RP, P<:AbstractReciprocalPath{RP}, L,UC,HB,H<:AbstractHamiltonian{L,UC,HB}, BS<:Union{Bandstructure, Bandstructure{P,H}}}

    # read attribute data
    attr_data = h5readattr(fn, group)

    # group for hamiltonian
    group_h = attr_data["hamiltonian"]
    # group for path
    group_p = attr_data["path"]
    # group for energy bands
    group_e = attr_data["bands"]

    # load number of segments
    segments = attr_data["segments"]

    # load hamiltonian
    h = loadHamiltonian(fn, group_h)
    # load path
    p = loadReciprocalPath(fn, group_p)

    # load all energy bands
    bds = Vector{Vector{Float64}}[]
    # iterate over all expected segments
    for s in 1:segments
        # load segment as matrix
        segmat = h5read(fn, group_e*"/segment_$(s)")
        # reformat segment and push to list
        segment = Vector{Float64}[
            Float64[segmat[i,b] for i in 1:size(segmat,1)]
            for b in 1:size(segmat,2)
        ]
        # push to the list
        push!(bds, segment)
    end

    # return the new Bandstructure
    bs = getBandstructure(h,p,recalculate=false)
    bs.bands = bds
    return bs
end

function loadBandstructure(
        fn :: AbstractString,
        group :: AbstractString = "bandstructure"
    )

    return loadBandstructure(Bandstructure, fn, group)
end

# recalculate the band structure (energy values)
function recalculate!(
            bs :: Bandstructure{P,H}
            ;
            resolution :: Union{Int64, Vector{Int64}} = 100,
            batch_size :: Integer = -1,
            batch_size_min :: Integer = 1,
            batch_size_max :: Integer = 100
        ) where {D,RP<:AbstractReciprocalPoint{D}, P<:AbstractReciprocalPath{RP}, L,UC,HB,H<:AbstractHamiltonian{L,UC,HB}}

    # calculate the number of segments
    N_segments = length(path(bs)) - 1
    # parse the segment resolution
    seg_resolution = (zeros(Int64,N_segments).+1).*resolution
    if typeof(resolution)==Vector{Int64}
        if length(resolution) == N_segments
            seg_resolution = resolution
        else
            error("given resolution \"$(resolution)\" has not enough components ($(N_segments) needed)")
        end
    end

    # calculate number of bands
    N_bands = dim(hamiltonian(bs))

    # make a list of all k values
    k_values = Vector{Vector{Vector{Float64}}}(undef, N_segments)
    # make a list for all energy values
    e_values = Vector{Vector{Vector{Float64}}}(undef, N_segments)

    # fill all k values
    @inbounds for i in 1:N_segments
        # fill the k values with points on the path segment
        k_values[i] = getPointsOnLine(path(bs)[i], path(bs)[i+1], seg_resolution[i])
        # fill energy values with zeros
        e_values[i] = Vector{Float64}[
            zeros(Float64, seg_resolution[i]) for j in 1:N_bands
        ]
    end

    # calculate the energy values
    @inbounds for i in 1:N_segments
        # the actual batch size for pmap
        batch_size_used = 1
        if batch_size > 0
            batch_size_used = batch_size
        else
            batch_size_used = min(batch_size_min, max(batch_size_max, round(Int64, length(k_values[i])/(10*length(workers())))))
        end
        # get the energy values
        e_values_band :: Vector{Vector{Float64}} = pmap(
            k -> sort(real.(eigvals!(Hermitian(matrixAtK(hamiltonian(bs), k))))),
            k_values[i],
            batch_size=batch_size_used
        )
        # format the energy values
        for j in 1:seg_resolution[i]
            # put the eigenvalues into the bands
            @simd for b in 1:N_bands
                e_values[i][b][j] = e_values_band[j][b]
            end
        end
    end

    # set the energy values within band structure
    bs.bands = e_values

    # return nothing
    return nothing
end

# export the function
export recalculate!

# Convinience function if path and hamiltonian are known
function getBandstructure(
            path :: P,
            h    :: H
            ;
            recalculate :: Bool = true
        ) :: Bandstructure{P,H} where {D,RP<:AbstractReciprocalPoint{D}, P<:AbstractReciprocalPath{RP}, L,UC,HB,H<:AbstractHamiltonian{L,UC,HB}}

    # create a new object
    bs = Bandstructure{P,H}(path, h, Vector{Vector{Vector{Float64}}}(undef, length(path)-1))

    # recalculate the numerical energy values
    if recalculate
        recalculate!(bs)
    end

    # return the new object
    return bs
end

function getBandstructure(
            h    :: H,
            path :: P
            ;
            recalculate :: Bool = true
        ) :: Bandstructure{P,H} where {D,RP<:AbstractReciprocalPoint{D}, P<:AbstractReciprocalPath{RP}, L,UC,HB,H<:AbstractHamiltonian{L,UC,HB}}

    # create a new object
    bs = Bandstructure{P,H}(path, h, Vector{Vector{Vector{Float64}}}(undef, length(path)-1))

    # recalculate the numerical energy values
    if recalculate
        recalculate!(bs)
    end

    # return the new object
    return bs
end

function getBandstructure(
            h      :: H,
            points :: Union{Symbol, Tuple{Symbol, Int64}} ...
            ;
            recalculate :: Bool = true
        ) :: Bandstructure{ReciprocalPath{ReciprocalPoint{D}},H} where {LS,D,S<:AbstractSite{LS,D},B,L,UC<:AbstractUnitcell{S,B},HB,H<:AbstractHamiltonian{L,UC,HB}}

    # create new path
    path = getReciprocalPath(unitcell(h), points...)
    # create a new object
    bs = Bandstructure{ReciprocalPath{ReciprocalPoint{D}},H}(path, h, Vector{Vector{Vector{Float64}}}(undef, length(path)-1))

    # recalculate the numerical energy values
    if recalculate
        recalculate!(bs)
    end

    # return the new object
    return bs
end

# Convinience function if only path, unitcell and the bond hamiltonian is known (uses concrete Hamiltonian)
function getBandstructure(
            uc   :: UC,
            path :: P,
            h    :: HB
            ;
            recalculate :: Bool = true
        ) :: Bandstructure{P,Hamiltonian{L,UC,HB}} where {D,RP<:AbstractReciprocalPoint{D}, P<:AbstractReciprocalPath{RP}, L,S,N,B<:AbstractBond{L,N},UC<:AbstractUnitcell{S,B},NS,HB<:AbstractBondHamiltonian{L,NS}}

    # create a new object
    bs = Bandstructure{P,Hamiltonian{L,UC,HB}}(path, Hamiltonian{L,UC,HB}(uc, h), Vector{Vector{Vector{Float64}}}(undef, length(path)-1))

    # recalculate the numerical energy values
    if recalculate
        recalculate!(bs)
    end

    # return the new object
    return bs
end

# Convinience function if only path and unitcell but not even the bond hamiltonian is known (uses concrete Hamiltonian and simple hopping hamiltonian)
function getBandstructure(
            uc   :: UC,
            path :: P
            ;
            recalculate :: Bool = true
        ) :: Bandstructure{P,Hamiltonian{L,UC,BondHoppingHamiltonianSimple{L}}} where {D,RP<:AbstractReciprocalPoint{D}, P<:AbstractReciprocalPath{RP}, L,S,N,B<:AbstractBond{L,N},UC<:AbstractUnitcell{S,B}}

    # create a new object
    bs = Bandstructure{P,Hamiltonian{L,UC,BondHoppingHamiltonianSimple{L}}}(path, Hamiltonian{L,UC,BondHoppingHamiltonianSimple{L}}(uc, getHoppingHamiltonianSimple(uc)), Vector{Vector{Vector{Float64}}}(undef, length(path)-1))

    # recalculate the numerical energy values
    if recalculate
        recalculate!(bs)
    end

    # return the new object
    return bs
end

# Convinience function if only unitcell and points are known
function getBandstructure(
            uc     :: UC,
            points :: Union{Symbol, Tuple{Symbol, Int64}} ...
            ;
            recalculate :: Bool = true
        ) :: Bandstructure{ReciprocalPath{ReciprocalPoint{D}},Hamiltonian{L,UC,BondHoppingHamiltonianSimple{L}}} where {LS,D,L,S<:AbstractSite{LS,D},N,B<:AbstractBond{L,N},UC<:AbstractUnitcell{S,B}}

    # create a new object
    bs = Bandstructure{ReciprocalPath{ReciprocalPoint{D}},Hamiltonian{L,UC,BondHoppingHamiltonianSimple{L}}}(getReciprocalPath(uc, points...), Hamiltonian{L,UC,BondHoppingHamiltonianSimple{L}}(uc, getHoppingHamiltonianSimple(uc)), Vector{Vector{Vector{Float64}}}(undef, length(points)-1))

    # recalculate the numerical energy values
    if recalculate
        recalculate!(bs)
    end

    # return the new object
    return bs
end

# export the convenience function
export getBandstructure


text:
module GymWrappers

import ..AbstractEnvironment,
    ..reset!,
    ..set_render_mode!,
    ..render,
    ..close,
    ..step!,
    ..observation_type,
    ..observation_extrema,
    ..observation_length,
    ..action_type,
    ..action_set,
    ..action_length,
    ..fps

using PyCall
using StaticArrays

export GymCartPoleV1,
    GymLunarLander,
    GymMountainCarDiscrete,
    GymAcrobot,
    GymBreakoutRAM,
    GymInvertedPendulum,
    GymHopperV2,
    GymCartPoleBulletV1,
    GymReacherBulletV0,
    GymHopperBulletV0,
    GymHalfCheetahBulletV0

const gym = PyNULL()
const pybullet_envs = PyNULL()

abstract type AbstractGymEnvironment{O <: AbstractVector,A} <:
              AbstractEnvironment{O,A} end

reset!(env::AbstractGymEnvironment) = 0.0, env.env.reset(), false
set_render_mode!(env::AbstractGymEnvironment, mode::Symbol) = nothing
function render(env::AbstractGymEnvironment, mode::Union{Nothing,Symbol}=nothing)
    isnothing(mode) && return env.env.render()

    env.env.render(mode="human")
end
close(env::AbstractGymEnvironment) = env.env.close()

# --------------------------------- CartPole --------------------------------- #
struct GymCartPoleV1 <: AbstractGymEnvironment{SizedVector{4,Float64},Int}
    env::PyObject

    GymCartPoleV1() = new(gym.make("CartPole-v1"))
end

observation_type(::Type{GymCartPoleV1}) = SizedVector{4,Float64}
observation_extrema(::Type{GymCartPoleV1}) =
    (SVector(-4.8, -4.0, -0.418, -3.0), SVector(4.8, 4.0, 0.418, 3.0))
observation_length(::Type{GymCartPoleV1}) = 4
action_type(::Type{GymCartPoleV1}) = Int
action_set(::Type{GymCartPoleV1}) = SA[-1, 1]

function step!(env::GymCartPoleV1, action::Int)
    (action != -1 && action != 1) &&
        throw(DomainError(action, "Invalid action"))

    obs, reward, done, _ = env.env.step(action == -1 ? 0 : 1)

    return reward, SizedVector{4}(obs), done
end

reset!(env::GymCartPoleV1) = 0.0, SizedVector{4}(env.env.reset()), false

# -------------------------------- LunarLander ------------------------------- #
struct GymLunarLander <: AbstractGymEnvironment{SizedVector{8,Float32},Int}
    env::PyObject

    GymLunarLander() = new(gym.make("LunarLander-v2"))
end

observation_type(::Type{GymLunarLander}) = SizedVector{8,Float32}
observation_extrema(::Type{GymLunarLander}) =
    ((@SVector fill(-1.0f0, 8)), (@SVector fill(1.0f0, 8)))
observation_length(::Type{GymLunarLander}) = 8
action_type(::Type{GymLunarLander}) = Int
action_set(::Type{GymLunarLander}) = SA[0,1,2,3]

function step!(env::GymLunarLander, action::Int)
    # (action != -1 && action != 1) &&
    #     throw(DomainError(action, "Invalid action"))

    obs, reward, done, _ = env.env.step(action)

    return Float64(reward), SizedVector{8}(obs), done
end

reset!(env::GymLunarLander) = 0.0, SizedVector{8}(env.env.reset()), false
fps(env::GymLunarLander) = 50

# -------------------------- MontainCar (Discrete) --------------------------- #
struct GymMountainCarDiscrete <: AbstractGymEnvironment{Vector{Float64},Int}
    env::PyObject

    function GymMountainCarDiscrete(steps::Int)
        env = gym.make("MountainCar-v0")
        env._max_episode_steps = steps

        new(env)
    end
end

# TODO: Validate
GymMountainCarDiscrete(; steps=200) = GymMountainCarDiscrete(steps)

observation_type(::GymMountainCarDiscrete) = Vector{Float64}
observation_extrema(::GymMountainCarDiscrete) = ([-1.2, -0.07], [0.6, 0.07])
observation_length(::GymMountainCarDiscrete) = 2
action_type(::GymMountainCarDiscrete) = Int
action_set(::GymMountainCarDiscrete) = SA[0,1,2]

function step!(env::GymMountainCarDiscrete, action::Int)
    (action != 0 && action != 1 && action != 2) &&
        throw(DomainError(action, "Invalid action"))

    obs, reward, done, _ = env.env.step(action)

    return reward, obs, done
end

# ------------------------- MontainCar (Continuous) -------------------------- #
struct GymMountainCarContinuous <: AbstractGymEnvironment{SizedVector{2,Float64},SVector{1,Float64}}
    env::PyObject

    function GymMountainCarContinuous()
        env = gym.make("MountainCarContinuous-v0")
        env._max_episode_steps = 999

        new(env)
    end
end

# TODO: Validate
# GymMountainCarContinuous() = GymMountainCarDiscrete(steps)

observation_type(::Type{GymMountainCarContinuous}) = SizedVector{2,Float64}
observation_extrema(::Type{GymMountainCarContinuous}) =
    (SVector(-1.2, -0.07), SVector(0.6, 0.07))
observation_length(::Type{GymMountainCarContinuous}) = 2
action_type(::Type{GymMountainCarContinuous}) = SVector{1,Float64}
# action_set(::GymMountainCarContinuous) = SA[0,1,2]

function step!(env::GymMountainCarContinuous, action::SVector{1,Float64})
    # (action != 0 && action != 1 && action != 2) &&
    #     throw(DomainError(action, "Invalid action"))

    obs, reward, done, _ = env.env.step(clamp.(action, -1.0, 1.0))

    return Float64(reward), SizedVector{2}(obs), done
end

reset!(env::GymMountainCarContinuous) = 0.0, SizedVector{2}(env.env.reset()), false

# --------------------------------- Acrobot ---------------------------------- #
struct GymAcrobot <: AbstractGymEnvironment{Vector{Float64},Int}
    env::PyObject

    function GymAcrobot(steps::Int)
        env = gym.make("Acrobot-v1")
        env._max_episode_steps = steps

        new(env)
    end
end

# TODO: Validate
GymAcrobot(; steps=500) = GymAcrobot(steps)

observation_type(::GymAcrobot) = Vector{Float64}
observation_extrema(::GymAcrobot) = (fill(-1, 6), fill(1, 6))
observation_length(::GymAcrobot) = 6
action_type(::GymAcrobot) = Int
action_set(::GymAcrobot) = SA[-1, 0, 1]

function step!(env::GymAcrobot, action::Int)
    (action != 0 && action != 1 && action != -1) &&
        throw(DomainError(action, "Invalid action"))

    obs, reward, done, _ = env.env.step(action)

    return reward, obs, done
end
# =============================== Atari ====================================== #
# ----------------------------- Breakout (RAM) ------------------------------- #
struct GymBreakoutRAM <: AbstractGymEnvironment{Vector{UInt8},Int}
    env::PyObject

    GymBreakoutRAM() = new(gym.make("Breakout-ram-v0"))
end

observation_type(::GymBreakoutRAM) = Vector{UInt8}
observation_extrema(::GymBreakoutRAM) = (typemin(UInt8), typemax(UInt8))
observation_length(::GymBreakoutRAM) = 128
action_type(::GymBreakoutRAM) = Int
action_set(::GymBreakoutRAM) = 0:3

function step!(env::GymBreakoutRAM, action::Int)
    !(action in action_set(env)) && throw(DomainError(action, "Invalid action"))

    obs, reward, done, _ = env.env.step(action)

    return reward, obs, done
end

# ================================== Mujoco ================================== #
# --------------------------- InvertedPendulum-v2 ---------------------------- #
struct GymInvertedPendulum <: AbstractGymEnvironment{Vector{Float64},Float64}
    env::PyObject

    GymInvertedPendulum() = new(gym.make("InvertedPendulum-v2"))
end

observation_type(::GymInvertedPendulum) = Vector{Float64}
observation_extrema(::GymInvertedPendulum) =
    ([-4.8, -4.0, -0.418, -3.0], [4.8, 4.0, 0.418, 3.0])
observation_length(::GymInvertedPendulum) = 4
action_type(::GymInvertedPendulum) = Float64
action_set(::GymInvertedPendulum) = -3.0:1:3.0

# -------------------------------- Hopper-v2 --------------------------------- #
struct GymHopperV2 <: AbstractGymEnvironment{SizedVector{11,Float32},SVector{3,Float32}}
    env::PyObject

    function GymHopperV2()
        env = gym.make("Hopper-v2")

        new(env)
    end
end

observation_type(::Type{GymHopperV2}) = SizedVector{11,Float32}
observation_extrema(::Type{GymHopperV2}) =
    ((@SVector fill(-1.0f0, 11)), (@SVector fill(1.0f0, 11)))

observation_length(::Type{GymHopperV2}) = 11
action_type(::Type{GymHopperV2}) = SVector{3,Float32}

function step!(env::GymHopperV2, action::SVector{3,Float32})
    obs, reward, done, _ = env.env.step(clamp.(action, -1.0, 1.0))

    return Float64(reward), SizedVector{11}(obs), done
end

reset!(env::GymHopperV2) = 0.0, SizedVector{11}(env.env.reset()), false

# ================================= pybullet ================================= #
# ---------------------------- CartPoleBulletEnv-v1 -------------------------- #
struct GymCartPoleBulletV1 <:
    AbstractGymEnvironment{SizedVector{4,Float64},Int}
    env::PyObject

    GymCartPoleBulletV1() = new(gym.make("CartPoleBulletEnv-v1"))
end

observation_type(::Type{GymCartPoleBulletV1}) = SizedVector{4,Float64}
# observation_extrema(::Type{GymCartPoleBulletV1}) =
#     (SA_F64[-4.8, -4.0, -0.418, -3.0], SA_F64[4.8, 4.0, 0.418, 3.0])
observation_extrema(::Type{GymCartPoleBulletV1}) = (fill(-0.05, 4), fill(0.05, 4))
observation_length(::Type{GymCartPoleBulletV1}) = 4
action_type(::Type{GymCartPoleBulletV1}) = Int
action_set(::Type{GymCartPoleBulletV1}) = SA[-1, 1]

set_render_mode!(env::GymCartPoleBulletV1, mode::Symbol=:human) = render(env, mode)

function step!(env::GymCartPoleBulletV1, action::Int)
    (action != -1 && action != 1) &&
        throw(DomainError(action, "Invalid action"))

    obs, reward, done, _ = env.env.step(action == -1 ? 0 : 1)
    # @show typeof(obs)
    return reward, SizedVector{4}(obs), done
end

reset!(env::GymCartPoleBulletV1) = 0.0, SizedVector{4}(env.env.reset()), false

# ---------------------------- ReacherBulletEnv-v0 -------------------------- #
struct GymReacherBulletV0 <:
    AbstractGymEnvironment{SizedVector{9,Float64},SVector{2,Float64}}
    env::PyObject

    GymReacherBulletV0() = new(gym.make("ReacherBulletEnv-v0"))
end

observation_type(::Type{GymReacherBulletV0}) = SizedVector{9,Float64}
observation_extrema(::Type{GymReacherBulletV0}) = (fill(-1.0, 9), fill(1.0, 9))
observation_length(::Type{GymReacherBulletV0}) = 9
action_type(::Type{GymReacherBulletV0}) = SVector{2,Float64}
# action_set(::Type{GymReacherBulletV0}) = SA[-1, 1]

set_render_mode!(env::GymReacherBulletV0, mode::Symbol=:human) = render(env, mode)

function step!(env::GymReacherBulletV0, action::StaticArray{Tuple{2},Float64,1})
    # (action != -1 && action != 1) &&
    #     throw(DomainError(action, "Invalid action"))

    obs, reward, done, _ = env.env.step(clamp.(action, -1.0f0, 1.0f0))

    return reward, SizedVector{9}(obs), done
end

reset!(env::GymReacherBulletV0) = 0.0, SizedVector{9}(env.env.reset()), false

# ---------------------------- HopperBulletEnv-v0 ---------------------------- #
struct GymHopperBulletV0 <:
       AbstractGymEnvironment{SizedVector{15,Float32},SVector{3,Float32}}
    env::PyObject

    GymHopperBulletV0() = new(gym.make("HopperBulletEnv-v0"))
end

observation_type(::Type{GymHopperBulletV0}) = SizedVector{15,Float32}
observation_extrema(::Type{GymHopperBulletV0}) = (SizedVector{15}(fill(0.0f0, 15)), SizedVector{15}(fill(1.5f0, 15)))
observation_length(::Type{GymHopperBulletV0}) = 15
action_type(::Type{GymHopperBulletV0}) = SVector{3,Float32} # Vector{Float32}
action_length(::Type{GymHopperBulletV0}) = 3
action_set(::Type{GymHopperBulletV0}) = error("Not implemented")
# fps(::Type{GymHopperBulletV0}) = 60

set_render_mode!(env::GymHopperBulletV0, mode::Symbol=:human) = render(env, mode)

function step!(env::GymHopperBulletV0, action::StaticArray{Tuple{3},Float32,1})
    # !env.env.action_space.contains(action) && throw(DomainError(action, "Invalid action"))

    obs, reward, done, _ = env.env.step(clamp.(action, -1.0f0, 1.0f0))

    return reward, SizedVector{15}(obs), done
end

close(env::GymHopperBulletV0) = nothing

reset!(env::GymHopperBulletV0) = 0.0, SizedVector{15}(env.env.reset()), false

# -------------------------- HalfCheetahBulletEnv-v0 ------------------------- #
struct GymHalfCheetahBulletV0 <:
    AbstractGymEnvironment{SizedVector{26,Float32},SVector{6,Float32}}
    env::PyObject

    GymHalfCheetahBulletV0() = new(gym.make("HalfCheetahBulletEnv-v0"))
end

observation_type(::Type{GymHalfCheetahBulletV0}) = SizedVector{26,Float32}
observation_extrema(::Type{GymHalfCheetahBulletV0}) = (SizedVector{26}(fill(-0.5f0, 26)), SizedVector{26}(fill(1.4f0, 26)))
observation_length(::Type{GymHalfCheetahBulletV0}) = 26
action_type(::Type{GymHalfCheetahBulletV0}) = SVector{6,Float32} # Vector{Float32}
action_length(::Type{GymHalfCheetahBulletV0}) = 6
action_set(::Type{GymHalfCheetahBulletV0}) = error("Not implemented")

set_render_mode!(env::GymHalfCheetahBulletV0, mode::Symbol=:human) = render(env, mode)

function step!(env::GymHalfCheetahBulletV0, action::StaticArray{Tuple{6},Float32,1})
    # !env.env.action_space.contains(action) && throw(DomainError(action, "Invalid action"))

    obs, reward, done, _ = env.env.step(clamp.(action, -1, 1))

    return reward, SizedVector{26}(obs), done
end

close(env::GymHalfCheetahBulletV0) = nothing

reset!(env::GymHalfCheetahBulletV0) = 0.0, SizedVector{26}(env.env.reset()), false

function __init__()
    copy!(gym, pyimport_conda("gym", "gym", "conda-forge"))
    copy!(pybullet_envs, pyimport_conda("pybullet_envs", "pybullet"))
end

end


text:
using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack
using DelimitedFiles
using DataFrames
using JLD2
using FileIO

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D

using SetupDG

push!(LOAD_PATH, "./examples/EntropyStableEuler.jl/src")
using EntropyStableEuler
#using EntropyStableEuler.Fluxes1D

const LIMITOPT = 2     # 1 if elementwise limiting lij, 2 if elementwise limiting li
const POSDETECT = 0    # 1 if turn on detection, 0 otherwise
const LBOUNDTYPE = 0.5   # 0 if use POSTOL as lower bound, if > 0, use LBOUNDTYPE*loworder
const POSTOL = 1e-14
const TOL = 5e-16
const CFL = 0.5

# Ignacio - Leblanc shocktube 2
const γ = 5/3
const Bl = 0.0
const Br = 1.0
const rhoL = 1.0
const rhoR = 1e-3
const uL = 0.0
const uR = 0.0
const pL = (γ-1)*1e-1
const pR = (γ-1)*1e-10
const xC = 0.33
T = 2/3

# const GIFINTERVAL = Int(floor((T-t0)/dt/60))

function wavespeed_1D(rho,rhou,E)
    p = pfun_nd(rho,(rhou,),E)
    cvel = @. sqrt(γ*p/rho)
    return @. abs(rhou/rho) + cvel
end
unorm(U) = sum(map((x->x.^2),U))
function pfun_nd(rho, rhoU, E)
    rhoUnorm2 = unorm(rhoU)./rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

function primitive_to_conservative_hardcode(rho,U,p)
    rhoU = rho.*U
    Unorm = unorm(U)
    E = @. p/(γ-1) + .5*rho*Unorm
    return (rho,rhoU,E)
end

function euler_fluxes_1D(rhoL,uL,betaL,rhologL,betalogL,
    rhoR,uR,betaR,rhologR,betalogR)

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))

    unorm = (@. uL*uR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. f4aux*uavg)

    return (FxS1,FxS2,FxS3)
end

function euler_fluxes(rhoL,uL,betaL,rhoR,uR,betaR)
    rhologL,betalogL,rhologR,betalogR = map(x->log.(x),(rhoL,betaL,rhoR,betaR))
    return euler_fluxes_1D(rhoL,uL,betaL,rhologL,betalogL,
                           rhoR,uR,betaR,rhologR,betalogR)
end

function exact_sol_Leblanc(x,t)
    xi = (x-0.33)/t
    rhoLstar = 5.4079335349316249*1e-2
    rhoRstar = 3.9999980604299963*1e-3
    vstar    = 0.62183867139173454
    pstar    = 0.51557792765096996*1e-3
    lambda1  = 0.49578489518897934
    lambda3  = 0.82911836253346982
    if xi <= -1/3
        return rhoL, 0.0, pL
    elseif xi <= lambda1
        return (0.75-0.75*xi)^3, 0.75*(1/3+xi), 1/15*(0.75-0.75*xi)^5
    elseif xi <= vstar
        return rhoLstar, vstar, pstar
    elseif xi <= lambda3
        return rhoRstar, vstar, pstar
    else
        return rhoR, 0.0, pR
    end
end

function exact_sol_Leblanc(x,t)
    xi = (x-0.33)/t
    rhoLstar = 5.4079335349316249*1e-2
    rhoRstar = 3.9999980604299963*1e-3
    vstar    = 0.62183867139173454
    pstar    = 0.51557792765096996*1e-3
    lambda1  = 0.49578489518897934
    lambda3  = 0.82911836253346982
    if xi <= -1/3
        return rhoL, 0.0, pL
    elseif xi <= lambda1
        return (0.75-0.75*xi)^3, 0.75*(1/3+xi), 1/15*(0.75-0.75*xi)^5
    elseif xi <= vstar
        return rhoLstar, vstar, pstar
    elseif xi <= lambda3
        return rhoRstar, vstar, pstar
    else
        return rhoR, 0.0, pR
    end
end

function limiting_param(U_low, P_ij, Lrho, Lrhoe)
    l = 1.0
    # Limit density
    if U_low[1] + P_ij[1] < Lrho
        l = max((Lrho-U_low[1])/P_ij[1],0.0)
    end

    # limiting internal energy (via quadratic function)
    a = P_ij[1]*P_ij[3]-1.0/2.0*P_ij[2]^2
    b = U_low[3]*P_ij[1]+U_low[1]*P_ij[3]-U_low[2]*P_ij[2]-P_ij[1]*Lrhoe
    c = U_low[3]*U_low[1]-1.0/2.0*U_low[2]^2-U_low[1]*Lrhoe

    l_eps_ij = 1.0
    if b^2-4*a*c >= 0
        r1 = (-b+sqrt(b^2-4*a*c))/(2*a)
        r2 = (-b-sqrt(b^2-4*a*c))/(2*a)
        if r1 > TOL && r2 > TOL
            l_eps_ij = min(r1,r2)
        elseif r1 > TOL && r2 < -TOL
            l_eps_ij = r1
        elseif r2 > TOL && r1 < -TOL
            l_eps_ij = r2
        end
    end

    l = min(l,l_eps_ij)
    return l
end

function flux_lowIDP(U_i,U_j,f_i,f_j,c_ij,wavespd)
    return c_ij*(f_i+f_j)-abs(c_ij)*wavespd*(U_j-U_i)
end

function flux_high(f_i,f_j,c_ij)
    return c_ij*(f_i+f_j)
end

function flux_ES(rho_i,u_i,beta_i,rho_j,u_j,beta_j,c_ij)
    return 2*c_ij.*euler_fluxes(rho_i,u_i,beta_i,rho_j,u_j,beta_j)
end

function flux_central(c_ij,f_i,f_j)
    return c_ij*(f_i+f_j)
end

function rhs_IDP(U,K,N,wq,S,S0,Mlump_inv,T,dt,t,in_s1,is_low_order)
    p = pfun_nd.(U[1],U[2],U[3])
    flux = zero.(U)
    @. flux[1] = U[2]
    @. flux[2] = U[2]^2/U[1]+p
    @. flux[3] = U[3]*U[2]/U[1]+p*U[2]/U[1]

    Ub = zero.(U)
    @. Ub[1] = U[1]
    @. Ub[2] = U[2]/U[1]
    @. Ub[3] = U[1]/(2*p)

    J = (Br-Bl)/K/2 # assume uniform interval

    # Low order and high order algebraic fluxes
    F_low       = [zeros(N+1,N+1),zeros(N+1,N+1),zeros(N+1,N+1)]
    F_high      = [zeros(N+1,N+1),zeros(N+1,N+1),zeros(N+1,N+1)]
    F_low_P     = [zeros(2),zeros(2),zeros(2)] # 1: left boundary, 2: right boundary
    F_high_P    = [zeros(2),zeros(2),zeros(2)]
    L           = zeros(N+1,N+1) # Array of limiting params
    L_P         = zeros(K-1) # limiting params at left/right boundary
    rhsU        = [zeros(N+1,K),zeros(N+1,K),zeros(N+1,K)]
    dii_arr     = zeros(N+1,K)

    wavespd_arr = zeros(N+1,K)
    for k = 1:K
        for i = 1:N+1
            wavespd_arr[i,k] = wavespeed_1D(U[1][i,k],U[2][i,k],U[3][i,k])
        end
    end

    for k = 1:K
        # Interior dissp coeff
        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    dij = abs(S0[i,j])*max(wavespd_arr[i,k],wavespd_arr[j,k])
                    if in_s1
                        dii_arr[i,k] = dii_arr[i,k] + dij
                    end
                end
            end
        end

        # Interface dissipation
        EL = pL/(γ-1)+0.5*rhoL*uL^2
        ER = pR/(γ-1)+0.5*rhoR*uR^2
        U_left   = (k == 1) ? [rhoL; rhoL*uL; EL]                 : [U[1][end,k-1]; U[2][end,k-1]; U[3][end,k-1]]
        Ub_left  = (k == 1) ? [rhoL; uL; rhoL/(2*pL)]             : [Ub[1][end,k-1]; Ub[2][end,k-1]; Ub[3][end,k-1]]
        f_left   = (k == 1) ? [rhoL*uL; rhoL*uL^2+pL; uL*(pL+EL)] : [flux[1][end,k-1]; flux[2][end,k-1]; flux[3][end,k-1]]
        U_right  = (k == K) ? [rhoR; rhoR*uR; ER]                 : [U[1][1,k+1]; U[2][1,k+1]; U[3][1,k+1]]
        Ub_right = (k == K) ? [rhoR; uR; rhoR/(2*pR)]             : [Ub[1][1,k+1]; Ub[2][1,k+1]; Ub[3][1,k+1]]
        f_right  = (k == K) ? [rhoR*uR; rhoR*uR^2+pR; uR*(pR+ER)] : [flux[1][1,k+1]; flux[2][1,k+1]; flux[3][1,k+1]]
        wavespd_l = max(wavespd_arr[1,k],wavespeed_1D(U_left[1],U_left[2],U_left[3]))
        wavespd_r = max(wavespd_arr[end,k],wavespeed_1D(U_right[1],U_right[2],U_right[3]))

        if in_s1
            dii_arr[1,k] = dii_arr[1,k] + 0.5*wavespd_l
            dii_arr[N+1,k] = dii_arr[N+1,k] + 0.5*wavespd_r
        end
    end

    if in_s1
        for k = 1:K
            for i = 1:N+1
                dt = min(dt,1.0*wq[i]*J/2.0/dii_arr[i,k])
            end
        end
        dt = min(CFL*dt,T-t)
    end

    for k = 1:K
        L = zeros(N+1,N+1)
        # Assemble matrix of low and high order algebraic fluxes
        # interior of the element
        for i = 1:N+1
            for j = 1:N+1
                if i != j # skip diagonal
                    wavespd = max(wavespd_arr[i,k],wavespd_arr[j,k])
                    fluxS = flux_ES(Ub[1][i,k],Ub[2][i,k],Ub[3][i,k],Ub[1][j,k],Ub[2][j,k],Ub[3][j,k],S[i,j])
                    for c = 1:3
                        F_low[c][i,j]  = flux_lowIDP(U[c][i,k],U[c][j,k],flux[c][i,k],flux[c][j,k],S0[i,j],wavespd)
                        F_high[c][i,j] = fluxS[c]
                    end
                end
            end
        end

        # Assemble matrix of low and high order algebraic fluxes
        # interface of the element
        EL = pL/(γ-1)+0.5*rhoL*uL^2
        ER = pR/(γ-1)+0.5*rhoR*uR^2
        U_left   = (k == 1) ? [rhoL; rhoL*uL; EL]                 : [U[1][end,k-1]; U[2][end,k-1]; U[3][end,k-1]]
        Ub_left  = (k == 1) ? [rhoL; uL; rhoL/(2*pL)]             : [Ub[1][end,k-1]; Ub[2][end,k-1]; Ub[3][end,k-1]]
        f_left   = (k == 1) ? [rhoL*uL; rhoL*uL^2+pL; uL*(pL+EL)] : [flux[1][end,k-1]; flux[2][end,k-1]; flux[3][end,k-1]]
        U_right  = (k == K) ? [rhoR; rhoR*uR; ER]                 : [U[1][1,k+1]; U[2][1,k+1]; U[3][1,k+1]]
        Ub_right = (k == K) ? [rhoR; uR; rhoR/(2*pR)]             : [Ub[1][1,k+1]; Ub[2][1,k+1]; Ub[3][1,k+1]]
        f_right  = (k == K) ? [rhoR*uR; rhoR*uR^2+pR; uR*(pR+ER)] : [flux[1][1,k+1]; flux[2][1,k+1]; flux[3][1,k+1]]
        wavespd_l = max(wavespd_arr[1,k],wavespeed_1D(U_left[1],U_left[2],U_left[3]))
        wavespd_r = max(wavespd_arr[end,k],wavespeed_1D(U_right[1],U_right[2],U_right[3]))

        for c = 1:3
            F_low_P[c][1] = flux_lowIDP(U[c][1,k],U_left[c],flux[c][1,k],f_left[c],-0.5,wavespd_l)
            F_low_P[c][2] = flux_lowIDP(U[c][end,k],U_right[c],flux[c][end,k],f_right[c],0.5,wavespd_r)

            F_high_P[c][1] = F_low_P[c][1]
            F_high_P[c][2] = F_low_P[c][2]
        end

        # Calculate limiting parameters over interior of the element
        P_ij = zeros(3,1)
        # TODO: redundant
        U_low =  [zeros(N+1,1),zeros(N+1,1),zeros(N+1,1)]
        for c = 1:3
            U_low[c] .= sum(-F_low[c],dims=2)
            U_low[c][1] -= F_low_P[c][1]
            U_low[c][end] -= F_low_P[c][2]
            U_low[c] .= U[c][:,k]+dt/J*Mlump_inv*U_low[c]
        end
        U_high =  [zeros(N+1,1),zeros(N+1,1),zeros(N+1,1)]
        for c = 1:3
            U_high[c] .= sum(-F_high[c],dims=2)
            U_high[c][1] -= F_high_P[c][1]
            U_high[c][end] -= F_high_P[c][2]
            U_high[c] .= U[c][:,k]+dt/J*Mlump_inv*U_high[c]
        end
        is_H_positive = true
        for i = 1:N+1
            rhoH_i  = U_high[1][i]
            rhouH_i = U_high[2][i]
            EH_i    = U_high[3][i]
            pH_i    = pfun_nd(rhoH_i,rhouH_i,EH_i)
            if pH_i < POSTOL || rhoH_i < POSTOL
                is_H_positive = false
            end
        end

        if (POSDETECT == 0)
            is_H_positive = false
        end

        if (LIMITOPT == 1)
            for i = 1:N+1
                #lambda_j = (i >= 2 && i <= N) ? 1/N : 1/(N+1)
                lambda_j = 1/N
                m_i = J*wq[i]
                for j = 1:N+1
                    if i != j
                        for c = 1:3
                            P_ij[c] = dt/(m_i*lambda_j)*(F_low[c][i,j]-F_high[c][i,j])
                        end
                        if !is_H_positive
                            Lrho  = POSTOL
                            Lrhoe = POSTOL
                            L[i,j] = limiting_param([U_low[1][i]; U_low[2][i]; U_low[3][i]],P_ij,Lrho,Lrhoe)
                        else
                            L[i,j] = 1.0
                        end
                    end
                end
            end
        elseif (LIMITOPT == 2)
            li_min = 1.0
            for i = 1:N+1
                lambda_j = 1/N
                m_i = J*wq[i]
                for c = 1:3
                    P_ij[c] = 0.0
                end
                for j = 1:N+1
                    for c = 1:3
                        P_ij[c] += dt/(m_i*lambda_j)*(F_low[c][i,j]-F_high[c][i,j])
                    end
                end
                if (LBOUNDTYPE == 0)
                    Lrho  = POSTOL
                    Lrhoe = POSTOL
                elseif (LBOUNDTYPE > 0)
                    Lrho  = LBOUNDTYPE*U_low[1][i]
                    Lrhoe = LBOUNDTYPE*pfun_nd(U_low[1][i],U_low[2][i],U_low[3][i])/(γ-1)
                end
                l = limiting_param([U_low[1][i]; U_low[2][i]; U_low[3][i]],P_ij,Lrho,Lrhoe)
                li_min = min(li_min,l)
            end

            for i = 1:N+1
                for j = 1:N+1
                    if i != j
                        L[i,j] = li_min
                        L[j,i] = li_min
                    end
                end
            end
        end

        # # Symmetrize limiting parameters
        # for i = 1:N+1
        #     for j = 1:N+1
        #         if i != j
        #             l_ij = min(L[i,j],L[j,i])
        #             L[i,j] = l_ij
        #             L[j,i] = l_ij
        #         end
        #     end
        # end

        # elementwise limiting
        l = 1.0
        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    l = min(l,L[i,j])
                end
            end
        end
        if is_low_order
            l = 0.0
        end

        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    L[i,j] = l
                    L[j,i] = l
                end
            end
        end

        # construct rhs
        for c = 1:3
            # With limiting
            rhsU[c][:,k] = sum((L.-1).*F_low[c] - L.*F_high[c],dims=2)

            if k > 1
                rhsU[c][1,k] += (L_P[k-1]-1)*F_low_P[c][1] - L_P[k-1]*F_high_P[c][1]
            else
                rhsU[c][1,k] += -F_high_P[c][1]
            end
            if k < K
                rhsU[c][N+1,k] += (L_P[k]-1)*F_low_P[c][2] - L_P[k]*F_high_P[c][2]
            else
                rhsU[c][N+1,k] += -F_high_P[c][2]
            end

            rhsU[c][:,k] .= 1/J*Mlump_inv*rhsU[c][:,k]
        end
    end

    return rhsU,dt
end

function rhs_IDPlow(U,K,N,Mlump_inv,wq)
    p = pfun_nd.(U[1],U[2],U[3])
    flux = zero.(U)
    @. flux[1] = U[2]
    @. flux[2] = U[2]^2/U[1]+p
    @. flux[3] = U[3]*U[2]/U[1]+p*U[2]/U[1]

    J = (Br-Bl)/K/2

    dt = Inf
    d_ii_arr = zeros(N+1,K)
    dfdx = (zeros(N+1,K),zeros(N+1,K),zeros(N+1,K))
    for i = 2:K*(N+1)-1
        for c = 1:3
            dfdx[c][i] = 1/2*(flux[c][mod1(i+1,K*(N+1))] - flux[c][mod1(i-1,K*(N+1))])
        end
    end
    dfdx[1][1] = 1/2*(flux[1][2] - 0.0)
    dfdx[2][1] = 1/2*(flux[2][2] - pL)
    dfdx[3][1] = 1/2*(flux[3][2] - 0.0)
    dfdx[1][end] = 1/2*(0.0 - flux[1][end-1])
    dfdx[2][end] = 1/2*(pR - flux[2][end-1])
    dfdx[3][end] = 1/2*(0.0 - flux[3][end-1])

    visc = (zeros(N+1,K),zeros(N+1,K),zeros(N+1,K))
    for i = 2:K*(N+1)-1
        wavespd_curr = wavespeed_1D(U[1][i],U[2][i],U[3][i])
        wavespd_R = wavespeed_1D(U[1][mod1(i+1,K*(N+1))],U[2][mod1(i+1,K*(N+1))],U[3][mod1(i+1,K*(N+1))])
        wavespd_L = wavespeed_1D(U[1][mod1(i-1,K*(N+1))],U[2][mod1(i-1,K*(N+1))],U[3][mod1(i-1,K*(N+1))])
        dL = 1/2*max(wavespd_curr,wavespd_L)
        dR = 1/2*max(wavespd_curr,wavespd_R)
        for c = 1:3
            visc[c][i] = dL*(U[c][mod1(i-1,K*(N+1))]-U[c][i]) + dR*(U[c][mod1(i+1,K*(N+1))]-U[c][i])
        end
        d_ii_arr[i] += .5*max(wavespd_curr,wavespd_L) + .5*max(wavespd_curr,wavespd_R)
    end

    # i = 1
    wavespd_curr = wavespeed_1D(U[1][1],U[2][1],U[3][1])
    wavespd_R = wavespeed_1D(U[1][2],U[2][2],U[3][2])
    wavespd_L = wavespeed_1D(rhoL,0.0,pL/(γ-1))
    dL = 1/2*max(wavespd_curr,wavespd_L)
    dR = 1/2*max(wavespd_curr,wavespd_R)
    visc[1][1] = dL*(rhoL-U[1][1]) + dR*(U[1][2]-U[1][1])
    visc[2][1] = dL*(0.0-U[2][1]) + dR*(U[2][2]-U[2][1])
    visc[3][1] = dL*(pL/(γ-1)-U[3][1]) + dR*(U[3][2]-U[3][1])

    d_ii_arr[1] += (dL+dR)

    # i = end
    wavespd_curr = wavespeed_1D(U[1][end],U[2][end],U[3][end])
    wavespd_R = wavespeed_1D(rhoR,0.0,pR/(γ-1))
    wavespd_L = wavespeed_1D(U[1][end-1],U[2][end-1],U[3][end-1])
    dL = 1/2*max(wavespd_curr,wavespd_L)
    dR = 1/2*max(wavespd_curr,wavespd_R)
    visc[1][end] = dL*(U[1][end-1]-U[1][end]) + dR*(rhoR-U[1][end])
    visc[2][end] = dL*(U[2][end-1]-U[2][end]) + dR*(0.0-U[2][end])
    visc[3][end] = dL*(U[3][end-1]-U[3][end]) + dR*(pR/(γ-1)-U[3][end])

    d_ii_arr[end] += (dL+dR)
    dt = minimum(J*wq./d_ii_arr./2.0)

    rhsU = (x->1/J*Mlump_inv*x).(.-dfdx.+visc)
    return rhsU,dt
end

function rhs_high(U,K,N,Mlump_inv,S)
    p = pfun_nd.(U[1],U[2],U[3])
    flux = zero.(U)
    @. flux[1] = U[2]
    @. flux[2] = U[2]^2/U[1]+p
    @. flux[3] = U[3]*U[2]/U[1]+p*U[2]/U[1]

    Ub = zero.(U)
    @. Ub[1] = U[1]
    @. Ub[2] = U[2]/U[1]
    @. Ub[3] = U[1]/(2*p)

    J = (Br-Bl)/K/2 # assume uniform interval

    # Low order and high order algebraic fluxes
    F_high      = [zeros(N+1,N+1),zeros(N+1,N+1),zeros(N+1,N+1)]
    F_high_P    = [zeros(2),zeros(2),zeros(2)]
    rhsU        = [zeros(N+1,K),zeros(N+1,K),zeros(N+1,K)]

    # Determine dt
    wavespd_arr = zeros(N+1,K)
    for k = 1:K
        for i = 1:N+1
            wavespd_arr[i,k] = wavespeed_1D(U[1][i,k],U[2][i,k],U[3][i,k])
        end
    end

    for k = 1:K
        # Assemble matrix of low and high order algebraic fluxes
        # interior of the element
        for i = 1:N+1
            for j = 1:N+1
                if i != j # skip diagonal
                    # wavespd = max(wavespd_arr[i,k],wavespd_arr[j,k])
                    # for c = 1:3
                    #     F_high[c][i,j] = flux_ES(Ub[c][i,k],Ub[c][j,k],S[i,j])
                    # end
                    tmp = flux_ES(Ub[1][i,k],Ub[2][i,k],Ub[3][i,k],Ub[1][j,k],Ub[2][j,k],Ub[3][j,k],S[i,j])
                    for c = 1:3
                        F_high[c][i,j] = tmp[c]
                    end
                end
            end
        end

        # Assemble matrix of low and high order algebraic fluxes
        # interface of the element
        U_left   = (k == 1) ? [rhoL; 0.0; pL/(γ-1)]    : [U[1][end,k-1]; U[2][end,k-1]; U[3][end,k-1]]
        Ub_left  = (k == 1) ? [rhoL; 0.0; rhoL/(2*pL)] : [Ub[1][end,k-1]; Ub[2][end,k-1]; Ub[3][end,k-1]]
        f_left   = (k == 1) ? [0.0; pL; 0.0]           : [flux[1][end,k-1]; flux[2][end,k-1]; flux[3][end,k-1]]
        U_right  = (k == K) ? [rhoR; 0.0; pR/(γ-1)]    : [U[1][1,k+1]; U[2][1,k+1]; U[3][1,k+1]]
        Ub_right = (k == K) ? [rhoR; 0.0; rhoR/(2*pR)] : [Ub[1][1,k+1]; Ub[2][1,k+1]; Ub[3][1,k+1]]
        f_right  = (k == K) ? [0.0; pR; 0.0]           : [flux[1][1,k+1]; flux[2][1,k+1]; flux[3][1,k+1]]
        wavespd_l = max(wavespd_arr[1,k],wavespeed_1D(U_left[1],U_left[2],U_left[3]))
        wavespd_r = max(wavespd_arr[end,k],wavespeed_1D(U_right[1],U_right[2],U_right[3]))

        tmp = .-euler_fluxes(Ub[1][1,k],Ub[2][1,k],Ub[3][1,k],Ub_left[1],Ub_left[2],Ub_left[3])
        tmp2 = euler_fluxes(Ub[1][end,k],Ub[2][end,k],Ub[3][end,k],Ub_right[1],Ub_right[2],Ub_right[3])
        for c = 1:3
            # F_high_P[c][1] = -1/2*(flux[c][1,k]+f_left[c])-wavespd_l/2*(U_left[c]-U[c][1,k])#tmp[c]-wavespd_l/2*(U_left[c]-U[c][1,k])#-euler_fluxes(Ub[c][1,k],Ub_left[c])#flux_high(Ub[c][1,k],Ub_left[c],-0.5)-wavespd_l/2*(Ub_left[c]-Ub[c][1,k])
            # F_high_P[c][2] = 1/2*(flux[c][end,k]+f_right[c])-wavespd_r/2*(U_right[c]-U[c][end,k])#tmp2[c]-wavespd_r/2*(U_right[c]-U[c][end,k])#euler_fluxes(Ub[c][end,k],Ub_right[c])#flux_high(Ub[c][end,k],Ub_right[c],0.5)-wavespd_r/2*(Ub_right[c]-Ub[c][end,k])
            F_high_P[c][1] = tmp[c]-wavespd_l/2*(U_left[c]-U[c][1,k])#-euler_fluxes(Ub[c][1,k],Ub_left[c])#flux_high(Ub[c][1,k],Ub_left[c],-0.5)-wavespd_l/2*(Ub_left[c]-Ub[c][1,k])
            F_high_P[c][2] = tmp2[c]-wavespd_r/2*(U_right[c]-U[c][end,k])#euler_fluxes(Ub[c][end,k],Ub_right[c])#flux_high(Ub[c][end,k],Ub_right[c],0.5)-wavespd_r/2*(Ub_right[c]-Ub[c][end,k])
        end

        # construct rhs
        for c = 1:3
            # With limiting
            rhsU[c][:,k] = -sum(F_high[c],dims=2)
            rhsU[c][1,k] += -F_high_P[c][1]
            rhsU[c][N+1,k] += -F_high_P[c][2]
            rhsU[c][:,k] .= 1/J*Mlump_inv*rhsU[c][:,k]
        end

    end

    return rhsU
end

exact_sol = exact_sol_Leblanc

Narr = [2;5]
Karr = [50;100;200;400;800]
low_order_arr = [false]#[true;false]

# Start N, K loop
for N in Narr
for K in Karr
for is_low_order in low_order_arr

"Approximation parameters"
# N = 2 # The order of approximation
# K = 800
t0 = 0.01
# dt = 0.0003

"Mesh related variables"
VX = LinRange(Bl,Br,K+1)
EToV = transpose(reshape(sort([1:K; 2:K+1]),2,K))

"Initialize reference element"
r,_ = gauss_lobatto_quad(0,0,N)   # Reference nodes
VDM = vandermonde_1D(N,r)         # modal to nodal
Dr = grad_vandermonde_1D(N,r)/VDM # nodal differentiation
V1 = vandermonde_1D(1,r)/vandermonde_1D(1,[-1;1]) # nodal linear interpolation

Nq = N
rq,wq = gauss_lobatto_quad(0,0,Nq)
Vq = vandermonde_1D(N,rq)/VDM
M = Vq'*diagm(wq)*Vq
Mlump = zeros(size(M))
Mlump_inv = zeros(size(M))
for i = 1:Nq+1
    Mlump[i,i] = sum(M[i,:])
    Mlump_inv[i,i] = 1.0/Mlump[i,i]
end

# operators
Qr = M*Dr
B = zeros(N+1,N+1)
B[1,1] = -1
B[end,end] = 1
L = Array(spdiagm(0=>-2*ones(N+1), 1=>ones(N), -1=>ones(N)))
L[1,1] = -1
L[end,end] = -1
psi = pinv(L)*-1/2*B*ones(N+1)
S0 = zeros(N+1,N+1)
for i = 1:N+1
    for j = 1:N+1
        if L[i,j] != 0
            S0[i,j] = psi[j] - psi[i]
        end
    end
end
Qr0 = S0+1/2*B

# Drop zeros
Qr = Matrix(droptol!(sparse(Qr),TOL))
Qr0 = Matrix(droptol!(sparse(Qr0),TOL))
B = Matrix(droptol!(sparse(B),TOL))
L = Matrix(droptol!(sparse(L),TOL))
S0 = Matrix(droptol!(sparse(S0),TOL))
S = Matrix(droptol!(sparse((Qr-Qr')/2),TOL))

rf = [-1.0;1.0]
nrJ = [-1.0;1.0]
Vf = vandermonde_1D(N,rf)/VDM

"""High order mesh"""
x = V1*VX[transpose(EToV)]
xf = Vf*x
mapM = reshape(1:2*K,2,K)
mapP = copy(mapM)
mapP[1,2:end] .= mapM[2,1:end-1]
mapP[2,1:end-1] .= mapM[1,2:end]

"""Geometric factors"""
J = repeat(transpose(diff(VX)/2),N+1,1)
nxJ = repeat([-1;1],1,K)
rxJ = 1.0

"""Geometric factors"""
J = repeat(transpose(diff(VX)/2),N+1,1)
nxJ = repeat([-1;1],1,K)
rxJ = 1.0

"""Initial condition"""
rho_x(x) = (x <= xC) ? rhoL : rhoR
u_x(x) = 0.0
p_x(x) = (x <= xC) ? pL : pR

U_init = @. exact_sol(x,t0)
rho_init = [x[1] for x in U_init]
u_init = [x[2] for x in U_init]
p_init = [x[3] for x in U_init]
U = primitive_to_conservative_hardcode.(rho_init,u_init,p_init)
rho = [x[1] for x in U]
rhou = [x[2] for x in U]
E = [x[3] for x in U]
U = (rho,rhou,E)

# Time stepping
"Time integration"
t = t0
U = collect(U)

Vp = vandermonde_1D(N,LinRange(-1,1,10))/VDM
# Nsteps = Int(ceil((T-t0)/dt))
resW = [zeros(size(x)),zeros(size(x)),zeros(size(x))]
resZ = [zeros(size(x)),zeros(size(x)),zeros(size(x))]
#for i = 1:Nsteps
#@gif for i = 1:Nsteps
while t < T

    # if abs(T-t) < dt
    #     global dt = T - t
    # end

    dt = Inf
    # SSPRK(3,3)
    rhsU,dt = rhs_IDP(U,K,N,wq,S,S0,Mlump_inv,T,dt,t,true,is_low_order)
    #rhsU = rhs_high(U,K,N,Mlump_inv,S)
    # rhsU,dt = rhs_IDPlow(U,K,N,Mlump_inv,wq)
    @. resW = U + dt*rhsU
    rhsU,_ = rhs_IDP(resW,K,N,wq,S,S0,Mlump_inv,T,dt,t,false,is_low_order)
    #rhsU = rhs_high(resW,K,N,Mlump_inv,S)
    # rhsU,_ = rhs_IDPlow(resW,K,N,Mlump_inv,wq)
    @. resZ = resW+dt*rhsU
    @. resW = 3/4*U+1/4*resZ
    rhsU,_ = rhs_IDP(resW,K,N,wq,S,S0,Mlump_inv,T,dt,t,false,is_low_order)
    #rhsU = rhs_high(resW,K,N,Mlump_inv,S)
    # rhsU,_ = rhs_IDPlow(resW,K,N,Mlump_inv,wq)
    @. resZ = resW+dt*rhsU
    @. U = 1/3*U+2/3*resZ

    t = t + dt
    println("Current time $t with time step size $dt, and final time $T")
    if t == T
        break
    end

    # if i % GIFINTERVAL == 0
    #     #plot(Vp*x,Vp*U[1])
    #     plot(Vp*x,Vp*U[1])
    #     # ptL = Bl+(Br-Bl)/K/(N+1)/2
    #     # ptR = Br-(Br-Bl)/K/(N+1)/2
    #     # hplot = (Br-Bl)/K/(N+1)
    #     # for k = 1:K
    #     #     plot!(ptL+(k-1)*hplot*(N+1):hplot:ptL+k*hplot*(N+1), 1 .-L_plot2[:,k],st=:bar,alpha=0.2)
    #     # end
    # end
end
#end every GIFINTERVAL
#end

@inline function Efun(rho,u,p)
    return p/(γ-1) + .5*rho*(u^2)
end

exact_U = @. exact_sol(x,T)
exact_rho = [x[1] for x in exact_U]
exact_u = [x[2] for x in exact_U]
exact_p = [x[3] for x in exact_U]
exact_rhou = exact_rho .* exact_u
exact_E = Efun.(exact_rho,exact_u,exact_p)

rho = U[1]
u = U[2]./U[1]
p = pfun_nd.(U[1],U[2],U[3])
rhou = U[2]
E = U[3]
J = (Br-Bl)/K/2

Linferr = maximum(abs.(exact_rho-rho))/maximum(abs.(exact_rho)) +
          maximum(abs.(exact_rhou-rhou))/maximum(abs.(exact_rhou)) +
          maximum(abs.(exact_E-E))/maximum(abs.(exact_E))

L1err = sum(J*wq.*abs.(exact_rho-rho))/sum(J*wq.*abs.(exact_rho)) +
        sum(J*wq.*abs.(exact_rhou-rhou))/sum(J*wq.*abs.(exact_rhou)) +
        sum(J*wq.*abs.(exact_E-E))/sum(J*wq.*abs.(exact_E))

L2err = sqrt(sum(J*wq.*abs.(exact_rho-rho).^2))/sqrt(sum(J*wq.*abs.(exact_rho).^2)) +
        sqrt(sum(J*wq.*abs.(exact_rhou-rhou).^2))/sqrt(sum(J*wq.*abs.(exact_rhou).^2)) +
        sqrt(sum(J*wq.*abs.(exact_E-E).^2))/sqrt(sum(J*wq.*abs.(exact_E).^2))

println("N = $N, K = $K")
println("L1 error is $L1err")
println("L2 error is $L2err")
println("Linf error is $Linferr")

# open("/data/yl184/1D-euler-conv/N=$N,K=$K,CFL=$CFL,LOW=$is_low_order,x,1D-euler-conv.txt","w") do io
#     writedlm(io,x)
# end
# open("/data/yl184/1D-euler-conv/N=$N,K=$K,CFL=$CFL,LOW=$is_low_order,rho,1D-euler-conv.txt","w") do io
#     writedlm(io,U[1])
# end

# df = DataFrame(N = Int64[], K = Int64[], T = Float64[], CFL = Float64[], LIMITOPT = Int64[], POSDETECT = Int64[], LBOUNDTYPE = Float64[], L1err = Float64[], L2err = Float64[], Linferr = Float64[])
df = load("dg1D_euler_convergence.jld2","convergence_data")
push!(df,(N,K,T,CFL,LIMITOPT,POSDETECT,LBOUNDTYPE,L1err,L2err,Linferr))
save("dg1D_euler_convergence.jld2","convergence_data",df)

end
end
end # End N, K loop

# println("Linf error is $Linferr")

# open("convergence.txt", "a") do io
#     write(io, "N = $N, K = $K\n")
#     write(io, "L1 error is $L1err\n")
#     write(io, "Linf error is $Linferr\n")
# end

# plot(Vp*x,Vp*U[1])
# savefig("~/Desktop/N=$N,K=$K,dt=$dt,modifiedESDG.png")

max_stars_repo_name: yiminllin/ESDG-PosLimit

max_issues_repo_name: yiminllin/ESDG-PosLimit

max_forks_repo_name: yiminllin/ESDG-PosLimit