text:
abstract type AbstractSpectralElementSpace <: AbstractSpace end

Topologies.nlocalelems(space::AbstractSpectralElementSpace) =
    Topologies.nlocalelems(Spaces.topology(space))

local_geometry_data(space::AbstractSpectralElementSpace) = space.local_geometry

eachslabindex(space::AbstractSpectralElementSpace) =
    1:Topologies.nlocalelems(Spaces.topology(space))

function Base.show(io::IO, space::AbstractSpectralElementSpace)
    typname = Base.typename(typeof(space)).name
    println(io, "$(typname):")
    println(io, "  topology: ", space.topology)
    println(io, "  quadrature: ", space.quadrature_style)
end

topology(space::AbstractSpectralElementSpace) = space.topology
quadrature_style(space::AbstractSpectralElementSpace) = space.quadrature_style

"""
    SpectralElementSpace1D <: AbstractSpace
A one-dimensional space: within each element the space is represented as a polynomial.
"""
struct SpectralElementSpace1D{T, Q, G, D, M} <: AbstractSpectralElementSpace
    topology::T
    quadrature_style::Q
    local_geometry::G
    dss_weights::D
    inverse_mass_matrix::M
end

function SpectralElementSpace1D(topology, quadrature_style)
    # TODO: we need to massage the coordinate points because the grid is assumed 2D
    CoordType = Topologies.coordinate_type(topology)
    CoordType1D = Geometry.coordinate_type(CoordType, 1)
    AIdx = Geometry.coordinate_axis(CoordType1D)
    FT = eltype(CoordType1D)
    nelements = Topologies.nlocalelems(topology)
    Nq = Quadratures.degrees_of_freedom(quadrature_style)

    Mxξ = Geometry.Axis2Tensor{
        FT,
        Tuple{Geometry.CartesianAxis{AIdx}, Geometry.CovariantAxis{AIdx}},
        SMatrix{1, 1, FT, 1},
    }
    Mξx = Geometry.Axis2Tensor{
        FT,
        Tuple{Geometry.ContravariantAxis{AIdx}, Geometry.CartesianAxis{AIdx}},
        SMatrix{1, 1, FT, 1},
    }
    LG = Geometry.LocalGeometry{CoordType1D, FT, Mxξ, Mξx}
    local_geometry = DataLayouts.IFH{LG, Nq}(Array{FT}, nelements)
    quad_points, quad_weights =
        Quadratures.quadrature_points(FT, quadrature_style)

    for elem in 1:nelements
        local_geometry_slab = slab(local_geometry, elem)
        for i in 1:Nq
            ξ = quad_points[i]
            # TODO: we need to massage the coordinate points because the grid is assumed 2D
            vcoords = Topologies.vertex_coordinates(topology, elem)
            vcoords1D = (
                Geometry.coordinate(vcoords[1], 1),
                Geometry.coordinate(vcoords[2], 1),
            )
            x = Geometry.interpolate(vcoords1D, ξ)
            ∂x∂ξ =
                (
                    Geometry.component(vcoords1D[2], 1) -
                    Geometry.component(vcoords1D[1], 1)
                ) / 2
            J = abs(∂x∂ξ)
            ∂ξ∂x = inv(∂x∂ξ)
            WJ = J * quad_weights[i]
            local_geometry_slab[i] = Geometry.LocalGeometry(
                x,
                J,
                WJ,
                Geometry.AxisTensor(
                    (
                        Geometry.CartesianAxis{AIdx}(),
                        Geometry.CovariantAxis{AIdx}(),
                    ),
                    ∂x∂ξ,
                ),
                Geometry.AxisTensor(
                    (
                        Geometry.ContravariantAxis{AIdx}(),
                        Geometry.CartesianAxis{AIdx}(),
                    ),
                    ∂ξ∂x,
                ),
            )
        end
    end
    dss_weights = copy(local_geometry.J)
    dss_weights .= one(FT)
    dss_1d!(dss_weights, dss_weights, topology, Nq)
    dss_weights = one(FT) ./ dss_weights

    inverse_mass_matrix = copy(local_geometry.WJ)
    dss_1d!(inverse_mass_matrix, inverse_mass_matrix, topology, Nq)
    inverse_mass_matrix = one(FT) ./ inverse_mass_matrix

    return SpectralElementSpace1D(
        topology,
        quadrature_style,
        local_geometry,
        dss_weights,
        inverse_mass_matrix,
    )
end

nlevels(space::SpectralElementSpace1D) = 1

"""
    SpectralElementSpace2D <: AbstractSpace

A two-dimensional space: within each element the space is represented as a polynomial.
"""
struct SpectralElementSpace2D{T, Q, G, D, IS, BS} <:
       AbstractSpectralElementSpace
    topology::T
    quadrature_style::Q
    local_geometry::G
    dss_weights::D
    internal_surface_geometry::IS
    boundary_surface_geometries::BS
end

"""
    SpectralElementSpace2D(topology, quadrature_style)

Construct a `SpectralElementSpace2D` instance given a `topology` and `quadrature`.
"""
function SpectralElementSpace2D(topology, quadrature_style)
    CoordType2D = Topologies.coordinate_type(topology)
    AIdx = Geometry.coordinate_axis(CoordType2D)
    FT = eltype(CoordType2D)
    nelements = Topologies.nlocalelems(topology)
    Nq = Quadratures.degrees_of_freedom(quadrature_style)

    Mxξ = Geometry.Axis2Tensor{
        FT,
        Tuple{Geometry.CartesianAxis{AIdx}, Geometry.CovariantAxis{AIdx}},
        SMatrix{2, 2, FT, 4},
    }
    Mξx = Geometry.Axis2Tensor{
        FT,
        Tuple{Geometry.ContravariantAxis{AIdx}, Geometry.CartesianAxis{AIdx}},
        SMatrix{2, 2, FT, 4},
    }
    LG = Geometry.LocalGeometry{CoordType2D, FT, Mxξ, Mξx}

    local_geometry = DataLayouts.IJFH{LG, Nq}(Array{FT}, nelements)
    quad_points, quad_weights =
        Quadratures.quadrature_points(FT, quadrature_style)

    for elem in 1:nelements
        local_geometry_slab = slab(local_geometry, elem)
        for i in 1:Nq, j in 1:Nq
            # this hard-codes a bunch of assumptions, and will unnecesarily duplicate data
            # e.g. where all metric terms are uniform over the space
            # alternatively: move local_geometry to a different object entirely, to support overintegration
            # (where the integration is of different order)
            ξ = SVector(quad_points[i], quad_points[j])
            x = Geometry.interpolate(
                CoordType2D.(Topologies.vertex_coordinates(topology, elem),),
                ξ[1],
                ξ[2],
            )
            ∂x∂ξ = ForwardDiff.jacobian(ξ) do ξ
                local x
                x = Geometry.interpolate(
                    CoordType2D.(
                        Topologies.vertex_coordinates(topology, elem),
                    ),
                    ξ[1],
                    ξ[2],
                )
                SVector(Geometry.component(x, 1), Geometry.component(x, 2))
            end
            J = det(∂x∂ξ)
            ∂ξ∂x = inv(∂x∂ξ)
            WJ = J * quad_weights[i] * quad_weights[j]

            local_geometry_slab[i, j] = Geometry.LocalGeometry(
                x,
                J,
                WJ,
                Geometry.AxisTensor(
                    (
                        Geometry.CartesianAxis{AIdx}(),
                        Geometry.CovariantAxis{AIdx}(),
                    ),
                    ∂x∂ξ,
                ),
                Geometry.AxisTensor(
                    (
                        Geometry.ContravariantAxis{AIdx}(),
                        Geometry.CartesianAxis{AIdx}(),
                    ),
                    ∂ξ∂x,
                ),
            )
        end
    end

    # dss_weights = J ./ dss(J)
    dss_weights = copy(local_geometry.J)
    dss_2d!(dss_weights, local_geometry.J, topology, Nq)
    dss_weights .= local_geometry.J ./ dss_weights

    # TODO: this assumes XYDomain, need to dispatch SG Vector type on domain axis type
    SG = Geometry.SurfaceGeometry{FT, Geometry.Cartesian12Vector{FT}}
    interior_faces = Topologies.interior_faces(topology)

    internal_surface_geometry =
        DataLayouts.IFH{SG, Nq}(Array{FT}, length(interior_faces))
    for (iface, (elem⁻, face⁻, elem⁺, face⁺, reversed)) in
        enumerate(interior_faces)
        internal_surface_geometry_slab = slab(internal_surface_geometry, iface)

        local_geometry_slab⁻ = slab(local_geometry, elem⁻)
        local_geometry_slab⁺ = slab(local_geometry, elem⁺)

        for q in 1:Nq
            sgeom⁻ = compute_surface_geometry(
                local_geometry_slab⁻,
                quad_weights,
                face⁻,
                q,
                false,
            )
            sgeom⁺ = compute_surface_geometry(
                local_geometry_slab⁺,
                quad_weights,
                face⁺,
                q,
                false,
            )

            @assert sgeom⁻.sWJ ≈ sgeom⁺.sWJ
            @assert sgeom⁻.normal ≈ -sgeom⁺.normal

            internal_surface_geometry_slab[q] = sgeom⁻
        end
    end

    boundary_surface_geometries =
        map(Topologies.boundaries(topology)) do boundarytag
            boundary_faces = Topologies.boundary_faces(topology, boundarytag)
            boundary_surface_geometry =
                DataLayouts.IFH{SG, Nq}(Array{FT}, length(boundary_faces))
            for (iface, (elem, face)) in enumerate(boundary_faces)
                boundary_surface_geometry_slab =
                    slab(boundary_surface_geometry, iface)
                local_geometry_slab = slab(local_geometry, elem)
                for q in 1:Nq
                    boundary_surface_geometry_slab[q] =
                        compute_surface_geometry(
                            local_geometry_slab,
                            quad_weights,
                            face,
                            q,
                            false,
                        )
                end
            end
            boundary_surface_geometry
        end

    return SpectralElementSpace2D(
        topology,
        quadrature_style,
        local_geometry,
        dss_weights,
        internal_surface_geometry,
        boundary_surface_geometries,
    )
end

nlevels(space::SpectralElementSpace2D) = 1

function compute_surface_geometry(
    local_geometry_slab,
    quad_weights,
    face,
    q,
    reversed = false,
)
    Nq = length(quad_weights)
    @assert size(local_geometry_slab) == (Nq, Nq, 1, 1, 1)
    i, j = Topologies.face_node_index(face, Nq, q, reversed)

    local_geometry = local_geometry_slab[i, j]
    @unpack J, ∂ξ∂x = local_geometry

    # surface mass matrix
    n = if face == 1
        -J * ∂ξ∂x[1, :] * quad_weights[j]
    elseif face == 2
        J * ∂ξ∂x[1, :] * quad_weights[j]
    elseif face == 3
        -J * ∂ξ∂x[2, :] * quad_weights[i]
    elseif face == 4
        J * ∂ξ∂x[2, :] * quad_weights[i]
    end
    sWJ = norm(n)
    n = n / sWJ
    return Geometry.SurfaceGeometry(sWJ, n)
end

function variational_solve!(data, space::AbstractSpace)
    data .= RecursiveApply.rdiv.(data, space.local_geometry.WJ)
end

"""
    SpectralElementSpaceSlab <: AbstractSpace

A view into a `SpectralElementSpace2D` for a single slab.
"""
struct SpectralElementSpaceSlab{Q, G} <: AbstractSpectralElementSpace
    quadrature_style::Q
    local_geometry::G
end

const SpectralElementSpaceSlab1D =
    SpectralElementSpaceSlab{Q, DL} where {Q, DL <: DataLayouts.DataSlab1D}

const SpectralElementSpaceSlab2D =
    SpectralElementSpaceSlab{Q, DL} where {Q, DL <: DataLayouts.DataSlab2D}

nlevels(space::SpectralElementSpaceSlab1D) = 1
nlevels(space::SpectralElementSpaceSlab2D) = 1

function slab(space::AbstractSpectralElementSpace, v, h)
    SpectralElementSpaceSlab(
        space.quadrature_style,
        slab(space.local_geometry, v, h),
    )
end


text:
# This file is currently required to get AD on KL(::MvNormal, ::MvNormal) to work

import ChainRulesCore

# opt-out of generic rrule for `Matrix` defined in ChainRules.jl.
# It's horrible that this code is needed. Given that it does, its ideal home is PDMats.jl.
# We're type-pirating for now.
ChainRulesCore.@opt_out ChainRulesCore.rrule(::Type{<:Matrix}, ::PDMats.PDMat)

ext: jl


text:
"""
The Modal Policy has an in-horizon and an out-horizon policy. If the context has no temporal nature,
the in-horizon-policy is null and only out-horizon policy is used.
"""
struct ModalHorizonPolicy{P <: Policy}
    in_horizon_policy::P
    out_horizon_policy::P
end

get_mdp(mhp::ModalHorizonPolicy) = mhp.in_horizon_policy.mdp

struct ModalFinInfHorPolicy{P <: Policy}
    fin_hor_policy::Union{ModalHorizonPolicy{P},P}
    inf_hor_policy::Union{Nothing,P}
end

"""
CMSSP State augmented with horizon value. Needed for LocalApproximationVI
"""
struct ModalStateAugmented{C}
    relative_state::C
    horizon::Int64
end

function ModalStateAugmented(state::C) where {C}
    return ModalStateAugmented(state, 0)
end

function ModalStateAugmented{C}(state::C) where {C}
    return ModalStateAugmented{C}(state, 0)
end

struct ModalAction{AC}
    action::AC
    action_idx::Int64
end

