text:
using BlockMatchingSOMs
using Distances
using Compose
using Colors

labels = ["red", "green", "blue", "yellow", "pink",
          "magenta", "skyblue", "darkblue", "orange"]
data = reduce(hcat,map(labels) do x
    c = parse(Colorant, x)
    Float64[c.r, c.g, c.b]
end)

epoch = 2000
mapsize = 20
dim = size(data,1)
som = BMSOM{Float64}(mapsize, dim, metric=Euclidean())
train!(som, data, epoch, umode=Batch())
labeling!(som, data, labels)

w=1.0/mapsize
c = [RGBA(som[:,i,j]...) for i=1:mapsize, j=1:mapsize]
img = SVG("map.svg", 6cm, 6cm)
draw(img, compose(context(),
[(context((i-1)/mapsize, (j-1)/mapsize, w, w),
rectangle(), fill(c[i, j])) for i=1:mapsize, j=1:mapsize]...))


text:
"Compute pressure levels, vmr, vcd for atmospheric profile, given p_half, T, q"
function compute_atmos_profile_fields(p_half::AbstractArray, T, q; g₀=9.8196)

    # Floating type to use
    FT = eltype(T)

    # Calculate full pressure levels
    p_full = (p_half[2:end] + p_half[1:end - 1]) / 2

    # Dry and wet mass
    dry_mass = 28.9647e-3  / Nₐ  # in kg/molec, weighted average for N2 and O2
    wet_mass = 18.01528e-3 / Nₐ  # just H2O
    ratio = dry_mass / wet_mass
    n_layers = length(T)

    # Also get a VMR vector of H2O (volumetric!)
    vmr_h2o = zeros(FT, n_layers, )
    vcd_dry = zeros(FT, n_layers, )
    vcd_h2o = zeros(FT, n_layers, )

    # Now actually compute the layer VCDs
    for i = 1:n_layers
        Δp = p_half[i + 1] - p_half[i]
        vmr_h2o[i] = q[i] * ratio
        vmr_dry = 1 - vmr_h2o[i]
        M  = vmr_dry * dry_mass + vmr_h2o[i] * wet_mass
        vcd_dry[i] = vmr_dry * Δp / (M * g₀ * 100.0^2)   # includes m2->cm2
        vcd_h2o[i] = vmr_h2o[i] * Δp / (M * g₀ * 100^2)
    end

    return p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o

end

"Validate an input atmospheric profile (YAML only)"
function validate_atmos_profile(params::Dict{Any,Any})

    @assert Set(keys(params)) == Set(["p_surf", "T", "q", "ak", "bk"]) ||
            Set(keys(params)) == Set(["p_half", "T", "q"]) ||
            Set(keys(params)) == Set(["p_half", "T"]) "Set of atmospheric profile fields must be one of (p_surf, T, q, ak, bk), (p_half, T, q), or (p_half, T)"

end

"From a yaml file, get the stored fields (psurf, T, q, ak, bk), calculate derived fields,
and return an AtmosphericProfile object"
function read_atmos_profile(file_path::String)

    # Make sure file is yaml type
    @assert endswith(file_path, ".yaml") "File must be yaml"

    # Read in the data and pass to compute fields
    params_dict = YAML.load_file(file_path)

    # Validate the parameters before doing anything else
    validate_atmos_profile(params_dict)

    T = convert.(Float64, params_dict["T"])

    # Calculate derived fields
    if ("ak" in keys(params_dict))
        psurf = convert(Float64, params_dict["p_surf"])
        q     = convert.(Float64, params_dict["q"])
        ak    = convert.(Float64, params_dict["ak"])
        bk    = convert.(Float64, params_dict["bk"])
        p_half = (ak + bk * psurf)
        p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o = compute_atmos_profile_fields(p_half, T, q)
    elseif ("q" in keys(params_dict))
        p_half = convert(Float64, params_dict["p_half"])
        q      = convert.(Float64, params_dict["q"])
        p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o = compute_atmos_profile_fields(p_half, T, q)
    else
        p_half = convert(Float64, params_dict["p_half"])
        q = zeros(length(T))
        p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o = compute_atmos_profile_fields(p_half, T, q)
    end

    # Return the atmospheric profile struct
    return AtmosphericProfile(nothing, nothing, psurf, T, q, p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o)

end

"Read atmospheric profile from netCDF file (ONLY works for FraLab test file)"
function read_atmos_profile(file::String, lat::Real, lon::Real, time_idx::Int)

    # Time index must be ∈ [1, 2, 3, 4]
    @assert 1 <= time_idx <= 4 "Time index must be ∈ [1, 2, 3, 4]"

    # Make sure file is nc type
    @assert endswith(file, ".nc4") "File must be nc4"

    # Load in the atmospheric profile
    @timeit "loading file" ds = Dataset(file)

    # See how easy it is to actually extract data?
    # Note the [:] in the end reads in ALL the data in one step
    file_lats, file_lons = ds["YDim"][:], ds["XDim"][:]

    # Convert the input lat/lon to right type
    FT = eltype(file_lats)
    lat, lon = FT(lat), FT(lon)

    # Find index (nearest neighbor, one could envision interpolation in space and time!):
    lat_idx, lon_idx = argmin(abs.(file_lats .- lat)), argmin(abs.(file_lons .- lon))

    # Temperature profile
    @timeit "getting T" T = convert(Array{FT,1}, ds["T"][lon_idx, lat_idx,  :, time_idx])

    # Specific humidity profile
    q = convert(Array{FT,1}, ds["QV"][lon_idx, lat_idx, :, time_idx])

    # Surface pressure
    psurf = convert(FT, ds["PS"][lon_idx, lat_idx, time_idx])

    # AK and BK global attributes (important to calculate pressure half-levels)
    ak, bk = ds.attrib["HDF_GLOBAL.ak"][:], ds.attrib["HDF_GLOBAL.bk"][:]

    # Close the file
    close(ds)

    # Calculate derived fields
    p_half = (ak + bk * psurf)
    p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o = compute_atmos_profile_fields(p_half, T, q)

    # Return the atmospheric profile struct
    return AtmosphericProfile(lat, lon, psurf, T, q, p_full, p_half, vmr_h2o, vcd_dry, vcd_h2o)
end

"Reduce profile dimensions by re-averaging to near-equidistant pressure grid"
function reduce_profile(n::Int, profile::AtmosphericProfile{FT}) where {FT}

    # Can only reduce the profile, not expand it
    @assert n < length(profile.T)

    # Unpack the profile
    @unpack lat, lon, psurf = profile

    # New rough half levels (boundary points)
    a = range(0, maximum(profile.p), length=n + 1)

    # Matrices to hold new values
    T = zeros(FT, n);
    q = zeros(FT, n);
    p_full = zeros(FT, n);
    p_levels = zeros(FT, n + 1);
    vmr_h2o  = zeros(FT, n);
    vcd_dry  = zeros(FT, n);
    vcd_h2o  = zeros(FT, n);

    # Loop over target number of layers
    for i = 1:n

        # Get the section of the atmosphere with the i'th section pressure values
        ind = findall(a[i] .< profile.p .<= a[i + 1]);

        # Set the pressure levels accordingly
        p_levels[i] = profile.p_levels[ind[1]]
        p_levels[i + 1] = profile.p_levels[ind[end]]

        # Re-average the other parameters to produce new layers
        p_full[i] = mean(profile.p_levels[ind])
        T[i] = mean(profile.T[ind])
        q[i] = mean(profile.q[ind])
        vmr_h2o[i] = mean(profile.vmr_h2o[ind])
        vcd_dry[i] = sum(profile.vcd_dry[ind])
        vcd_h2o[i] = sum(profile.vcd_h2o[ind])
    end

    return AtmosphericProfile(lat, lon, psurf, T, q, p_full, p_levels, vmr_h2o, vcd_dry, vcd_h2o)
end;

"""
$(FUNCTIONNAME)(psurf, λ, depol_fct, vcd_dry)
Returns the Rayleigh optical thickness per layer at reference wavelength `λ` (N₂,O₂ atmosphere, i.e. terrestrial)
Input:
- `psurf` surface pressure in `[hPa]`
- `λ` wavelength in `[μm]`
- `depol_fct` depolarization factor
- `vcd_dry` dry vertical column (no water) per layer
"""
function getRayleighLayerOptProp(psurf, λ, depol_fct, vcd_dry)
    FT = eltype(λ)
    # Total vertical Rayleigh scattering optical thickness
    tau_scat = FT(0.00864) * (psurf / FT(1013.25)) * λ^(-FT(3.916) - FT(0.074) * λ - FT(0.05) / λ)
    tau_scat = tau_scat * (FT(6.0) + FT(3.0) * depol_fct) / (FT(6.0)- FT(7.0) * depol_fct)
    Nz = length(vcd_dry)
    τRayl = zeros(FT,Nz)
    k = tau_scat / sum(vcd_dry)
    for i = 1:Nz
        τRayl[i] = k * vcd_dry[i]
    end

    return convert.(FT, τRayl)
end

"""
$(FUNCTIONNAME)(total_τ, p₀, σp, p_half)
Returns the aerosol optical depths per layer using a Gaussian distribution function with p₀ and σp on a pressure grid
"""
function getAerosolLayerOptProp(total_τ, p₀, σp, p_half)

    # Need to make sure we can also differentiate wrt σp (FT can be Dual!)
    FT = eltype(p₀)
    Nz = length(p_half)
    ρ = zeros(FT,Nz)

    for i = 2:Nz
        dp = p_half[i] - p_half[i - 1]
        ρ[i] = (1 / (σp * sqrt(2π))) * exp(-(p_half[i] - p₀)^2 / (2σp^2)) * dp
    end
    Norm = sum(ρ)
    τAer  =  (total_τ / Norm) * ρ
    return convert.(FT, τAer)
