text:
#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#    SatelliteToolbox orbit propagator API: init_orbit_propagator
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Hoots, F. R., Roehrich, R. L (1980). Models for Propagation of NORAD
#       Elements Set. Spacetrack Report No. 3.
#
#   [2] Vallado, D. A., Crawford, P., Hujsak, R., Kelso, T. S (2006). Revisiting
#       Spacetrack Report #3: Rev1. AIAA.
#
#   [3] SGP4 Source code of STRF: https://github.com/cbassa/strf
#       The SGP4 C code available on STRF was converted by Paul. S. Crawford and
#       Andrew R. Brooks.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export init_orbit_propagator

"""
    init_orbit_propagator(T, epoch::Number, a_0::Number, e_0::Number, i_0::Number, Ω_0::Number, ω_0::Number, f_0::Number, ...)
    init_orbit_propagator(T, orb_0::Orbit, ...)

Initialize the orbit propagator `T` using the initial mean orbital elements. The
propagator type `T` can be:

* `Val(:J2)`: J2 orbit propagator;
* `Val(:J4)`: J4 orbit propagator; or
* `Val(:twobody)`: Two-body orbit propagator.

The mean orbital elements can be passed individually of using an instance of the
structure `Orbit`.

# Args

* `epoch`: Initial orbit epoch [Julian Day].
* `a_0`: Initial semi-major axis [m].
* `e_0`: Initial eccentricity.
* `i_0`: Initial inclination [rad].
* `Ω_0`: Initial right ascension of the ascending node [rad].
* `ω_0`: Initial argument of perigee [rad].
* `f_0`: Initial true anomaly [rad].
* `n_0`: Initial angular velocity [rad/s].
* `M_0`: Initial mean anomaly [rad].
* `orb_0`: Instance of the structure `Orbit` with the initial mean orbital
           elements [SI].

## Additional optional arguments for the J2 orbit propagator

The initialization function of the J2 orbit propagator can receive the
following optional parameters:

* `dn_o2`: (OPTIONAL) First time derivative of mean motion divided by 2
           \\[rad/s²] (**Default** = 0).
* `ddn_o6`: (OPTIONAL) Second time derivative of mean motion divided by 6
            \\[rad/s³] (**Default** = 0).
* `j2_gc`: (OPTIONAL) J2 orbit propagator gravitational constants
           (**Default** = `j2_gc_egm08`).

## Additional optional arguments for the J4 orbit propagator

The initialization function of the J4 orbit propagator can receive the
following optional parameters:

* `dn_o2`: (OPTIONAL) First time derivative of mean motion divided by 2
           \\[rad/s²] (**Default** = 0).
* `ddn_o6`: (OPTIONAL) Second time derivative of mean motion divided by 6
            \\[rad/s³] (**Default** = 0).
* `j4_gc`: (OPTIONAL) J4 orbit propagator gravitational constants
           (**Default** = `j4_gc_egm08`).

## Additional optional arguments for the two body orbit propagator

The initialization function of the two body orbit propagator can receive the
following optional parameter:

* `μ`: (OPTIONAL) Standard gravitational parameter of the central body
       \\[m^3/s^2] (**Default** = `m0`).

# Returns

A new instance of the orbit propagator structure that stores the information of
the orbit propagator.

# Remarks

If the orbit is defined in terms of the angular velocity (mean motion) instead
of the semi-major axis, then it is possible to use the function `angvel_to_a` to
convert.

"""
function init_orbit_propagator(::Val{:J2}, epoch::Number, a_0::Number,
                               e_0::Number, i_0::Number, Ω_0::Number,
                               ω_0::Number, f_0::Number, dn_o2::Number = 0,
                               ddn_o6::Number = 0,
                               j2_gc::J2_GravCte{T} = j2_gc_egm08) where T

    # Create the new Two Body propagator structure.
    j2d = j2_init(j2_gc, epoch, a_0, e_0, i_0, Ω_0, ω_0, f_0, dn_o2, ddn_o6)

    # Create the `Orbit` structure.
    orb_0 = Orbit(epoch, j2d.al_0*j2_gc.R0, e_0, i_0, Ω_0, ω_0, j2d.f_k)

    # Create and return the orbit propagator structure.
    OrbitPropagatorJ2(orb_0, j2d)
end

function init_orbit_propagator(::Val{:J4}, epoch::Number, a_0::Number,
                               e_0::Number, i_0::Number, Ω_0::Number,
                               ω_0::Number, f_0::Number, dn_o2::Number = 0,
                               ddn_o6::Number = 0,
                               j4_gc::J4_GravCte{T} = j4_gc_egm08) where T

    # Create the new J4 propagator structure.
    j4d = j4_init(j4_gc, epoch, a_0, e_0, i_0, Ω_0, ω_0, f_0, dn_o2, ddn_o6)

    # Create the `Orbit` structure.
    orb_0 = Orbit(epoch, j4d.al_0*j4_gc.R0, e_0, i_0, Ω_0, ω_0, j4d.f_k)

    # Create and return the orbit propagator structure.
    return OrbitPropagatorJ4(orb_0, j4d)
end

function init_orbit_propagator(::Val{:twobody}, epoch::Number,
                               a_0::Number, e_0::Number, i_0::Number,
                               Ω_0::Number, ω_0::Number, f_0::Number,
                               μ::T = m0) where T

    # Create the new Two Body propagator structure.
    tbd = twobody_init(epoch, a_0, e_0, i_0, Ω_0, ω_0, f_0, μ)

    # Create the `Orbit` structure.
    orb_0 = Orbit(epoch, a_0, e_0, i_0, Ω_0, ω_0, f_0)

    # Create and return the orbit propagator structure.
    return OrbitPropagatorTwoBody(orb_0, tbd)
end

init_orbit_propagator(::Val{:J2}, orb_0::Orbit, dn_o2::Number = 0,
                      ddn_o6::Number = 0, j2_gc::J2_GravCte = j2_gc_egm08) =
    init_orbit_propagator(Val(:J2), orb_0.t, orb_0.a, orb_0.e, orb_0.i, orb_0.Ω,
                          orb_0.ω, orb_0.f, dn_o2, ddn_o6, j2_gc)

init_orbit_propagator(::Val{:J4}, orb_0::Orbit, dn_o2::Number = 0,
                      ddn_o6::Number = 0, j4_gc::J4_GravCte = j4_gc_egm08) =
    init_orbit_propagator(Val(:J4), orb_0.t, orb_0.a, orb_0.e, orb_0.i, orb_0.Ω,
                          orb_0.ω, orb_0.f, dn_o2, ddn_o6, j4_gc)

init_orbit_propagator(::Val{:twobody}, orb_0::Orbit, μ::Number = m0) =
    init_orbit_propagator(Val(:twobody), orb_0.t, orb_0.a, orb_0.e, orb_0.i,
                          orb_0.Ω, orb_0.ω, orb_0.f, μ)

"""
    init_orbit_propagator(T, tle::TLE, ...)

Initialize the orbit propagator `T` using the TLE `tle`. The propagator type `T`
can be:

* `Val(:J2)`: J2 orbit propagator;
* `Val(:J4)`: J4 orbit propagator;
* `Val(:twobody)`: Two-body orbit propagator; or
* `Val(:sgp4)`: SGP4 orbit propagator.

## Additional optional arguments for the J2 orbit propagator

The initialization function of the J2 orbit propagator can receive the
following optional parameter:

* `j2_gc`: (OPTIONAL) J2 orbit propagator gravitational constants
           (**Default** = `j2_gc_egm08`).

## Additional optional arguments for the J4 orbit propagator

The initialization function of the J4 orbit propagator can receive the
following optional parameter:

* `j4_gc`: (OPTIONAL) J4 orbit propagator gravitational constants
           (**Default** = `j4_gc_egm08`).

## Additional optional arguments for the two body orbit propagator

The initialization function of the two body orbit propagator can receive the
following optional parameter:

* `μ`: (OPTIONAL) Standard gravitational parameter of the central body
       \\[m^3/s^2] (**Default** = `m0`).

## Additional optional arguments for the SGP4 orbit propagator

The initialization function of the SGP4 orbit propagator can receive the
following optional parameter:

* `sgp4_gc`: (OPTIONAL) Gravitational constants (**Default** = `sgp4_gc_wgs84`).

# Returns

A new instance of the orbit propagator structure that stores the information of
the orbit propagator.

# Remarks

The SGP4 implementation includes also the deep space perturbations, which was
originally called SDP4 algorithm. Modern approaches, such as [2] and [3],
identifies if the selected orbit must be propagated using the deep space
perturbations and automatically applied them. This is sometimes called SGDP4
algorithm.

"""
function init_orbit_propagator(::Val{:J2}, tle::TLE,
                               j2_gc::J2_GravCte = j2_gc_egm08)

    # Unpack the gravitational constants to improve code readability.
    @unpack_J2_GravCte j2_gc

    # Constants.
    revday2radsec = 2π/86400    # Revolutions per day to radians per second.
    d2r           =  π/180      # Degrees to radians.

    # Obtain the data from the TLE.
    n_0 = tle.n*revday2radsec
    e_0 = tle.e
    i_0 = tle.i*d2r
    Ω_0 = tle.Ω*d2r
    ω_0 = tle.ω*d2r
    M_0 = tle.M*d2r

    # Auxiliary variables.
    e_0²    = e_0^2
    cos_i_0 = cos(i_0)

    # Recover the original mean motion (nll_0) and semi-major axis (all_0) from
    # the input elements and the same algorithm of SGP4.
    aux   = 1/2*J2*(3*cos_i_0^2-1)/(1-e_0²)^(3/2)
    a_1   = (μm/n_0)^(2/3)
    δ_1   = 3/2*aux/a_1^2
    a_0   = a_1*(1 - 1/3*δ_1 - δ_1^2 - 134/81*δ_1^3)
    δ_0   = 3/2*aux/a_0^2
    nll_0 = n_0/(1 + δ_0)
    a_0   = (μm/nll_0)^(2/3)*R0
    f_0   = M_to_f(e_0, M_0)

    return init_orbit_propagator(Val(:J2), tle.epoch, a_0, e_0, i_0, Ω_0, ω_0,
                          f_0, tle.dn_o2*2π/(86400)^2, tle.ddn_o6*2π/(86400)^3,
                          j2_gc_egm08)
end