"""
Define modal MDP type which will be used with approximate VI to get local layer policy

Attributes:
    - `mode::D` The specific mode that defines the smaller regional MDP
    - `actions::Vector{ModalAction{AC}}` The control actions available in this region
    - `beta_threshold::Float64` The interrupt threshold for this MDP
    - `horizon_limit::Int64` The horizon up to which the value function is computed
    - `min_value_per_horizon::Vector{Float64}` The worst cost or min value from each horizon value
    - `terminal_cost_penalty::Float64` The phi_CF value
    - `terminal_costs_set::Bool` Flag used for out-horizon vs in-horizon VI
    - `params::P` For tracking any other parameters used by the solver.
"""
mutable struct ModalFinHorMDP{D,C,AC,P} <: POMDPs.MDP{ModalStateAugmented{C},ModalAction{AC}}
    mode::D
    actions::Vector{ModalAction{AC}}
    beta_threshold::Float64
    horizon_limit::Int64
    min_value_per_horizon::Vector{Float64}
    terminal_cost_penalty::Float64
    terminal_costs_set::Bool
    params::P
end

function ModalFinHorMDP{D,C,AC,P}(mode::D, params::P, actions::Vector{AC}, beta::Float64=1.0, horizon_limit::Int64=0) where {D,C,AC,P}
    modal_actions = [ModalAction(a,i) for (i,a) in enumerate(actions)]
    return ModalFinHorMDP{D,C,AC,P}(mode, modal_actions, beta, horizon_limit, Inf*ones(horizon_limit), Inf, false, params)
end

function ModalFinHorMDP{D,C,AC,P}(mode::D, params::P, modal_actions::Vector{ModalAction{AC}}, beta::Float64=1.0, horizon_limit::Int64=0) where {D,C,AC,P}
    return ModalFinHorMDP{D,C,AC,P}(mode, modal_actions, beta, horizon_limit, Inf*ones(horizon_limit), Inf, false, params)
end

function ModalFinHorMDP{D,C,AC,P}(mode::D, params::P) where {D,C,AC,P}
    return ModalFinHorMDP{D,C,AC,P}(mode, params, Vector{ModalAction{AC}}(undef,0))
end

function set_horizon_limit!(mdp::ModalFinHorMDP, h::Int64) where {D,C,AC}
    mdp.horizon_limit = h
    mdp.min_value_per_horizon = Inf*ones(h)
end

POMDPs.actions(mdp::ModalFinHorMDP) = mdp.actions
POMDPs.n_actions(mdp::ModalFinHorMDP) = length(mdp.actions)
POMDPs.discount(mdp::ModalFinHorMDP) = 1.0 # SSP - Undiscounted
POMDPs.actionindex(mdp::ModalFinHorMDP, a::ModalAction) = a.action_idx

## Inf hor stuff
mutable struct ModalInfHorMDP{D,C,AC,P} <: POMDPs.MDP{C,ModalAction{AC}}
    mode::D
    actions::Vector{ModalAction{AC}}
    params::P
    terminal_reward::Float64
end

function ModalInfHorMDP{D,C,AC,P}(mode::D, actions::Vector{AC}, params::P, terminal_reward::Float64=0.0) where {D,C,AC,P}
    modal_actions = [ModalAction(a,i) for (i,a) in enumerate(actions)]
    return ModalInfHorMDP{D,C,AC,P}(mode, modal_actions, params, terminal_reward)
end

function ModalInfHorMDP{D,C,AC,P}(mode::D, params::P, terminal_reward::Float64=0.0) where {D,C,AC,P}
    return ModalInfHorMDP{D,C,AC,P}(mode, Vector{ModalAction{AC}}(undef,0), params, terminal_reward)
end

POMDPs.actions(mdp::ModalInfHorMDP) = mdp.actions
POMDPs.n_actions(mdp::ModalInfHorMDP) = length(mdp.actions)
POMDPs.discount(mdp::ModalInfHorMDP) = 1.0 # SSP - Undiscounted
POMDPs.actionindex(mdp::ModalInfHorMDP, a::ModalAction) = a.action_idx

get_mdp(pol::LocalApproximationValueIterationPolicy) = pol.mdp

"""
Convert augmented state to a vector by calling the underlying convert_s function
for non-augmented state. IMPORTANT: The underlying dynamics should be for RELATIVE STATE
"""
function POMDPs.convert_s(::Type{Vector{Float64}}, s::ModalStateAugmented,
                          mdp::ModalFinHorMDP)
    v = convert_s(Vector{Float64}, s.relative_state, mdp)
    push!(v, convert(Float64, s.horizon))
    return v
end

"""
Convert horizon-augmented vector to ModalStateAugmented
"""
function POMDPs.convert_s(::Type{ModalStateAugmented{C}}, v::AbstractVector{Float64},
                          mdp::ModalFinHorMDP) where {C}
    state = convert_s(C, v[1:end-1], mdp)
    horizon = convert(Int64, v[end])
    s = ModalStateAugmented{C}(state, horizon)
    return s
end