end

"""
$(FUNCTIONNAME)(τRayl, τAer,  aerosol_optics, Rayl𝐙⁺⁺, Rayl𝐙⁻⁺, Aer𝐙⁺⁺, Aer𝐙⁻⁺, τ_abs, arr_type)
Computes the composite layer single scattering parameters (τ, ϖ, Z⁺⁺, Z⁻⁺)

Returns:
- `τ`, `ϖ`   : only Rayleigh scattering and aerosol extinction, no gaseous absorption (no wavelength dependence)
- `τ_λ`,`ϖ_λ`: Rayleigh scattering + aerosol extinction + gaseous absorption (wavelength dependent)
- `Z⁺⁺`,`Z⁻⁺`: Composite Phase matrix (weighted average of Rayleigh and aerosols)

# Arguments
- `τRay` layer optical depth for Rayleigh
- `τAer` layer optical depth for Aerosol(s) (vector)
- `aerosol_optics` array of aerosol optics struct
- `Rayl𝐙⁺⁺` Rayleigh 𝐙⁺⁺ phase matrix (2D)
- `Rayl𝐙⁻⁺` Rayleigh 𝐙⁻⁺ phase matrix (2D)
- `Aer𝐙⁺⁺` Aerosol 𝐙⁺⁺ phase matrix (3D)
- `Aer𝐙⁻⁺` Aerosol 𝐙⁻⁺ phase matrix (3D)
- `τ_abs` layer absorption optical depth array (per wavelength) by gaseous absorption
"""
function construct_atm_layer(τRayl, τAer,  aerosol_optics, Rayl𝐙⁺⁺, Rayl𝐙⁻⁺, Aer𝐙⁺⁺, Aer𝐙⁻⁺, τ_abs, arr_type)
    FT = eltype(τRayl)
    nAer = length(aerosol_optics)

    # Fix Rayleigh SSA to 1
    ϖRayl = FT(1)

    @assert length(τAer) == nAer "Sizes don't match"

    τ = FT(0)
    ϖ = FT(0)
    A = FT(0)
    Z⁺⁺ = similar(Rayl𝐙⁺⁺);
    Z⁻⁺ = similar(Rayl𝐙⁺⁺);

    if (τRayl + sum(τAer)) < eps(FT)
        fill!(Z⁺⁺, 0); fill!(Z⁻⁺, 0);
        return FT(0), FT(1), Z⁺⁺, Z⁻⁺
    end

    τ += τRayl
    ϖ += τRayl * ϖRayl
    A += τRayl * ϖRayl

    Z⁺⁺ = τRayl * ϖRayl * Rayl𝐙⁺⁺
    Z⁻⁺ = τRayl * ϖRayl * Rayl𝐙⁻⁺

    for i = 1:nAer
        τ   += τAer[i]
        ϖ   += τAer[i] * aerosol_optics[i].ω̃
        A   += τAer[i] * aerosol_optics[i].ω̃ * (1 - aerosol_optics[i].fᵗ)
        Z⁺⁺ += τAer[i] * aerosol_optics[i].ω̃ * (1 - aerosol_optics[i].fᵗ) * Aer𝐙⁺⁺[:,:,i]
        Z⁻⁺ += τAer[i] * aerosol_optics[i].ω̃ * (1 - aerosol_optics[i].fᵗ) * Aer𝐙⁻⁺[:,:,i]
    end

    Z⁺⁺ /= A
    Z⁻⁺ /= A
    A /= ϖ
    ϖ /= τ

    # Rescaling composite SSPs according to Eqs. A.3 of Sanghavi et al. (2013) or Eqs.(8) of Sanghavi & Stephens (2015)
    τ *= (FT(1) - (FT(1) - A) * ϖ)
    ϖ *= A / (FT(1) - (FT(1) - A) * ϖ)#Suniti

    # Adding absorption optical depth / albedo:
    τ_λ = τ_abs .+ τ
    ϖ_λ = (τ .* ϖ) ./ τ_λ

    return Array(τ_λ), Array(ϖ_λ), τ, ϖ, Array(Z⁺⁺), Array(Z⁻⁺)
end

"When performing RT_run, this function pre-calculates properties for all layers, before any Core RT is performed"
function construct_all_atm_layers(FT, nSpec, Nz, NquadN, τRayl, τAer, aerosol_optics, Rayl𝐙⁺⁺, Rayl𝐙⁻⁺, Aer𝐙⁺⁺, Aer𝐙⁻⁺, τ_abs, arr_type, qp_μ, μ₀, m)

    # Empty matrices to hold all values
    τ_λ_all   = zeros(FT, nSpec, Nz)
    ϖ_λ_all   = zeros(FT, nSpec, Nz)
    τ_all     = zeros(FT, Nz)
    ϖ_all     = zeros(FT, Nz)
    Z⁺⁺_all   = zeros(FT, NquadN, NquadN, Nz)
    Z⁻⁺_all   = zeros(FT, NquadN, NquadN, Nz)

    dτ_max_all  = zeros(FT, Nz)
    dτ_all      = zeros(FT, Nz)
    ndoubl_all  = zeros(Int64, Nz)
    dτ_λ_all    = zeros(FT, nSpec, Nz)
    expk_all    = zeros(FT, nSpec, Nz)
    scatter_all = zeros(Bool, Nz)

    Threads.@threads for iz=1:Nz

        # Construct atmospheric properties
        τ_λ_all[:, iz], ϖ_λ_all[:, iz], τ_all[iz], ϖ_all[iz], Z⁺⁺_all[:,:,iz], Z⁻⁺_all[:,:,iz] = construct_atm_layer(τRayl[iz], τAer[:,iz], aerosol_optics, Rayl𝐙⁺⁺, Rayl𝐙⁻⁺, Aer𝐙⁺⁺, Aer𝐙⁻⁺, τ_abs[:,iz], arr_type)

        # Compute doubling number
        dτ_max_all[iz] = minimum([τ_all[iz] * ϖ_all[iz], FT(0.001) * minimum(qp_μ)])
        dτ_all[iz], ndoubl_all[iz] = doubling_number(dτ_max_all[iz], τ_all[iz] * ϖ_all[iz]) #Suniti

        # Compute dτ vector
        dτ_λ_all[:, iz] = (τ_λ_all[:, iz] ./ (FT(2)^ndoubl_all[iz]))
        #@show maximum(dτ_λ_all[:,iz])
        expk_all[:, iz] = exp.(-dτ_λ_all[:, iz] /μ₀) #Suniti

        # Determine whether there is scattering
        scatter_all[iz] = (  sum(τAer[:,iz]) > 1.e-8 ||
                          (( τRayl[iz] > 1.e-8 ) && (m < 3))) ?
                            true : false
    end

    # Compute sum of optical thicknesses of all layers above the current layer
    τ_sum_all = accumulate(+, τ_λ_all, dims=2)

    # First start with all zeros
    # At the bottom of the atmosphere, we have to compute total τ_sum (bottom of lowest layer), for the surface interaction
    τ_sum_all = hcat(zeros(FT, size(τ_sum_all[:,1])), τ_sum_all)

    # Starting scattering interface (None for both added and composite)
    scattering_interface = ScatteringInterface_00()
    scattering_interfaces_all = []

    for iz = 1:Nz
        # Whether there is scattering in the added layer, composite layer, neither or both
        scattering_interface = get_scattering_interface(scattering_interface, scatter_all[iz], iz)
        push!(scattering_interfaces_all, scattering_interface)
    end

    return ComputedAtmosphereProperties(τ_λ_all, ϖ_λ_all, τ_all, ϖ_all, Z⁺⁺_all, Z⁻⁺_all, dτ_max_all, dτ_all, ndoubl_all, dτ_λ_all, expk_all, scatter_all, τ_sum_all, scattering_interfaces_all)
end

"Given the CrossSectionModel, the grid, and the AtmosphericProfile, fill up the τ_abs array with the cross section at each layer
(using pressures/temperatures) from the profile"
function compute_absorption_profile!(τ_abs::Array{FT,2},
                                     model::AbstractCrossSectionModel,
                                     grid,
                                     profile::AtmosphericProfile,
                                     ) where FT <: AbstractFloat

    # The array to store the cross-sections must be same length as number of layers
    @assert size(τ_abs,2) == length(profile.p)

    for iz in 1:length(profile.p)

        # Pa -> hPa
        p = profile.p[iz] / 100
        T = profile.T[iz]

        # Changed index order
        τ_abs[:,iz] = Array(absorption_cross_section(model, grid, p, T)) * profile.vcd_dry[iz] * model.vmr
    end

    return nothing

end

ext: jl


text:
#=
sampler.jl:

Author: Peter Wolf (opus111@gmail.com)

Port of the Pytorch Sampler API to Julia

Utility methods for selecting various subsets of Datasets

The original source is here

https://pytorch.org/docs/stable/_modules/torch/utils/data/sampler.html#Sampler

The documentation is copied from here

https://pytorch.org/docs/stable/data.html
=#

"""
Abstract type for all Samplers.

Samplers are iterable types, with integer elements.  Those elements
are indexes into a data source.  Where data sources are any indexable type e.g. MapDataset

In Julia duck typing, implementing an interface just requires
implementing a set of required fuctions.

Required functions for Samplers are:

iterate(s<:Sampler)		    Returns either a tuple of the first index into data source
                            and initial state, or nothing if empty
iterate(s<:Sampler, state)	Returns either a tuple of the next index and next state,
                            or nothing if no items remain

Optional function:

length(s<:Sampler)          returns the number of indexes in this Sampler.

See Julia iteration documentation for more information

    https://docs.julialang.org/en/v1/manual/interfaces/
"""
abstract type Sampler end