function init_orbit_propagator(::Val{:J4}, tle::TLE,
                               j4_gc::J4_GravCte = j4_gc_egm08)

    # Unpack the gravitational constants to improve code readability.
    @unpack_J2_GravCte j4_gc

    revday2radsec = 2π/86400    # Revolutions per day to radians per second.
    # Constants.
    d2r           =  π/180      # Degrees to radians.

    # Obtain the data from the TLE.
    n_0 = tle.n*revday2radsec
    e_0 = tle.e
    i_0 = tle.i*d2r
    Ω_0 = tle.Ω*d2r
    ω_0 = tle.ω*d2r
    M_0 = tle.M*d2r

    # Auxiliary variables.
    e_0²    = e_0^2
    cos_i_0 = cos(i_0)

    # Recover the original mean motion (nll_0) and semi-major axis (all_0) from
    # the input elements and the same algorithm of SGP4.
    aux   = 1/2*J2*(3*cos_i_0^2-1)/(1-e_0²)^(3/2)
    a_1   = (μm/n_0)^(2/3)
    δ_1   = 3/2*aux/a_1^2
    a_0   = a_1*(1 - 1/3*δ_1 - δ_1^2 - 134/81*δ_1^3)
    δ_0   = 3/2*aux/a_0^2
    nll_0 = n_0/(1 + δ_0)
    a_0   = (μm/nll_0)^(2/3)*R0
    f_0   = M_to_f(e_0, M_0)

    return init_orbit_propagator(Val(:J4), tle.epoch, a_0, e_0, i_0, Ω_0, ω_0,
                                 f_0, tle.dn_o2*2π/(86400)^2,
                                 tle.ddn_o6*2π/(86400)^3, j4_gc_egm08)
end

function init_orbit_propagator(::Val{:twobody}, tle::TLE, μ::Number = m0)
    # Constants.
    revday2radsec = 2π/86400    # Revolutions per day to radians per second.
    d2r           =  π/180      # Degrees to radians.

    # Obtain the data from the TLE.
    n_0 = tle.n*revday2radsec
    e_0 = tle.e
    i_0 = tle.i*d2r
    Ω_0 = tle.Ω*d2r
    ω_0 = tle.ω*d2r
    M_0 = tle.M*d2r

    # Since we do not have in the Two Body algorithms the coefficient `J2`, then
    # the semi-major axis will be recovered considering a Keplerian orbit.

    a_0 = (μ/n_0^2)^(1/3)
    f_0 = M_to_f(e_0, M_0)

    init_orbit_propagator(Val(:twobody), tle.epoch, a_0, e_0, i_0, Ω_0, ω_0,
                          f_0, μ)
end

function init_orbit_propagator(::Val{:sgp4}, tle::TLE,
                               sgp4_gc::SGP4_GravCte{T} = sgp4_gc_wgs84) where T

    # Constants.
    d2r           =  π/180      # Degrees to radians.
    revday2radmin = 2π/1440     # Revolutions per day to radians per minute.

    # Obtain the data from the TLE.
    epoch = tle.epoch
    n_0   = tle.n*revday2radmin
    e_0   = tle.e
    i_0   = tle.i*d2r
    Ω_0   = tle.Ω*d2r
    ω_0   = tle.ω*d2r
    M_0   = tle.M*d2r
    bstar = tle.bstar

    # Create the new SGP4 structure.
    sgp4d = sgp4_init(sgp4_gc, tle.epoch, n_0, e_0, i_0, Ω_0, ω_0, M_0,
                      tle.bstar)

    # Create the `Orbit` structure.
    orb_0 = Orbit(epoch, 1000sgp4d.a_k*sgp4_gc.R0, e_0, i_0, Ω_0, ω_0,
                  M_to_f(e_0, M_0))

    # Create and return the orbit propagator structure.
    return OrbitPropagatorSGP4(orb_0, sgp4_gc, sgp4d)
end


text:
@doc raw"""
    InviscidBurgersEquation1D

The inviscid Burgers' equation
```math
\partial_t u + \frac{1}{2} \partial_1 u^2 = 0
```
in one space dimension.
"""
struct InviscidBurgersEquation1D <: AbstractInviscidBurgersEquation{1, 1} end

get_name(::InviscidBurgersEquation1D) = "InviscidBurgersEquation1D"
varnames(::typeof(cons2cons), ::InviscidBurgersEquation1D) = ("scalar", )
varnames(::typeof(cons2prim), ::InviscidBurgersEquation1D) = ("scalar", )

# Set initial conditions at physical location `x` for time `t`
"""
    initial_condition_constant(x, t, equations::InviscidBurgersEquation1D)

A constant initial condition to test free-stream preservation.
"""
function initial_condition_constant(x, t, equation::InviscidBurgersEquation1D)
  return SVector(2.0)
end

"""
    initial_condition_convergence_test(x, t, equations::InviscidBurgersEquation1D)

A smooth initial condition used for convergence tests.
"""
function initial_condition_convergence_test(x, t, equation::InviscidBurgersEquation1D)
  c = 2.0
  A = 1.0
  L = 1
  f = 1/L
  omega = 2 * pi * f
  scalar = c + A * sin(omega * (x[1] - t))

  return SVector(scalar)
end

"""
    source_terms_convergence_test(u, x, t, equations::InviscidBurgersEquation1D)

Source terms used for convergence tests in combination with
[`initial_condition_convergence_test`](@ref).
"""
@inline function source_terms_convergence_test(u, x, t, equations::InviscidBurgersEquation1D)
  # Same settings as in `initial_condition`
  c = 2.0
  A = 1.0
  L = 1
  f = 1/L
  omega = 2 * pi * f
  du = omega * cos(omega * (x[1] - t)) * (1 + sin(omega * (x[1] - t)))

  return SVector(du)
end

# Pre-defined source terms should be implemented as
# function source_terms_WHATEVER(u, x, t, equations::InviscidBurgersEquation1D)

# Calculate 1D flux in for a single point
@inline function flux(u, orientation, equation::InviscidBurgersEquation1D)
  return SVector(0.5 * u[1]^2)
end

function flux_lax_friedrichs(u_ll, u_rr, orientation, equation::InviscidBurgersEquation1D)
  u_L = u_ll[1]
  u_R = u_rr[1]

  return SVector(0.5 * ( 0.5 * (u_L^2 + u_R^2) - max(abs(u_L), abs(u_R)) * (u_R - u_L) ))
end

@inline function max_abs_speeds(u, equation::InviscidBurgersEquation1D)
  return (abs(u[1]),)
end

function flux_ec(u_ll, u_rr, orientation, equation::InviscidBurgersEquation1D)
  u_L = u_ll[1]
  u_R = u_rr[1]

  return SVector((u_L^2 + u_L * u_R + u_R^2) / 6)
end

# Convert conservative variables to primitive
@inline cons2prim(u, equation::InviscidBurgersEquation1D) = u

# Convert conservative variables to entropy variables
@inline cons2entropy(u, equation::InviscidBurgersEquation1D) = u

# Calculate entropy for a conservative state `cons`
@inline entropy(u::Real, ::InviscidBurgersEquation1D) = 0.5 * u^2
@inline entropy(u, equation::InviscidBurgersEquation1D) = entropy(u[1], equation)

# Calculate total energy for a conservative state `cons`
@inline energy_total(u::Real, ::InviscidBurgersEquation1D) = 0.5 * u^2
@inline energy_total(u, equation::InviscidBurgersEquation1D) = energy_total(u[1], equation)

ext: jl


text:
module CompositeGates

export xx,yy,zz,swap,iswap

using QXZoo.GateOps
using QXZoo.DefaultGates
using QXZoo.GateMap
using StaticArrays

module GateSymbols
    using QXZoo.GateOps

    """Two qubit"""
    swap = GateOps.GateSymbol(:swap)
    iswap = GateOps.GateSymbol(:iswap)

    xx(θ) = GateOps.GateSymbolP(:xx, 0, false, θ)
    yy(θ) = GateOps.GateSymbolP(:yy, 0, false, θ)
    zz(θ) = GateOps.GateSymbolP(:zz, 0, false, θ)

    """Three qubit"""
    c_c_x = GateOps.GateSymbol(:c_c_x)
    c_swap = GateOps.GateSymbol(:c_swap)
end

"""
    xx(q_target1::Int, q_target2::Int, theta::Real)

Generate a XX parametric gate-equivalent circuit, composed of fundamental gate operations (H, CX, Rz)

# Examples
```julia-repl
```
"""
function xx(q_target1::Int, q_target2::Int, θ::Number)
    if ~haskey(GateMap.gates, GateSymbols.xx(θ)) # Build XX matrix if not available
        GateMap.cache_gate!(GateSymbols.xx(θ), (θ)->exp(1im*θ*(kron(GateMap.x(),GateMap.x()))))
    end
    return GateOps.GateCall2(GateSymbols.xx(θ), q_target1, q_target2, nothing)
end

"""
    yy(q_target1::Int, q_target2::Int, theta::Real)

Generate a YY parametric gate, which can be constructed from fundamental 1 & 2 qubit operations

# Examples
```julia-repl
```
"""
function yy(q_target1::Int, q_target2::Int, θ::Number)
    if ~haskey(GateMap.gates, GateSymbols.yy(θ)) # Build YY matrix if not available
        GateMap.cache_gate!(GateSymbols.yy(θ), (θ)->exp(1im*θ*(kron(GateMap.y(),GateMap.y()))))
    end
    return GateOps.GateCall2(GateSymbols.yy(θ), q_target1, q_target2, nothing)
end

"""
    zz(q_target1::Int, q_target2::Int, theta::Real)

Generate a ZZ parametric gate, which can be constructed from fundamental 1 & 2 qubit operations

# Examples
```julia-repl
```
"""
function zz(q_target1::Int, q_target2::Int, θ::Number)
    if ~haskey(GateMap.gates, GateSymbols.zz(θ)) # Build YY matrix if not available
        GateMap.cache_gate!(GateSymbols.zz(θ), (θ)->exp(1im*θ*(kron(GateMap.z(),GateMap.z()))))
    end
    return GateOps.GateCall2(GateSymbols.zz(θ), q_target1, q_target2, nothing)
end

"""
    swap(q_target1::Int, q_target2::Int)

Generate a SWAP gate equivalent circuit, composed of fundamental gate operations (CX)

# Examples
```julia-repl
```
"""
function swap(q_target1::Int, q_target2::Int)
    if ~haskey(GateMap.gates, GateSymbols.swap) # Build SWAP matrix if not available
        GateMap.cache_gate!(GateSymbols.swap, ()->
            @SMatrix ([1 0 0 0; 0 0 1.0+0im 0; 0 1.0+0im 0 0; 0 0 0 1])
        )
    end
    return GateOps.GateCall2(GateSymbols.swap, q_target1, q_target2, nothing)
end