function POMDPs.generate_sr(mdp::ModalFinHorMDP, s::ModalStateAugmented, a::ModalAction,
                            rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    # Get next state and underlying cost for dynamics
    # TODO : This should be the same regardless of relative or absolute state
    (sp,r) = generate_sr(mdp, s.relative_state, a.action, rng)
    cost = -1.0*r

    # If at end of horizon, increment with terminal cost penalty
    # if sp is not terminal
    if s.horizon == 1
        if mdp.terminal_costs_set
            if POMDPs.isterminal(mdp, sp) == false
                cost += mdp.terminal_cost_penalty
            end
        end
    end

    return ModalStateAugmented(sp, s.horizon-1), -cost
end

function POMDPs.generate_sr(mdp::ModalInfHorMDP, s::S, a::ModalAction,
                            rng::RNG=Random.GLOBAL_RNG) where {S, RNG <: AbstractRNG}
    (sp, r) = generate_sr(mdp, s, a.action, rng)
    return sp, r
end

function POMDPs.isterminal(mdp::ModalFinHorMDP, s::ModalStateAugmented)
    return s.horizon == 0
end

"""
Compute the terminal cost penalty for the modal region, i.e. phi_CF from paper
This version uses a local function approximator.

Arguments:
    - `mdp::ModalFinHorMDP{D,C,AC}`
    - `lfa::LFA` A LocalFunctionApproximator object for the value iteration
"""
function compute_terminalcost_localapprox!(mdp::ModalFinHorMDP{D,C,AC,P}, lfa::LFA) where {D,C,AC,P,LFA <: LocalFunctionApproximator}
    max_contr_cost = 0.0
    max_switch_cost = 0.0

    # First iterate over continuous actions
    @info "Getting max control cost"
    for a in actions(mdp)
        for point in get_all_interpolating_points(lfa)
            state = convert_s(C, point, mdp)
            cost = -1.0*expected_reward(mdp,state,a.action)
            if cost > max_contr_cost
                max_contr_cost = cost
            end
        end
    end
    @show max_contr_cost
    mdp.terminal_cost_penalty = max_contr_cost * mdp.horizon_limit
end

"""
Perform finite horizon value approximation on the regional MDP

Arguments:
    - `mdp::ModalFinHorMDP{D,C,AC}`
    - `lfa::LFA`
    - `is_mdp_generative::Bool`
    - `n_generative_samples::Int64`
    - `rng::RNG`
"""
function finite_horizon_VI_localapprox!(mdp::ModalFinHorMDP, lfa::LFA,
                                       is_mdp_generative::Bool=false, n_generative_samples::Int64=0,
                                       rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG,LFA<:LocalFunctionApproximator}

    # Setup an out-horizon-approximator
    # Terminal costs are 0 in this case
    @assert mdp.horizon_limit > 1
    mdp.terminal_costs_set = false
    outhor_lfa = finite_horizon_extension(lfa, 0:1:mdp.horizon_limit)
    outhor_solver = LocalApproximationValueIterationSolver(outhor_lfa, max_iterations = 1,
                                                           verbose=true, rng=rng, is_mdp_generative=is_mdp_generative,
                                                           n_generative_samples=n_generative_samples)
    out_horizon_policy_augmented = solve(outhor_solver, mdp)

    # Create the true out-horizon LFA that has just one time axis
    outhor_lfa_true = finite_horizon_extension(lfa, mdp.horizon_limit+1:1:mdp.horizon_limit+1)
    all_interp_values = get_all_interpolating_values(out_horizon_policy_augmented.interp)
    all_interp_states = get_all_interpolating_points(out_horizon_policy_augmented.interp)

    outhor_interp_values = Vector{Float64}(undef,0)
    # copy over value for horizon = K
    for (v,s) in zip(all_interp_values, all_interp_states)
        if s[end] == mdp.horizon_limit
            push!(outhor_interp_values, v)
        end
    end

    # Check that there are indeed that many points as true out-horizon-LFA
    @assert length(outhor_interp_values) == n_interpolating_points(outhor_lfa_true)
    out_horizon_policy = LocalApproximationValueIterationPolicy(outhor_lfa_true, ordered_actions(mdp),
                                                                mdp, is_mdp_generative, n_generative_samples, rng)

    ## Now do for in-horizon
    mdp.terminal_costs_set = true
    inhor_lfa = finite_horizon_extension(lfa, 0:1:mdp.horizon_limit)
    inhor_solver = LocalApproximationValueIterationSolver(inhor_lfa, max_iterations = 1,
                                                           verbose=true, rng=rng, is_mdp_generative=is_mdp_generative,
                                                           n_generative_samples=n_generative_samples)
    in_horizon_policy = solve(inhor_solver, mdp)

    return ModalHorizonPolicy(in_horizon_policy, out_horizon_policy)
end

function infinite_horizon_VI_localapprox(mdp::ModalInfHorMDP, lfa::LFA, max_iterations::Int64=50,
                                 is_mdp_generative::Bool=false, n_generative_samples::Int64=0,
                                 rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG, LFA <: LocalFunctionApproximator}
    infhor_solver = LocalApproximationValueIterationSolver(lfa, max_iterations=max_iterations,
                                                           verbose=true, rng=rng, is_mdp_generative=is_mdp_generative,
                                                           n_generative_samples=n_generative_samples)
    infhor_policy = solve(infhor_solver, mdp)
    return infhor_policy
end

"""
Compute the worst cost or minimum value per horizon and update mdp object in place.

Arguments:
    - `modal_policy::ModalHorizonPolicy{P}`
"""
function compute_min_value_per_horizon_localapprox!(modal_policy::ModalHorizonPolicy)

    # Access the underlying value function approximator
    all_interp_values_inhor = get_all_interpolating_values(modal_policy.in_horizon_policy.interp)
    all_interp_states_inhor = get_all_interpolating_points(modal_policy.in_horizon_policy.interp)

    for (v,s) in zip(all_interp_values_inhor, all_interp_states_inhor)

        hor_val = convert(Int64,s[end])

        if hor_val == 0
            continue
        end

        # Compute lowest value
        if v < modal_policy.in_horizon_policy.mdp.min_value_per_horizon[hor_val]
            modal_policy.in_horizon_policy.mdp.min_value_per_horizon[hor_val] = v
        end
    end
end

"""
Compute the negative value of the relative state between current and target. The value is weighted by the
distribution over horizons.

Arguments:
    - `modal_policy::ModalHorizonPolicy`
    - `curr_timestep::Int64` The current system time
    - `tp_dist::TPDistribution` The distribution over GLOBAL times for reaching target state
    - `curr_state::C`
    - `target_state::C`
"""
function horizon_weighted_value(modal_policy::ModalHorizonPolicy, curr_timestep::Int64,
                                tp_dist::TPDistribution, curr_state::C, target_state::C) where C

    weighted_value = 0.0
    weighted_minvalue = 0.0

    mdp = modal_policy.in_horizon_policy.mdp

    # Iterate over horizon values and weight
    # Subtract the current time-step to get the true relative time
    for (h_, p) in tp_dist

        h = h_ - curr_timestep

        if h <= 0
            continue
        end

        temp_relative_state = get_relative_state(mdp, curr_state, target_state)

        if h > mdp.horizon_limit
            temp_state = ModalStateAugmented(temp_relative_state, mdp.horizon_limit+1)
            # weighted_value += p*POMDPs.value(modal_policy.out_horizon_policy,temp_state)

            weighted_value += p*(-Inf)
            # IMP - For any prob of out-of-horizon, never abort
            weighted_minvalue += p*(-Inf)
        else
            temp_state = ModalStateAugmented(temp_relative_state, h)
            weighted_value += p*value(modal_policy.in_horizon_policy,temp_state)
            weighted_minvalue += p*mdp.min_value_per_horizon[h]
        end
    end

    return weighted_value, weighted_minvalue
end

"""
Compute the negative value of the relative state between current and target. The value is weighted by the
distribution over horizons.

Arguments:
    - `modal_policy::ModalHorizonPolicy`
    - `curr_timestep::Int64` The current system time
    - `tp_dist::TPDistribution` The distribution over GLOBAL times for reaching target state
    - `curr_state::C`
    - `target_state::C`
    - `a::ModalAction{AC}`
"""
function horizon_weighted_actionvalue(modal_policy::ModalHorizonPolicy, curr_timestep::Int64,
                                tp_dist::TPDistribution, curr_state::C, target_state::C, a::ModalAction{AC}) where {C,AC}

    total_value = 0.0

    mdp = modal_policy.in_horizon_policy.mdp

    # Iterate over horizon values and weight
    for (h_,p) in tp_dist

        h = h_ - curr_timestep

        if h <= 0
            continue
        end

        temp_relative_state = get_relative_state(mdp, curr_state, target_state)

        if h > mdp.horizon_limit
            temp_state = ModalStateAugmented(temp_relative_state, mdp.horizon_limit+1)
            total_value += p*value(modal_policy.out_horizon_policy, temp_state, a)
        else
            temp_state = ModalStateAugmented(temp_relative_state, h)
            total_value += p*value(modal_policy.in_horizon_policy, temp_state, a)
        end
    end

    return total_value
end

"""
Returns action with lowest average cost (i.e. highest average value, since value is negative cost)

Arguments:
    - `modal_policy::ModalHorizonPolicy`
    - `curr_timestep::Int64` The current system time
    - `tp_dist::TPDistribution` The distribution over GLOBAL times for reaching target state
    - `curr_state::C`
    - `target_state::C`
"""
function get_best_intramodal_action(modal_policy::ModalHorizonPolicy, curr_timestep::Int64,
                               tp_dist::TPDistribution, curr_state::C, target_state::C) where C

    mdp = modal_policy.in_horizon_policy.mdp

    # First check if abort needed
    (weighted_value, weighted_minvalue) = horizon_weighted_value(modal_policy, curr_timestep,
                                                                 tp_dist, curr_state, target_state)

    if weighted_value <= mdp.beta_threshold*weighted_minvalue
        return nothing # closed-loop interrupt
    end

    best_action = mdp.actions[1]
    best_action_val = -Inf

    for a in mdp.actions
        action_val = horizon_weighted_actionvalue(modal_policy, curr_timestep, tp_dist, curr_state, target_state, a)
        if action_val > best_action_val
            best_action_val = action_val
            best_action = a
        end
    end

    return best_action
end

function inf_hor_value(policy::P, curr_state::C, target_state::C) where {P <: Policy, C}

    # This is for an inf horizon policy
    mdp = get_mdp(policy)

    temp_relative_state = get_relative_state(mdp, curr_state, target_state)
    value = POMDPs.value(policy, temp_relative_state) - mdp.terminal_reward

    return value
end

function get_best_intramodal_action(policy::LocalApproximationValueIterationPolicy, curr_timestep::Int64,
                               tp_dist::TPDistribution, curr_state::C, target_state::C) where C

    mdp = get_mdp(policy)

    temp_relative_state = get_relative_state(mdp, curr_state, target_state)
    best_action = action(policy, temp_relative_state)

    return best_action

end


text:
module CorrDCA

using Distributed: @distributed
using SharedArrays: SharedArray, sdata
using FastaIO: FastaReader

export read_fasta_alignment,
    remove_duplicate_sequences,
    compute_weights,
    compute_frequencies,
    add_pseudocount,
    covariance_matrix

include("readfasta.jl")
include("reweighting.jl")
include("frequencies.jl")
end

max_stars_repo_name: pagnani/CorrDCA

max_issues_repo_name: pagnani/CorrDCA

max_forks_repo_name: pagnani/CorrDCA


text:
# This file is pat of RLinearAlgebra.jl
# 1. Specifies type
# 2. Implements sample function
# 3. Exports Type

# using Random

"""
    LinSysVecRowDetermCyclic <: LinSysVecRowSelect

An immutable structure without any fields. Specifies deterministic cycling through the
equations of a linear system.
"""
struct LinSysVecRowDetermCyclic <: LinSysVecRowSelect end

# Common sample interface for linear systems
function sample(
    type::LinSysVecRowDetermCyclic,
    A::AbstractArray,
    b::AbstractVector,
    x::AbstractVector,
    iter::Int64
)

    eqn_ind = mod(iter, 1:length(b))
    return A[eqn_ind, :], b[eqn_ind]
end

#export LinSysVecRowDetermCyclic


text:
mutable struct MiniMaxConfigType{T}
      tol_converged::T # used to compare the score; this is the convergence criterion in step 4 of algorithm 2 from Filip's thesis.
      tol_𝓧_spacing::T # minimum distance between reference positions.
      tol_derivative_zero::T
      tol_no_update::T
      #tol_h::T # new reference points are allowed to be this much
               #  smaller than the absolute level error.

      N_samples_interval::Int # number of samples to take for each error function's interval.
      N_test_positions::Int # number of samples used to evaluate the convergence score.
      max_iters::Int
      L_interval::Int

      verbose_flag::Bool
      plot_flag::Bool # if true, plots error function and new reference positions every iteration.
 end

max_stars_repo_name: RoyCCWang/FIRRemez

max_issues_repo_name: RoyCCWang/FIRRemez

max_forks_repo_name: RoyCCWang/FIRRemez


text:
using MultivariateFunctions
using DataFrames
using Dates
using Random
using Distributions
tol = 1e-14
fun1 = PE_Function(1.0,1.0,1.0,1)
fun2 = PE_Function(1.2,1.0,1.0,1)
fun3 = PE_Function(1.5,1.0,2.0,3)
fun4 = PE_Function(3.0,1.0,3.0,2)
fun5 = PE_Function(0.0,2.0,1.0,1)
fun6 = PE_Function(2.0,0.0,1.0,1)
fun7 = PE_Function(3.0,0.0,2.0,4)
fun8 = PE_Function(3.0,0.0,2.0,2)
fun9 = PE_Function(3.0,0.0,Date(2015,1,1),4)
fun10 = PE_Function(3.0,0.0,Date(2016,1,1),2)

limits = Dict{Symbol, Tuple{Union{Symbol,Float64},Union{Symbol,Float64}}}( [:x,:y] .=> [(:x_left, 5.6), (:y_left, :y_right) ]   )

mfun1 = PE_Function(1.8, Dict([:x, :y, :z] .=> [PE_Unit(1.0,1.0,1), PE_Unit(1.0,1.0,1), PE_Unit(1.0,2.0,3)]))
abs(mfun1.multiplier_ - 1.2*1.5) < 1e-10
mfun2 = PE_Function(4.5, Dict([:x, :y] .=> [PE_Unit(1.0,2.0,3), PE_Unit(1.0,3.0,2)]))
mfun3 = PE_Function(3.0, Dict([:z, :y] .=> [ PE_Unit(1.0,3.0,2), PE_Unit(1.0,1.0,1)]))
mfun4 = PE_Function(6.0, Dict([:x, :y] .=> [PE_Unit(0.0,1.0,1), PE_Unit(0.0,2.0,4)]))
mfun5 = PE_Function(6.0, Dict([:x, :z] .=> [PE_Unit(0.0,1.0,1), PE_Unit(0.0,2.0,2)]))
mfun6 = PE_Function(18.0, Dict([:x, :y, :z] .=> [PE_Unit(0.0,Date(2015,1,1),4), PE_Unit(0.0,1.0,1), PE_Unit(0.0,Date(2016,1,1),2)]))

mSumFunction1 = Sum_Of_Functions([mfun1, mfun2])
mSumFunction2 = Sum_Of_Functions([mSumFunction1, mfun3])
mSumFunction3 = Sum_Of_Functions([mSumFunction1, mSumFunction2])

coordinates = Dict([:x, :y, :z] .=> [0.5, 2.0, 3.0])

# Conversion
typeof(convert(MultivariateFunctions.Sum_Of_Functions, mfun1)) == MultivariateFunctions.Sum_Of_Functions

function test_result(func, eval_to)
    return abs(evaluate(func, coordinates) - eval_to) < 1e-05
end

# Testing numerical results
mfun1_value = evaluate(fun1, coordinates[:x]) * evaluate(fun2, coordinates[:y]) * evaluate(fun3, coordinates[:z])
test_result(mfun1, mfun1_value)
mfun2_value = evaluate(fun3, coordinates[:x]) * evaluate(fun4, coordinates[:y])
test_result(mfun2, mfun2_value)
test_result(mSumFunction1, mfun1_value + mfun2_value)
mfun3_value = evaluate(fun4, coordinates[:z]) * evaluate(fun1, coordinates[:y])
test_result(mSumFunction3, 2* mfun1_value + 2*mfun2_value + mfun3_value)
mfun6_value = evaluate(fun4, coordinates[:z]) * evaluate(fun1, coordinates[:y])
# Testing partial evaluations
mfun1_partiallyevaluated = evaluate(mfun1, Dict{Symbol,Float64}([:x, :y] .=> [2.0,2.0]))
abs(mfun1_partiallyevaluated.multiplier_ - 13.30030097807517) < tol
underlying_dimensions(mfun1_partiallyevaluated) == Set(Symbol[:z])

# Testing partial evaluations of Sum_Of_Functions
mSumFunction1_partiallyevaluated = evaluate(mSumFunction1, Dict{Symbol,Float64}([:x, :y] .=> [3.0,2.0]))
abs(mSumFunction1_partiallyevaluated.functions_[1].multiplier_ - 72.3079329234756) < tol
underlying_dimensions(mSumFunction1_partiallyevaluated) == Set(Symbol[:z])

# Additions and subtractions
test_result( ((7 - (5.0 - mfun1))  + 5.0) + 7        , 14 + mfun1_value                )
test_result( ((7 - (5.0 - mSumFunction1))  + 5.0) + 7, 14 +  mfun1_value + mfun2_value )
# multiplications
test_result( (5 * mfun1) * 7.0        , 35* mfun1_value                  )
test_result( (5 * mSumFunction1) * 7.0, 35*  (mfun1_value + mfun2_value) )
# Divisions
test_result( (mfun1/2) / 3.0       ,  mfun1_value/6                  )
test_result( (mSumFunction1/2) /3.0, (mfun1_value + mfun2_value)/6 )
# powers
test_result( mfun1^0, 1.0)
test_result( mSumFunction1^0, 1)
test_result( mfun1^1, mfun1_value)
test_result( mSumFunction1^1, (mfun1_value + mfun2_value))
test_result( mfun1^2, mfun1_value^2)
test_result( mSumFunction1^2, (mfun1_value + mfun2_value)^2)
# multiplying functions
test_result( mfun1 * mfun2, mfun1_value * mfun2_value)
# Using Dates
dateCoordinates = Dict([:x, :y, :z] .=> [Date(2017,1,1), 2, 17.0])
dateCoordinates2 = Dict([:x, :y, :z] .=> [years_from_global_base(dateCoordinates[:x]), 2.0, 17.0])
abs(evaluate(mfun6, dateCoordinates) - evaluate(mfun6, dateCoordinates2) ) < 1e-14

# Derivatives
test_result(derivative(mfun4, Dict{Symbol,Int}(:x => 2)) , 0.0)
test_result(derivative(mfun4, Dict{Symbol,Int}(:y => 4)), 24 * 6 * (-0.5) )
test_result(derivative(mfun5, Dict([:x, :z] .=> [1,1])), 12 * (3-2)^1 )
# Integration
int_limits = Dict{Symbol,Tuple{Float64,Float64}}([:x, :y] .=> [(2.0,2.5), (3.0,3.5)])
abs(integral(mfun4, int_limits) - 0.6 * ((int_limits[:x][2] - 1)^2 - (int_limits[:x][1]-1)^2)*((int_limits[:y][2]-2)^5 - (int_limits[:y][1]-2)^5)) < tol
abs(integral(Sum_Of_Functions([mfun4, mfun4]), int_limits) - 1.2 * ((int_limits[:x][2] - 1)^2 - (int_limits[:x][1]-1)^2)*((int_limits[:y][2]-2)^5 - (int_limits[:y][1]-2)^5)) < tol

# Integration of constant function
constant_func = PE_Function(5.0)
abs(integral(constant_func, int_limits) - 5.0/4 ) < tol

## Dataframe evaluations
Random.seed!(1992)
nObs = 100
dd = DataFrame()
dd[:x] = rand( Normal(),nObs) + 0.1 .* rand( Normal(),nObs)
dd[:z] = rand( Normal(),nObs) + 0.1 .* rand( Normal(),nObs)
dd[:y] = (dd[:x] ) .* (dd[:z]) .+ dd[:x] + rand( Normal(),nObs)
sum(abs.(evaluate(mfun1, dd) + evaluate(mfun2, dd) - evaluate(mSumFunction1, dd))) < tol


text:
const RGBAf0 = RGBA{Float32}
# Solution 1: strictly type fields
const TorVector{T} = Union{Vector{T}, T}

@enum Space Pixel Data

mutable struct Camera
    pixel_space::Observable{Mat4f0}
    view::Observable{Mat4f0}
    projection::Observable{Mat4f0}
    projectionview::Observable{Mat4f0}
    resolution::Observable{Vec2f0}
    eyeposition::Observable{Vec3f0}
    steering_nodes::Vector{Any}
end

function Camera()
    Camera(
        Observable(Mat4f0(I)),
        Observable(Mat4f0(I)),
        Observable(Mat4f0(I)),
        Observable(Mat4f0(I)),
        Observable(Vec2f0(0)),
        Observable(Vec3f0(1)),
        []
    )
end

function connect!(camera::Camera, pixel_area::Observable{FRect2D})
    on(pixel_area) do window_size
        nearclip = -10_000f0
        farclip = 10_000f0
        w, h = Float32.(widths(window_size))
        camera.pixel_space[] = orthographicprojection(0f0, w, 0f0, h, nearclip, farclip)
        camera.resolution[] = Vec2f0(w, h)
    end
end

struct Transformation
    parent::Base.RefValue{Any}
    translation::Observable{Vec3f0}
    scale::Observable{Vec3f0}
    rotation::Observable{Quaternionf0}
    model::Observable{Mat4f0}
    # data conversion node, for e.g. log / log10 etc
    transform_func::Observable{Any}
    function Transformation(translation, scale, rotation, model, transform_func)
        return new(
            Base.RefValue{Any}(),
            translation, scale, rotation, model, transform_func
        )
    end
end

function transformationmatrix(translation, scale)
    T = eltype(translation)
    T0, T1 = zero(T), one(T)
    Mat{4}(
        scale[1],T0,  T0,  T0,
        T0,  scale[2],T0,  T0,
        T0,  T0,  scale[3],T0,
        translation[1],translation[2],translation[3], T1
    )
end
function transformationmatrix(translation, scale, rotation::Quaternion)
    trans_scale = transformationmatrix(translation, scale)
    return trans_scale * Mat4f0(rotation)
end

function Transformation(transform_func=identity;
                        translation=Observable(Vec3f0(0)),
                        scale=Observable(Vec3f0(1)),
                        rotation=Observable(Quaternionf0(0, 0, 0, 1)))
    model = map(transformationmatrix, translation, scale, rotation)
    return Transformation(
        translation,
        scale,
        rotation,
        model,
        Observable{Any}(transform_func)
    )
end


text:
module GrowingRegions

using Printf
import SearchBall, GeometryTypes
const sb = SearchBall
const gt = GeometryTypes

function init_strategies(config; kwargs...)
    # myPlayers
    ws = config.ws

    xlen, ylen = Int(ceil(config.field_size[1]))*10, Int(ceil(config.field_size[2]))*10
    config.global_communication["sight_regions"] = fill(0.0, (xlen, ylen))
    config.global_communication["x_raster"] = sb.raster(-config.field_size[1]/2, config.field_size[1]/2, xlen)
    config.global_communication["y_raster"] = sb.raster(-config.field_size[2]/2, config.field_size[2]/2, ylen)

    if length(ws.myPlayers) > 0
        # sb.add_check_button!(config, "draw function", "draw background")do widget
        #     if sb.get_value(widget, Bool)
        #         run_strategies(config, 0.0, :reset)
        #     else
        #         sb.clear_background!(config)
        #     end
        # end

    # radiobuttons = []
    #     for myIdx in eachindex(ws.myPlayers)
    #         myPlayer = ws.myPlayers[myIdx]
    #         rb = sb.add_radio_button!(config, "rb$myIdx", string(myIdx))do widget
    #             # if sb.get_manip_value(config, "draw function")
    #                 run_strategies(config, 0.0, :reset)
    #             # end
    #         end
    #         push!(radiobuttons, rb)
    #     end
    #     sb.add_radio_button_group!(config, "show function for", radiobuttons)
    end
    sb.add_check_button!(config, "sight_regions", "show sight_regions") do widget
        run_strategies(config, 0.0, :reset)
    end
    sb.add_check_button!(config, "with hfov", "with hfov") do widget
        run_strategies(config, 0.0, :reset)
    end

    add_shadowcheckbuttons(ws.myPlayers, "myPlayer", config)
    add_shadowcheckbuttons(ws.opponents, "opponent", config)

    return
end

function update_sightregions!(config::sb.Config, dt::Number)
    x_raster = config.global_communication["x_raster"]
    y_raster = config.global_communication["y_raster"]
    regs = config.global_communication["sight_regions"]
    for (idx, x) in enumerate(x_raster), (idy, y) in enumerate(y_raster)
        position = [x, y]
        # @warn position
        # @warn config.ws.myPlayers[1]
        # @warn sb.is_in_line_of_sight(config.ws, config.ws.myPlayers[1], position)
        if any(map(p -> sb.iswithin_hfov(p, position) && sb.is_in_line_of_sight(config.ws, p, position, excludes=[p]), config.ws.myPlayers))
            # @warn "inside"
            regs[idx, idy] = 0.0
        elseif regs[idx, idy] ≤ 5.0
            regs[idx, idy] += dt/10
        end
    end
    # @warn regs[1, 1]
    # @warn regs[1, end]
    # @warn regs[end, 1]
    # @warn regs[end, end]
    return
end

function add_shadowcheckbuttons(player::AbstractVector{<:sb.Robot}, name::AbstractString, config::sb.Config)
    for myplayer_idx in eachindex(player)
        sb.add_check_button!(config, string(name, "Shadow", myplayer_idx), string(name, " ", myplayer_idx)) do widget
            if !sb.get_value(widget, Bool)
                sb.show_text!(config, string(name, " ", myplayer_idx, " area:"), "")
            end
            run_strategies(config)
        end
        sb.init_text!(config, string(name, " ", myplayer_idx, " area:"), "")
    end
end

function handle_shadow(player::AbstractVector{<:sb.Robot}, name::AbstractString, config::sb.Config)
    for myplayer_idx in eachindex(player)
        sb.remove_graphics_filter!(config) do s
            startswith(s, string(name, " ", myplayer_idx, " area:"))
        end

        if !sb.get_manip_value(config, string(name, "Shadow", myplayer_idx))
            continue
        end
        shadows = sb.get_shadow_polygons(player[myplayer_idx], config, sb.get_manip_value(config, "with hfov"))

        shadow_area = sum([sb.polygon_area(s) for s in shadows])

        sb.show_text!(config, string(name, " ", myplayer_idx, " area:"), @sprintf("%.2fm²",shadow_area))

        for (idx, shadow) in enumerate(shadows)
            if isempty(shadow)
                continue
            end
            sb.show_graphic!(config, string(name, " ", myplayer_idx, " area:", idx), shadow[vcat(1:end, 1)], sb.GraphicLine(fill=true, fillcolor=[0.0, 0.0, 0.0, 0.2]))
        end
    end
end

function calc_myPlayers(config::sb.Config, dt::Number, set_type::Symbol, moved_object::Union{sb.Object{<:Number}, Nothing}=nothing)
    ws = config.ws

    handle_shadow(ws.myPlayers, "myPlayer", config)
    handle_shadow(ws.opponents, "opponent", config)

    if length(ws.myPlayers) > 0
        if dt > 0
            update_sightregions!(config, dt)
        end
        if sb.get_manip_value(config, "sight_regions")
            sb.send_background(config.global_communication["sight_regions"], [true], config, (0.0, 5.0), allways_repaint=true)
            # @warn "sended"
        else
            sb.clear_background!(config)
        end
        # myforfunction = parse(Int, sb.get_manip_value(config, "show function for"))

        # if set_type ∈ [:reset, :release] && sb.get_manip_value(config, "draw function") && moved_object!=ws.myPlayers[myforfunction] && (sb.get_manip_value(config, "with hfov") || length(ws.myPlayers) > 1 || length(ws.opponents) > 0)
        #     vals, ins = sb.sample(config, myforfunction, sb.get_manip_value(config, "with hfov"))
        #     sb.send_background(vals, ins, config)
        # end
    end

    return
end

function calc_opponend(config::sb.Config, dt::Number, set_type::Symbol, moved_object::Union{sb.Object{<:Number}, Nothing}=nothing)
    ws = config.ws

    # do stuff

    return
end

function run_strategies(config::sb.Config, dt::Number=0, set_type::Symbol=:default, moved_object::Union{sb.Object{<:Number}, Nothing}=nothing)
    calc_myPlayers(config, dt, set_type, moved_object)
    calc_opponend(config, dt, set_type, moved_object)
    return
end

end


text:
using ConstrainedDynamics

# Parameters
ex = [1.;0.;0.]
ey = [0.;1.;0.]
ez = [0.;0.;1.]

l1 = 1.0
x, y = .1, .1
b1 = Box(x, y, l1, l1, color = RGBA(1., 1., 0.))

vert11 = [0.;0.;l1 / 2]
vert12 = -vert11

# Initial orientation
phi1, phi2, phi3 = pi / 4, -pi / 4, pi / 2
q1, q2, q3 = Quaternion(RotX(phi1)), Quaternion(RotX(phi2)), Quaternion(RotX(phi3))

# Links
origin = Origin{Float64}()
link1 = Body(b1)
link2 = Body(b1)

# Constraints
joint0to12 = EqualityConstraint(CylindricalFree(origin, link1, zeros(3), vert11, ey), Spherical(origin, link2, zeros(3), vert11))
joint1to2 = EqualityConstraint(CylindricalFree(link1, link2, zeros(3), zeros(3), ex))

links = [link1; link2]
constraints = [joint0to12;joint1to2]
shapes = [b1]

mech = Mechanism(origin, links, constraints, shapes = shapes)
setPosition!(mech,origin,link1,p1 = [0;-0.5 * sqrt(2);0],p2 = vert11,Δq = q1)
setPosition!(mech,origin,link2,p2 = vert11,Δq = q2)

simulate!(mech,save = true)
visualize!(mech)


text:
using InvariantMeasures
using ValidatedNumerics

using Plots
using LaTeXStrings
using StatsPlots

pgfplotsx()

function onegrid(T, Btype, size)

    time_assembling = @elapsed begin
        D = Mod1Dynamic(T)
        # different backend, a tad slower
        # D = mod1_dynamic(x -> x->2*x+0.5*x*(1-x))
        B = Btype(size)
        Q = DiscretizedOperator(B, D)
    end

    time_norms = @elapsed norms = powernormbounds(B, D; Q=Q)
    time_eigen = @elapsed w = invariant_vector(B, Q)
    time_error = @elapsed error = distance_from_invariant(B, D, Q, w, norms)

    times = [time_error time_eigen time_norms time_assembling 0]

    return times, error, (B, D, norms, time_assembling+time_norms)
end

function twogrid(Btype, size, (B, D, norms, time_coarse))
    time_assembling_fine = @elapsed begin
        B_fine = Btype(size)
        Q_fine = DiscretizedOperator(B_fine, D)
    end

    time_norms_fine = @elapsed norms_fine = finepowernormbounds(B, B_fine, D, norms; Q_fine=Q_fine)
    time_eigen_fine = @elapsed w_fine = invariant_vector(B_fine, Q_fine)
    time_error_fine = @elapsed error_fine = distance_from_invariant(B_fine, D, Q_fine, w_fine, norms_fine)

    times_fine = [time_error_fine time_eigen_fine time_norms_fine time_assembling_fine time_coarse]
    return times_fine, error_fine
end

function time_convergence_plot(T, Btype, k_onegrid, k_twogrid)
    n_onegrid = 2 .^ k_onegrid

    times_onegrid = fill(NaN, length(n_onegrid), 5)
    errors_onegrid = fill(NaN, size(n_onegrid))
    onegrid(T, Btype, n_onegrid[1]) #warmup
    for i in 1:length(n_onegrid)
        times_onegrid[i,:], errors_onegrid[i], _ = onegrid(T, Btype, n_onegrid[i])
    end

    n_twogrid = 2 .^ k_twogrid
    _, _, coarse_data = onegrid(T, Btype, 1024)
    times_twogrid = fill(NaN, length(n_twogrid), 5)
    errors_twogrid = fill(NaN, size(n_twogrid))

    twogrid(Btype, n_twogrid[1], coarse_data) #warmup
    for i in 1:length(n_twogrid)
        times_twogrid[i,:], errors_twogrid[i] = twogrid(Btype, n_twogrid[i], coarse_data)
    end

    p1 = groupedbar(
        times_onegrid,
        bar_position = :stack,
        legend = false,
        label = ["err" "eigen" "norm" "matrix" "coarse"],
        title = "CPU time breakdown (s)",
        xticks = (1:length(n_onegrid), LaTeXString.(raw"2^{" .* string.(k_onegrid) .* raw"}")),
        link = :y,
    )

    p2 = groupedbar(
        times_twogrid,
        bar_position = :stack,
        legend = false,
        label = ["err" "eigen" "norm" "matrix" "coarse"],
        title = "CPU time breakdown (s)",
        xticks = (1:length(n_twogrid), LaTeXString.(raw"2^{" .* string.(k_twogrid) .* raw"}")),
        link = :y,
    )

    p3 = plot(
        n_onegrid,
        errors_onegrid,
        title = "Error",
        mark = :dot,
        yscale = :log10,
        xscale = :log10,
        xticks = (n_onegrid, LaTeXString.(raw"2^{" .* string.(k_onegrid) .* raw"}" )),
        label = "One-grid strategy",
        legend = :bottomleft,
        link = :y,
    )

    p4 = plot(
        n_twogrid,
        errors_twogrid,
        title = "Error",
        mark = :dot,
        yscale = :log10,
        xscale = :log10,
        color = :red,
        xticks = (n_twogrid, LaTeXString.(raw"2^{" .* string.(k_twogrid) .* raw"}")),
        label = "Two-grid strategy",
        legend = :bottomleft,
        link = :y,
    )

    p = plot(p1, p2, p3, p4)

end

p = time_convergence_plot(x->2x+0.5x*(1-x), Ulam, 8:11, 11:14)

# large-scale version:
# p = time_convergence_plot(x->2x+0.5x*(1-x), Ulam, 8:14, 13:26); savefig(p, "time_convergence_plot.tikz"); savefig(p, "time_convergence_plot.pdf");

# savefig(p, "time_convergence_plot.tikz")


text:
using MPI
using Test
using LinearAlgebra
using Random
using StaticArrays
using ClimateMachine
using ClimateMachine.SystemSolvers
using ClimateMachine.MPIStateArrays
using CUDAapi
using Random
using KernelAbstractions
using CuArrays

let
    Random.seed!(1235)
    if CUDAapi.has_cuda_gpu()
        arrays = [Array, CuArray]
    else
        arrays = [Array]
    end
    # Initialize
    ClimateMachine.init()
    T = Float64
    mpicomm = MPI.COMM_WORLD
    # set the error threshold
    err_thresh = 1e-8

    @kernel function multiply_A_kernel!(x, A, y, n1, n2)
        I = @index(Global)
        for i in 1:n1
            tmp = zero(eltype(x))
            for j in 1:n2
                tmp += A[i, j, I] * y[j, I]
            end
            x[i, I] = tmp
        end
    end

    function multiply_by_A!(
        x,
        A,
        y,
        n1,
        n2;
        ndrange = size(x[1, :]),
        cpu_threads = Threads.nthreads(),
        gpu_threads = 256,
    )
        if isa(x, Array)
            kernel! = multiply_A_kernel!(CPU(), cpu_threads)
        else
            kernel! = multiply_A_kernel!(CUDA(), gpu_threads)
        end
        return kernel!(x, A, y, n1, n2, ndrange = ndrange)
    end
    # for defining linear_operator
    function closure_linear_operator_multi!(A, n1, n2, n3)
        function linear_operator!(x, y)
            event = multiply_by_A!(x, A, y, n1, n2, ndrange = n3)
            wait(event)
            return nothing
        end
    end

    # Test 1: Basic Functionality
    for array in arrays
        ArrayType = array
        n = 100  # size of vector space
        ni = 10 # number of independent linear solves
        b = ArrayType(randn(n, ni)) # rhs
        x = ArrayType(randn(n, ni)) # initial guess
        A = ArrayType(randn((n, n, ni)) ./ sqrt(n) .* 1.0)
        for i in 1:n
            A[i, i, :] .+= 1.0
        end
        ss = size(b)[1]
        gmres = BatchedGeneralizedMinimalResidual(
            b,
            ArrayType = ArrayType,
            subspace_size = ss,
        )
        for i in 1:ni
            x[:, i] = A[:, :, i] \ b[:, i]
        end
        sol = copy(x)
        x += ArrayType(randn(n, ni) * 0.01 * maximum(abs.(x)))
        y = copy(x)
        linear_operator! = closure_linear_operator_multi!(A, size(A)...)
        iters = linearsolve!(linear_operator!, gmres, x, b; max_iters = ss)
        linear_operator!(y, x)
        test_name = "(" * string(ArrayType) * ") Basic Test"
        @testset "$test_name" begin
            @test norm(y - b) / norm(b) < err_thresh
        end
    end
    ###
    # Test 2: MPI State Array
    for array in arrays
        ArrayType = array
        Random.seed!(1235)
        n1 = 8
        n2 = 3
        n3 = 10
        mpi_b = MPIStateArray{T}(mpicomm, ArrayType, n1, n2, n3)
        mpi_x = MPIStateArray{T}(mpicomm, ArrayType, n1, n2, n3)
        mpi_A = ArrayType(randn(n1 * n2, n1 * n2, n3))

        # need to make sure that mpi_b and mpi_x are reproducible
        mpi_b.data[:] .= ArrayType(randn(n1 * n2 * n3))
        mpi_x.data[:] .= ArrayType(randn(n1 * n2 * n3))
        mpi_y = copy(mpi_x)

        # for defining linear_operator
        function closure_linear_operator_mpi!(A, n1, n2, n3)
            function linear_operator!(x, y)
                alias_x = reshape(x.data, (n1, n3))
                alias_y = reshape(y.data, (n1, n3))
                event =
                    multiply_by_A!(alias_x, A, alias_y, n1, n2, ndrange = n3)
                wait(event)
                return nothing
            end
        end

        gmres = BatchedGeneralizedMinimalResidual(
            mpi_b,
            ArrayType = ArrayType,
            m = n1 * n2,
            n = n3,
        )

        # Now define the linear operator
        linear_operator! = closure_linear_operator_mpi!(mpi_A, size(mpi_A)...)
        iters = linearsolve!(
            linear_operator!,
            gmres,
            mpi_x,
            mpi_b;
            max_iters = n1 * n2,
        )
        linear_operator!(mpi_y, mpi_x)
        # check one in the batch
        sol = mpi_A[:, :, 1] \ mpi_b.data[:, :, 1][:]
        test_name = "(" * string(ArrayType) * ") MPIStaterray Test"
        @testset "$test_name" begin
            @test norm(mpi_y.data - mpi_b.data) / norm(mpi_b) < err_thresh
            @test norm(sol - mpi_x.data[:, :, 1][:]) < err_thresh
        end
    end

    ###
    # Test 3: Columnwise test (CPU ONLY)
    Random.seed!(1235)
    # ## More Complex Example
    function closure_linear_operator!(A, tup)
        function linear_operator!(y, x)
            alias_x = reshape(x, tup)
            alias_y = reshape(y, tup)
            for i6 in 1:tup[6]
                for i4 in 1:tup[4]
                    for i2 in 1:tup[2]
                        for i1 in 1:tup[1]
                            tmp = alias_x[i1, i2, :, i4, :, i6][:]
                            tmp2 = A[i1, i2, i4, i6] * tmp
                            alias_y[i1, i2, :, i4, :, i6] .=
                                reshape(tmp2, (tup[3], tup[5]))
                        end
                    end
                end
            end
        end
    end

    tup = (3, 4, 7, 6, 5, 2)

    B = [
        randn(tup[3] * tup[5], tup[3] * tup[5])
        for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]
    ]
    columnwise_A = [
        B[i1, i2, i4, i6] + 10I
        for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]
    ]
    # taking the inverse of A isn't great, but it is convenient
    columnwise_inv_A = [
        inv(columnwise_A[i1, i2, i4, i6])
        for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]
    ]
    columnwise_linear_operator! = closure_linear_operator!(columnwise_A, tup)
    columnwise_inverse_linear_operator! =
        closure_linear_operator!(columnwise_inv_A, tup)

    mpi_tup = (tup[1] * tup[2] * tup[3], tup[4], tup[5] * tup[6])
    b = randn(mpi_tup)
    x = copy(b)
    columnwise_inverse_linear_operator!(x, b)
    x += randn((tup[1] * tup[2] * tup[3], tup[4], tup[5] * tup[6])) * 0.1

    reshape_tuple_f = tup
    permute_tuple_f = (5, 3, 1, 4, 2, 6)

    gmres = BatchedGeneralizedMinimalResidual(
        b,
        m = tup[3] * tup[5],
        n = tup[1] * tup[2] * tup[4] * tup[6],
        reshape_tuple_f = reshape_tuple_f,
        permute_tuple_f = permute_tuple_f,
        atol = eps(T),
        rtol = eps(T),
    )

    x_exact = copy(x)
    iters = linearsolve!(
        columnwise_linear_operator!,
        gmres,
        x,
        b,
        max_iters = tup[3] * tup[5],
    )

    @testset "(Array) Columnwise Test" begin
        columnwise_inverse_linear_operator!(x_exact, b)
        @test norm(x - x_exact) / norm(x_exact) < err_thresh
        columnwise_linear_operator!(x_exact, x)
        @test norm(x_exact - b) / norm(b) < err_thresh
    end