"Samples elements sequentially, always in the same order."
struct SequentialSampler <: Sampler
    n::Int
end
SequentialSampler(data_source) = SequentialSampler(length(data_source))

iterate(s::SequentialSampler) = if s.n < 1 nothing else (1,2) end
iterate(s::SequentialSampler, state) = if state > length(s) nothing else (state,state+1) end
Base.length(s::SequentialSampler) = s.n

"Shared code between Random and Weighted Samplers"
abstract type IterableSampler <: Sampler end
function iterate(s::IterableSampler)
    if length(s) < 1
        return nothing
    else
        i = ids(s)
        return (i[1],(i,2))
    end
end
iterate(s::IterableSampler, state) = nothing ? state[2] > length(s) : (state[1][state[2]],state[2]+1)

"""
Samples elements randomly without replacement. Sample from a shuffled dataset.

Arguments:
        data_source: dataset to sample from
"""
struct RandomSamplerWithoutReplacement <: Sampler
    n::Int
end
RandomSamplerWithoutReplacement(data_source) = RandomSamplerWithoutReplacement(length(data_source))
ids(s::RandomSamplerWithoutReplacement) = randperm(length(s))
Base.length(s::RandomSamplerWithoutReplacement) = s.n

struct RandomSamplerWithReplacement <: Sampler
    n::Int
    num_samples::Int
end

"""
Samples elements randomly with replacement.

Arguments:
        data_source: dataset to sample from
        num_samples: number of samples to draw, default=`length(dataset)`.
"""
function RandomSamplerWithReplacement(data_source,num_samples::Int)
    if length(data_source) < num_samples
        @error "Data source must have at least as many samples as num_samples"
    end
    if num_samples < 1
        @error "num_samples must be a positive integer"
    end
    return RandomSamplerWithReplacement(length(data_source),num_samples)
end
ids(s::RandomSamplerWithReplacement)=rand(1:length(s),s.num_samples)
Base.length(s::RandomSamplerWithReplacement) = s.num_samples

RandomSampler(data_source,num_samples=nothing,replacement::Bool =false) = if(replacement) RandomSamplerWithReplacement(data_source,num_samples) else RandomSamplerWithReplacement(data_source) end
"""
SubsetRandomSampler

Samples elements randomly from a given list of indices, without replacement.

Arguments:
        indices: an indexable sequence of indices
"""
struct SubsetRandomSampler <: Sampler
    indicies::AbstractArray{Int}
end

ids(s::SubsetRandomSampler) = shuffle(s.indicies)
Base.length(s::SubsetRandomSampler) = length(s.indicies)

"""
WeightedRandomSampler
    r'Samples elements from ``[0,..,len(weights)-1]`` with given probabilities (weights).

    Args:
        weights (sequence)   : a sequence of weights, not necessary summing up to one
        num_samples (int): number of samples to draw
        replacement (bool): if ``True``, samples are drawn with replacement.
            If not, they are drawn without replacement, which means that when a
            sample index is drawn for a row, it cannot be drawn again for that row.

    Example:
        >>> list(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True))
        [4, 4, 1, 4, 5]
        >>> list(WeightedRandomSampler([0.9, 0.4, 0.05, 0.2, 0.3, 0.1], 5, replacement=False))
        [0, 1, 4, 3, 2]
"""
struct WeightedRandomSamplerWithReplacement <: Sampler
    weights::AbstractArray{Real}
    num_samples::Int
end
ids(s::WeightedRandomSamplerWithReplacement) = sample(1:s.num_samples,s.weights)
Base.length(s::WeightedRandomSamplerWithReplacement) = s.num_samples

struct WeightedRandomSamplerWithoutReplacement <: Sampler
    weights::AbstractArray{Real}
    num_samples::Int
end
ids(s::WeightedRandomSamplerWithoutReplacement) = sample(1:s.num_samples,s.weights,replace=false)
Base.length(s::WeightedRandomSamplerWithoutReplacement) = s.num_samples

"""
BatchSampler
    Wraps another sampler to yield a mini-batch of indices.

    Args:
        sampler: Base sampler.
        batch_size: Size of mini-batch.
        drop_last: If ``true``, the sampler will drop the last batch if
            its size would be less than ``batch_size``

    Example:
        >>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False))
        [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
        >>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True))
        [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
"""

struct BatchSampler <: Sampler
    sampler :: Sampler
    batch_size :: Int
    drop_last :: Bool
end

function batch(s::BatchSampler, offset::Int)
    batch = []
    for i in 1:s.batch_size
        if offset+i > length(s.sampler)
            return batch,offset+i
        else
            push!(batch,s.sampler[offset+i])
        end
    end
    return batch,offset+i
end

iterate(s::BatchSampler) =
    if s.drop_last
        if length(s.sampler) < s.batch_size
            return nothing
        else
            return batch(s,0)
        end
    else
        return batch(s,0)
    end

iterate(s::BatchSampler, state) =
    if s.drop_last
        if length(s.sampler) < s.batch_size + state[2]
            return nothing
        else
            return batch(s,state[2])
        end
    else
        return batch(s,0)
    end

Base.length(s::BatchSampler) =
    if s.drop_last
        length(s.sampler) ÷ s.batch_size
    else
        (length(s.sampler) + s.batch_size - 1) ÷ s.batch_size
    end


text:
"""
DG Balance Law Discretizations module.
Attempts to provide a reasonable implementation of a discontinuous Galerkin
method (in weak form) on tensor product quadrilateral (2D) and hexahedral (3D)
elements for balance laws of the form

```math
q_{,t} + Σ_{i=1,...d} F_{i,i} = s
```

where ``q`` is the state vector, ``F`` is the flux function, and ``s`` is the
source function. ``F`` includes both the "inviscid" and "viscous" fluxes. Note
that this is a space only discretization, time must be advanced using some
ordinary differential equations methods; see [`ODESolvers`](@ref).

The flux function `F_{i}` is taken to be of the form:
```math
F_{i} := F_{i}(q, σ; a)
σ = H(q, ∇G(q; a); a)
```
where ``a`` is a set of parameters and viscous terms enter through ``σ``

The source term is of the form:
```math
s := s(q; a)
```

In the code and docs the following terminology is used:
- ``q`` is referred to as the state
- ``σ`` is the viscous state
- ``a`` is the auxiliary state
- ``F`` is the physical flux
- ``H`` is the viscous transform
- ``G`` is the gradient transform

Much of the notation used in this module follows Hesthaven and Warburton (2008).

!!! note

    Currently all the functions take the same parameters and the gradient
    transform can take a user-specified subset of the state vector.

!!! note

    We plan to switch to a skew-symmetric formulation (at which time this note
    will be removed)

!!! references

    ```
    @BOOK{HesthavenWarburton2008,
      title = {Nodal Discontinuous {G}alerkin Methods: {A}lgorithms, Analysis,
               and Applications},
      publisher = {Springer-Verlag New York},
      year = {2008},
      author = {Hesthaven, Jan S. and Warburton, Tim},
      volume = {54},
      series = {Texts in Applied Mathematics},
      doi = {10.1007/978-0-387-72067-8}
    }
    ```
"""
module DGBalanceLawDiscretizations

using MPI
using ..Grids
using ..MPIStateArrays
using StaticArrays
using ..SpaceMethods
using DocStringExtensions
using ..Topologies
using GPUifyLoops

export DGBalanceLaw

include("DGBalanceLawDiscretizations_kernels.jl")
include("NumericalFluxes.jl")

"""
    DGBalanceLaw <: AbstractDGMethod

This contains the necessary information for a discontinuous Galerkin method for
balance laws.

See also: Outer constructor [`DGBalanceLaw`](@ref)

# Fields

$(DocStringExtensions.FIELDS)

"""
struct DGBalanceLaw <: AbstractDGMethod
  "computational grid / mesh"
  grid::DiscontinuousSpectralElementGrid

  "number of state"
  nstate::Int

  "physical flux function"
  flux!::Function

  "numerical flux function"
  numerical_flux!::Function

  "numerical boundary flux function"
  numerical_boundary_flux!::Union{Nothing, Function}

  "storage for the viscous state"
  Qvisc::MPIStateArray

  "number of out states for gradient_transform!"
  number_gradient_states::Int

  "number of out states for the viscous_transform!"
  number_viscous_states::Int

  "tuple of states going into gradient_transform!"
  states_for_gradient_transform::Tuple

  "transform from state to variables to take gradient of"
  gradient_transform!::Union{Nothing, Function}

  "transform from Q and gradient state to viscous states"
  viscous_transform!::Union{Nothing, Function}

  "penalty for the viscous state computation"
  viscous_penalty!::Union{Nothing, Function}

  "boundary penalty for the viscous state computation (e.g., Dirichlet)"
  viscous_boundary_penalty!::Union{Nothing, Function}

  "auxiliary state array"
  auxstate::MPIStateArray

  "source function"
  source!::Union{Nothing, Function}
end