"""
    iswap(q_target1::Int, q_target2::Int)

Generate a iSWAP gate equivalent circuit, generatable using
`exp(1im*(pi/4)*(kron(GateMap.x(),GateMap.x()) + kron(GateMap.y(),GateMap.y())))`

# Examples
```julia-repl
```
"""
function iswap(q_target1::Int, q_target2::Int)
    if ~haskey(GateMap.gates, GateSymbols.iswap) # Build iSWAP matrix if not available
        GateMap.cache_gate!(GateSymbols.iswap, ()->
            @SMatrix ([1 0 0 0; 0 0 1.0im 0; 0 1.0im 0 0; 0 0 0 1])
        )
    end
    return GateOps.GateCall2(GateSymbols.iswap, q_target1, q_target2, nothing)
end

end


text:
@defcomp co2emissions begin
    region = Index()

    e_globalCO2emissions = Variable(index=[time], unit="Mtonne/year")
    e0_baselineCO2emissions = Parameter(index=[region], unit="Mtonne/year")
    e_regionalCO2emissions = Variable(index=[time,region], unit="Mtonne/year")
    er_CO2emissionsgrowth = Parameter(index=[time,region], unit="%")

    # read in counterfactual GDP in absence of growth effects (gdp_leveleffects) and actual GDP
    gdp = Parameter(index=[time, region], unit="\$M")
    gdp_leveleffect   = Parameter(index=[time, region], unit="\$M")
    emfeed_emissionfeedback = Parameter(unit="none", default=1.)

    function run_timestep(p, v, d, t)

        # eq.4 in Hope (2006) - regional CO2 emissions as % change from baseline
        for r in d.region
            v.e_regionalCO2emissions[t,r] = p.er_CO2emissionsgrowth[t,r] * p.e0_baselineCO2emissions[r] / 100

            # rescale emissions based on GDP deviation from original scenario pathway
            if p.emfeed_emissionfeedback == 1.
                v.e_regionalCO2emissions[t,r] = v.e_regionalCO2emissions[t,r] * (p.gdp[t,r] / p.gdp_leveleffect[t,r])
            end
        end
        # eq. 5 in Hope (2006) - global CO2 emissions are sum of regional emissions
        v.e_globalCO2emissions[t] = sum(v.e_regionalCO2emissions[t,:])
    end
end


text:
# Timing methods for the whole function
function gettime(fex::nGraph.FluxExecutable; timeout = Second(5), min_calls = 3)
    start = now()
    mintime = typemax(Float64)
    times = 1
    while (now() < start + timeout) || (times <= min_calls)
        @info "Running Function"
        runtime = @elapsed(nGraph.fetch(fex()))
        @info "Done Running Function"
        mintime = min(mintime, runtime)
        times += 1
    end
    return mintime
end

function get_baseline_allocation(backend, f)
    allocation_ref = Ref{Int}(0)
    io_ref = Ref{Int}(0)
    callbacks = CallbackChain()

    # For the first callback, we want to set all intermediate algorithms to their fastest
    # setting.
    function set(f::nGraph.NFunction)
        data = profile(f, backend)
        for node in nodes(data)
            if nGraph.Lib.can_select_algo(nGraph.getpointer(unx(node)))
                # Find the minimum runtime of the algorithms
                runtimes = get_times(gettime(data, node))
                _, ind = findmin(runtimes)
                enum = get_enums(gettime(data, node))[ind]

                nGraph.Lib.set_algo(
                    nGraph.getpointer(node),
                    convert(UInt, enum),
                    convert(UInt, get_bytes(gettime(data, node), enum)),
                )
            end
        end
    end

    function ret(f::nGraph.NFunction)
        allocation_ref[] = nGraph.get_temporary_pool_size(f)
        io_ref[] = sum(sizeof, input_tensors(f)) + sum(sizeof, output_tensors(f))
        throw(CompilerExit())
    end

    callback!(callbacks, set)
    callback!(callbacks, ret)

    try
        return actualize(backend, f; callback = callbacks)
    catch e
        isa(e, CompilerExit) || rethrow(e)
    end

    return allocation_ref[], io_ref[]
end

_base_stats() = (
    io_size = Ref(0),
    default_alloc_size = Ref(0),
    # GPU runtime usinc CUDA Malloc Managed
    gpu_managed_runtime = Ref(0.0),
    runs = Vector{Dict{Symbol,Any}}(),
)

"""
$(SIGNATURES)

* `f`: A constructor function `f() -> fex, args` returning a `FluxExecutable` and a tuple
    of arguments to be passed to the executable.

* `opt_iter`: An iterator returning optimization arguments that will be passed to `factory`.

Keywords
--------

* `cache`: A cache for kernel timings. Defaults to `CPUKernelCache(BASE_CACHE_PATH)`. That
    is, the default cache.

* `statspath`: An optional file path to saved stats. If this is given, any DRAM limits
    already in the cached stats will be skipped on this profiling run.
"""
function compare(
        func,
        opt,
        backend::nGraph.Backend;
        skip_base_check = false,
        statspath = nothing,
        kw...
    )

    if (isnothing(statspath) || !ispath(statspath))
        stats = _base_stats()
        if !skip_base_check
            initialize!(stats, func, backend)
        end
    else
        stats = deserialize(statspath)
    end

    GC.gc()
    _compare!(
        stats,
        func,
        opt,
        backend;
        kw...
    )

    if !isempty(stats.runs)
        @info """
        Predicted Run Time: $(last(stats.runs)[:predicted_runtime])
        Actual Run Time: $(last(stats.runs)[:actual_runtime])
        """
    end

    sort!(stats.runs; rev = true, by = x -> x[:dram_limit])
    isnothing(statspath) || serialize(statspath, stats)

    return stats
end

function initialize!(stats, func, backend::nGraph.Backend{nGraph.CPU})
    # Instantiate the function
    fex = actualize(backend, func)

    stats.io_size[] = sum(sizeof, input_tensors(fex)) + sum(sizeof, output_tensors(fex))
    stats.default_alloc_size[] = nGraph.get_temporary_pool_size(fex.ex.ngraph_function)
    return nothing
end

# To initialize the GPU stuff, we turn on the Managed Memory flag and compile the function
function initialize!(stats, func, backend::nGraph.Backend{nGraph.GPU})
    fex = withenv("NGRAPH_GPU_CUDA_MALLOC_MANAGED" => true) do
        actualize(backend, func)
    end

    stats.io_size[] = sum(sizeof, input_tensors(fex)) + sum(sizeof, output_tensors(fex))
    stats.default_alloc_size[] = nGraph.get_temporary_pool_size(fex.ex.ngraph_function)
    stats.gpu_managed_runtime[] = gettime(fex)
end

# Overload point for optimizers
_compare!(args...; kw...) = error("_compare! not implemented for $(typeof.(args))")

max_stars_repo_name: darchr/AutoTM

max_issues_repo_name: darchr/AutoTM

max_forks_repo_name: darchr/AutoTM


text:
# import CPEG

p = CPEG.CPEGDensityParameters()

# this data came from a marsgramm output
hs = [5*1e3, 30e3, 50e3, 75e3, 100e3]
ρs = [8.277E-03, 7.555E-04, 7.081E-05, 5.867E-06, 3.612E-07]

for i = 1:length(hs)

    @test log(ρs[i]) ≈ log(CPEG.density(p,hs[i])) rtol = 1e-2

end


text:
using .Pathfinder

"""
    `pathfinder(post::Posterior, ndraws::Int; init_params=nothing, kwargs...)`

Estimates the posterior using the path finder algorithm from [Zhang 2021](https://arxiv.org/abs/2108.03782).

# Arguments

- post::Posterior: is the posterior you want to fit
- ndraws::Int: is the number of draws from the fit to the posterior
- init_params: Initial location to start pathfinder. If nothing the starting location
               will be a random draw from the prior.

# Returns

- q::Distributions::MvNormal: ELBO maximizing variational approximation to the posterior.
                              **Note** this is the approximation in the transformed space,
                              not the original model parameter space.
- ϕ::TupleVector: `ndraws` from th variational estimated draws from the target distribution
- logqϕ::Vector{Int}: logdensity of draws `ϕ`

# Notes
For a list of potential keyword arguments please the the [Pathfinder.jl](https://sethaxen.github.io/Pathfinder.jl/dev/)
documentation.
"""
function Pathfinder.pathfinder(post::Posterior,
                               ndraws::Int;
                               init_params=nothing,
                               kwargs...)
    # transform to unit space
    tpost = asflat(post)

    # set initial location, if nothing provided grab a draw from the prior
    p0 = init_params
    if isnothing(p0)
        p0 = rand(post.prior)
    end
    θ0 = HypercubeTransform.inverse(tpost, p0)
    ℓ(x) = logdensityof(tpost, x)

    q, ϕ, logqϕ = pathfinder(ℓ, θ0, ndraws; kwargs...)
    out = transform_and_logjac.(Ref(tpost.transform), eachcol(ϕ))
    tϕ = first.(out)
    tlogqϕ = logqϕ .+ last.(out)
    return q, TupleVector(tϕ), tlogqϕ
end

"""
    `mulitpathfinder(post::Posterior, ndraws::Int; init_params=50, kwargs...)`

Estimates the posterior using the path finder algorithm from [Zhang 2021](https://arxiv.org/abs/2108.03782).

# Arguments

- post::Posterior: is the posterior you want to fit
- ndraws::Int: is the number of draws from the fit to the posterior
- init_params: If an integers this is the number of pathfinder algorithms you want to run.
               Otherwise, it must be a Vector where each element is a starting location of
               the model.

# Returns

- q::Distributions::MixtureModel: Uniformly weights mixture of ELBO-maximizing MvNormals.
                                  **Note** this normal approximation is done in the transformed
                                  space, not the model parameter space!
- ϕ::TupleVector: `ndraws` from th variational estimated draws from the target distribution
- component_inds::Vector{Int}: Indices `k` of components in `q` from which each element in
                               ϕ is drawn.

# Notes
For a list of potential keyword arguments please the the [Pathfinder.jl](https://sethaxen.github.io/Pathfinder.jl/dev/)
documentation.
"""
function Pathfinder.multipathfinder(post::Posterior, ndraws::Int; init_params=50, kwargs...)
    tpost = asflat(post)
    ℓ(x) = logdensityof(tpost, x)

    # if init_params is an integer then we draw init_params from the prior
    if init_params isa Integer
        θ0 = HypercubeTransform.inverse.(Ref(tpost.transform), rand(post.prior, init_params))
    else
        θ0 = HypercubeTransform.inverse.(Ref(tpost.transform), init_params)
    end

    q, ϕ, inds = multipathfinder(ℓ, θ0, ndraws; kwargs...)
    out = transform.(Ref(tpost), eachcol(ϕ))
    return q, TupleVector(out), inds