end

max_stars_repo_name: leios/CLIMA

max_issues_repo_name: leios/CLIMA

max_forks_repo_name: leios/CLIMA


text:
function GTest()
    return GTest(())
end

function g(obj::GTest, arg0::Vector{jdouble}, arg1::Vector{jlong})
    return jcall(obj, "g", jdouble, (Vector{jdouble}, Vector{jlong}), arg0, arg1)
end

function g_data_sets_comparison(obj::GTest, arg0::Vector{jlong}, arg1::Vector{jlong})
    return jcall(obj, "gDataSetsComparison", jdouble, (Vector{jlong}, Vector{jlong}), arg0, arg1)
end

function g_test(obj::GTest, arg0::Vector{jdouble}, arg1::Vector{jlong})
    return jcall(obj, "gTest", jdouble, (Vector{jdouble}, Vector{jlong}), arg0, arg1)
end

function g_test(obj::GTest, arg0::Vector{jdouble}, arg1::Vector{jlong}, arg2::jdouble)
    return jcall(obj, "gTest", jboolean, (Vector{jdouble}, Vector{jlong}, jdouble), arg0, arg1, arg2)
end

function g_test_data_sets_comparison(obj::GTest, arg0::Vector{jlong}, arg1::Vector{jlong})
    return jcall(obj, "gTestDataSetsComparison", jdouble, (Vector{jlong}, Vector{jlong}), arg0, arg1)