"""
    DGBalanceLaw(;grid::DiscontinuousSpectralElementGrid,
                 length_state_vector,
                 flux!,
                 numerical_flux!,
                 numerical_boundary_flux! = nothing,
                 states_for_gradient_transform = (),
                 number_gradient_states = 0,
                 number_viscous_states = 0,
                 gradient_transform! = nothing,
                 viscous_transform! = nothing,
                 viscous_penalty! = nothing,
                 viscous_boundary_penalty! = nothing,
                 auxiliary_state_length = 0,
                 auxiliary_state_initialization! = nothing,
                 source! = nothing)

Constructs a `DGBalanceLaw` spatial discretization type for the physics defined
by `flux!` and `source!`. The computational domain is defined by `grid`. The
number of state variables is defined by `length_state_vector`. The user may also
specify an auxiliary state which will be unpacked by the compute kernel passed
on to the user-defined flux and numerical flux functions. The source function
`source!` is optional.

The flux function is called with data from a degree of freedom (DOF) as
```
flux!(F, Q, V, aux, t)
```
where
- `F` is an `MArray` of size `(dim, length_state_vector)` to be filled (note
  that this is uninitialized so the user must set to zero if is this desired)
- `Q` is the state to evaluate (`MArray`)
- `V` is the viscous state to evaluate (`MArray`)
- `aux` is the user-defined auxiliary state (`MArray`)
- `t` is the current simulation time
Warning: Modifications to `Q` or `aux` may cause side effects and should be
avoided.

The numerical flux function is called with data from two DOFs as
```
numerical_flux!(F, nM, QM, VM, auxM, QP, VP, auxP, t)
```
where
- `F` is an `MVector` of length `length_state_vector` to be filled with the
  numerical flux across the face (note that this is uninitialized so user must
  set to zero if is this desired)
- `nM` is the unit outward normal to the face with respect to the minus side
  (`MVector` of length `3`)
- `QM` and `QP` are the minus and plus side states (`MArray`)
- `VM` and `VP` are the minus and plus viscous side states (`MArray`)
- `auxM` and `auxP` are the auxiliary states (`MArray`)
- `t` is the current simulation time
Warning: Modifications to `nM`, `QM`, `auxM`, `QP`, or `auxP` may cause side
effects and should be avoided.

If `grid.topology` has a boundary then the function
`numerical_boundary_flux!` must be specified. This function is called
with the data from the neighbouring DOF as
```
numerical_boundary_flux!(F, nM, QM, VM, auxM, QP, VP, auxP, bctype, t)
```
where
- `F` is an `MArray` of size `(dim, length_state_vector)` to be filled with the
  numerical flux across the face (note that this is uninitialized so user must
  set to zero is this desired)
- `nM` is the unit outward normal to the face with respect to the minus side
  (`MVector` of length `3`)
- `QM` and `QP` are the minus and plus side states (`MArray`)
- `VM` and `VP` are the minus and plus viscous side states (`MArray`)
- `auxM` and `auxP` are the auxiliary states (`MArray`)
- `bctype` is the boundary condition flag for the connected face and element of
   `grid.elemtobndy`
- `t` is the current simulation time
Note: `QP` and `auxP` are filled with values based on degrees of freedom
referenced in `grid.vmapP`; `QP` and `auxP` may be modified by the calling
function.

Warning: Modifications to `nM`, `QM`, or `auxM` may cause side effects and
should be avoided.

If present the source function is called with data from a DOF as
```
source!(S, Q, aux, t)
```
where `S` is an `MVector` of length `length_state_vector` to be filled; other
arguments are the same as `flux!` and the same warning concerning `Q` and `aux`
applies.

When `auxiliary_state_initialization! !== nothing` then this is called on the
auxiliary state (assuming `auxiliary_state_length > 0`) as
```
auxiliary_state_initialization!(aux, x, y, z)
```
where `aux` is an `MArray` to fill with the auxiliary state for a DOF located at
Cartesian coordinate locations `(x, y, z)`; see also
[`grad_auxiliary_state!`](@ref) allows the user to take the gradient of a field
stored in the auxiliary state.

When viscous terms are needed, the user must specify values for the following
keyword arguments:
- `states_for_gradient_transform` (`Tuple`)
- `number_gradient_states` (`Int`)
- `number_viscous_states` (`Int`)
- `gradient_transform!` (`Function`)
- `viscous_transform!` (`Function`)
- `viscous_penalty!` (`Function`)
- `viscous_boundary_penalty!` (`Function`); only required if the topology has a
  boundary

The function `gradient_transform!` is the implementation of the function `G` in
the module docs; see [`DGBalanceLawDiscretizations`](@ref). It transforms the
elements of the components of state vector specified by
`states_for_gradient_transform` into the values that should have their gradient
taken. It is called on each DOF as:
```
gradient_transform!(G, Q, aux, t)
```
where `G` is an `MVector` of length `number_gradient_states` to be filled, `Q`
is an `MVector` containing only the states specified by
`states_for_gradient_transform`, `aux` is the full auxiliary state at the DOF,
and `t` is the simulation time.Q

The function `viscous_transform!` is the implementation of the function `H` in
the module docs; see [`DGBalanceLawDiscretizations`](@ref). It transforms the
gradient ``∇G`` and ``q`` into the viscous state ``σ``. It is called on each DOF
as:
```
viscous_transform!(V, gradG, Q, aux, t)
```
where `V` is an `MVector` of length `number_viscous_states` to be filled,
`gradG` is an `MMatrix` containing the DG-gradient of ``G``, `Q` is an `MVector`
containing only the states specified by `states_for_gradient_transform`, `aux`
is the full auxiliary state at the DOF, and `t` is the simulation time. Note
that `V` is a vector not a matrix so that minimal storage can be used if
symmetry can be exploited.

The function `viscous_penalty!` is the penalty terms to be used for the
DG-gradient calculation. It is called with data from two neighbouring degrees of
freedom as
```
viscous_penalty!(V, nM, HM, QM, auxM, HP, QP, auxP, t)
```
where:
- `V` is an `MVector` of length `number_viscous_states` to be filled with the
  numerical penalty across the face; see below.
- `nM` is the unit outward normal to the face with respect to the minus side
  (`MVector` of length `3`)
- `HM` and `HP` are the minus and plus evaluation of `gradient_transform!` on
  either side of the face
- `QM` and `QP` are the minus and plus side states (`MArray`); filled only with
  `states_for_gradient_transform` states.
- `auxM` and `auxP` are the auxiliary states (`MArray`)
- `t` is the current simulation time
The viscous penalty function is should compute on the faces
```math
n^{-} \\cdot H^{*} - n^{-} \\cdot H^{-}
```
where ``n^{-} \\cdot H^{*}`` is the "numerical-flux" for the viscous state
computation and ``H^{-}`` is the value of `viscous_transform!` evaluated on the
minus side ``n^{-} \\cdot G^{-}`` as an argument.

If `grid.topology` has a boundary then the function `viscous_boundary_penalty!`
must be specified. This function is called with the data from the neighbouring
DOF as
```
viscous_boundary_penalty!(V, nM, HM, QM, auxM, HP, QP, auxP, bctype, t)
```
where the required behaviour mimics that of `viscous_penalty!` and
`numerical_boundary_flux!`.

!!! note

    If `(x, y, z)`, or data derived from this such as spherical coordinates, is
    needed in the flux or source the user is responsible to storing this in the
    auxiliary state

"""
function DGBalanceLaw(;grid::DiscontinuousSpectralElementGrid,
                      length_state_vector, flux!,
                      numerical_flux!,
                      numerical_boundary_flux! = nothing,
                      states_for_gradient_transform=(),
                      number_gradient_states=0,
                      number_viscous_states=0,
                      gradient_transform! = nothing,
                      viscous_transform! = nothing,
                      viscous_penalty! = nothing,
                      viscous_boundary_penalty! = nothing,
                      auxiliary_state_length=0,
                      auxiliary_state_initialization! = nothing,
                      source! = nothing)

  topology = grid.topology
  Np = dofs_per_element(grid)
  h_vgeo = Array(grid.vgeo)
  DFloat = eltype(h_vgeo)
  DA = arraytype(grid)

  (Topologies.hasboundary(topology) &&
   numerical_boundary_flux! === nothing &&
   error("no `numerical_boundary_flux!` given when topology "*
         "has boundary"))

  if number_viscous_states > 0 || number_gradient_states > 0 ||
    length(states_for_gradient_transform) > 0

    # These should all be true in this case
    @assert number_viscous_states > 0
    @assert number_gradient_states > 0
    @assert length(states_for_gradient_transform) > 0
    @assert gradient_transform! !== nothing
    @assert viscous_transform! !== nothing
    @assert viscous_penalty! !== nothing
    (Topologies.hasboundary(topology)) && (@assert viscous_boundary_penalty! !==
                                           nothing)
  end

  # TODO: Clean up this MPIStateArray interface...
  Qvisc = MPIStateArray{Tuple{Np, number_viscous_states},
                     DFloat, DA
                    }(topology.mpicomm,
                      length(topology.elems),
                      realelems=topology.realelems,
                      ghostelems=topology.ghostelems,
                      sendelems=topology.sendelems,
                      nabrtorank=topology.nabrtorank,
                      nabrtorecv=topology.nabrtorecv,
                      nabrtosend=topology.nabrtosend,
                      weights=view(h_vgeo, :, grid.Mid, :),
                      commtag=111)

  auxstate = MPIStateArray{Tuple{Np, auxiliary_state_length}, DFloat, DA
                          }(topology.mpicomm,
                            length(topology.elems),
                            realelems=topology.realelems,
                            ghostelems=topology.ghostelems,
                            sendelems=topology.sendelems,
                            nabrtorank=topology.nabrtorank,
                            nabrtorecv=topology.nabrtorecv,
                            nabrtosend=topology.nabrtosend,
                            weights=view(h_vgeo, :, grid.Mid, :),
                            commtag=222)

  if auxiliary_state_initialization! !== nothing
    @assert auxiliary_state_length > 0
    dim = dimensionality(grid)
    N = polynomialorder(grid)
    vgeo = grid.vgeo
    device = typeof(auxstate.Q) <: Array ? CPU() : CUDA()
    nrealelem = length(topology.realelems)
    @launch(device, threads=(Np,), blocks=nrealelem,
            initauxstate!(Val(dim), Val(N), Val(auxiliary_state_length),
                          auxiliary_state_initialization!, auxstate.Q, vgeo,
                          topology.realelems))
    MPIStateArrays.start_ghost_exchange!(auxstate)
    MPIStateArrays.finish_ghost_exchange!(auxstate)
  end

  DGBalanceLaw(grid, length_state_vector, flux!,
               numerical_flux!, numerical_boundary_flux!,
               Qvisc, number_gradient_states, number_viscous_states,
               states_for_gradient_transform, gradient_transform!,
               viscous_transform!, viscous_penalty!,
               viscous_boundary_penalty!, auxstate, source!)