end

"""
    `mulitpathfinder(post::TransformedPosterior, ndraws::Int; init_params=50, kwargs...)`

Estimates the posterior using the path finder algorithm from [Zhang 2021](https://arxiv.org/abs/2108.03782).

# Arguments

- post::Posterior: is the posterior you want to fit
- ndraws::Int: is the number of draws from the fit to the posterior
- init_params: If an integers this is the number of pathfinder algorithms you want to run.
               Otherwise, it must be a Vector where each element is a starting location of
               the model.

# Returns

- q::Distributions::MixtureModel: Uniformly weights mixture of ELBO-maximizing MvNormals.
                                  **Note** this normal approximation is done in the transformed
                                  space, not the model parameter space!
- ϕ::TupleVector: `ndraws` from th variational estimated draws from the target distribution
- component_inds::Vector{Int}: Indices `k` of components in `q` from which each element in
                               ϕ is drawn.

# Notes
For a list of potential keyword arguments please the the [Pathfinder.jl](https://sethaxen.github.io/Pathfinder.jl/dev/)
documentation.
"""
function Pathfinder.multipathfinder(tpost::TransformedPosterior, ndraws::Int; init_params=50, kwargs...)
    ℓ(x) = logdensityof(tpost, x)

    # if init_params is an integer then we draw init_params from the prior
    if init_params isa Integer
        θ0 = inverse.(Ref(tpost.transform), rand(post.prior, init_params))
    else
        θ0 = init_params
    end

    #println(θ0)
    q, ϕ, inds = multipathfinder(ℓ, θ0, ndraws; kwargs...)
    out = collect(eachcol(ϕ))
    return q, out, inds
end


text:
realified_coefficient_type(::Type{Complex{T}}) where T<:Real = promote_type(Float64, T)
realified_coefficient_type(::Type{T}) where T<:Real = promote_type(Float64, T)

"""
    realify(p::AbstractPolynomial)

Compute the realification of `p`. This is ``p(z) = p(x+iy) = p^{re}(x, y) + i p^{im}(x,y)``.
"""
function realify(p::MP.AbstractPolynomial, reals, imags)
    reals_imags = map((x, y) -> x + im*y, reals, imags)
    tmp = MP.subs(p, MP.variables(p) => reals_imags)

    parts = MP.map(MP.terms(tmp)) do term
        c = MP.coefficient(term)
        c_re = real(c)
        c_im = imag(c)
        (c_re * MP.monomial(term), c_im * MP.monomial(term))
    end

    f_re = MP.polynomial(first.(parts))
    f_im = MP.polynomial(last.(parts))

    f_re, f_im
end

function realify(p::MP.AbstractPolynomial)
    realify(p, realifyvariables(p)...)
end

function realifyvariables(p::MP.AbstractPolynomial)
    if MP.nvariables(p) == 2
        MP.@similarvariable p x1
        MP.@similarvariable p x2
        MP.@similarvariable p y1
        MP.@similarvariable p y2

        [x1, x2], [y1, y2]
    elseif MP.nvariables(p) == 3
        MP.@similarvariable p x1
        MP.@similarvariable p x2
        MP.@similarvariable p x3
        MP.@similarvariable p y1
        MP.@similarvariable p y2
        MP.@similarvariable p y3

        [x1, x2, x3], [y1, y2, y3]
    end
end


text:
abstract type AbstractRenewableFormulation <: AbstractDeviceFormulation end
abstract type AbstractRenewableDispatchFormulation <: AbstractRenewableFormulation end
struct RenewableFullDispatch <: AbstractRenewableDispatchFormulation end
struct RenewableConstantPowerFactor <: AbstractRenewableDispatchFormulation end

########################### renewable generation variables #################################
function activepower_variables!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{R},
) where {R <: PSY.RenewableGen}
    add_variable(
        psi_container,
        devices,
        variable_name(ACTIVE_POWER, R),
        false,
        :nodal_balance_active;
        lb_value = x -> 0.0,
        ub_value = x -> PSY.get_rating(x),
    )
    return
end

function reactivepower_variables!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{R},
) where {R <: PSY.RenewableGen}
    add_variable(
        psi_container,
        devices,
        variable_name(REACTIVE_POWER, R),
        false,
        :nodal_balance_reactive,
    )
    return
end

####################################### Reactive Power Constraints #########################
function reactivepower_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{R},
    model::DeviceModel{R, RenewableFullDispatch},
    system_formulation::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {R <: PSY.RenewableGen}
    constraint_data = Vector{DeviceRange}(undef, length(devices))
    for (ix, d) in enumerate(devices)
        name = PSY.get_name(d)
        if isnothing(PSY.get_reactivepowerlimits(d))
            lims = (min = 0.0, max = 0.0)
            @warn("Reactive Power Limits of $(lims) are nothing. Q_$(lims) is set to 0.0")
        else
            lims = PSY.get_reactivepowerlimits(d)
        end
        constraint_data[ix] = DeviceRange(name, lims)
    end
    device_range(
        psi_container,
        constraint_data,
        constraint_name(REACTIVE_RANGE, R),
        variable_name(REACTIVE_POWER, R),
    )
    return
end

function reactivepower_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{R},
    model::DeviceModel{R, RenewableConstantPowerFactor},
    system_formulation::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {R <: PSY.RenewableGen}
    names = (PSY.get_name(d) for d in devices)
    time_steps = model_time_steps(psi_container)
    p_var = get_variable(psi_container, ACTIVE_POWER, R)
    q_var = get_variable(psi_container, REACTIVE_POWER, R)
    constraint_val = JuMPConstraintArray(undef, names, time_steps)
    assign_constraint!(psi_container, REACTIVE_RANGE, R, constraint_val)
    for t in time_steps, d in devices
        name = PSY.get_name(d)
        pf = sin(acos(PSY.get_powerfactor(d)))
        constraint_val[name, t] =
            JuMP.@constraint(psi_container.JuMPmodel, q_var[name, t] == p_var[name, t] * pf)
    end
    return
end

function activepower_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{R},
    model::DeviceModel{R, <:AbstractRenewableDispatchFormulation},
    system_formulation::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {R <: PSY.RenewableGen}
    parameters = model_has_parameters(psi_container)
    use_forecast_data = model_uses_forecasts(psi_container)

    if !parameters && !use_forecast_data
        constraint_data = Vector{DeviceRange}(undef, length(devices))
        for (ix, d) in enumerate(devices)
            name = PSY.get_name(d)
            ub = PSY.get_activepower(d)
            limits = (min = 0.0, max = ub)
            range_data = DeviceRange(name, limits)
            add_device_services!(range_data, d, model)
            constraint_data[ix] = range_data
        end
        device_range(
            psi_container,
            constraint_data,
            constraint_name(ACTIVE_RANGE, R),
            variable_name(ACTIVE_POWER, R),
        )
        return
    end

    forecast_label = "get_rating"
    constraint_data = Vector{DeviceTimeSeries}(undef, length(devices))
    for (ix, d) in enumerate(devices)
        ts_vector = get_time_series(psi_container, d, forecast_label)
        timeseries_data = DeviceTimeSeries(d, x -> PSY.get_rating(x), ts_vector)
        add_device_services!(timeseries_data, d, model)
        constraint_data[ix] = timeseries_data
    end

    if parameters
        device_timeseries_param_ub(
            psi_container,
            constraint_data,
            constraint_name(ACTIVE, R),
            UpdateRef{R}(ACTIVE_POWER, forecast_label),
            variable_name(ACTIVE_POWER, R),
        )
    else
        device_timeseries_ub(
            psi_container,
            constraint_data,
            constraint_name(ACTIVE, R),
            variable_name(ACTIVE_POWER, R),
        )
    end
    return
end

########################## Addition of to the nodal balances ###############################
function nodal_expression!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{R},
    system_formulation::Type{<:PM.AbstractPowerModel},
) where {R <: PSY.RenewableGen}
    nodal_expression!(psi_container, devices, PM.AbstractActivePowerModel)
    parameters = model_has_parameters(psi_container)
    use_forecast_data = model_uses_forecasts(psi_container)
    if use_forecast_data
        forecast_label = "get_rating"
        peak_value_function = x -> PSY.get_rating(x) * sin(acos(PSY.get_powerfactor(x)))
    else
        forecast_label = ""
        peak_value_function = x -> PSY.get_reactivepower(x)
    end
    constraint_data = Vector{DeviceTimeSeries}(undef, length(devices))
    for (ix, d) in enumerate(devices)
        ts_vector = get_time_series(psi_container, d, forecast_label)
        timeseries_data = DeviceTimeSeries(d, peak_value_function, ts_vector)
        constraint_data[ix] = timeseries_data
    end

    if parameters
        include_parameters(
            psi_container,
            constraint_data,
            UpdateRef{R}(REACTIVE_POWER, forecast_label),
            :nodal_balance_active,
        )
        return
    else
        for t in model_time_steps(psi_container)
            for device in constraint_data
                add_to_expression!(
                    psi_container.expressions[:nodal_balance_reactive],
                    device.bus_number,
                    t,
                    device.multiplier * device.timeseries[t],
                )
            end
        end
    end
    return
end

function nodal_expression!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{R},
    system_formulation::Type{<:PM.AbstractActivePowerModel},
) where {R <: PSY.RenewableGen}
    parameters = model_has_parameters(psi_container)
    use_forecast_data = model_uses_forecasts(psi_container)
    if use_forecast_data
        forecast_label = "get_rating"
        peak_value_function = x -> PSY.get_rating(x) * PSY.get_powerfactor(x)
    else
        forecast_label = ""
        peak_value_function = x -> PSY.get_activepower(x)
    end
    constraint_data = Vector{DeviceTimeSeries}(undef, length(devices))
    for (ix, d) in enumerate(devices)
        ts_vector = get_time_series(psi_container, d, forecast_label)
        timeseries_data = DeviceTimeSeries(d, peak_value_function, ts_vector)
        constraint_data[ix] = timeseries_data
    end

    if parameters
        include_parameters(
            psi_container,
            constraint_data,
            UpdateRef{R}(ACTIVE_POWER, forecast_label),
            :nodal_balance_active,
        )
        return
    else
        for t in model_time_steps(psi_container)
            for device in constraint_data
                add_to_expression!(
                    psi_container.expressions[:nodal_balance_active],
                    device.bus_number,
                    t,
                    device.multiplier * device.timeseries[t],
                )
            end
        end
    end
    return
end