end

function g_test_data_sets_comparison(obj::GTest, arg0::Vector{jlong}, arg1::Vector{jlong}, arg2::jdouble)
    return jcall(obj, "gTestDataSetsComparison", jboolean, (Vector{jlong}, Vector{jlong}, jdouble), arg0, arg1, arg2)
end

function g_test_intrinsic(obj::GTest, arg0::Vector{jdouble}, arg1::Vector{jlong})
    return jcall(obj, "gTestIntrinsic", jdouble, (Vector{jdouble}, Vector{jlong}), arg0, arg1)
end

function root_log_likelihood_ratio(obj::GTest, arg0::jlong, arg1::jlong, arg2::jlong, arg3::jlong)
    return jcall(obj, "rootLogLikelihoodRatio", jdouble, (jlong, jlong, jlong, jlong), arg0, arg1, arg2, arg3)
end


text:
module Z3

using CxxWrap
using z3_jll

import Base: +, -, *, /, ^, ==, !=, !, <=, >=, <, >, xor, rem, mod, &, |, ~
import Base: min, max, abs, sqrt, sum
import Base: string, getindex, size, length, push!, isequal, hash
import Base: numerator, denominator

if !isdefined(z3_jll, :libz3jl)
    error("Platform not supported")
end

@wrapmodule(z3_jll.libz3jl)

function __init__()
    @initcxx
    __main_ctx[] = Context()
end

const __main_ctx = Ref{Context}()
main_ctx() = __main_ctx[]

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

const AstVector      = AstVectorTpl{Ast}
const ExprVector     = AstVectorTpl{Expr}
const SortVector     = AstVectorTpl{Sort}
const FuncDeclVector = AstVectorTpl{FuncDecl}

Base.iterate(x::AstVectorTpl, s = 1) = s <= length(x) ? (x[s], s + 1) : nothing

# ------------------------------------------------------------------------------
# Model

consts(m::Model) = (get_const_decl(m, i)() => get_const_interp(m, get_const_decl(m, i)) for i in 0:num_consts(m) - 1)
funcs(m::Model) = (get_func_decl(m, i) => get_func_interp(m, get_func_decl(m, i)) for i in 0:num_funcs(m) - 1)

Base.length(m::Model) = size(m)

# ------------------------------------------------------------------------------
# Expr

or(xs::ExprVector) = mk_or(xs)
function or(xs::Union{AbstractArray, Tuple})
    @assert length(xs) > 1
    or(ExprVector(ctx(first(xs)), xs))
end
or(xs...) = or(xs)

and(xs::ExprVector) = mk_and(xs)
function and(xs::Union{AbstractArray, Tuple})
    @assert length(xs) > 1
    and(ExprVector(ctx(first(xs)), xs))
end
and(xs...) = and(xs)

# ------------------------------------------------------------------------------
# ExprVector

CxxWrap.@cxxdereference function ExprVector(ctx::Context, xs)
    vec = ExprVector(ctx)
    for x in xs
        push!(vec, x)
    end
    vec
end

Base.promote(x::S, y::T) where {S <: Expr, T <: Bool}                = (x, bool_val(ctx(x), y))
Base.promote(x::S, y::T) where {S <: Expr, T <: Integer}             = (x, num_val(ctx(x), y, get_sort(x)))
Base.promote(x::S, y::T) where {S <: Expr, T <: AbstractFloat}       = promote(x, rationalize(y))
Base.promote(x::S, y::T) where {S <: Expr, T <: Rational{<:Integer}} = (x, real_val(ctx(x), numerator(y), denominator(y)))
Base.promote(x::T, y::S) where {S <: Expr, T <: Number}              = reverse(promote(y, x))

for op in [:+, :-, :*, :/, :^, :(==), :(!=), :(<=), :(>=), :(<), :(>), :&, :|, :xor]
    @eval begin
        $(op)(x::S, y::T) where {S <: Expr, T <: Number} = $(op)(promote(x, y)...)
        $(op)(x::T, y::S) where {S <: Expr, T <: Number} = $(op)(promote(x, y)...)
    end
end

# ------------------------------------------------------------------------------
# Solver

check(s::Solver, xs::AbstractVector{<:Expr}) = check(s, ExprVector(ctx(s), xs))

# ------------------------------------------------------------------------------
# Model

eval(m::Model, e::Expr, model_completion::Bool = false) = __eval(m, e, model_completion)

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

for T in [Ast, AstVectorTpl, Solver, Model, Symbol, Fixedpoint, Optimize, Params, ParamDescrs, Goal]
    Base.show(io::IO, x::T) = print(io, string(x))
end

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

# function __with_main_ctx(name::Core.Symbol, m::Method)
#     try
#         arg_types = fieldtypes(m.sig)[2:end]
#         params = [:($(Core.Symbol(string("x", i)))::$T) for (i, T) in enumerate(arg_types) if !(Context <: T)]
#         length(arg_types) == length(params) && return nothing
#         call_args = [Context <: T ? :(main_ctx()) : :($(Core.Symbol(string("x", i)))) for (i, T) in enumerate(arg_types)]
#         :($(name)($(params...)) = $(name)($(call_args...)))
#     catch
#         # TODO: varargs not supported yet
#         return nothing
#     end
# end

function __should_be_exported(s::Core.Symbol)
    str = string(s)
    # Private
    startswith(str, "__")      && return false
    startswith(str, "#")       && return false
    # Auxiliary types created by CxxWrap
    endswith(str, "Allocated") && return false
    endswith(str, "Ref")       && return false
    # Own module
    isequal(s, :Z3)            && return false
    # Methods from Base.MainInclude
    isequal(s, :include)       && return false
    isequal(s, :eval)          && return false
    # Enums
    eval(s) isa CheckResult    && return false
    eval(s) isa RoundingMode   && return false
    # Conflicting types of Z3 and Base/Core
    isequal(s, :Expr)          && return false
    isequal(s, :Symbol)        && return false
    # Export otherwise
    return true
end

for name in names(Z3, all = true, imported = false)
    if __should_be_exported(name)
        @eval export $name

        # for m in methods(eval(name))
        #     @eval $(__with_main_ctx(name, m))
        # end
    end
end

end # module


text:
# Generator ops have no input.  There are four types:
# Par: writes output only at initialization.  Later modified using gradients by update!.
# Arr: writes output only at initialization.  Never changed, used for constant arrays.
# Rnd: generates random output every call.
# Input: copies the next network input to its output every call.

# Fields:
# dims: gives information about the output size, used by infersize
# init: an array or a random generator, used to fill the output
# initialized: used by Arr and Par to make sure output written only at initialization

abstract GenOp <: Op
type Par <: GenOp; dims; init; initialized; Par(; dims=nothing, init=nothing, o...)=new(dims,init,false); end
type Arr <: GenOp; dims; init; initialized; Arr(; dims=nothing, init=nothing, o...)=new(dims,init,false); end
type Rnd <: GenOp; dims; init; Rnd(; dims=nothing, init=nothing, o...)=new(dims,init); end
type Input <: GenOp; Input(;o...)=new(); end

forw(p::Par,y;o...)=(p.initialized ? y : (p.initialized=true; genforw(p.init,y)))
forw(p::Arr,y;o...)=(p.initialized ? y : (p.initialized=true; genforw(p.init,y)))
forw(p::Rnd,y;o...)=genforw(p.init,y)
forw(::Input,y;o...)=error("Input has no forward function, forw(::Net) handles it.")
back(p::GenOp,y;o...)=error("Back should not be called on inputless $(typeof(p))")

genforw(init::BaseArray,y)=copysync!(y,init)
genforw(init::Rgen,y)=rgen!(init,y)
genforw(init,y)=error("Unknown init type: $init")