end

"""
    MPIStateArray(disc::DGBalanceLaw; nstate=disc.nstate, commtag=888)

Given a discretization `disc` constructs an `MPIStateArrays` for holding a
solution state. The optional 'nstate' arguments allows the user to specify a
specific number of states. The optional `commtag` allows the user to set the tag
to use for communication with this `MPIStateArray`.
"""
function MPIStateArrays.MPIStateArray(disc::DGBalanceLaw; nstate=disc.nstate,
                                      commtag=888)
  grid = disc.grid
  topology = disc.grid.topology
  # FIXME: Remove after updating CUDA
  h_vgeo = Array(disc.grid.vgeo)
  DFloat = eltype(h_vgeo)
  Np = dofs_per_element(grid)
  DA = arraytype(grid)
  MPIStateArray{Tuple{Np, nstate}, DFloat, DA}(topology.mpicomm,
                                               length(topology.elems),
                                               realelems=topology.realelems,
                                               ghostelems=topology.ghostelems,
                                               sendelems=topology.sendelems,
                                               nabrtorank=topology.nabrtorank,
                                               nabrtorecv=topology.nabrtorecv,
                                               nabrtosend=topology.nabrtosend,
                                               weights=view(h_vgeo, :,
                                                            disc.grid.Mid, :),
                                               commtag=commtag)
end

"""
    MPIStateArray(disc::DGBalanceLaw, initialization!::Function; commtag=888)

Given a discretization `disc` constructs an `MPIStateArrays` for holding a
solution state. The optional `commtag` allows the user to set the tag to use
for communication with this `MPIStateArray`.

After allocation the `MPIStateArray` is initialized using the function
`initialization!` which will be called as:
```
initialization!(Q, x, y, z, [aux])
```
where `Q` is an `MArray` with the solution state at a single degree of freedom
(DOF) to initialize and `(x,y,z)` is the coordinate point for the allocation. If
`disc` contains an auxiliary data the values of this at the DOF are passed
through as an `MArray` through the `aux` argument

!!! note

    `Q` is `undef` at start the function (i.e., not initialized to zero)

!!! note

    Modifications of the `aux` array will be discarded.

!!! todo

    GPUify this function to remove `host` and `device` data transfers

"""
function MPIStateArrays.MPIStateArray(disc::DGBalanceLaw,
                                      ic!::Function; commtag=888)
  Q = MPIStateArray(disc; commtag=commtag)

  nvar = disc.nstate
  grid = disc.grid
  vgeo = grid.vgeo
  Np = dofs_per_element(grid)
  auxstate = disc.auxstate
  nauxstate = size(auxstate, 2)

  # FIXME: GPUify me
  host_array = Array ∈ typeof(Q).parameters
  (h_vgeo, h_Q, h_auxstate) = host_array ? (vgeo, Q, auxstate) :
                                       (Array(vgeo), Array(Q), Array(auxstate))
  Qdof = MArray{Tuple{nvar}, eltype(h_Q)}(undef)
  auxdof = MArray{Tuple{nauxstate}, eltype(h_Q)}(undef)
  @inbounds for e = 1:size(Q, 3), i = 1:Np
    (x, y, z) = (h_vgeo[i, grid.xid, e], h_vgeo[i, grid.yid, e],
                 h_vgeo[i, grid.zid, e])
    if nauxstate > 0
      for s = 1:nauxstate
        auxdof[s] = h_auxstate[i, s, e]
      end
      ic!(Qdof, x, y, z, auxdof)
    else
      ic!(Qdof, x, y, z)
    end
    for n = 1:nvar
      h_Q[i, n, e] = Qdof[n]
    end
  end
  if !host_array
    Q .= h_Q
  end

  Q
end

"""
    MPIStateArray(initialization!::Function, disc::DGBalanceLaw; commtag=888)

Wrapper function to allow for calls of the form

```
MPIStateArray(disc) do  Q, x, y, z
  # fill Q
end
```

See also [`MPIStateArray`](@ref)
"""
MPIStateArrays.MPIStateArray(f::Function,
                             d::DGBalanceLaw; commtag=888
                            ) = MPIStateArray(d, f; commtag=commtag)

"""
    odefun!(disc::DGBalanceLaw, dQ::MPIStateArray, Q::MPIStateArray, t)

Evaluates the right-hand side of the discontinuous Galerkin semi-discretization
defined by `disc` at time `t` with state `Q`. The result is added into
`dQ`. Namely, the semi-discretization is of the form
```math
Q̇ = F(Q, t)
```
and after the call `dQ += F(Q, t)`
"""
function SpaceMethods.odefun!(disc::DGBalanceLaw, dQ::MPIStateArray,
                              Q::MPIStateArray, t)

  device = typeof(Q.Q) <: Array ? CPU() : CUDA()

  grid = disc.grid
  topology = grid.topology

  dim = dimensionality(grid)
  N = polynomialorder(grid)
  Nq = N + 1
  Nqk = dim == 2 ? 1 : Nq
  Nfp = Nq * Nqk
  nrealelem = length(topology.realelems)

  Qvisc = disc.Qvisc
  auxstate = disc.auxstate

  nstate = disc.nstate
  nviscstate = disc.number_viscous_states
  ngradstate = disc.number_gradient_states
  nauxstate = size(auxstate, 2)
  states_grad = disc.states_for_gradient_transform

  Dmat = grid.D
  vgeo = grid.vgeo
  sgeo = grid.sgeo
  vmapM = grid.vmapM
  vmapP = grid.vmapP
  elemtobndy = grid.elemtobndy

  ########################
  # Gradient Computation #
  ########################
  MPIStateArrays.start_ghost_exchange!(Q)

  if nviscstate > 0

    @launch(device, threads=(Nq, Nq, Nqk), blocks=nrealelem,
            volumeviscterms!(Val(dim), Val(N), Val(nstate), Val(states_grad),
                             Val(ngradstate), Val(nviscstate), Val(nauxstate),
                             disc.viscous_transform!, disc.gradient_transform!,
                             Q.Q, Qvisc.Q, auxstate.Q, vgeo, t, Dmat,
                             topology.realelems))

    MPIStateArrays.finish_ghost_recv!(Q)

    @launch(device, threads=Nfp, blocks=nrealelem,
            faceviscterms!(Val(dim), Val(N), Val(nstate), Val(states_grad),
                           Val(ngradstate), Val(nviscstate), Val(nauxstate),
                           disc.viscous_penalty!,
                           disc.viscous_boundary_penalty!,
                           disc.gradient_transform!, Q.Q, Qvisc.Q, auxstate.Q,
                           vgeo, sgeo, t, vmapM, vmapP, elemtobndy,
                           topology.realelems))

    MPIStateArrays.start_ghost_exchange!(Qvisc)
  end

  ###################
  # RHS Computation #
  ###################

  @launch(device, threads=(Nq, Nq, Nqk), blocks=nrealelem,
          volumerhs!(Val(dim), Val(N), Val(nstate), Val(nviscstate),
                     Val(nauxstate), disc.flux!, disc.source!, dQ.Q, Q.Q,
                     Qvisc.Q, auxstate.Q, vgeo, t, Dmat, topology.realelems))

  MPIStateArrays.finish_ghost_recv!(nviscstate > 0 ? Qvisc : Q)

  # The main reason for this protection is not for the MPI.Waitall!, but the
  # make sure that we do not recopy data to the GPU
  nviscstate > 0 && MPIStateArrays.finish_ghost_recv!(Qvisc)
  nviscstate == 0 && MPIStateArrays.finish_ghost_recv!(Q)

  @launch(device, threads=Nfp, blocks=nrealelem,
          facerhs!(Val(dim), Val(N), Val(nstate), Val(nviscstate),
                   Val(nauxstate), disc.numerical_flux!,
                   disc.numerical_boundary_flux!, dQ.Q, Q.Q, Qvisc.Q,
                   auxstate.Q, vgeo, sgeo, t, vmapM, vmapP, elemtobndy,
                   topology.realelems))

  # Just to be safe, we wait on the sends we started.
  MPIStateArrays.finish_ghost_send!(Qvisc)
  MPIStateArrays.finish_ghost_send!(Q)
end