##################################### renewable generation cost ############################
function cost_function(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{PSY.RenewableDispatch},
    ::Type{D},
    ::Type{<:PM.AbstractPowerModel},
) where {D <: AbstractRenewableDispatchFormulation}
    add_to_cost(
        psi_container,
        devices,
        variable_name(ACTIVE_POWER, PSY.RenewableDispatch),
        :fixed,
        -1.0,
    )
    return
end


text:
using Distributed
using StatsPlots

addprocs(4)

@everywhere using CSV
@everywhere using DataFrames
@everywhere using StatsFuns: logistic
@everywhere using Turing

"""
    get_data(path, :y_label)

Returns data for the logistic model split by X, y structure.
Retruns included breast cancer data by default.

# Arguments:
- `path::String="data/breast_cancer.csv"` path to csv file.
- `y_label::Symbol=:diagnosis` y label for target/dependent variable

# Returns:
- `X` Matrix containing X data/exogenous variables.
- `y` Vector containing y label data/endogenous variables.
"""
@everywhere function get_data(
        path::String="data/breast_cancer.csv",
        y_label::Symbol=:diagnosis
    )
    df = DataFrame(CSV.File(path))
    y = df[!, y_label]
    X = convert(Matrix, select!(df, Not(y_label)))
    X = (X .- mean(X, dims=1)) ./ std(X, dims=1)
    return X, y
end

"""
    logistic_regression(x::Matrix{Float64}, y::Vector{Int64})

Creates a logistic model to be used as an input argument in a sampler.

# Arguments:
- `X::Matrix{Float64}`: Exogenous input data X in matrix format.
- `y::Vector{Int64}`: Endogenous/dependent variables vector.
"""
@everywhere @model function logistic_regression(
        X::Matrix{Float64},
        y::Vector{Int64}
    )
    intercept ~ Normal(0, sqrt(25))

    nfeatures = size(X, 2)
    n = size(X, 1)
    slopes ~ MvNormal(nfeatures, sqrt(25))
    p = intercept .+ X * slopes
    for i = 1:n

        y[i] ~ Bernoulli(logistic(p[i]))

    end
end

"""
    profiler(X::Matrix{Float64}, y::Vector{Int64}, max_iters::Int64=10)

# Arguments:
- `X::Matrix{Float64}`: Exogenous variables
- `y::Matrix{Int64}`: Endogenous variables
- `max_iters::Int64`: The specified number of iterations to execute the sampler for logistic regression.

# Returns:
- `times::Vector{Any}`: Execution times for each iteration.
- `trace::Chains{Any}`: The trace generated in the final sample.
"""
@everywhere function profiler(
        X::Matrix{Float64},
        y::Vector{Int64},
        max_iters::Int64=10
    )
    times = []
    for iteration in 1:max_iters

        time = @elapsed global trace = sample(
                logistic_regression(X, y),
                NUTS(1000, 0.90),
                MCMCDistributed(),
                1000,
                4
            )
        append!(times, time)
    end
    return times, trace
end

@everywhere function main()
    X, y = get_data()
    times, traces = profiler(X, y)
    println(times)
    return times, trace
end

times, trace = main()

max_stars_repo_name: mtaung/bazaar

max_issues_repo_name: mtaung/bazaar

max_forks_repo_name: mtaung/bazaar


text:
module NBodyUnitTests

using Test
using UnitfulAstrodynamics

@testset "NBody" begin

    r̅₁ = [0, 0, 0]u"km"
    v̅₁ = [0, 0, 0]u"km/s"
    m₁ = uconvert(u"kg", 1u"Mearth")
    myEarth = Body(r̅₁, v̅₁, m₁)

    r̅₂ = [0, 11681, 0]u"km"
    v̅₂ = [5.134, 4.226, 2.787]u"km/s"
    m₂ = 150.0u"kg"
    mySatellite = Body(r̅₂, v̅₂, m₂)

    sys1 = NBodySystem([myEarth, mySatellite])

    # No errors!
    @test true

end

end


text:
"""
3D grid with the following configuration. x and y are periodic coordinates,
uniformly spaced. The u coordinate is decomposed into several u-domains, and in
each of these domains the u-direction uses Gauss-Lobatto points.
"""
Base.@kwdef struct SpecCartGrid3D{T<:Real}
    x_min            :: T
    x_max            :: T
    x_nodes          :: Int
    y_min            :: T
    y_max            :: T
    y_nodes          :: Int
    u_min            :: T       = 0.0
    u_max            :: T
    u_domains        :: Int     = 1
    u_nodes          :: Int # number of points per domain
    fd_order         :: Int     = 4
end

function Jecco.Atlas(grid::SpecCartGrid3D{T}) where {T}
    Nsys = grid.u_domains
    delta_udom = (grid.u_max - grid.u_min) / Nsys

    ucoords = [GaussLobatto{1}("u", grid.u_min + (i-1)*delta_udom, grid.u_min +
                               i*delta_udom, grid.u_nodes) for i in 1:Nsys]

    xcoord  = Cartesian{2}("x", grid.x_min, grid.x_max, grid.x_nodes, endpoint=false)
    ycoord  = Cartesian{3}("y", grid.y_min, grid.y_max, grid.y_nodes, endpoint=false)

    charts  = [Chart(ucoords[i], xcoord, ycoord) for i in 1:Nsys]

    Atlas(charts)
end

struct System{Cu,Cx,Cy,Du,Dx,Dy}
    ucoord :: Cu
    xcoord :: Cx
    ycoord :: Cy
    Du     :: Du
    Duu    :: Du
    Dx     :: Dx
    Dxx    :: Dx
    Dy     :: Dy
    Dyy    :: Dy
end

function System(ucoord::GaussLobattoCoord, xcoord::CartesianCoord,
                ycoord::CartesianCoord, ord::Int)
    Du  = ChebDeriv{1}(1, ucoord.min, ucoord.max, ucoord.nodes)
    Duu = ChebDeriv{1}(2, ucoord.min, ucoord.max, ucoord.nodes)

    Dx  = CenteredDiff{2}(1, ord, Jecco.delta(xcoord), xcoord.nodes)
    Dxx = CenteredDiff{2}(2, ord, Jecco.delta(xcoord), xcoord.nodes)

    Dy  = CenteredDiff{3}(1, ord, Jecco.delta(ycoord), ycoord.nodes)
    Dyy = CenteredDiff{3}(2, ord, Jecco.delta(ycoord), ycoord.nodes)

    System{typeof(ucoord), typeof(xcoord), typeof(ycoord), typeof(Du), typeof(Dx),
           typeof(Dy)}(ucoord, xcoord, ycoord, Du, Duu, Dx, Dxx, Dy, Dyy)
end

Base.size(sys::System) = (sys.ucoord.nodes, sys.xcoord.nodes, sys.ycoord.nodes)

struct SystemPartition{N,A} <: AbstractPartition{N,A}
    x :: NTuple{N,A}
end

function SystemPartition(x_::AbstractVector{S}) where {S<:System}
    x = Tuple(x_)
    N = length(x)
    SystemPartition{N,eltype(x)}(x)
end

"""
    SystemPartition(grid::SpecCartGrid3D)

Create a `SystemPartition`, which is a `Tuple` of `System` where each element
corresponds to a different u-domain
"""
function SystemPartition(grid::SpecCartGrid3D{T}) where {T}
    Nsys = grid.u_domains
    delta_udom = (grid.u_max - grid.u_min) / Nsys

    ucoords = [GaussLobatto{1}("u", grid.u_min + (i-1)*delta_udom, grid.u_min +
                               i*delta_udom, grid.u_nodes) for i in 1:Nsys]

    xcoord  = Cartesian{2}("x", grid.x_min, grid.x_max, grid.x_nodes, endpoint=false)
    ycoord  = Cartesian{3}("y", grid.y_min, grid.y_max, grid.y_nodes, endpoint=false)

    systems = [System(ucoords[i], xcoord, ycoord, grid.fd_order) for i in 1:Nsys]

    SystemPartition(systems)
end

"""
    Boundary(grid::SpecCartGrid3D)

Create a `Boundary` struct with arrays of `size = (1,grid.x_nodes,grid.y_nodes)`
"""
function Boundary(grid::SpecCartGrid3D{T}) where {T}
    Nx = grid.x_nodes
    Ny = grid.y_nodes
    Boundary{T}(undef, Nx, Ny)
end

"""
    BulkEvolvedPartition(grid::SpecCartGrid3D)

Returns `BulkPartition`, of `length = grid.u_domains`, with
elements of type `BulkEvolved`.
"""
function BulkEvolvedPartition(grid::SpecCartGrid3D{T}) where {T}
    Nx   = grid.x_nodes
    Ny   = grid.y_nodes
    Nu   = grid.u_nodes
    Nsys = grid.u_domains

    bulk  = [BulkEvolved{T}(undef, Nu, Nx, Ny) for i in 1:Nsys]

    BulkPartition(bulk...)
end

"""
    BulkConstrainedPartition(grid::SpecCartGrid3D)

Returns `BulkPartition`, of `length = grid.u_domains`, with
elements of type `BulkConstrained`.
"""
function BulkConstrainedPartition(grid::SpecCartGrid3D{T}) where {T}
    Nx   = grid.x_nodes
    Ny   = grid.y_nodes
    Nu   = grid.u_nodes
    Nsys = grid.u_domains

    bulk  = [BulkConstrained{T}(undef, Nu, Nx, Ny) for i in 1:Nsys]

    BulkPartition(bulk...)
end

function BulkDerivPartition(grid::SpecCartGrid3D{T}) where {T}
    Nx   = grid.x_nodes
    Ny   = grid.y_nodes
    Nu   = grid.u_nodes
    Nsys = grid.u_domains

    bulk = [BulkDeriv{T}(undef, Nu, Nx, Ny) for i in 1:Nsys]

    BulkPartition(bulk...)
end


text:
module LDScore
    include("ldsc.jl")
    include("LDScore/LDScoreRegression.jl")
    include("LDScore/Hsq.jl")
    include("LDScore/regression.jl")
    include("LDScore/sumstats.jl")
    include("LDScore/parse.jl")

    export hsq_weights
    export aggregate
    export estimate_h2
    export _read_sumstats
    export parse_sumstats
    export args
end


text:
#cfg title = "Learning a Control Policy"
#cfg weight = 12
#cfg active = true

#md # ## Overview

# In this example we walk through the process of setting up an experiment
# that runs [Natural Policy Gradient](https://papers.nips.cc/paper/2073-a-natural-policy-gradient.pdf)
# (or more recently [in this work](https://arxiv.org/pdf/1703.02660.pdf)).
# This is an on-policy reinforcement learning method that is comparable to TRPO,
# PPO, and other policy gradient methods. See the documentation for `NaturalPolicyGradient`
# for full implementation details.