ninputs(::GenOp)=0
canoverwrite(::GenOp)=false
back_reads_x(::GenOp)=false
back_reads_y(::GenOp)=false

# We have three sources of information to infer the output size.
# ysize is the size of the output register.
# p.dims may be specified by the user.
# the size of p.init if it is an array (could be a random distribution)
# We need to check their consistency and possibly fill in missing information.
function infersize(p::GenOp,ysize)
    dims = inferhelp(nothing, ysize)
    isdefined(p,:dims) && (dims = inferhelp(dims, p.dims))
    isdefined(p,:init) && isa(p.init,BaseArray) && (dims = inferhelp(dims, size(p.init)))
    return (dims == nothing ? nothing : tuple(dims))
end

function inferhelp(a::Dims,b::Dims)
    length(a)==length(b) || throw(DimensionMismatch())
    map(a,b) do ai,bi
        ai == bi ? ai :
        ai == 0  ? bi :
        bi == 0  ? ai :
        throw(DimensionMismatch())
    end
end

inferhelp(::Void,::Void)=nothing
inferhelp(a::Dims,::Void)=a
inferhelp(::Void,b::Dims)=b


text:
using GLVisualize, GLAbstraction, GeometryTypes, Reactive, Colors, GLFW
function circlescale(v0, pressed)
    scale = v0
    if isempty(pressed)
        for i in eachindex(scale)
            scale[i] = Vec2f0(20)
        end
    else
        ls = linspace(1.0, 2.0, length(scale))
        for i in eachindex(scale)
            scale[i] = Vec2f0(20*ls[i])
        end
    end
    scale
end
function to_color(button)
    if button == GLFW.MOUSE_BUTTON_LEFT
        return GLVisualize.default(RGBA)
    elseif button == GLFW.MOUSE_BUTTON_MIDDLE
        return RGBA{Float32}(0.9,0.1, 0.3, 1.0)
    elseif button == GLFW.MOUSE_BUTTON_MIDDLE
        return RGBA{Float32}(0.9,0.1, 0.9, 1.0)
    end
    RGBA{Float32}(0.6,0.6,0.6,1.0)
end
function mouse_color(pressed)
    if isempty(pressed)
        return RGBA{Float32}(0,0,0,0.8)
    else
        return to_color(first(pressed))
    end
end

function transmat(x, pressed)
    c = !isempty(pressed)
    translationmatrix(Vec3f0(x, 0)) * scalematrix(Vec3f0(c ? 0.7 : 1.0))
end

function add_mouse(window)
    @materialize mouseposition, mouse_buttons_pressed = window.inputs
    model_matrix = map(transmat, mouseposition, mouse_buttons_pressed)
    cursor = map(RGBA{N0f8}, loadasset("cursor.png"))
    w,h = size(cursor)
    ratio = Float32(w/h)
    w, h = ratio*15f0, 15f0
    _view(visualize(
        cursor,
        model=model_matrix,
        primitive=SimpleRectangle(-1f0, -(h+1f0), w, h)
    ), window, camera=:fixed_pixel)
end


text:
#This file will contain some demo functions for DDPG algorithms: https://arxiv.org/abs/1509.02971

#Function optimising the Q-function depending on the replay buffer
function q_training_all(buffer::Traj_Data,Q,Q_mod,θa,θa_mod, opt_para_a, opt_para_q, params::Params)
    train_size  = buffer.n_traj
    shuffle_arr = shuffle(1:train_size)
    #Sample from Batch:
    for i_true = 1:train_size
        i_traj = shuffle_arr[i_true]
        traj   = buffer.traj[i_traj]

        #get the local gradients
        ĝ_q = grad_Q(Q,Q_mod,θa,θa_mod,traj,params)
        ĝ_a = grad_DP(θa,θa_mod,Q,Q_mod,traj,params)

        #Update the evaluated networks:
        Knet.update!(θa, ĝ_a, opt_para_a)
        Knet.update!(Q,  ĝ_q, opt_para_q)

        #Update the target networks as well: (but slower)
        Q_mod  = params.τ .* Q  + (1 - params.τ) .* Q_mod
        θa_mod = params.τ .* θa + (1 - params.τ) .* θa_mod
    end
end

#Function calculationg the q-loss for updating
@everywhere function loss_Q(Q,Q_mod,θa,θa_mod,traj::Trajectory,params::Params)
    #
    len = traj.len
    lossQ = 0
    γ = params.γ_disc

    #Over the whole minibatch:
    for i_t = 1:len
        #Take from struct and precalculate
        r_t       = traj.r_t[i_t]
        s_t       = traj.s_t[i_t,:]
        s_t2      = traj.s_t2[i_t,:]
        a_t       = traj.a_t[i_t,:]

        out_t     = mlp_nlayer_policy_norm(θa_mod,s_t2,params)
        in_q_mod  = cat(s_t2,out_t; dims = 1)
        in_q      = cat(s_t,a_t; dims = 1)
        #Get both values
        y_i       = r_t + γ * mlp_nlayer_value_norm(Q_mod,in_q_mod,params)
        Q_out     = mlp_nlayer_value_norm(Q,in_q,params)
        #accumulate loss
        lossQ   += (y_i - Q_out)^2
    end

    #Return the final quadratic loss per timestep
    return (lossQ / len)
end

@everywhere grad_Q = grad(loss_Q,1)

#Function calculationg the q-loss for updating
@everywhere function loss_DP2(θa,θa_mod,Q,Q_mod,traj::Trajectory,params::Params)
    #general setup
    len = traj.len
    loss = 0

    #Over the whole minibatch:
    for i_t = 1:len
        #Take from struct and precalculate
        r_t    = traj.r_t[i_t]
        s_t    = traj.s_t[i_t,:]
        a_t    = traj.a_t[i_t,:]

        #get inputs
        out_p  = mlp_nlayer_policy_norm(θa,s_t,params)
        #in_q   = cat(s_t,out_p; dims = 1)
        #out_q  = mlp_nlayer_value_norm(Q,in_q,params)

        x  = cat(s_t,out_p; dims = 1) #stuff the input together

        #Here comes the evaluation as autograd has problems with calling the combined version
        #----------------------------------------
        for i = 1:3:length(Q)-2
            #1. Get the new summed input
            x = Q[i] * x

            #2. Calculate the mean and std of the layer
            x_mean, x_std = get_norm(x)

            #3. Normalize activation by applying the norm and mean
            x_norm = (x .- x_mean) / x_std

            #4. Reparametrize the output
            x = x_norm .* Q[i+1] + Q[i+2]

            #5. Apply the non-linearity
            x = relu.(x)
        end
        out_q = (Q[end-1] * x .+ Q[end])[1]
        #-----------------------------------------

        loss  += out_q
    end

    #Return the final loss
    return (loss / len)
end

@everywhere grad_DP2 = grad(loss_DP2,1)

#Function calculationg the q-loss for updating
@everywhere function grad_DP(θa,θa_mod,Q,Q_mod,traj::Trajectory,params::Params)
    #general setup
    len = traj.len
    γ = params.γ_disc
    ĝ_DP = net_to_vec(θa)*0

    #Over the whole minibatch:
    for i_t = 1:len
        #Take from struct and precalculate
        r_t    = traj.r_t[i_t]
        s_t    = traj.s_t[i_t,:]
        a_t    = traj.a_t[i_t,:]

        #get inputs
        out_t  = mlp_nlayer_policy_norm(θa,s_t,params)
        in_q   = cat(s_t, out_t; dims = 1)

        #Get the derivation of the Q-function depending on action:
        ∇Q_a = simple_grad_q_norm(out_t,s_t,Q,params)

        #Get the gradient of the policy depending on state
        ∇θ_s = zeros(length(net_to_vec(θa)),params.l_out)
        for d = 1:params.l_out
            ∇θ_s[:,d] = vec(net_to_vec(simple_grad_policy(θa,s_t,params,d)))
        end

        #accumulate loss
        ĝ_DP += ∇θ_s * ∇Q_a
    end

    #Return the final gradient
    return vec_to_net(ĝ_DP ./ len, θa)
end

#Get the gradient of the Policy towards input
@everywhere function eval_p(θa,s_t,params::Params,dim)
    out_t = mlp_nlayer_policy_norm(θa,s_t,params)[dim]
    return sum(out_t) #after some thinking...
end
@everywhere simple_grad_policy = grad(eval_p,1)

#Get the gradient of the q function after a
@everywhere function eval_Q(a_t,s_t,Q,params::Params)
    x  = cat(s_t,a_t; dims = 1) #stuff the input together

    #Here comes the evaluation as autograd has problems with calling the combined version
    #----------------------------------------
    for i = 1:2:length(Q)-2
        x = relu.(Q[i] * x .+ Q[i+1])
    end
    value    = (Q[end-1] * x .+ Q[end])[1]
    #----------------------------------------

    return value
end
@everywhere simple_grad_q = grad(eval_Q,1)

#Get the gradient of the q function after a
@everywhere function eval_Q_norm(a_t,s_t,Q,params::Params)
    x  = cat(s_t,a_t; dims = 1) #stuff the input together

    #Here comes the evaluation as autograd has problems with calling the combined version
    #----------------------------------------
    for i = 1:3:length(Q)-2
        #1. Get the new summed input
        x = Q[i] * x

        #2. Calculate the mean and std of the layer
        x_mean, x_std = get_norm(x)

        #3. Normalize activation by applying the norm and mean
        x_norm = (x .- x_mean) / x_std

        #4. Reparametrize the output
        x = x_norm .* Q[i+1] + Q[i+2]

        #5. Apply the non-linearity
        x = relu.(x)
    end

    value    = (Q[end-1] * x .+ Q[end])[1]
    #----------------------------------------
    return value
end
@everywhere simple_grad_q_norm = grad(eval_Q_norm,1)

"""
Function enabling parallel execution of the DPG gradients
"""
function parallel_DDPG_update(buffer::Traj_Data,Q,Q_mod,θa,θa_mod, opt_para_a, opt_para_q, params::Params)
        #Take all partial traj and shuffle their order!
        n_traj = buffer.n_traj
        shuffle_arr = shuffle(1:n_traj)

        active_workers = workers()
        n_workers = length(active_workers)
        ĝ_a = copy(θa)
        ĝ_q = copy(Q)

        # Create Channels storing results and starting values
        input_channel  = RemoteChannel(()->Channel{Tuple}(length(workers())+1)) # will contain networks and traj
        output_channel = RemoteChannel(()->Channel{Tuple}(length(workers())+1)) # will contain networks

        # Create problem on every worker
        begin
            for p in workers()
                remote_do(distributed_gradients_dpg, p, params, input_channel, output_channel)
            end
        end

        # Start optimizations
        n_parallel = Int(ceil(n_traj / (2*n_workers)))
        prog = Progress(n_traj,1)

        for i_parallel = 1 : n_parallel

            #1. Send the work tasks:
            for i_w = 1 : n_workers
                i_true = ((i_parallel-1)*n_workers + i_w) #get true order
                if i_true <= n_traj
                    i_traj = shuffle_arr[i_true] #shuffle order
                    traj_local = buffer.traj[i_traj]
                    @async put!(input_channel, (traj_local, Q, Q_mod, θa, θa_mod) )
                end
            end

            #2. Take the results
            ĝ_q *= 0
            ĝ_a *= 0
            for i_w = 1: n_workers
                i_true = ((i_parallel-1)*n_workers + i_w)
                if i_true <= n_traj
                    #Take from Channel
                    gradients = take!(output_channel)
                    #Reorder
                    ĝ_q += gradients[1]
                    ĝ_a += gradients[2]
                end
            end
            ĝ_q /= n_workers
            ĝ_a /= n_workers

            #Update the evaluated networks:
            Knet.update!(θa, ĝ_a, opt_para_a)
            Knet.update!(Q,  ĝ_q, opt_para_q)

            #Update the target networks as well: (but slower)
            Q_mod  = params.τ .* Q  + (1 - params.τ) .* Q_mod
            θa_mod = params.τ .* θa + (1 - params.τ) .* θa_mod
        end

        # Tell workers that its done
        NaN_mat = zeros(length(params.batch_size))
        fill!(NaN_mat,NaN)
        ending_traj = init_traj(params,params.batch_size)
        ending_traj.r_t = NaN_mat

        for j = 1 : n_workers
            @async put!(input_channel, (ending_traj, Q, Q_mod, θa, θa_mod))
        end

        #Close Channels
        sleep(0.1)
        close(input_channel)
        close(output_channel)

        #Return the updated Networks
        return (Q, Q_mod, θa, θa_mod)
end

# Function running on the worker
@everywhere function distributed_gradients_dpg(params::Params, input_channel::RemoteChannel, output_channel::RemoteChannel)
    # Start loop where work is done
    while true
        # break loop if NaN state received
        try
            #Take the new Input
            input = take!(input_channel)

            #Reorder data:
            traj_local = input[1]
            Q          = input[2]
            Q_mod      = input[3]
            θa         = input[4]
            θa_mod     = input[5]

            #Check if u reached final
            if (sum(isnan, traj_local.r_t) > 0)
                GC.gc()
                break
            end

            #DO THE WORK:
            #-------------------------------------------------------------------
            ĝ_q  = grad_Q(Q,Q_mod,θa,θa_mod,traj_local,params)
            #ĝ_a  = grad_DP(θa,θa_mod,Q,Q_mod,traj_local,params)
            ĝ_a = grad_DP2(θa,θa_mod,Q,Q_mod,traj_local,params)
            #-------------------------------------------------------------------

            put!(output_channel, (ĝ_q , ĝ_a))

        catch
            println("Some problem occured, worker $(myid()) exiting trajectory optimization function.")
            break
        end
    end

    return nothing
end

max_stars_repo_name: NikonPic/REINFORCEMENT-JULIA

max_issues_repo_name: NikonPic/REINFORCEMENT-JULIA

max_forks_repo_name: NikonPic/REINFORCEMENT-JULIA


text:
using KSVM

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

msg("Test interface stability")

a = hcat(rand(100)+0, rand(100)+2)
b = hcat(rand(100)+2, rand(100)+0)
X = vcat(a,b)'
Y = vcat(ones(100), ones(100)*-1)