"""
    grad_auxiliary_state!(disc, i, (ix, iy, iz))

Computes the gradient of a the field `i` of the constant auxiliary state of
`disc` and stores the `x, y, z` compoment in fields `ix, iy, iz` of constant
auxiliary state.

!!! note

    This only computes the element gradient not a DG gradient. If your constant
    auxiliary state is discontinuous this may or may not be what you want!
"""
function grad_auxiliary_state!(disc::DGBalanceLaw, id, (idx, idy, idz))
  grid = disc.grid
  topology = grid.topology

  dim = dimensionality(grid)
  N = polynomialorder(grid)

  auxstate = disc.auxstate

  nauxstate = size(auxstate, 2)

  @assert nauxstate >= max(id, idx, idy, idz)
  @assert 0 < min(id, idx, idy, idz)
  @assert allunique((idx, idy, idz))

  Dmat = grid.D
  vgeo = grid.vgeo

  device = typeof(auxstate.Q) <: Array ? CPU() : CUDA()

  nelem = length(topology.elems)
  Nq = N + 1
  Nqk = dim == 2 ? 1 : Nq

  @launch(device, threads=(Nq, Nq, Nqk), blocks=nelem,
          elem_grad_field!(Val(dim), Val(N), Val(nauxstate), auxstate.Q, vgeo,
                           Dmat, topology.elems, id, idx, idy, idz))
end

"""
    dof_iteration!(dof_fun!::Function, R::MPIStateArray, disc::DGBalanceLaw,
                   Q::MPIStateArray)

Iterate over each dof to fill `R` using the `dof_fun!`. The syntax of the
`dof_fun!` is
```
dof_fun!(l_R, l_Q, l_Qvisc, l_aux)
```
where `l_R`, `l_Q`, `l_Qvisc`, and `l_aux` are of type `MArray` filled initially
with the values at a single degree of freedom. After the call the values in
`l_R` will be written back to the degree of freedom of `R`.
"""
function dof_iteration!(dof_fun!::Function, R::MPIStateArray, disc::DGBalanceLaw,
                        Q::MPIStateArray)
  grid = disc.grid
  topology = grid.topology

  @assert size(R)[end] == size(Q)[end] == size(disc.auxstate)[end]
  @assert size(R)[1] == size(Q)[1] == size(disc.auxstate)[1]

  dim = dimensionality(grid)
  N = polynomialorder(grid)

  Qvisc = disc.Qvisc
  auxstate = disc.auxstate

  nstate = size(Q, 2)
  nviscstate = size(Qvisc, 2)
  nauxstate = size(auxstate, 2)

  nRstate = size(R, 2)

  Dmat = grid.D
  vgeo = grid.vgeo

  device = typeof(auxstate.Q) <: Array ? CPU() : CUDA()

  nelem = length(topology.elems)
  Nq = N + 1
  Nqk = dim == 2 ? 1 : Nq
  Np = Nq * Nq * Nqk

  nrealelem = length(topology.realelems)

  @launch(device, threads=(Np,), blocks=nrealelem,
          knl_dof_iteration!(Val(dim), Val(N), Val(nRstate), Val(nstate),
                             Val(nviscstate), Val(nauxstate), dof_fun!, R.Q,
                             Q.Q, Qvisc.Q, auxstate.Q, topology.realelems))
end

end

max_stars_repo_name: jkozdon/CLIMA

max_issues_repo_name: jkozdon/CLIMA

max_forks_repo_name: jkozdon/CLIMA


text:
# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

@with_kw mutable struct ProgressiveHedgingData{T <: AbstractFloat}
    Q::T = 1e10
    δ₁::T = 1.0
    δ₂::T = 1.0
    iterations::Int = 0
end

@with_kw mutable struct ProgressiveHedgingParameters{T <: AbstractFloat}
    ϵ₁::T = 1e-5
    ϵ₂::T = 1e-5
    time_limit::T = Inf
    log::Bool = true
    keep::Bool = true
    offset::Int = 0
    indent::Int = 0
end

"""
    ProgressiveHedgingAlgorithm

Functor object for the progressive-hedging algorithm. Create using the `ProgressiveHedgingSolver` factory function and then pass to a `StochasticPrograms.jl` model.

...
# Parameters
- `τ::AbstractFloat = 1e-6`: Relative tolerance for convergence checks.
- `log::Bool = true`: Specifices if progressive-hedging procedure should be logged on standard output or not.
...
"""
struct ProgressiveHedgingAlgorithm{T <: AbstractFloat,
                                   A <: AbstractVector,
                                   ST <: ScenarioDecompositionStructure,
                                   E <: AbstractProgressiveHedgingExecution,
                                   P <: AbstractPenalization,
                                   PT <: AbstractPenaltyTerm} <: AbstractProgressiveHedging
    structure::ST
    data::ProgressiveHedgingData{T}
    parameters::ProgressiveHedgingParameters{T}

    # Estimate
    ξ::A
    decisions::DecisionMap
    Q_history::A
    primal_gaps::A
    dual_gaps::A

    # Execution
    execution::E

    penalization::P

    # Params
    progress::ProgressThresh{T}

    function ProgressiveHedgingAlgorithm(structure::ScenarioDecompositionStructure,
                                         x₀::AbstractVector,
                                         executer::AbstractExecution,
                                         penalizer::AbstractPenalizer,
                                         penaltyterm::AbstractPenaltyTerm; kw...)
        # Sanity checks
        length(x₀) != num_decisions(structure) && error("Incorrect length of starting guess, has ", length(x₀), " should be ", num_decisions(structure))
        num_subproblems == 0 && error("No subproblems in stochastic program. Cannot run progressive-hedging procedure.")
        n = num_subproblems(structure, 2)
        # Float types
        T = promote_type(eltype(x₀), Float32)
        x₀_ = convert(AbstractVector{T}, copy(x₀))
        A = typeof(x₀_)
        # Structure
        ST = typeof(structure)
        # Penalty term
        PT = typeof(penaltyterm)
        # Execution policy
        execution = executer(T, A, PT)
        E = typeof(execution)
        # Penalization policy
        penalization = penalizer()
        P = typeof(penalization)
        # Algorithm parameters
        params = ProgressiveHedgingParameters{T}(; kw...)

        ph = new{T,A,ST,E,P,PT}(structure,
                                ProgressiveHedgingData{T}(),
                                params,
                                x₀_,
                                structure.decisions[1],
                                A(),
                                A(),
                                A(),
                                execution,
                                penalization,
                                ProgressThresh(T(1.0), 0.0, "$(indentstr(params.indent))Progressive Hedging"))
        # Initialize solver
        initialize!(ph, penaltyterm)
        return ph
    end
end

function show(io::IO, ph::ProgressiveHedgingAlgorithm)
    println(io, typeof(ph).name.name)
    println(io, "State:")
    show(io, ph.data)
    println(io, "Parameters:")
    show(io, ph.parameters)
end

function show(io::IO, ::MIME"text/plain", ph::ProgressiveHedgingAlgorithm)
    show(io, ph)
end

function (ph::ProgressiveHedgingAlgorithm)()
    # Reset timer
    ph.progress.tinit = ph.progress.tlast = time()
    # Start workers (if any)
    start_workers!(ph)
    # Start procedure
    while true
        status = iterate!(ph)
        if status !== nothing
            close_workers!(ph)
            return status
        end
    end
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

struct SecurityRulesApi <: SwaggerApi
    client::Swagger.Client
end