#md # ## The Code

# First, let's go head and grab all the dependencies:
using LinearAlgebra, Random, Statistics # From Stdlib
using LyceumAI                          # For the NPG controller
using LyceumMuJoCo                      # For the Hopper environment
using Flux                              # For our neural networks needs
using UniversalLogger                   # For logging experiment data
using Plots                             # For plotting the results
using LyceumBase.Tools                  # Miscellaneous utilities

# We first instantiate a `HopperV2` environment to grab useful
# environment-specific values, such as the size of the observation and action vectors:
env = LyceumMuJoCo.HopperV2();
dobs, dact = length(obsspace(env)), length(actionspace(env));

# We'll also seed the per-thread global RNGs:
seed_threadrngs!(1)

#md # ### Policy Gradient Components

# Policy Gradient methods require a policy: a function that takes in the state/observations
# of the agent, and outputs an action i.e. `action = π(obs)`. In much of Deep RL, the
# policy takes the form of a neural network which can be built on top of the
# [Flux.jl](https://github.com/FluxML/Flux.jl) library. We utilize a stochastic policy in
# this example. Specifically, our policy is represented as a multivariate Gaussian
# distribution of the form:
# ```math
# \pi(a | o) = \mathcal{N}(\mu_{\theta_1}(o), \Sigma_{\theta_2})
# ```
# where ``\mu_{\theta_1}`` is a neural network, parameterized by ``\theta_1``, that maps
# an observation to a mean action and ``\Sigma_{\theta_2}`` is a diagonal
# covariance matrix parameterized by ``\theta_2``, the diagonal entries of the matrix.
# For ``\mu_{\theta_1}`` we utilize a 2-layer neural network, where each layer has a "width"
# of 32. We use tanh activations for each hidden layer and initialize the network weights
# with Glorot Uniform initializations. Rather than tracking ``\Sigma_{\theta_2}`` directly,
# we track the log standard deviations, which are easier to learn. We initialize
# ``\log \text{diag}(\Sigma_{\theta_2})`` as `zeros(dact)`, i.e. a `Vector` of length `dact`,
# initialized to 0. Both ``\theta_1`` and ``\theta_2`` are learned in this example.
# Note that ``\mu_{\theta_1}`` is a _state-dependent_ mean while ``\Sigma_{\theta_2}``
# is a _global_ covariance.
const policy = DiagGaussianPolicy(
    multilayer_perceptron(
        dobs,
        32,
        32,
        dact;
        σ = tanh,
        initb = Flux.glorot_uniform,
        initb_final = Flux.glorot_uniform,
        dtype = Float32,
    ),
    zeros(Float32, dact),
);

# This NPG implementation uses [Generalized Advantaged Estimation](https://arxiv.org/pdf/1506.02438.pdf),
# which requires an estimate of the value function, `value(state)`, which we
# represent using a 2-layer, feedforward neural network where each layer has a width of
# 128 and uses the ReLU activation function. The model weights are initialized using
# Glorot Uniform initialization as above.
const value = multilayer_perceptron(
    dobs,
    128,
    128,
    1;
    σ = Flux.relu,
    initb = Flux.glorot_uniform,
    initb_final = Flux.glorot_uniform,
    dtype = Float32,
);

# Next, we set up the optimization pipeline for `value`. We use a mini-batch size of 64
# and the [ADAM](https://arxiv.org/pdf/1412.6980.pdf) optimizer. `FluxTrainer` is an
# iterator that loops on the model provided, performing a single step of gradient
# descent at each iteration. The result at each loop is passed to `stopcb` below, so you
# can quit after a number of epochs, convergence, or other criteria; here it's capped at
# two epochs. See the documentation for `FluxTrainer` for more information.
valueloss(bl, X, Y) = Flux.mse(vec(bl(X)), vec(Y))
stopcb(x) = x.nepochs > 2
const valuetrainer = FluxTrainer(
    optimiser = ADAM(1e-3),
    szbatch = 64,
    lossfn = valueloss,
    stopcb = stopcb
);

# The `NaturalPolicyGradient` iterator is a type that pre-allocates all necesary
# data structures and performs one gradient update to `policy` at each iteration.
# We first pass in a constructor that given `n` returns `n` instances of
# `LyceumMuJoCo.HopperV2`, all sharing the same `jlModel`, to allow `NaturalPolicyGradient`
# to allocate per-thread environments and enable performant, parallel sampling from
# `policy`. We then pass in the `policy`, `value`, and `valuetrainer` instances
# constructed above and override a few of the default `NaturalPolicyGradient` parameters:
# `gamma`, `gaelambda`, and `norm_step_size`. Finally, we set the max trajectory length
# `Hmax` and total number of samples per iteration, `N`. Under the hood,
# `NaturalPolicyGradient` will use approximately `div(N, Hmax)` threads to perform the
# sampling.
const npg = NaturalPolicyGradient(
    n -> tconstruct(LyceumMuJoCo.HopperV2, n),
    policy,
    value,
    valuetrainer;
    gamma = 0.995,
    gaelambda = 0.97,
    norm_step_size = 0.05,
    Hmax = 1000,
    N = 10240,
);

#md # ### Running Experiments

# Finally, let's spin on our iterator 200 times, plotting every 20 iterations.
# This lets us break out of the loop if certain conditions are met, or re-start training
# manually if needed. We of course wish to track results, so we create a `ULogger` and
# `Experiment` to which we can save data. We also have useful timing information displayed
# every 20 iterations to better understand the performance of our algorithm and identify
# any potential bottlenecks. Rather than iterating on `npg` at the global scope, we'll
# do it inside of a function to avoid the performance issues associated with global
# variables as discussed in the
# [Julia performance tips](https://docs.julialang.org/en/v1/manual/performance-tips/).
# Note, to keep the Markdown version of this tutorial readable, we skip the plots
# and performance statistics. To enable them, simply call `hopper_NPG(npg, true)`.
function hopper_NPG(npg::NaturalPolicyGradient, plot::Bool)
    exper = Experiment("/tmp/hopper_example.jlso", overwrite = true)
    ## Walks, talks, and acts like a Julia logger. See the UniversalLogger.jl docs for more info.
    lg = ULogger()
    for (i, state) in enumerate(npg)
        if i > 200
            ## serialize some stuff and quit
            exper[:policy] = npg.policy
            exper[:value] = npg.value
            exper[:etype] = LyceumMuJoCo.HopperV2
            exper[:meanstates] = state.meanbatch
            exper[:stocstates] = state.stocbatch
            break
        end

        ## log everything in `state` except meanbatch and stocbatch
        push!(lg, :algstate, filter_nt(state, exclude = (:meanbatch, :stocbatch)))

        if plot && mod(i, 20) == 0
            x = lg[:algstate]
            ## The following are helper functions for plotting to the terminal.
            ## The first plot displays the `geteval` function for our stochastic
            ## and mean policy rollouts.
            display(expplot(
                Line(x[:stocterminal_eval], "StocLastE"),
                Line(x[:meanterminal_eval], "MeanLastE"),
                title = "Evaluation Score, Iter=$i",
                width = 60,
                height = 8,
            ))
            ## While the second one similarly plots `getreward`.
            display(expplot(
                Line(x[:stoctraj_reward], "StocR"),
                Line(x[:meantraj_reward], "MeanR"),
                title = "Reward, Iter=$i",
                width = 60,
                height = 8,
            ))

            ## The following are timing values for various parts of the Natural Policy Gradient
            ## algorithm at the last iteration, useful for finding performance bottlenecks
            ## in the algorithm.
            println("elapsed_sampled  = ", state.elapsed_sampled)
            println("elapsed_gradll   = ", state.elapsed_gradll)
            println("elapsed_vpg      = ", state.elapsed_vpg)
            println("elapsed_cg       = ", state.elapsed_cg)
            println("elapsed_valuefit = ", state.elapsed_valuefit)
        end
    end
    exper, lg
end
exper, lg = hopper_NPG(npg, false);

# Let's go ahead and plot the final reward trajectory for our stochastic and mean policies
# to see how we did:
plot(
    [lg[:algstate][:meantraj_reward] lg[:algstate][:stoctraj_reward]],
    labels = ["Mean Policy" "Stochastic Policy"],
    title = "HopperV2 Reward",
    legend = :bottomright,
)

# We'll also plot the evaluations:
plot(
    [lg[:algstate][:meantraj_eval] lg[:algstate][:stoctraj_eval]],
    labels = ["Mean Policy" "Stochastic Policy"],
    title = "HopperV2 Eval",
    legend = :bottomright,
)

# Finally, we save the logged results to `exper` for later review:
exper[:logs] = get(lg)
finish!(exper); # flushes everything to disk

using Test #src
@test lg[:algstate][:stoctraj_reward][end] > 1000 #src


text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã ``2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

"""
Probabilistic Context-Free Grammar
Wraps a CFG and adds a parallel data structure for the probabilities
"""
module PCFGs

export PCFG, fit_mle!, weighted_sum!

using GrammaticalEvolution
using LinearDerivTrees
using AbstractTrees
using RLESUtils, StatUtils
import Base: rand, rand!, getindex, setindex!, normalize!, fill!, copy

typealias CFG Grammar
typealias ProbDict Dict{Symbol,Vector{Float64}}

type PCFG
    cfg::CFG
    probs::ProbDict #probability vector for each Or rule
end

"""
PCFG constructor that initializes all probabilities to uniform
"""
function PCFG(cfg::CFG)
    probs = ProbDict()
    for (sym,rule) in cfg.rules
        if isa(rule, OrRule) || isa(rule, RangeRule)
            probs[sym] = uniform_probs(length(rule))
        end
    end
    PCFG(cfg, probs)
end

getindex(pcfg::PCFG, index) = pcfg.probs[index]
setindex!(pcfg::PCFG, index, val) = pcfg.probs[index] = val

"""
Draw a random sample from pcfg.  The output is in-place
into tree.
"""
function rand!(tree::LinearDerivTree, pcfg::PCFG)
    rand!(Base.GLOBAL_RNG, tree, pcfg)
end
function rand!(rng::AbstractRNG, tree::LinearDerivTree, pcfg::PCFG)
    probs = pcfg.probs
    initialize!(tree)
    while !isdone(tree)
        as = actionspace(tree)
        sym = get_sym(tree)
        step!(tree, weighted_rand(rng, as, probs[sym]))
    end
    iscomplete(tree)
end