@test_throws ArgumentError ret = svm(X, Y, regtype = L1Penalty)

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

msg("Native interface using the svm function")

a = hcat(rand(100)+2, rand(100)+2)
b = hcat(rand(100)+0, rand(100)+0)
X = vcat(a,b)'
Y = vcat(ones(100), ones(100)*-1)

ret = svm(X, Y, bias = 0., loss = HingeLoss(), dual = false)
@test _accuracy(sign(predict(ret, X)), Y) < 1.0

ret = svm(X, Y, bias = 1., loss = HingeLoss(), dual = false)
@test _accuracy(sign(predict(ret, X)), Y) == 1.0

ret = svm(X, Y, bias = 0., loss = L2HingeLoss(), dual = false)
@test _accuracy(sign(predict(ret, X)), Y) < 1.0

ret = svm(X, Y, bias = 1., loss = L2HingeLoss(), dual = false)
@test _accuracy(sign(predict(ret, X)), Y) == 1.0

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

msg("Converting dual to primal solution")

model1 = svm(X, Y, bias = 1., dual = true)
model2 = convert(KSVM.PrimalSVM, model1)
@test typeof(model1) <: DualSVM
@test typeof(model2) <: PrimalSVM
@test all(round(predict(model1, X), 4) .== round(predict(model2, X), 4))

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

msg("Test callback function for primal solution")

ret = svm(X, Y, bias = 1., dual = false) do t, w, v, g
  @test length(w) == size(X,1) + 1
  @test length(g) == size(X,2) # can't provide grad of w ...
  t % 50 == 0 && msg2("$t : $v", newline = true)
end
@test _accuracy(sign(predict(ret, X)), Y) == 1.0

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

msg("Test callback function for primal solution (sparse)")

ret = svm(sparse(X), Y, bias = 1., dual = false) do t, w, v, g
  @test length(w) == size(X,1) + 1
  @test length(g) == size(X,2) # can't provide grad of w ...
  t % 50 == 0 && msg2("$t : $v", newline = true)
end
@test _accuracy(sign(predict(ret, X)), Y) == 1.0

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

# msg("Test callback function for primal solution (Regression.jl)")

# ret = svm(X, Y, bias = 1., dual = false, verbosity = :final, solver = Regression.BFGS()) do t, w, v, g
#   @test length(w) == size(X,1) + 1
#   @test length(g) == size(X,1) + 1
#   t % 10 == 0 && msg2("$t : $v", newline = true)
# end
# @test _accuracy(sign(predict(ret, X)), Y) == 1.0

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

msg("Test callback function for dual solution")

ret = svm(X, Y, bias = 1., dual = true) do t, alpha, v, g
  @test length(alpha) == size(X,2)
  @test length(g) == size(X,2)
  t % 50 == 0 && msg2("$t : $v", newline = true)
end
@test _accuracy(sign(predict(ret, X)), Y) == 1.0

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

msg("Test callback function for dual solution (sparse)")

ret = svm(sparse(X), Y, bias = 1., dual = true) do t, alpha, v, g
  @test length(alpha) == size(X,2)
  @test length(g) == size(X,2)
  t % 50 == 0 && msg2("$t : $v", newline = true)
end
@test _accuracy(sign(predict(ret, X)), Y) == 1.0


text:
using Test
using GeoParams
@testset "EnergyParameters.jl" begin

# This tests the MaterialParameters structure
CharUnits_GEO   =   GEO_units(viscosity=1e19, length=10km);

# Heat capacity ---------

# Constant heat capacity
cp1      =  ConstantHeatCapacity()
info     =  param_info(cp1)
@test isbits(cp1)
@test cp1.cp.val == 1050.0

cp1     = nondimensionalize(cp1,CharUnits_GEO)
@test cp1.cp.val ≈ 1.3368075000000002e22
@test compute_heatcapacity(cp1, 10.0, 100.0) ≈ 1.3368075000000002e22  # compute

# Temperature-dependent heat capacity
# dimensional
T        =   250.0:100:1250;
cp2      =   T_HeatCapacity_Whittacker()
Cp       =   similar(T)
@test isbits(cp2)
compute_heatcapacity!(Cp, cp2, T)
@test sum(Cp) ≈ 11667.035717418683

# nondimensional
cp2_nd   =   T_HeatCapacity_Whittacker()
cp2_nd   =   nondimensionalize(cp2_nd,CharUnits_GEO)
T_nd     =   Float64.(T*K/CharUnits_GEO.Temperature)
Cp_nd    =   similar(T)
compute_heatcapacity!(Cp_nd, cp2_nd, T_nd)
@test sum(Cp_nd) ≈ 1.4853886523631602e23

# Dimensionalize again and double-check the results
@test sum(abs.(ustrip.(Cp_nd*CharUnits_GEO.heatcapacity) - Cp)) < 1e-11

# Test with arrays
T_array     =  T*ones(10)'
Cp_array    =  similar(T_array)
P_array     =  similar(T_array)
compute_heatcapacity!(Cp_array, cp1, P_array, T_array)
@test Cp_array[1] ≈ 1.3368075000000002e22

Cp_array    =  similar(T_array)
P_array     =  similar(T_array)
compute_heatcapacity!(Cp_array, cp2, P_array, T_array)
@test sum(Cp_array[:,1]) ≈ 11667.035717418683

T_array     =  T*ones(10)'
Cp_array    =  zeros(size(T_array))
P_array     =  zeros(size(T_array))
compute_heatcapacity!(Cp_array, cp2, T_array)
@test sum(Cp_array[:,1]) ≈ 11667.035717418683

# Check that it works if we give a phase array
MatParam    =   Array{MaterialParams, 1}(undef, 2);
MatParam[1] =   SetMaterialParams(Name="Mantle", Phase=1,
                    HeatCapacity  = ConstantHeatCapacity());

MatParam[2] =   SetMaterialParams(Name="Crust", Phase=2,
                    HeatCapacity  = T_HeatCapacity_Whittacker());

Mat_tup = Tuple(MatParam)

# test computing material properties
n = 100;
Phases              = ones(Int64,n,n,n);
Phases[:,:,20:end] .= 2

Cp = zeros(size(Phases))
T  =  ones(size(Phases))*1500
P  =  zeros(size(Phases))

compute_heatcapacity!(Cp, Mat_tup, Phases, P, T)    # computation routine w/out P (not used in most heat capacity formulations)
num_alloc = @allocated compute_heatcapacity!(Cp, Mat_tup, Phases, P, T)
@test sum(Cp[1,1,:]) ≈ 121399.0486067196
@test num_alloc == 0

compute_heatcapacity!(Cp, Mat_tup, Phases, T)       # computation routine w/out P (not used in most heat capacity formulations)
@test sum(Cp[1,1,:]) ≈ 121399.0486067196

# test if we provide phase ratios
PhaseRatio  = zeros(n,n,n,3);
for i in CartesianIndices(Phases)
    iz = Phases[i]
    I = CartesianIndex(i,iz)
    PhaseRatio[I] = 1.0
end
compute_heatcapacity!(Cp, Mat_tup, PhaseRatio, P, T)
num_alloc = @allocated compute_heatcapacity!(Cp, Mat_tup, PhaseRatio, P, T)
@test sum(Cp[1,1,:]) ≈ 121399.0486067196
@test num_alloc == 0

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

# Conductivity ----------

# Constant

# Constant conductivity
cond      =   ConstantConductivity()
@test isbits(cond)
@test NumValue(cond.k) == 3.0
@test cond.k.unit==u"W"/K/m

cond = nondimensionalize(cond,CharUnits_GEO)
@test  NumValue(cond.k) ≈ 3.8194500000000007

@test compute_conductivity(cond, 100.0) ≈ 3.8194500000000007 # compute

# Temperature-dependent conductivity
# dimensional
T        =   Vector{Float64}(250:100:1250);
cond2    =   T_Conductivity_Whittacker()
k        =   compute_conductivity(cond2, T)
@test isbits(cond2)
@test sum(k) ≈ 27.503366436682285

# nondimensional
cond2_nd =   T_Conductivity_Whittacker()
cond2_nd =   nondimensionalize(cond2_nd,CharUnits_GEO)
T_nd     =   Float64.(ustrip.(T/CharUnits_GEO.Temperature))
k_nd     =   compute_conductivity(cond2_nd, T_nd)
@test sum(k_nd) ≈ 35.01591097886205

# Dimensionalize again and double-check the results
@test sum(abs.(ustrip.(k_nd*CharUnits_GEO.conductivity) - k)) < 1e-11

# Check if we use arrays
T_array     =  ustrip.(T)*ones(100)'
k_array     =  copy(T_array)
P_array     =  copy(T_array)

compute_conductivity!(k_array,cond,P_array,T_array)
@test k_array[1] ≈ 3.8194500000000007

compute_conductivity!(k_array,cond2,P_array,T_array)
@test sum(k_array) ≈ 2750.3366436682285

k_TP    =   Set_TP_Conductivity("LowerCrust")
compute_conductivity!(k_array, k_TP, P_array, T_array)
@test sum(k_array) ≈ 2055.7129327367625

# Check that it works if we give a phase array
MatParam    =   Array{MaterialParams, 1}(undef, 3);
MatParam[1] =   SetMaterialParams(Name="Mantle", Phase=1,
                    Conductivity  = ConstantConductivity());

MatParam[2] =   SetMaterialParams(Name="Crust", Phase=2,
                    Conductivity  = T_Conductivity_Whittacker());

MatParam[3] =   SetMaterialParams(Name="MantleLithosphere", Phase=3,
                    Conductivity = Set_TP_Conductivity("Mantle"));

Mat_tup = Tuple(MatParam)
# test computing material properties
n = 100;
Phases              = ones(Int64,n,n,n);
Phases[:,:,20:end] .= 2
Phases[:,:,60:end] .= 3

PhaseRatio  = zeros(n,n,n,3);
for i in CartesianIndices(Phases)
    iz = Phases[i]
    I = CartesianIndex(i,iz)
    PhaseRatio[I] = 1.0
end

k   =  zeros(size(Phases))
T   =  ones(size(Phases))*1500
P   =  zeros(size(Phases))

compute_conductivity!(k, Mat_tup, Phases, P, T)
num_alloc = @allocated compute_conductivity!(k, Mat_tup, Phases, P, T)
@test sum(k) ≈ 1.9216938849389635e6
@test num_alloc == 0

compute_conductivity!(k, Mat_tup, PhaseRatio, P, T)
num_alloc = @allocated compute_conductivity!(k, Mat_tup, PhaseRatio, P, T)
@test num_alloc == 0
@test sum(k) ≈ 1.9216938849389635e6

#

######

# TP-dependent conductivity for different predefines cases
T       =   Vector{Float64}(250:100:1250);
P       = 1e6*ones(size(T))/ustrip(uconvert(Pa,1MPa))  # must be in MPa!
List    = ["LowerCrust"   "Mantle"        "OceanicCrust"  "UpperCrust"]
Sol_kT  = [20.55712932736763 28.700405819019323 20.55712932736763 19.940302462417037]
for i=1:length(List)
    k_TP    =   Set_TP_Conductivity(List[i])
    k       =   compute_conductivity(k_TP,P,T)           # note that P must be in MPa
    @test sum(k) ≈ Sol_kT[i]

    k_TP_nd  =   deepcopy(k_TP)
    k_TP_nd  =   nondimensionalize(k_TP_nd,CharUnits_GEO)
    T_nd     =   Float64.(ustrip.(T/CharUnits_GEO.Temperature))
    P_nd     =   Float64.(ustrip(P/CharUnits_GEO.stress))
    k_nd     =   compute_conductivity(k_TP_nd,P_nd,T_nd)

    @test ustrip(sum(abs.(ustrip.(k_nd*CharUnits_GEO.conductivity) - k))) < 1e-11

end

T = [200. 300.; 400. 500.]
k1        =   compute_conductivity(cond2,T)

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

# Latent heat -----------
a = ConstantLatentHeat()
Q_L = compute_latent_heat(a)
@test isbits(a)
@test Q_L == 400

a   = nondimensionalize(a,CharUnits_GEO)
Q_L = compute_latent_heat(a)
@test Q_L ≈ 4e21
# -----------------------

# Radioactive heat ------
a = ConstantRadioactiveHeat()
H_r = compute_radioactive_heat(a)
@test isbits(a)
@test H_r ≈ 1.0e-6

a = nondimensionalize(a,CharUnits_GEO)
H_r = compute_radioactive_heat(a)
@test H_r == 0.1
# -----------------------

# Shear heating -------
Χ       = ConstantShearheating(1.0)
@test isbits(Χ)

# Define parameters as vectors
τ       = [1 2 3 4]*1e6
ε       = [1 0.1 0.1 1]
ε_el    = [0.01 0.01 0.01 0.01]

τ_2D       = [1 2; 3 4]*1e6
ε_2D       = [1 0.1; 0.1 1]
ε_el_2D    = [0.01 0.01; 0.01 0.01]

# With elasticity
H_s1 = compute_shearheating(Χ, τ,   ε,    ε_el   )
H_s2 = compute_shearheating(Χ, τ_2D,ε_2D, ε_el_2D)
@test H_s1 ≈ 5.4e6
@test H_s2 ≈ 5.4e6

# No elasticity
H_s3 = compute_shearheating(Χ, τ,   ε   )
H_s4 = compute_shearheating(Χ, τ_2D,ε_2D)
@test H_s3 ≈ 5.5e6
@test H_s4 ≈ 5.5e6

# Now in non-dimensional units
τ       = [1 2 3 4]
ε       = [1 0.1 0.1 1]
ε_el    = [0.01 0.01 0.01 0.01]

τ_2D       = [1 2; 3 4]
ε_2D       = [1 0.1; 0.1 1]
ε_el_2D    = [0.01 0.01; 0.01 0.01]
Χ       = nondimensionalize(Χ,CharUnits_GEO)
info    = param_info(Χ)