"""

Creates or updates a security rule in the specified network security group.
Param: resourceGroupName::String (required)
Param: networkSecurityGroupName::String (required)
Param: securityRuleName::String (required)
Param: securityRuleParameters::SecurityRule (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: SecurityRule
"""
function securityRulesCreateOrUpdate(_api::SecurityRulesApi, resourceGroupName::String, networkSecurityGroupName::String, securityRuleName::String, securityRuleParameters, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "PUT", SecurityRule, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/networkSecurityGroups/{networkSecurityGroupName}/securityRules/{securityRuleName}", ["azure_auth"], securityRuleParameters)
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "networkSecurityGroupName", networkSecurityGroupName)  # type String
    Swagger.set_param(_ctx.path, "securityRuleName", securityRuleName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Deletes the specified network security rule.
Param: resourceGroupName::String (required)
Param: networkSecurityGroupName::String (required)
Param: securityRuleName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: Nothing
"""
function securityRulesDelete(_api::SecurityRulesApi, resourceGroupName::String, networkSecurityGroupName::String, securityRuleName::String, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "DELETE", Nothing, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/networkSecurityGroups/{networkSecurityGroupName}/securityRules/{securityRuleName}", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "networkSecurityGroupName", networkSecurityGroupName)  # type String
    Swagger.set_param(_ctx.path, "securityRuleName", securityRuleName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Get the specified network security rule.
Param: resourceGroupName::String (required)
Param: networkSecurityGroupName::String (required)
Param: securityRuleName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: SecurityRule
"""
function securityRulesGet(_api::SecurityRulesApi, resourceGroupName::String, networkSecurityGroupName::String, securityRuleName::String, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", SecurityRule, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/networkSecurityGroups/{networkSecurityGroupName}/securityRules/{securityRuleName}", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "networkSecurityGroupName", networkSecurityGroupName)  # type String
    Swagger.set_param(_ctx.path, "securityRuleName", securityRuleName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Gets all security rules in a network security group.
Param: resourceGroupName::String (required)
Param: networkSecurityGroupName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: SecurityRuleListResult
"""
function securityRulesList(_api::SecurityRulesApi, resourceGroupName::String, networkSecurityGroupName::String, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", SecurityRuleListResult, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/networkSecurityGroups/{networkSecurityGroupName}/securityRules", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "networkSecurityGroupName", networkSecurityGroupName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

export securityRulesCreateOrUpdate, securityRulesDelete, securityRulesGet, securityRulesList


text:
function subfor!(l::Matrix{Float64}, b::Vector{Float64})
# Forward-substitution on a Lower Triangle
  n = size(l, 1)
  for i in 1:n
    total = b[i]
    if i > 1
      for j in 1:i-1
        total -= l[i, j] * b[j]
      end
    end
    b[i] = total / l[i, i]
  end
end

# Forward-substitution
function subfor(l::Matrix{Float64}, b_in::Vector{Float64})
  b = deepcopy(b_in)
  n = size(l, 1)
  for i in 1:n
    total = b[i]
    if i > 1
      for j in 1:i-1
        total -= l[i, j] * b[j]
      end
    end
    b[i] = total / l[i, i]
  end
  b
end

export
  subfor!,
  subfor


text:
# ---
# title: 1206. Design Skiplist
# id: problem1206
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Design
# link: <https://leetcode.com/problems/design-skiplist/description/>
# hidden: true
# ---
#
# Design a Skiplist without using any built-in libraries.
#
# _A Skiplist is a data structure that takes  O(log(n)) time to `add`, `erase`
# and `search`. Comparing with treap and red-black tree which has the same
# function and performance, the code length of Skiplist can be comparatively
# short and the idea behind Skiplists are just simple linked lists._
#
# _For example:  we have a Skiplist containing `[30,40,50,60,70,90]` and we want
# to add `80` and `45` into it. The Skiplist works this way:_
#
# ![](https://assets.leetcode.com/uploads/2019/09/27/1506_skiplist.gif)
# Artyom Kalinin [CC BY-SA 3.0], via [Wikimedia
# Commons](https://commons.wikimedia.org/wiki/File:Skip_list_add_element-en.gif
# "Artyom Kalinin \[CC BY-SA 3.0 \(https://creativecommons.org/licenses/by-
# sa/3.0\)\], via Wikimedia Commons")
#
# _You can see there are many layers in the Skiplist. Each layer is a sorted
# linked list. With the help of the top layers,`add` , `erase` and `search `can
# be faster than O(n). It can be proven that the average time complexity for
# each operation is O(log(n)) and space complexity is O(n)._
#
# To be specific, your design should include these functions:
#
#   * `bool search(int target)` : Return whether the `target` exists in the Skiplist or not.
#   * `void add(int num)`: Insert a value into the SkipList.
#   * `bool erase(int num)`: Remove a value in the Skiplist. If `num` does not exist in the Skiplist, do nothing and return false. If there exists multiple `num` values, removing any one of them is fine.
#
# See more about Skiplist : <https://en.wikipedia.org/wiki/Skip_list>
#
# Note that duplicates may exist in the Skiplist, your code needs to handle this
# situation.
#
#
#
# **Example:**
#
#
#
#     Skiplist skiplist = new Skiplist();
#
#     skiplist.add(1);
#     skiplist.add(2);
#     skiplist.add(3);
#     skiplist.search(0);   // return false.
#     skiplist.add(4);
#     skiplist.search(1);   // return true.
#     skiplist.erase(0);    // return false, 0 is not in skiplist.
#     skiplist.erase(1);    // return true.
#     skiplist.search(1);   // return false, 1 has already been erased.
#
#
#
# **Constraints:**
#
#   * `0 <= num, target <= 20000`
#   * At most `50000` calls will be made to `search`, `add`, and `erase`.
#
#
## @lc code=start
using LeetCode

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


text:
# this source will be included by MLStyle.jl
import Base: convert
@case Fun{T, R}(fn :: Function)

function (typed_fn :: Fun{T, R})(arg :: T) :: R where {T, R}
    typed_fn.fn(arg)
end

function convert(::Type{Fun{T, R}}, fn :: Function) where {T, R}
    Func{T, R}(fn)
end

function convert(::Type{Fun{T, R}}, fn :: Fun{C, D}) where{T, R, C <: T, D <: R}
    Fun{T, R}(fn.fn)
end

(⇒)(::Type{T}, ::Type{R}) where {T, R} = Fun{T, R}


text:
#===============================================================================

    TmCheck1w.jl

    Author: Tom Price
    Date:   July 2019

===============================================================================#

struct TmCheck1w
    Tm                  ::Tuple{Float_T, Bool}
    area                ::Float_T
end

max_stars_repo_name: MakerButt/chaipcr

max_issues_repo_name: MakerButt/chaipcr

max_forks_repo_name: MakerButt/chaipcr


text:
#also too slow
using CategoricalArrays
import CategoricalArrays.CategoricalVector

function sumby_cate(by, val)
 byc = CategoricalVector(by)
 sumby(byc, val)
end

@time id6c = CategoricalVector(id6)
@which sumby(id6c, v1)
@time sumby_cate(id6, v1);
@time sumby(id6, v1);


text:
# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   Modia3D.Composition (Modia3D/Composition/_module.jl)
#

mutable struct AssemblyInternal <: ModiaMath.AbstractComponentInternal
   name::Symbol                                   # Instance name of assembly within parent assembly (is set by parent assembly)
   within::Union{Modia3D.AbstractAssembly,NOTHING}   # Parent assembly (if within==nothing, there is no parent assembly yet)
   scene::Union{Scene,NOTHING}                       # Scene in which the Assembly is present (if scene==nothing, Assembly is not yet in a scene)
   sceneOptions::Union{SceneOptions,NOTHING}         # SceneOptions defined for the Assembly (if sceneOptions==nothing, no sceneOptions are defined)
   referenceObject3D::Union{Object3D,NOTHING}        # Reference frame of assembly (if referenceObject3D==nothing, Assembly has not yet a reference frame)

   uniqueSignals::Array{Modia3D.AbstractSignal}
   uniqueForceTorques::Array{Modia3D.AbstractForceTorque}
   potentialVarInput::Array{ModiaMath.RealScalar}
   potentialVarOutput::Array{ModiaMath.RealScalar}
   flowVarInput::Array{ModiaMath.RealScalar}
   flowVarOutput::Array{ModiaMath.RealScalar}
   AssemblyInternal(name,sceneOptions=nothing) = new(Symbol(name), nothing, nothing, sceneOptions, nothing,
                    Array{Modia3D.AbstractSignal,1}(),Array{Modia3D.AbstractForceTorque,1}(),Array{ModiaMath.RealScalar,1}(),Array{ModiaMath.RealScalar,1}(),Array{ModiaMath.RealScalar,1}(),Array{ModiaMath.RealScalar,1}())
end

function initAssemblyComponent!(within::Modia3D.AbstractAssembly, component::Any, name)::NOTHING
    setfield!(within, Symbol(name), component)
    return nothing
end

function initAssemblyComponent!(within::Modia3D.AbstractAssembly, component::Modia3D.AbstractAssemblyComponent, name)::NOTHING
   component._internal.within = within
   component._internal.name   = Symbol(name)

   # within.name = component
   setfield!(within, Symbol(name), component)
   return nothing
end

function initAssemblyComponent!(within::Modia3D.AbstractAssembly, component::Object3D, name)::NOTHING
   component._internal.within = within
   component._internal.name   = Symbol(name)

   if component === component.parent && typeof(within._internal.referenceObject3D) == NOTHING
      within._internal.referenceObject3D = component
   end

   # within.name = component
   setfield!(within, Symbol(name), component)
   return nothing
end

function initAssemblyComponent!(within::Modia3D.AbstractAssembly, component::AbstractVector, name)::NOTHING
   for i in eachindex(component)
      component[i]._internal.within = within
      component[i]._internal.name   = Symbol( string(name) * "[" * string(i) * "]" )
   end

   # within.name = component
   setfield!(within, Symbol(name), component)
   return nothing
end

#-------------------------  Print assemblies -----------------------------
function Base.show(io::IO, a::Modia3D.AbstractAssemblyComponent)
   #if typeof(a._internal.scene) == NOTHING
      print(io, typeof(a), "(\n")
   #else
   #   print(io, typeof(a), "(scene=[...],")
   #end

   for c in fieldnames(typeof(a))
      field = getfield(a,c)
      if typeof(field) <: Modia3D.AbstractAssembly
         println(io,"\n   ", c, " = ", field)
      elseif typeof(field) <: AbstractVector
         for i in 1:length(field)
            println(io,"   ", c, "[", i, "] = ", field[i])
         end
      elseif typeof(field) != AssemblyInternal
         # Print fields, but not "_internal::Modia3D.Assembly"
         println(io,"   ", c, " = ", field)
      end
   end
   print(io, "   )")
end

#=
"""    printObject3DAndAllChildren(frame) - print frame and all of its children"""
function printObject3DAndAllChildren(frame::Object3D)::NOTHING
   twoObject3Dobjects = ObjectIdDict()
   if hasParent(frame)
      stack = frame.scene.stack
      empty!(stack)
   else
      stack = Object3D[]
   end
   push!(stack, frame)

   while length(stack) > 0
      frame1 = pop!(stack)
      println(frame1)
      for i in length(frame1.children):-1:1
         push!(stack, frame1.children[i])
      end
      for obj in frame1.twoObject3Dobject
         twoObject3Dobjects[obj] = obj
      end
   end

   if length(twoObject3Dobjects) > 0
      println("TwoObject3DObjects = [")
      first = true
      for (key,value) in twoObject3Dobjects
         if first
            first = false
         else
            println(",")
         end
         println("  ", value)
      end
      println("]")
   end

   return nothing
end

function JSON.show_json(io::JSON.StructuralContext,
                         s::JSON.CommonSerialization, scene::Scene)
   dictJoints = ObjectIdDict()
   stackObj   = Modia3D.AbstractObject3Ddata[]

   JSON.begin_object(io)
      JSON.show_pair(io, s, "name", scene.name)

      # Store all Object3Ds in sequence
      stack = scene.stack
      push!(stack, scene.world)
      JSON.show_key(io, "frames")
      JSON.begin_object(io)
         while length(stack) > 0
            frame = pop!(stack)
            JSON.show_key(io, frame.name)
            JSON.show_json(io, s, frame)

            # Put children on stack
            append!(stack, frame.children)

            # Put joint on dictJoints, if available
            if hasJoint(frame)
               dictJoints[frame.joint] = frame.joint
            end

            # Put object3D on stackObj, if available
            if hasData(frame) && typeof(frame.data) != Scene
               push!(stackObj, frame.data)
            end

            # Put twoObject3Dobject joints on dictJoints, if available
            for obj in frame.twoObject3Dobject
               if typeof(obj) <: Modia3D.AbstractJoint
                  dictJoints[obj] = obj
               end
            end
         end
      JSON.end_object(io)

      # Store all objects in sequence (sorted by name)
      perm = sortperm( [obj.name for obj in stackObj] )
      JSON.show_key(io, "objects3D")
      JSON.begin_object(io)
         for obj in stackObj[perm]
            JSON.show_key(io, obj.name)
            JSON.show_json(io, s, obj)
         end
      JSON.end_object(io)

      # Store all joints in sequence (sorted by name)
      JSON.show_key(io, "joints")
      joints = collect(keys(dictJoints))
      perm   = sortperm( [j.name for j in joints] )
      JSON.begin_object(io)
         for joint in joints[perm]
            JSON.show_key(io, joint.name)
            JSON.show_json(io, s, joint)
         end
      JSON.end_object(io)
   JSON.end_object(io)
end

function JSON.show_json(io::JSON.StructuralContext,
                         s::JSON.CommonSerialization, frame::Object3D)
   JSON.begin_object(io)
      JSON.show_pair(io, s, "scene"      , inScene(frame)   ? frame.scene.name : nothing)
      JSON.show_pair(io, s, "parent"     , hasParent(frame) ? frame.parent.name : nothing)
      JSON.show_pair(io, s, "children"   , [child.name for child in frame.children] )
      JSON.show_pair(io, s, "joint"      , hasJoint(frame) ? frame.joint.name : nothing)
      if frame.r_rel != ModiaMath.ZeroVector3D
         JSON.show_pair(io, s, "r_rel", frame.r_rel)
      end
      if frame.R_rel != ModiaMath.NullRotation
         JSON.show_pair(io, s, "R_rel", frame.R_rel)
      end
      if frame.r_abs != ModiaMath.ZeroVector3D
         JSON.show_pair(io, s, "r_abs", frame.r_abs)
      end
      if frame.R_abs != ModiaMath.NullRotation
         JSON.show_pair(io, s, "R_abs", frame.R_abs)
      end

      JSON.show_pair(io, s, "object3D"      , hasData(frame) ? frame.data.name : nothing)
      JSON.show_pair(io, s, "twoObject3Dobject", [obj.name for obj in frame.twoObject3Dobject])
   JSON.end_object(io)
end

"""
    writeObject3DAndAllChildrenOnJsonFile(frame; file=frame.name * ".json")

Write frame and all of its children in JSON format on file
"""
function writeObject3DAndAllChildrenOnJsonFile(frame::Object3D; file = frame.name*".json")::NOTHING
   if notInScene(frame)
      error("\nError from writeObject3DAndAllChildrenOnJsonFile:\n",
            "Can currently only write frames on file that are in a scene, but ", frame.name, " is not.")
   end
   open(file, "w") do io
      JSON.print(io, frame.scene, 3)
   end
end

=#

using Base.Meta:quot, isexpr

"""
    @assembly AssemblyName(arguments) begin ... end

Return the constructor for a new struct `AssemblyName`
consisting of Object3Ds that are connected together.
The new struct consists of all left-hand-side (scalar or vector) symbols
present between `begin ... end`.

# Examples
```julia
using Modia3D

@assembly Bar(;Lx = 0.1, Ly=Lx/5, Lz=Ly) begin
   frame0 = Modia3D.Object3D(Modia3D.Solid(Modia3D.SolidBeam(Lx,Ly,Lz), nothing, vmat1))
   frame1 = Modia3D.Object3D(frame0; r=[-Lx/2, 0.0, 0.0])
   frame2 = Modia3D.Object3D(frame0; r=[ Lx/2, 0.0, 0.0])
   cyl1   = Modia3D.Object3D(frame1, cyl)
   cyl2   = Modia3D.Object3D(frame2, cyl)
end
bar = Bar(;Lx=1.0)
Modia3D.visualizeAssembly!( bar )
```
"""
macro assembly(head, top_ex)
   # @assembly Name(arguments) begin ... end
   #    is mapped to:
   # mutable struct Name <: Modia3D.AbstractAssembly
   #    < declarations >
   #    function Name(sceneOptions=nothing; <arguments>)
   #       < equations top_ex >
   #       _internal = Modia3D.AssemblyInternal(:Name, sceneOptions)
   #       a = new(_internal, <declarations>)
   #       Modia3D.initAssemblyComponent(_internal, name1, :name1)
   #         ...
   #       return a
   #    end
   # end
   #
   # Todo: use  "gensym()" instead of "a" for the
   #
   # names:
   #
   if typeof(head) == Symbol
      assembly_symbol = head
      fc = Expr(:call, head)
   else
      @assert isexpr(head, :call)
      assembly_symbol = head.args[1]::Symbol
      fc = head
   end

   # Build:  internalAssemblyName = Modia3D.Assembly(assembly_name)
   # internalAssemblyName = gensym()
   # push!(eq, :( a = Modia3D.Assembly($assembly_name)) )

   # Inspect top_ex find
   #   symbol = < statements >
   # and store "symbol" in "names"
   this = gensym("this")
   eqInternal = []
   eq = []
   push!(eqInternal,:( $this = new( Modia3D.Composition.AssemblyInternal($(quot(assembly_symbol)),sceneOptions) ) ))
   namesSet = DataStructures.OrderedSet{Symbol}() # Set of left-hand side symbols
   for ex in top_ex.args
     # dump(ex)
      push!(eq, ex)
      if isexpr(ex, :(=), 2) && typeof(ex.args[1]) == Symbol
         symbol = ex.args[1]
         if in(symbol, namesSet)
            error("\nfrom @assembly ", assembly_symbol, ": Left hand side symbol in",
                  "\n   ", ex,
                  "\nis defined twice.")
         end
         push!(namesSet, symbol)
         push!(eq, :( Modia3D.Composition.initAssemblyComponent!($this,$symbol,$(quot(symbol)) ) ) )

 #        str = String(symbol)
 #        push!(eq, :( Modia3D.Composition.initAssemblyComponent!($this,$symbol,$str) ) )
     end
   end
   names = collect(namesSet)

   #helpArray = []
   #for name in names
   #  push!(helpArray, :(me.$name = $name))
   #end

   # Build:
   #   mutable struct assembly_symbol <: Modia3D.AbstractAssembly
   #      _internal::Modia3D.AssemblyInternal
   #      names[1]
   #      names[2]
   #       ...
   #      function assembly_symbol(...)
   #         $(eq...)
   #         Modia3D.Composition.initAssemblyInternal!( new(...) )
   #      end
   #   end
   code = :(mutable struct $assembly_symbol <: Modia3D.AbstractAssembly; end)
   code.args[3].args = vcat( :( _internal::Modia3D.Composition.AssemblyInternal ), names )

   code2 = code.args[3].args
   push!(code2, :(function fc(); end) )
   i = length(code2)
   code2[i].args[1] = fc

   # Add keyword argument sceneOptions to the arguments of the constructor function
   fcp = code2[i].args[1]
   if length(fcp.args) <= 1
      # No arguments; add new keyword argument sceneOptions
      push!(fcp.args, Expr(:parameters, Expr(:kw, :sceneOptions, :nothing) ))
   elseif typeof(fc.args[2]) == Symbol
      # Positional arguments present, but no keyword arguments yet; insert new kewword argument
      insert!(fcp.args, 2, Expr(:parameters, Expr(:kw, :sceneOptions, :nothing) ))
   elseif fc.args[2].head == :parameters
      # Keyword arguments already present; add sceneOptions
      push!(fcp.args[2].args, Expr(:kw, :sceneOptions, :nothing))
   else
      error("Unknown error in assembly.jl: typeof(fc.args[2]) = ", typeof(fc.args[2]))
   end

   # Generate constructor function
   code2[i].args[2] = quote
      $(eqInternal...)
 #     me = new($this)
      $(eq...)
#      $(helpArray...)
      return $this
   end

   # Print generated function (without "#....\n")
   # println( replace(sprint(showcompact,code) , r"# .*\n", "\n") )
   # println( sprint(showcompact,code) )

   return esc(code)
end


text:
using Random

const Torus = Int32

# Float64の値を[-0.5,0.5)に射影
FLoat2Torus05(x::Float64) = x - floor(x-0.5) - 1.0

# Float64の値を[0.0, 1.0)に射影
Float2Torus01(x::Float64) = x - floor(x)

# [-0.5,0.5)をTorus32に変換
dtot32(d::Float64) = trunc(Int32, d*2^32)

function gaussian_float(rng::AbstractRNG, sigma::Float64, dims...)
    randn(rng, dims...) .* sigma
end

# モジュラー正規分布に従う乱数生成
function modular_gaussian(rng::AbstractRNG, sigma::Float64, dims...)
    dtot32(Float2Torus05(gaussian_float(rng, sigma, dims...)))
end

# Torus32上の一様分布に従う乱数生成
function rand_uniform_Torus32(rng::AbstractRNG, dims...)
    rand(rng, Torus32, dims...)
end

max_stars_repo_name: bbtatt/BBTTFHE

max_issues_repo_name: bbtatt/BBTTFHE

max_forks_repo_name: bbtatt/BBTTFHE


text:
using Test, ScientificTypesBase, Tables
import ScientificTypesBase: scitype
const ST = ScientificTypesBase

include("convention.jl")

include("scitype.jl")