"""
Draw random samples according to pcfg and output them in-place into samples.
"""
function rand!(samples::Vector{LinearDerivTree}, pcfg::PCFG)
    rand!(Base.GLOBAL_RNG, samples, pcfg)
end
function rand!(rng::AbstractRNG, samples::Vector{LinearDerivTree}, pcfg::PCFG)
    for i = 1:length(samples)
        rand!(rng, samples[i], pcfg)
    end
end

"""
Draw N random samples from pcfg.  Derivation trees are limited to maxsteps.
Outputs a vector of LinearDerivTrees.
"""
function rand(pcfg::PCFG, N::Int64, maxsteps::Int64)
    rand(Base.GLOBAL_RNG, pcfg, N, maxsteps)
end
function rand(rng::AbstractRNG, pcfg::PCFG, N::Int64, maxsteps::Int64)
    params = LDTParams(pcfg.cfg, maxsteps)
    samples = Array(LinearDerivTree, N)
    for i = 1:N
        s = LinearDerivTree(params)
        rand!(rng, s, pcfg)
        samples[i] = s
    end
    samples
end

"""
Compute the maximum likelihood estimator from all the transitions in the
trees in samples and store the output in-lce into pcfg.
"""
function fit_mle!(pcfg::PCFG, samples::Vector{LinearDerivTree})
    fill!(pcfg, 0.0)
    for s in samples
        if iscomplete(s) #don't consider incomplete elite samples
            for node in PreOrderDFS(s.derivtree.root)
                if isa(node.rule, DecisionRule)
                    sym = Symbol(node.rule.name)
                    a = node.action
                    pcfg[sym][a] += 1.0
                end
            end
        end
    end
    normalize!(pcfg)
end

"""
Computes the weighted sum of two pcfgs and returns the output in-place
pcfg = w1 * pcfg + w2 * pcfg2
"""
function weighted_sum!(pcfg::PCFG, w1::Float64, pcfg2::PCFG, w2::Float64)
    for (k,v1) in pcfg.probs
        v2 = pcfg2.probs[k]
        v1[:] = w1*v1 + w2*v2
    end
end

"""
Fill all probabilities in pcfg to value x
"""
function fill!(pcfg::PCFG, x)
    for v in values(pcfg.probs)
        fill!(v, x)
    end
end

"""
Normalize the sum of each probability vector in pcfg
"""
function normalize!(pcfg::PCFG)
    for v in values(pcfg.probs)
        if sum(v) > 0.0 #make sure valid denom
            normalize!(v, 1)
        else #set to uniform, catches NaNs and 0.0
            len = length(v)
            fill!(v, 1.0/len)
        end
    end
end

function copy(pcfg::PCFG)
    PCFG(pcfg.cfg, deepcopy(pcfg.probs))
end

end #module

max_stars_repo_name: sisl/ExprSearch

max_issues_repo_name: sisl/ExprSearch


text:
using Andes

@testset "Andes.jl run power flow" begin
    kundur = Andes.py.utils.paths.get_case("kundur/kundur_full.xlsx")
    system = Andes.py.run(kundur, no_output=true)

    @test system.PFlow.converged == true

    # test sparse matrix conversion
    @test size(system.PFlow.A) == (28, 28)

end


text:
@testset "Test utility functions" begin
    concrete_types = IS.get_all_concrete_subtypes(IS.InfrastructureSystemsType)
    @test length([x for x in concrete_types if isconcretetype(x)]) == length(concrete_types)
end

@testset "Test strip_module_name" begin
    @test IS.strip_module_name("PowerSystems.HydroDispatch") == "HydroDispatch"

    @test IS.strip_module_name("InfrastructureSystems.Deterministic{PowerSystems.HydroDispatch}") ==
          "Deterministic{PowerSystems.HydroDispatch}"

    @test IS.strip_module_name("Deterministic{PowerSystems.HydroDispatch}") ==
          "Deterministic{PowerSystems.HydroDispatch}"
end

@testset "Test strip_parametric_type" begin
    @test IS.strip_parametric_type("Deterministic{PowerSystems.HydroDispatch}") ==
          "Deterministic"

    @test IS.strip_parametric_type("InfrastructureSystems.Deterministic{PowerSystems.HydroDispatch}") ==
          "InfrastructureSystems.Deterministic"
end

@testset "Test serialization utility functions" begin
    struct Foo{Float64, Int} end
    text = "Foo{Float64,Int64}"
    symbol = IS.parse_serialized_type(text)
    @test eval(symbol) == Foo{Float64, Int64}

    text = "Float64"
    symbol = IS.parse_serialized_type(text)
    @test eval(symbol) == Float64

    @test_throws ErrorException IS.parse_serialized_type("foo()")
end

@testset "Test checking of parametric expression characters" begin
    illegal_chars = (
        '+',
        '=',
        '!',
        '@',
        '#',
        '$',
        '%',
        '^',
        '&',
        '*',
        '(',
        ')',
        '-',
        '\\',
        '"',
        '\'',
        '[',
        ']',
        '|',
        '~',
        '<',
        '>',
        ';',
        ':',
    )
    for char in illegal_chars
        text = "foo" * char
        @test_throws ErrorException IS._check_expression_characters(text)
    end
end

@testset "Test checking of parametric expressions" begin
    IS._check_parametric_expression(Meta.parse("Foo{Float64,Int64}"))
    @test_throws ErrorException IS._check_parametric_expression(Meta.parse("foo()"))
    expr = Meta.parse("Foo{Float64, print(\"hello\")}")
    @test_throws ErrorException IS._check_parametric_expression(expr)
end

@testset "Test exported names" begin
    @test IS.validate_exported_names(IS)
end

IS.@scoped_enum Fruit begin
    APPLE
    ORANGE
end

@testset "Test scoped_enum" begin
    @test Fruits.APPLE isa Fruits.Fruit
    @test Fruits.ORANGE isa Fruits.Fruit
end


text:
using Rotations
using Documenter

DocMeta.setdocmeta!(Rotations, :DocTestSetup, :(using Rotations); recursive=true)

makedocs(;
    modules=[Rotations],
    repo="https://github.com/JuliaGeometry/Rotations.jl/blob/{commit}{path}#{line}",
    sitename="Rotations.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://JuliaGeometry.github.io/Rotations.jl",
        assets = ["assets/custom.css"],
    ),
    pages=[
        "Home" => "index.md",
        "Rotation Types" => "rotation_types.md",
        "2D Rotations" => [
            "Matrix" => "2d_matrix.md",
            "Angle" => "2d_angle.md",
            ],
        "3D Rotations" => [
            "Matrix" => "3d_matrix.md",
            "Euler Angles" => "3d_euler.md",
            "Angle and Axis" => "3d_angleaxis.md",
            "Quaternion and Related Parameters" => "3d_quaternion.md",
            ],
        "Rotation Generators" => [
            "Rotation Generator Types" => "rotation_generator_types.md",
            "2D Rotation Generators" => "2d_rotation_generator.md",
            "3D Rotation Generators" => "3d_rotation_generator.md",
            ],
        "General Dimensional Rotations" => "general_dimensional_rotations.md",
        "Common Methods for Rotations" => "functions.md",
        "Visualizing Rotations" => "visualizing.md",
        "Function Reference" => "functionreference.md",
        "References" => "reference.md",
    ],
)

deploydocs(;
    repo="github.com/JuliaGeometry/Rotations.jl",
)


text:
begin
    using Revise
    using BayesNegativeBinomial
    using BNPRegressionGGA2021
    using Distributions
    using Gadfly
    using Random
    using RCall
    using Statistics
    using StatsBase
    using Test
    using LinearAlgebra
end

#=================================#
# Example 1 - 1 trivial predictor #
#=================================#

function simulate_sample_01(N0, N1)
    dy = MixtureModel(Normal, [(1.0, 1.0), (-1.0, 0.25)], [0.5, 0.5])
    X0 = ones(N0, 1)
    X1 = ones(N1, 1)
    y0 = rand(dy, N0)
    y1 = LinRange(-3, 3, N1) |> collect
    mean_y0, std_y0 = mean_and_std(y0, 1)
    mean_X0, std_X0 = mean_and_std(X0, 1)
    y0 = (y0 .- mean_y0) ./ std_y0
    X0 = (X0 .- mean_X0) ./ std_X0
    X1 = (X1 .- mean_X0) ./ std_X0
    X0[:, 1] .= 1
    X1[:, 1] .= 1
    return y0, X0, y1, X1
end;