H_s1 = compute_shearheating(Χ, τ,   ε,    ε_el,  )
H_s2 = compute_shearheating(Χ, τ_2D,ε_2D, ε_el_2D)
H_s3 = compute_shearheating(Χ, τ,   ε)
H_s4 = compute_shearheating(Χ, τ_2D,ε_2D)
@test H_s1 ≈ 5.4
@test H_s2 ≈ 5.4
@test H_s3 ≈ 5.5
@test H_s4 ≈ 5.5
# -----------------------

end


text:
module AnsiColor
export colorize, black, red, green, yellow, blue, magenta, cyan, white, uncolor

COLORS = {
  "black"          => 0,
  "red"            => 1,
  "green"          => 2,
  "yellow"         => 3,
  "blue"           => 4,
  "magenta"        => 5,
  "cyan"           => 6,
  "white"          => 7,
  "default"        => 9,

  "light_black"    => 60,
  "light_red"      => 61,
  "light_green"    => 62,
  "light_yellow"   => 63,
  "light_blue"     => 64,
  "light_magenta"  => 65,
  "light_cyan"     => 66,
  "light_white"    => 67
  }

MODES = {
  "default"        => 0,
  "bold"           => 1,
  "underline"      => 4,
  "blink"          => 5,
  "swap"           => 7,
  "hide"           => 8
  }

COLOR_OFFSET       = 30
BACKGROUND_OFFSET  = 40

function colorize(color::String, str::String; background::String = "default", mode::String="default")
  "\033[$(MODES[mode]);$(COLOR_OFFSET + COLORS[color]);$(BACKGROUND_OFFSET + COLORS[background])m$(str)\033[0m"
end

function colorize(color::Symbol, str::String; background::String = "default", mode::String="default")
  colorize(string(color), str; background=background, mode=mode)
end

black(str::String)    = colorize(:black, str)
red(str::String)      = colorize(:red, str)
green(str::String)    = colorize(:green, str)
yellow(str::String)   = colorize(:yellow, str)
blue(str::String)     = colorize(:blue, str)
magenta(str::String)  = colorize(:magenta, str)
cyan(str::String)     = colorize(:cyan, str)
white(str::String)    = colorize(:white, str)
uncolor(str::String)  = colorize(:default, str)

light_black(str::String)    = colorize(:light_black, str)
light_red(str::String)      = colorize(:light_red, str)
light_green(str::String)    = colorize(:light_green, str)
light_yellow(str::String)   = colorize(:light_yellow, str)
light_blue(str::String)     = colorize(:light_blue, str)
light_magenta(str::String)  = colorize(:light_magenta, str)
light_cyan(str::String)     = colorize(:light_cyan, str)
light_white(str::String)    = colorize(:light_white, str)

end


text:
export
    Filter,
    init_filter,
    filter_step

"""
    abstract type Filter <: Algorithm end

General type of filtering algorithms.

Must implement init_filter and filter_step methods.
"""
abstract type Filter <: Algorithm end

"""
    init_filter!(::Filter, ::DynamicRuntime)

An interface for intializing the filter for a dynamic runtime.
"""
function init_filter!(::Filter, ::DynamicRuntime) end

"""
    filter_step(filter::Filter, runtime::Runtime, variables::Vector{Variable}, time::T, evidence::Dict{Symbol, Score})

Run one step of the filter by instantiating the given variables at the given time and passing in the given evidence.
"""
function filter_step(::Filter, ::DynamicRuntime{T}, ::Vector{Variable}, ::T, ::Dict{Symbol, Score}) where T end


text:
#####################
# Logical circuits
#####################

"""
Root of the logical circuit node hierarchy
"""
abstract type LogicalΔNode <: ΔNode end

"""
Root of the unstructured logical circuit node hierarchy
"""
abstract type UnstLogicalΔNode <: LogicalΔNode end

"""
A logical leaf node
"""
abstract type LogicalLeafNode <: UnstLogicalΔNode end

"""
A logical inner node
"""
abstract type LogicalInnerNode <: UnstLogicalΔNode end

"""
A logical literal leaf node, representing the positive or negative literal of its variable
"""
mutable struct LiteralNode <: LogicalLeafNode
    literal::Lit
    data
    bit::Bool
    LiteralNode(l) = new(l, nothing, false)
end

"""
A logical constant leaf node, representing true or false
"""
abstract type ConstantNode <: LogicalInnerNode end

"""
Constant True node
"""
mutable struct TrueNode <: ConstantNode
    data
    bit::Bool
    TrueNode() = new(nothing, false)
end

"""
Constant False node
"""
mutable struct FalseNode <: ConstantNode
    data
    bit::Bool
    FalseNode() = new(nothing, false)
end

"""
A logical conjunction node (And node)
"""
mutable struct ⋀Node <: LogicalInnerNode
    children::Vector{LogicalΔNode}
    data
    bit::Bool
    ⋀Node(c) = new(c, nothing, false)
end

"""
A logical disjunction node (Or node)
"""
mutable struct ⋁Node <: LogicalInnerNode
    children::Vector{LogicalΔNode}
    data
    bit::Bool
    ⋁Node(c) = new(c, nothing, false)
end

"""
A logical circuit represented as a bottom-up linear order of nodes
"""
const LogicalΔ = AbstractVector{<:LogicalΔNode}

"""
A unstructured logical circuit represented as a bottom-up linear order of nodes
"""
const UnstLogicalΔ = AbstractVector{<:UnstLogicalΔNode}

#####################
# traits
#####################

"""
Returns GateType of a node (Literal, Constant, And, Or)
"""
@inline GateType(::Type{<:LiteralNode}) = LiteralGate()
@inline GateType(::Type{<:ConstantNode}) = ConstantGate()
@inline GateType(::Type{<:⋀Node}) = ⋀Gate()
@inline GateType(::Type{<:⋁Node}) = ⋁Gate()

#####################
# methods
#####################

@inline node_type(::Type{<:UnstLogicalΔNode}) = UnstLogicalΔNode

"Get the logical literal in a given literal leaf node"
@inline literal(n::LiteralNode)::Lit = n.literal

"Get the logical constant in a given constant leaf node"
@inline constant(n::TrueNode)::Bool = true
@inline constant(n::FalseNode)::Bool = false

"Get the children of a given inner node"
@inline children(n::LogicalInnerNode) = n.children

@inline function conjoin_like(example::UnstLogicalΔNode, arguments::Vector)
    if isempty(arguments)
        TrueNode()
    # it's unclear if we want to also optimize the following
    # elseif length(arguments) == 1
    #     arguments[1]
    elseif example isa ⋀Node && children(example) == arguments
        example
    else
        ⋀Node(arguments)
    end
end

"Disjoin nodes in the same way as the example"
@inline function disjoin_like(example::UnstLogicalΔNode, arguments::Vector)
    if isempty(arguments)
        FalseNode()
    # it's unclear if we want to also optimize the following
    # elseif length(arguments) == 1
    #     arguments[1]
    elseif example isa ⋁Node && children(example) == arguments
        example
    else
        ⋁Node(arguments)
    end
end

"Construct a new literal node like the given node's type"
literal_like(::UnstLogicalΔNode, lit::Lit) = LiteralNode(lit)

"Generate a fully factorized (logistic regression) circuit over `n` variables"
function fully_factorized_circuit(n)
    lin = LogicalΔNode[]
    ors = map(1:n) do v
        v = Var(v)
        pos = LiteralNode( var2lit(v))
        push!(lin, pos)
        neg = LiteralNode(-var2lit(v))
        push!(lin, neg)
        or = ⋁Node([pos,neg])
        push!(lin, or)
        or
    end
    and = ⋀Node(ors)
    push!(lin, and)
    bias = ⋁Node([and])
    push!(lin, bias)
    lin
end

import Base.copy
function copy(n::ΔNode, depth::Int64)
    old2new = Dict{ΔNode, ΔNode}()
    copy_rec(n, depth, old2new)
end

function copy_rec(n::ΔNode, depth::Int64, old2new::Dict{ΔNode, ΔNode})
    if depth == 0 || isliteralgate(n) || isconstantgate(n)
        n
    else
        get!(old2new, n) do
            cns = map(children(n)) do c
                copy_rec(c, depth - 1, old2new)
            end
            copy_node(n, cns)
        end
    end
end

@inline copy_node(n::ΔNode, cns) = @assert false "TODO"
@inline copy_node(n::⋀Node, cns) = ⋀Node(cns)
@inline copy_node(n::⋁Node, cns) = ⋁Node(cns)


text:
module MissingsAsFalse

using MacroTools

export @mfalse

"""
    @mfalse(e)

Recursively walks through an expression and transforms
`missing` values to `false` in select operations.

* Checks for equality: `==`. `@mfalse missing == 1` returns
`false`. Without `@mfalse`, the expression returns `missing`.

* Greater than and less than comparisons: `>`, `>=`, `<`, and
`<=`. `missing < 1` returns `false`. Without `@mfalse` the
expression returns `missing`. The same holds for `>=`, `<`, and
`<=`.

* Broadcasted comparisons with arrays: `.==`, `.>`, `.>=`, `.<`, and `.<=`.
With `@mfalse`, all broadcasted comparisons with `missing` inside
collections return `false`. `@mfalse` only applies to comparisons
with `<:AbstractArray{Union{Missing, Bool}}`.

* Control flow: `if` and `ifelse`. With `@mfalse`, `missing`
values in `if` or `elseif` conditions are set to `false`. This
also applies to ternary operations such as `missing == 1 ? 1 : 2`.
Without `@mfalse`, `if-else` conditions error on missing values.

* Short circuiting. With `@mfalse`, `missing` acts as false in
`&&` and `||` operations. Without `@mfalse`, these comparisons
error.

* Indexing operations: `y[x]`. In literal indexing operations,
i.e. `y[x]` (but not `getindex`), if `x` is `<:AbstractVector{Union{Missing, Bool}}`
`missing` values are set to `false` and such indexes
are not returned. Without `@mfalse` such
operations error. Note: `@mfalse` only applies when `x` is
`<:AbstractArray{Union{Missing, Bool}}`.

!!! warning
    For array operaions, including braodcasted comparisons (like `.==`)
    and boolean indexing `y[x]` where x is a boolean array with missing
    values, `@mfalse` will allocate a new array without missing values.
    This will allocate memory and may cause noticeable slow-downs.

## Examples

```
julia> x = missing;

julia> @mfalse x == 1
false

julia> x = [1, missing]; y = [1, 200];

julia> @mfalse x .== y
2-element BitVector:
 1
 0

julia> @mfalse if x
           100
       else
           200
       end
200

julia> @mfalse (missing && true)
false

julia> @mfalse (missing || true)
true

julia> x = [true, missing, true]; y = [500, 600, 700];

julia> @mfalse y[x]
2-element Vector{Int64}:
 500
 700
```
"""
macro mfalse(e)
    esc(mfalse_helper(e))
end

wrap_coalesce_false(x) = :(coalesce($x, false))
wrap_fix_missing_arr(x) = :(MissingsAsFalse.fix_missing_arr($x))

function fix_missing_arr(v::AbstractArray{Union{Missing, Bool}})
    isequal.(v, true)
end
fix_missing_arr(v) = v

function mfalse_helper(e)
    t = MacroTools.postwalk(e) do x
        !(x isa Expr) && return x
        if x.head == :call
            a1 = x.args[1]
            if (a1 === :(==) ||
                a1 === :> ||
                a1 === :>= ||
                a1 === :< ||
                a1 === :<=)

                wrap_coalesce_false(x)
            elseif (a1 === :.== ||
                    a1 === :.> ||
                    a1 === :.>= ||
                    a1 === :.< ||
                    a1 === :.<=)

                wrap_fix_missing_arr(x)
            else
                x
            end
        elseif (x.head == :if) || (x.head == :elseif)
            x.args[1] = wrap_coalesce_false(x.args[1])
            x
        elseif (x.head === :(&&)) || x.head === :(||)
            x.args[1] = wrap_coalesce_false(x.args[1])
            x.args[2] = wrap_coalesce_false(x.args[2])
            x
        elseif x.head === :ref
            x.args[2:end] .= wrap_fix_missing_arr.(x.args[2:end])
            x
        else
            x
        end
    end
end

end


text:
"""
A test for nearly_zero().

Arguments
----------
number:AbstractFloat
    input number

threshold:AbstractFloat
    threshold for being nearly zero

solution:Bool
    answer key

msg="":AbstractString
    (optional) test message
"""
function test_nearly_zero(number::AbstractFloat, threshold::AbstractFloat, solution::Bool, msg::AbstractString="")
    println("---------------------------------")
    print_with_color(:blue, "Test nearly_zero()\n\n")
    if msg != ""
        print_with_color(:blue, "$(msg)\n\n")
    end
    answer = nearly_zero(number, threshold)
    @test answer == solution
    print_with_color(:green, "VERIFIED!\n")
    println("---------------------------------")
    return true
end
test_nearly_zero(1e-5,1e-7,false)
test_nearly_zero(1e-8,1e-7,true)


text:
using SpirvCross
using Test

const spv = SpirvCross

function test_ubo_decorations(ctx::Context, shader, descriptor_sets, bindings)
    ir = IR(ctx, shader)
    comp = Compiler(ctx, ir)
    res = ShaderResources(comp)
    list = get_resource_list(res)

    _descriptor_sets = spv.spvc_compiler_get_decoration.(comp, getproperty.(list, :id), spv.SpvDecorationDescriptorSet)
    _bindings = spv.spvc_compiler_get_decoration.(comp, getproperty.(list, :id), spv.SpvDecorationBinding)

    @test all(_descriptor_sets .== descriptor_sets)
    @test all(_bindings .== _bindings)
end

resource(filename) = joinpath(@__DIR__, "resources", filename)

ctx = Context()

@testset "Uniform Buffer Object decorations" begin
    for (file, ds, b) ∈ zip(
        ["vert.spv", "frag.spv"],
        [[0], []],
        [[0], []],
        )
        test_ubo_decorations(ctx, resource(file), ds, b)
    end
end

finalize(ctx)