begin
    Random.seed!(1)
    N0, N1 = 1000, 50;
    rng = MersenneTwister(1);
    y0, X0, y1, X1 = simulate_sample_01(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

f1 = mean(chainf);
plot(
    layer(x = y1, y = f1, Geom.line, color=["bnp"]),
    layer(x = y0, Geom.density, color=["kden"]),
    layer(x = y0, Geom.histogram(density = true), color=["hist"])
)

#========================================================#
# Example 2 - 1 discrete predictor (simple distribution) #
#========================================================#

# Simulate a sample
function simulate_sample(N0, N1)
    x0 = [-0.5, 0.5]
    X0 = [ones(N0) rand(x0, N0)]
    y0 = X0 * ones(2) + randn(N0)
    y1 = LinRange(minimum(y0), maximum(y0), N1) |> collect |> x -> repeat(x, 2)
    x1 = [x0[1] * ones(N1); x0[2] * ones(N1)]
    X1 = [ones(2 * N1) x1]
    return y0, X0, y1, X1
end;

begin
    Random.seed!(1)
    N0, N1 = 1000, 50;
    y0, X0, y1, X1 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

fb = mean(chainf);
dy(x) = Normal(1 + x, 1);
plot(x = y0, color = X0[:, 2], Geom.density)
plot(
    layer(x = y1, y = fb, color = X1[:, 2], Geom.line),
    layer(x = y1, y = pdf.(dy.(X1[:, 2]), y1), color = X1[:, 2], Geom.line)
)
mean(chainβ)

#=========================================================#
# Example 3 - 1 discrete predictor (mixture distribution) #
#=========================================================#

# Simulate a sample
function simulate_sample(N0, N1)
    dϵ = MixtureModel(Normal, [(1.0, 1.0), (-1.0, 0.25)], [0.5, 0.5])
    x0 = [-0.5, 0.5]
    X0 = [ones(N0) rand(x0, N0)]
    y0 = X0 * ones(2) + rand(dϵ, N0)
    y1 = LinRange(minimum(y0), maximum(y0), N1) |> collect |> x -> repeat(x, 2)
    x1 = [x0[1] * ones(N1); x0[2] * ones(N1)]
    X1 = [ones(2 * N1) x1]
    return y0, X0, y1, X1
end

begin
    Random.seed!(1)
    N0, N1 = 1000, 50
    y0, X0, y1, X1 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

fb = mean(chainf);
plot(x = y0, color = X0[:, 2], Geom.histogram(density = true, bincount = 50))
plot(x = y0, color = X0[:, 2], Geom.density)
plot(x = y1, y = fb, color = X1[:, 2], Geom.line)
mean(chainβ)

#=================================================#
# Example 4 - 2 predictors (mixture distribution) #
#=================================================#

# Simulate a sample
function simulate_sample(N0, N1)
    dϵ = MixtureModel(Normal, [(1.0, 1.0), (-1.0, 0.25)], [0.5, 0.5])
    x0 = [-0.5, 0.5]
    X0 = [ones(N0) rand(x0, N0, 2)]
    y0 = X0 * [ones(2); 0] + rand(dϵ, N0)
    y1 = LinRange(minimum(y0), maximum(y0), N1) |> collect |> x -> repeat(x, 4)
    x1 = kron([x0[1], x0[2], x0[1], x0[2]], ones(N1))
    x2 = kron([x0[1], x0[1], x0[2], x0[2]], ones(N1))
    X1 = [ones(4 * N1) x1 x2]
    return y0, X0, y1, X1
end

begin
    Random.seed!(1)
    N0, N1 = 1000, 50
    y0, X0, y1, X1 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

fb = mean(chainf);
X0concat = @. string(X0[:, 2]) * " " * string(X0[:, 3])
X1concat = @. string(X1[:, 2]) * " " * string(X1[:, 3])
plot(x = y0, color = X0concat, Geom.density)
plot(x = y1, y = fb, color = X1concat, Geom.line)
mean(chainβ)

#==========================================================#
# Example 5 - 1 continuous predictor (simple distribution) #
#==========================================================#

function simulate_sample(N0, N1)
    dϵ = Normal(0.0, 1.0)
    x0 = randn(N0)
    X0 = rcopy(R"cbind(1, splines::bs($x0, df = 6, Boundary.knots = c(-4, 4)))")
    y0 = x0 + rand(dϵ, N0)
    x_grid = [-0.5, 0.5]
    y_grid = LinRange(-3, 3, N1)
    y1, x1 = Iterators.product(y_grid, x_grid) |> x -> zip(x...) .|> collect
    X1 = rcopy(R"cbind(1, predict(splines::bs($x0, df = 6, Boundary.knots = c(-4, 4)), $x1))")
    # Standardization
    mean_y0, std_y0 = mean_and_std(y0, 1)
    mean_X0, std_X0 = mean_and_std(X0, 1)
    y0 = (y0 .- mean_y0) ./ std_y0
    X0 = (X0 .- mean_X0) ./ std_X0
    X1 = (X1 .- mean_X0) ./ std_X0
    X0[:, 1] .= 1
    X1[:, 1] .= 1
    return y0, x0, X0, y1, x1, X1, mean_y0, std_y0
end

begin
    Random.seed!(1)
    N0, N1 = 1000, 50
    y0, x0, X0, y1, x1, X1, mean_y0, std_y0 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

fb = mean(chainf);
dy1(x) = Normal((x - mean_y0[1]) / std_y0[1], 1 / std_y0[1])
plot(
    layer(x = y1, y = fb, color = x1, Geom.line),
    layer(x = y1, y = pdf.(dy1.(x1), y1), color = x1, Geom.line),
    # layer(x = y0, Geom.histogram)
)

#===========================================================#
# Example 6 - 1 continuous predictor (mixture distribution) #
#===========================================================#

function simulate_sample(N0, N1)
    dϵ = MixtureModel(Normal, [(-1.5, 0.5), (1.5, 0.5)], [0.5, 0.5])
    dx = Normal()
    x0 = rand(dx, N0)
    X0 = rcopy(R"cbind(1, splines::bs($x0, df = 8, Boundary.knots = c(-5, 5)))")
    y0 = x0 + rand(dϵ, N0)
    x_grid = [-0.5, 0.5]
    y_grid = LinRange(-3, 3, N1)
    y1, x1 = Iterators.product(y_grid, x_grid) |> x -> zip(x...) .|> collect
    X1 = rcopy(R"cbind(1, predict(splines::bs($x0, df = 8, Boundary.knots = c(-5, 5)), $x1))")
    # Standardization
    mean_y0, std_y0 = mean_and_std(y0, 1)
    mean_X0, std_X0 = mean_and_std(X0, 1)
    y0 = (y0 .- mean_y0) ./ std_y0
    X0 = (X0 .- mean_X0) ./ std_X0
    X1 = (X1 .- mean_X0) ./ std_X0
    X0[:, 1] .= 1
    X1[:, 1] .= 1
    return y0, x0, X0, y1, x1, X1
end

begin
    Random.seed!(1)
    N0, N1 = 1000, 50
    y0, x0, X0, y1, x1, X1 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

fb = mean(chainf);
plot(x = y1, y = fb, color = x1, Geom.line)
plot(x = y0[-1.5 .< x0 .< -0.5], Geom.density)
plot(x = y0[+0.5 .< x0 .< +1.5], Geom.density)


text:
function get_bb(obj::DSSTZonalContext)
    return jcall(obj, "getBB", jdouble, ())
end

function get_bo_a(obj::DSSTZonalContext)
    return jcall(obj, "getBoA", jdouble, ())
end

function get_bo_a_bpo(obj::DSSTZonalContext)
    return jcall(obj, "getBoABpo", jdouble, ())
end

function get_cxo2n2a2(obj::DSSTZonalContext)
    return jcall(obj, "getCXO2N2A2", jdouble, ())
end

function get_hk(obj::DSSTZonalContext)
    return jcall(obj, "getHK", jdouble, ())
end

function get_k2mh2(obj::DSSTZonalContext)
    return jcall(obj, "getK2MH2", jdouble, ())
end

function get_k2mh2o2(obj::DSSTZonalContext)
    return jcall(obj, "getK2MH2O2", jdouble, ())
end

function get_m2ao_a(obj::DSSTZonalContext)
    return jcall(obj, "getM2aoA", jdouble, ())
end

function get_m_co2ab(obj::DSSTZonalContext)
    return jcall(obj, "getMCo2AB", jdouble, ())
end

function get_mean_motion(obj::DSSTZonalContext)
    return jcall(obj, "getMeanMotion", jdouble, ())
end

function get_muoa(obj::DSSTZonalContext)
    return jcall(obj, "getMuoa", jdouble, ())
end

function get_oo_ab(obj::DSSTZonalContext)
    return jcall(obj, "getOoAB", jdouble, ())
end

function get_oon2a2(obj::DSSTZonalContext)
    return jcall(obj, "getOON2A2", jdouble, ())
end

function get_roa(obj::DSSTZonalContext)
    return jcall(obj, "getRoa", jdouble, ())
end

function get_x(obj::DSSTZonalContext)
    return jcall(obj, "getX", jdouble, ())
end

function get_x2on2a2xp1(obj::DSSTZonalContext)
    return jcall(obj, "getX2ON2A2XP1", jdouble, ())
end

function get_x3on2a(obj::DSSTZonalContext)
    return jcall(obj, "getX3ON2A", jdouble, ())
end

function get_xon2a2(obj::DSSTZonalContext)
    return jcall(obj, "getXON2A2", jdouble, ())
end

function get_xx(obj::DSSTZonalContext)
    return jcall(obj, "getXX", jdouble, ())
end

function get_xxx(obj::DSSTZonalContext)
    return jcall(obj, "getXXX", jdouble, ())
end


text:
### CONCRETE TYPE: ITERATIVE PROX EVALUATION

using LinearAlgebra
using IterativeSolvers

struct QuadraticIterative{R <: Real, M, V <: AbstractVector{R}} <: Quadratic
  Q::M
  q::V
  temp::V
end

is_prox_accurate(f::QuadraticIterative) = false

function QuadraticIterative(Q::M, q::V) where {R <: Real, M, V <: AbstractVector{R}}
  if size(Q, 1) != size(Q, 2) || length(q) != size(Q, 2)
    error("Q must be squared and q must be compatible with Q")
  end
  QuadraticIterative{R, M, V}(Q, q, similar(q))
end

function (f::QuadraticIterative{R, M, V})(x::AbstractArray{R}) where {R, M, V}
  mul!(f.temp, f.Q, x)
  return 0.5*dot(x, f.temp) + dot(x, f.q)
end

function prox!(y::AbstractArray{R}, f::QuadraticIterative{R, M, V}, x::AbstractArray{R}, gamma::R=one(R)) where {R, M, V}
  f.temp .= x./gamma .- f.q
  op = Shift(f.Q, one(R)/gamma)
  IterativeSolvers.cg!(y, op, f.temp)
  mul!(f.temp, f.Q, y)
  fy = 0.5*dot(y, f.temp) + dot(y, f.q)
  return fy
end

function gradient!(y::AbstractArray{R}, f::QuadraticIterative{R, M, V}, x::AbstractArray{R}) where {R, M, V}
  mul!(y, f.Q, x)
  y .+= f.q
  return 0.5*(dot(x, y) + dot(x, f.q))
end

function prox_naive(f::QuadraticIterative, x, gamma=1.0)
  y = IterativeSolvers.cg(gamma*f.Q + I, x - gamma*f.q)
  fy = 0.5*dot(y, f.Q*y) + dot(y, f.q)
  return y, fy
end


text:
# --- PRIVATE METHODS --------------------------------------------------------------------------------------- #
function _obj_function_breakeven(x,contractSet::Set{PSAbstractAsset})

    # what is the test price?
    underlyingPrice = x[1]

    # what is the PL value at this price?
    result = compute_profit_loss_value(contractSet,underlyingPrice)
    if (isa(result.value,Exception) == true)
        return result
    end
    pl_value = result.value

    # compute the objective function -
    obj_value = (pl_value^2)+10000000*max(0.0,-1*underlyingPrice)

    # return -
    return obj_value
end
# ----------------------------------------------------------------------------------------------------------- #

# --- PUBLIC METHODS ---------------------------------------------------------------------------------------- #
function compute_breakeven_price(contractSet::Set{PSAbstractAsset}; initialPrice::Float64=1.0)::PSResult

    # setup initial price -
    xinitial = [initialPrice]

    # setup the objective function -
    OF(p) = _obj_function_breakeven(p, contractSet)

    # call the optimizer -
    opt_result = optimize(OF,xinitial,BFGS())

    # breakeven price -
    breakeven_price = Optim.minimizer(opt_result)[1]

    # return -
    return PSResult(breakeven_price)
end
# ----------------------------------------------------------------------------------------------------------- #