text:
"""
    Spline{B<:BSplineBasis, C<:AbstractVector{<:Real}}

Type for a spline based on a B-spline basis of type `B` and coefficient vector of type `C`.
"""
struct Spline{B<:BSplineBasis,C<:AbstractVector{<:Real}}
    basis::B
    coeffs::C

    function Spline{B,C}(basis, coeffs) where {B<:BSplineBasis,C<:AbstractVector{<:Real}}
        @boundscheck if eachindex(basis) != eachindex(coeffs)
            throw(DimensionMismatch("basis has indices $(eachindex(basis)), coefficient vector " *
                                    "has indices $(eachindex(coeffs))."))
        end
        new(basis, coeffs)
    end
end

"""
    Spline(basis, coeffs)

Create a spline from a B-spline basis and a vector of coefficients.
"""
Spline(basis, coeffs) = Spline{typeof(basis),typeof(coeffs)}(basis, coeffs)

Base.:(==)(x::Spline, y::Spline) = basis(x) == basis(y) && coeffs(x) == coeffs(y)

Base.hash(x::Spline, h::UInt) =
    hash(coeffs(x), hash(breakpoints(basis(x)), hash(order(x), hash(:Spline, h))))

function Base.show(io::IO, ::MIME"text/plain", s::Spline)
    cio = IOContext(io, :compact=>true)
    summary(io, s); println(io, ':')
    print(io, " basis: "); summary(io, basis(s)); println(io, ':')
    println(io, "  order: ", order(basis(s)))
    print(io, "  breakpoints: ")
    show(cio, breakpoints(basis(s)))
    if s isa BSpline
        print(io, "\n index: ", coeffs(s).index, " (knots: ")
        show(cio, bsplineknots(s))
        print(io, ')')
    else
        print(io, "\n coeffs: ")
        show(cio, coeffs(s))
    end
end

# Splines act as scalars for broadcasting
Base.Broadcast.broadcastable(s::Spline) = Ref(s)

for op in (:+, :-)
    @eval Base.$op(x::Spline) = Spline(basis(x), $op(coeffs(x)))
    @eval function Base.$op(x::Spline, y::Spline)
        basis(x) == basis(y) || throw(ArgumentError("splines must be defined on the same basis."))
        Spline(basis(x), $op(coeffs(x), coeffs(y)))
    end
end

Base.:*(x::Spline, y::Real) = Spline(basis(x), coeffs(x)*y)
Base.:/(x::Spline, y::Real) = Spline(basis(x), coeffs(x)/y)

Base.:*(x::Real, y::Spline) = Spline(basis(y), x*coeffs(y))
Base.:\(x::Real, y::Spline) = Spline(basis(y), x\coeffs(y))

LinearAlgebra.lmul!(x::Real, y::Spline) = (lmul!(x, coeffs(y)); y)
LinearAlgebra.rmul!(x::Spline, y::Real) = (rmul!(coeffs(x), y); x)

@static if VERSION ≥ v"1.2"
    LinearAlgebra.ldiv!(x::Real, y::Spline) = (ldiv!(x, coeffs(y)); y)
    LinearAlgebra.rdiv!(x::Spline, y::Real) = (rdiv!(coeffs(x), y); x)
end

(s::Spline)(args...; kwargs...) = splinevalue(s, args...; kwargs...)

check_intervalindex(::Type{Bool}, s::Spline, x, index) =
    check_intervalindex(Bool, basis(s), x, index)

"""
    basis(spline::Spline)

Return the B-spline basis on which the spline is defined.
"""
basis(s::Spline) = s.basis

"""
    coeffs(spline::Spline)

Return the coefficient vector of `spline`.
"""
coeffs(s::Spline) = s.coeffs

order(s::Spline) = order(basis(s))

"""
    BSpline{B} <: Spline{B}

Type for a B-spline from a B-spline basis of type `B`. `BSpline`s can be obtained by
indexing into a B-spline basis.
"""
const BSpline = Spline{B, StandardBasisVector{Bool}} where B<:BSplineBasis

@propagate_inbounds BSpline{B}(basis, index::Integer) where B<:BSplineBasis =
    BSpline{B}(basis, StandardBasisVector(length(basis), index))

"""
    BSpline(basis::BSplineBasis, index)

Return the `index`-th B-spline of `basis`.

# Example

```jldoctest
julia> basis = BSplineBasis(5, 1:10);

julia> BSpline(basis, 3)
BSpline{BSplineBasis{UnitRange{Int64}}}:
 basis: 13-element BSplineBasis{UnitRange{Int64}}:
  order: 5
  breakpoints: 1:10
 index: 3 (knots: [1, 1, 1, 2, 3, 4])
```
"""
@propagate_inbounds BSpline(basis::B, index::Integer) where B<:BSplineBasis =
    BSpline{B}(basis, index)

Base.show(io::IO, x::BSpline) =
    (summary(io, x); print(io, '(', basis(x), ", ", coeffs(x).index, ')'))

@static if VERSION < v"1.6.0-DEV.347"
    Base.summary(io::IO, x::BSpline{B}) where B = print(io, "BSpline{", B, '}')
end

bsplineknots(x::BSpline) = bsplineknots(basis(x), coeffs(x).index)
bsplineknots(basis, index) = knots(basis)[index:index+order(basis)]

"""
    support(spline::Spline) -> a, b

If `spline` is a `BSpline`, return the interval ``[a,b]`` on which the B-spline is non-zero.
Otherwise, return `support(basis(spline))`.

# Examples

```jldoctest
julia> basis = BSplineBasis(3, 0:5);

julia> spline = Spline(basis, ones(7));

julia> zerospline = Spline(basis, zeros(7));

julia> support(spline)
(0, 5)

julia> support(zerospline) # even though the spline is zero everywhere
(0, 5)

julia> support(basis[4]) # for BSplines, return their actual support
(1, 4)
```
"""
support(s::Spline) = support(basis(s))
support(b::BSpline) = (t = knots(basis(b)); i = coeffs(b).index; (t[i], t[i+order(b)]))

"""
    splinevalue(spline::Spline, x[, ::Derivative{N}]; kwargs...)

Calculate the value of `spline`, or its `N`-th derivative, at `x`.

Two optional keyword arguments can be used to increase performance:
* `workspace`: By default, the function allocates a vector of length `order(spline)` in
  which the calculation is performed. To avoid this, a pre-allocated vector can be supplied
  with the `workspace` keyword. In this case, the returned value is always of type
  `eltype(workspace)`.
* `leftknot`: If the index of the relevant interval (i.e.,
  `intervalindex(basis(spline), x)`) is already known, it can be supplied with the
  `leftknot` keyword.

Instead of calling `splinevalue`, a spline object can be called directly:
`spline(x[, Derivative(N)]; kwargs...)` is equivalent to
`splinevalue(spline, x[, Derivative(N)]; kwargs...)`.

# Examples

```jldoctest
julia> spl = Spline(BSplineBasis(4, 0:5), 1:8);

julia> splinevalue(spl, 1.7)
3.69775

julia> splinevalue(spl, 1.7, Derivative(1))
1.0225

julia> splinevalue(spl, 3.6, leftknot=7)
5.618

julia> spl(18//5)
2809//500

julia> spl(3.6, Derivative(3), leftknot=7)
0.5
```
"""
function splinevalue(spline::Spline, x, drv::Derivative{N}=NoDerivative();
                     leftknot=intervalindex(basis(spline), x),
                     workspace=nothing) where N
    check_intervalindex(spline, x, leftknot)
    T = workspace === nothing ? bspline_returntype(spline, x) : eltype(workspace)
    leftknot === nothing && return isnan(x) ? T(NaN) : zero(T)
    if spline isa BSpline
        iszero_bspline(spline, leftknot) && return zero(T)
    end
    N ≥ order(spline) && return zero(T)
    if workspace === nothing
        workspace = Vector{T}(undef, order(spline))
    else
        exp = (Base.OneTo(order(spline)),)
        got = axes(workspace)
        if exp != got
            throw(DimensionMismatch("workspace has wrong axes: expected $exp, got $got"))
        end
    end
    @inbounds _splinevalue(workspace, basis(spline), coeffs(spline), x, leftknot, drv)
end

# Return true iff `spline` is zero on the interval given by `leftknot`
function iszero_bspline(spline::BSpline, leftknot::Integer)
    i = coeffs(spline).index
    (i < leftknot-order(spline)+1) | (i > leftknot)
end

# The implementation of this method (and the methods iterate_splinevalue_derivatives! and
# iterate_splinevalue_bsplines! it calls) is adapted from the Fortran subroutine BVALUE from
# Carl de Boor’s book *A practical Guide to Splines* [^deBoor1978].
#
# [^deBoor1978]:
#     Carl de Boor, *A Practical Guide to Splines*, New York, N.Y.: Springer-Verlag, 1978.
@propagate_inbounds function _splinevalue(A, basis::BSplineBasis, coeffs, x, leftknot::Integer, ::Derivative{N}) where N
    t = knots(basis)
    k = order(basis)
    xtyped = convert(eltype(A), x)
    A .= @view coeffs[leftknot-k+1:leftknot]
    # Difference the coefficients (stored in A) N times
    for j = 1:N
        iterate_splinevalue_derivatives!(A, t, k, j, leftknot)
    end
    # Compute value of N-th derivative at x from its B-spline coefficients in A[1:k-N]
    for j = N+1:k-1
        iterate_splinevalue_bsplines!(A, t, k, j, xtyped, leftknot)
    end
    A[1]
end

@propagate_inbounds function iterate_splinevalue_bsplines!(coeffs::AbstractArray{T}, knots, k, j, x::T, leftknot) where T
    kmj = k-j
    for i = 1:kmj
        rindex = leftknot + i
        lindex = rindex - kmj
        tl = convert(T, knots[lindex])
        tr = convert(T, knots[rindex])
        coeffs[i] = (coeffs[i+1]*(x-tl) + coeffs[i]*(tr-x)) / (tr - tl)
    end
end

@propagate_inbounds function iterate_splinevalue_derivatives!(coeffs::AbstractArray{T}, knots, k, j, leftknot) where T
    kmj = k-j
    for i = 1:kmj
        rindex = leftknot + i
        lindex = rindex - kmj
        coeffs[i] = (coeffs[i+1] - coeffs[i]) / (convert(T, knots[rindex]) - convert(T, knots[lindex])) * kmj
    end
end


text:
#=
    Scenario generation methods
=#
abstract type AbstractScenariosGenerator end

mutable struct MarkovChain
    states
    transition_matrices
end

# Compute markov chain
function MarkovChain(data, nstate::Int64, algo::String)
    # Parameters
    nk = size(data,1)
    nh = 24
    nm = size(data[1],1)

    # Pre-allocate
    states = [zeros(nh, nk, nstate) for m in 1:nm]
    transition_matrices = [zeros(nstate, nstate, nh - 1) for m in 1:nm]

    # For each month, we compute the markov states and transition matrices
    for m in 1:nm
        # For each hour, we extract markov states using k-means from aggregated data to keep the synchronicity
        states[m], sequences = compute_markov_states([data[k][m] for k in 1:nk], nstate, algo)
        # We compute the transition matrices between consecutive hours using the sequences
        transition_matrices[m] = compute_transition_matrices_between_hours(sequences, nstate)
    end

    return MarkovChain(states, transition_matrices)
end
# Compute states using clustering algorithm
function compute_markov_states(data, nstate::Int64, algo::String)
    # Parameters
    nk = size(data,1)
    nh = size(data[1], 1)
    nd = size(data[1], 2)

    # Pre-allocate
    states = zeros(nh, nk, nstate)
    sequences = zeros(Int64, nh+1, nd)

    # Normalization
    data_n = replace!.(data ./ maximum.(data), NaN => 0.)

    # Clustering
    for h in 1:nh
        if algo == "kmeans"
            # Aggregate normalized data
            data_agg = permutedims(hcat([data_n[k][h,:] for k in 1:nk]...))
            # Clustering
            clusters = kmeans(data_agg, nstate)
            # Store the denormalized states and sequences
            states[h,:,:] = clusters.centers .* maximum.(data)
            sequences[h,:] = permutedims(clusters.assignments)
        elseif algo == "kmedoids"
            # Aggregate normalized data
            data_agg = permutedims(hcat([data_n[k][h,:] for k in 1:nk]...))
            # Compute euclidean distance matrix
            dist = pairwise(Euclidean(), data_agg, dims = 2 )
            # Clustering
            clusters = kmedoids(dist, nstate)
            # Store the states and sequences
            states[h,:,:] = data_agg[:, clusters.medoids] .* maximum.(data)
            sequences[h,:] = permutedims(clusters.assignments)
        end
    end
    # The sequence of the last hour is equal to the first one
    sequences[end, :] = sequences[1, :]

    return states, sequences
end
# Compute transition matrix between consecutive hours
function compute_transition_matrices_between_hours(sequence, nstate::Int64)
    # Parameters
    nh = size(sequence, 1)

    # Pre-allocate
    transition_matrix = zeros(nstate, nstate, nh - 1)

    for h in 1:nh-1
        for (i,j) in zip(sequence[h, :], sequence[h+1, :])
            transition_matrix[i, j, h] += 1
        end

        # Convert into probabilities
        for i in 1:size(transition_matrix,1)

            sum_row = sum(transition_matrix[i, :, h])

            if sum_row > 0 # to avoid zero division...
                transition_matrix[i, :, h] = transition_matrix[i, :, h] ./ sum_row
            elseif sum_row == 0
                transition_matrix[i, i, h] = 1 # to avoid bug in Categorical...
            else
                print("Error! The sum couldn't be negative")
                return
            end
        end
    end

    return transition_matrix
end
# Clustering data by week for each month
function clustering_month_week(data, t; flag::String="")
    # Parameters
    nm = 12 # number of month in one year

    # Pre_allocate
    data_cluster= Vector{Array{Float64,2}}(undef, nm)

    for m in 1:nm
        # We cluster week days using the "not" operator
        if flag == "odd" # odd indices
            data_cluster[m] = reshape(data[(Dates.month.(t) .== m) .& .~isweekend(t)], 24, :)[:, isodd.(1:end)]
        elseif flag == "even" # even indices
            data_cluster[m] = reshape(data[(Dates.month.(t) .== m) .& .~isweekend(t)], 24, :)[:, iseven.(1:end)]
        else
            data_cluster[m] = reshape(data[(Dates.month.(t) .== m) .& .~isweekend(t)], 24, :)
        end
    end

    return data_cluster
end
# Clustering data by weekend for each month
function clustering_month_weekend(data, t; flag::String="")
    # Parameters
    nm = 12 # number of month in one year

    # Pre_allocate
    data_cluster = Vector{Array{Float64,2}}(undef, nm)

    for m in 1:nm
        # We cluster weekend days
        if flag == "odd" # odd indices
            data_cluster[m] = reshape(data[(Dates.month.(t) .== m) .& isweekend(t)], 24, :)[:, isodd.(1:end)]
        elseif flag == "even" # even indices
            data_cluster[m] = reshape(data[(Dates.month.(t) .== m) .& isweekend(t)], 24, :)[:, iseven.(1:end)]
        else
            data_cluster[m] = reshape(data[(Dates.month.(t) .== m) .& isweekend(t)], 24, :)
        end
    end

    return data_cluster
end
# Clustering data by month
function clustering_month(data, t; flag::String="")
    # Parameters
    nm = 12 # number of month in one year

    # Pre_allocate
    data_cluster = Vector{Array{Float64,2}}(undef, nm)

    for m in 1:nm
        # We cluster weekend days
        if flag == "odd" # odd indices
            data_cluster[m] = reshape(data[(Dates.month.(t) .== m)], 24, :)[:, isodd.(1:end)]
        elseif flag == "even" # even indices
            data_cluster[m] = reshape(data[(Dates.month.(t) .== m)], 24, :)[:, iseven.(1:end)]
        else
            data_cluster[m] = reshape(data[(Dates.month.(t) .== m)], 24, :)
        end
    end

    return data_cluster
end

mutable struct MarkovGenerator <: AbstractScenariosGenerator
    nstate::Int64
    algo::String
    markovchain::NamedTuple{(:wk, :wkd), Tuple{MarkovChain, MarkovChain}}

    MarkovGenerator(; nstate = 10, algo = "kmeans") = new(nstate, algo)
end

function initialize_generator!(generator::MarkovGenerator, data...)
    # TODO: generic function without .power and .t or specify the input type
    # Clustering by week and weekend for each month
    wk = [clustering_month_week(d.power, d.t) for d in data]
    wkd = [clustering_month_weekend(d.power, d.t) for d in data]

    # The number of states must be greater than 2 and lower than the minimum number of week and weekend days
    min_wk, min_wkd =  minimum([size(wk[k][m])[2] for k in 1:size(wk,1), m in 1:12]), minimum([size(wkd[k][m])[2] for k in 1:size(wk,1), m in 1:12])
    1 < generator.nstate < min(min_wk, min_wkd) ? nothing : generator.nstate = min(min_wk, min_wkd) - 1

    # Compute markov chains
    generator.markovchain = (wk = MarkovChain(wk, generator.nstate, generator.algo),
                             wkd = MarkovChain(wkd, generator.nstate, generator.algo))

    return generator
end

# Generate scenario from markov chains
function generate(generator::MarkovGenerator, s0, t0::DateTime, nstep::Int64; ny::Int64=1, ns::Int64=1)
    # Initialize state with the closest value
    mc = chose(generator, t0)
    Δ_s0 = [s0 .- mc.states[Dates.month(t0)][Dates.hour(t0)+1, :, :][:,state] for state in 1:generator.nstate]
    idx_0 = findmin(norm.(Δ_s0))[2]

    # Timestamp
    t = t0:Hour(1):t0+Hour(nstep-1)
    hour, month = Dates.hour.(t) .+ 1, Dates.month.(t) # hour(t[1:24]) = 0:23...

    # Pre-allocate
    ω_generated = [zeros(nstep, ny, ns) for _ in 1:size(mc.states[1],2)]
    probabilities = ones(nstep, ny, ns)

    # Initialization
    for j in 1:size(mc.states[1],2)
        ω_generated[j][1,:,:] .= s0[j]
    end

    # Generate scenarios
    for s in 1:ns, y in 1:ny, h in 1:nstep-1
        # Test to chose the appropriate markov chain
        mc = chose(generator, t[h])
        # Build categorical distribution with transition matrices - we make the assumption that the probabilities from 24h to 1h are equal whatever day transitions
        distribution = Categorical(mc.transition_matrices[month[h]][idx_0, :, hour[h]])
        # Sample state index from the distribution
        idx_1 = rand(distribution)
        # Retrieve the associated values
        for j in 1:size(mc.states[1],2)
            ω_generated[j][h+1,y,s] = mc.states[month[h]][hour[h+1], j, idx_1]
        end
        # Store probabilities
        probabilities[h+1,y,s] = probs(distribution)[idx_1]
        # Update idx
        idx_0 = idx_1
    end

    return ω_generated, prod(probabilities, dims=1)[1,:,:] / sum(prod(probabilities, dims=1)[1,:,:])
end

# Perfect foresight generator
mutable struct AnticipativeGenerator <: AbstractScenariosGenerator
    forecast
    AnticipativeGenerator() = new()
end

function initialize_generator!(generator::AnticipativeGenerator, data...)
    # Store anticipative forecast
    generator.forecast = [d for d in data]
    return generator
end

# Generate perfect forecast
function generate(generator::AnticipativeGenerator, s0, t0::DateTime, nstep::Int64; ny::Int64=1, ns::Int64=1, h::Int64=1)
    # Current index
    if length(generator.forecast[1].t[h:end]) < nstep
        # Windows
        windows = h : h + length(generator.forecast[1].t[h:end]) - 1
        # Add zeros to have a constant size
        n_zeros = nstep - length(windows)
        # Forecast
        forecast = [vcat(d.power[windows], zeros(n_zeros)) for d in generator.forecast]
    else
        # Windows
        windows = h : h + nstep - 1
        # Forecast
        forecast = [d.power[windows] for d in generator.forecast]
    end

    return forecast, [1.]
end

ext: jl


text:
# bands

using CairoMakie

f = Figure(resolution = (600, 400), backgroundcolor = :honeydew,)
Axis(f[1, 1])

xs = 1:0.2:10
ys_low = -0.2 .* sin.(xs) .- 0.25
ys_high = 0.2 .* sin.(xs) .+ 0.25

band!(xs, ys_low, ys_high, color=:maroon1)
band!(xs, ys_low .- 1, ys_high .-1, color = :purple1)

# bands ribbon

let
    x = LinRange(-10, 10, 200)
    fig = Figure(resolution = (600, 400))
    ax = Axis(fig[1, 1], xlabel = "x", ylabel = "y")
    band!(x, sin.(x), sin.(x) .+ 1; color = (:cadetblue2, 0.4))
    band!(x, cos.(x), 1 .+ cos.(x); color = (:red, 0.2))
    display(fig)
end;

max_stars_repo_name: Ellon-M/visualizations

max_issues_repo_name: Ellon-M/visualizations

max_forks_repo_name: Ellon-M/visualizations


text:
# Penalty Vector

PenaltyVector() = PenaltyVector(Vector{Penalty}());
PenaltyVector(n :: Integer) = PenaltyVector(Vector{Penalty}(undef, n));

Base.length(pv :: PenaltyVector) = length(pv.pv);
Base.isempty(pv :: PenaltyVector) = isempty(pv.pv);
Base.getindex(pv :: PenaltyVector, idx) = getindex(pv.pv, idx);
Base.setindex!(pv :: PenaltyVector, v :: Penalty, idx) = setindex!(pv.pv, v, idx);
Base.push!(pv :: PenaltyVector, v :: Penalty) = push!(pv.pv, v);

Base.iterate(pv :: PenaltyVector) = Base.iterate(pv, 1)

function Base.iterate(pv :: PenaltyVector, idx)
    if idx > length(pv)
        return nothing
    else
        return (pv.pv[idx], idx+1)
    end
end

use(pv :: PenaltyVector, pName :: Symbol) = set_use(pv, pName, true);
no_use(pv :: PenaltyVector, pName :: Symbol) = set_use(pv, pName, false);

function set_use(pv :: PenaltyVector, pName :: Symbol, doUse :: Bool)
    p = retrieve(pv, pName);
    @assert !isnothing(p)  "Penalty $pName not found"
    set_use(p, doUse);
end

function is_used(pv :: PenaltyVector, pName :: Symbol)
    p = retrieve(pv, pName);
    @assert !isnothing(p)  "Penalty $pName not found"
    return is_used(p);
end

"""
	$(SIGNATURES)

Retrieve a penalty by name. Returns nothing if not found.
"""
function retrieve(pv :: PenaltyVector, pName :: Symbol)
    pOut = nothing;
    for p in pv
        if name(p) == pName
            pOut = p;
            break;
        end
    end
    return pOut
end

"""
	$(SIGNATURES)

Compute all (used) scalar penalties.
"""
function compute_all_penalties(pv :: PenaltyVector, x)
    for p in pv
        if is_used(p)
            compute_penalty(p, x);
        end
    end
    return nothing
end

"""
	$(SIGNATURES)

Return scalar penalties as Dict mapping names into penalties.
"""
function penalty_dict(pv :: PenaltyVector)
    d = Dict{Symbol, Float64}();
    for p in pv
        if is_used(p)
            d[name(p)] = scalar_value(p);
        end
    end
    return d
end

"""
	$(SIGNATURES)

Computer sum of all used penalties.
Optionally recompute all penalties using the `data` provided.
"""
function total_penalty(pv :: PenaltyVector; data = nothing)
    pTotal = 0.0;
    for p in pv
        if is_used(p)
            if !isnothing(data)
                compute_penalty(p, data);
            end
            pTotal += scalar_value(p);
        end
    end
    return pTotal
end

"""
	$(SIGNATURES)

Report all (positive) penalties. Optionally recompute with provided data.
"""
function report_penalties(pv :: PenaltyVector;  io = stdout,
    data = nothing)

    nPen = length(pv);
    pTotal = total_penalty(pv; data = data);
    pStr = format_number(pTotal);
    println(io, "$nPen Penalties:   total = $pStr")
    for p in pv
        if is_used(p)
            pStr = format_number(scalar_value(p));
            println("  $(name(p)):  $pStr")
        end
    end
    return nothing
end

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


text:
"""
    VariationalDense(in, out, σ=identity;
                    weight_init=TrainableDistribution,
                    bias_init=TrainableDistribution,
                    bias=true)
    VariationalDense(weight_sampler, bias_sampler, act)

Creates a variational dense layer. Computes variational bayesian
approximation to the distribution over the parameters of the dense layer.
The stochasticity is during the forward pass, instead of using point
estimates for weights and biases, we sample from the distribution over
weights and biases. Gradients of the distribution's learnable parameters
are trained using the reparameterization trick.

Reference - https://arxiv.org/abs/1505.05424
We use DistributionsAD - https://github.com/TuringLang/DistributionsAD.jl
to help us with backprop.

# Fields
- `weight_sampler`: A trainable distribution from which weights are sampled
                    in every forward pass
- `bias_sampler`: A trainable distribution from which biases are sampled in
                    every forward pass
- `act`: Activation function, applies to logits after layer transformation

# Arguents
- `in::Integer`: Input dimension size
- `out::Integer`: Output dimension size
- `σ`: Acivation function, defaults to identity
- `init`: Distribution parameters Initialization, defaults to glorot_normal
- `weight_dist`: Weight distribution, defaults to a trainable multivariate normal
- `bias_dist`: Bias distribution, defaults to trainable multivariate normal

"""
struct VariationalDense{WS,BS,F}
    weight_sampler::WS
    bias_sampler::BS
    act::F
end

function VariationalDense(
    in::Integer,
    out::Integer,
    σ = identity;
    init = glorot_normal,
    weight_dist = TrainableMvNormal,
    bias_dist = TrainableMvNormal,
    device = cpu,
    bias = true,
)
    # Initialize alpha and gamma samplers
    weight_sampler = weight_dist((out, in), init = init, device = device)
    if bias
        bias_sampler = bias_dist((out,), init = init, device = device)
    else
        bias_sampler = () -> Flux.Zeros()
    end

    return VariationalDense(weight_sampler, bias_sampler, σ)
end

@functor VariationalDense

function (dv::VariationalDense)(x)
    weight = dv.weight_sampler()
    bias = dv.bias_sampler()
    logits = weight * x .+ bias
    return dv.act.(logits)
end


text:
charNum = parse(Int, ARGS[1])
charS = lpad(charNum, 2, "0")

##
cd(@__DIR__)
using Pkg
Pkg.activate(".")
Pkg.instantiate()

##
using LinearAlgebra
using StatsBase
using CSV
using DataFrames
using Pipe
using ProgressMeter
using StatsPlots
using MCPhylo
using Base.Threads
using StatsFuns
using Optim
using Distributions
using Random
Random.seed!(6897415724436294598)

##

wals = CSV.read("../data/charMtx.csv", DataFrame)
d = CSV.read("../data/fpairMtx.csv", DataFrame)

famFreqs = combine(groupby(wals, :glot_fam), nrow)
sort!(famFreqs, :nrow, rev=true)

families = filter(x -> x.nrow>1, famFreqs).glot_fam
isolates = filter(x -> x.nrow==1, famFreqs).glot_fam

isoDict = @pipe wals |>
      filter(x -> x.glot_fam ∈ isolates, _) |>
      zip(_.glot_fam, _.longname) |>
      Dict

##

fm2trees = Dict()
@showprogress for fm in families
      fm2trees[fm] = MCPhylo.ParseNewick("../data/posteriorTrees/$fm.posterior.tree")
end

for fm in isolates
      fm2trees[fm] = repeat([Node(isoDict[fm])], 1000)
end

##

function renumberTree(tree::GeneralNode)
      tips = [nd for nd in post_order(tree) if nd.nchild==0]
      nonTips = [nd for nd in post_order(tree) if nd.nchild > 0]
      sort!(tips, lt= (x,y) -> x.name < y.name)
      for (i,nd) in enumerate(tips)
            nd.num = i
      end
      for (i,nd) in enumerate(nonTips)
            nd.num = i+length(tips)
      end
      tree
end

##

function unnorm_log_post(mc::MCPhylo.ModelChains, θ_vec::Array)
      mod = deepcopy(mc.model)
      relist!(mod, θ_vec)
      logpdf(mod)
end

##

for fm in keys(fm2trees)
      fm2trees[fm] = renumberTree.(fm2trees[fm])
end

##

lineages = vcat(families, isolates)
nLineages = length(lineages)

##

ttrees = [[fm2trees[fm][i] for fm in lineages] for i in 1:1000]

nnodes = [length(post_order(t)) for t in ttrees[1]]

##

states = ["a", "b", "c", "d"]

nsites = 1

nbase = 4

rates = [1.0]

##
taxa = Vector{String}[]
for t in ttrees[1]
      push!(taxa, sort([nd.name for nd in pre_order(t) if nd.nchild==0]))
end

##

include("bridge_sampling.jl")
##

char = Dict(zip(d.taxon, d[:,charNum+1]))

data = zeros(nbase, nsites, maximum(nnodes), length(ttrees[1]))

for i = 1:length(taxa)
      for l in taxa[i]
            s = char[l]
            t_ind = find_by_name(ttrees[1][i], l).num
            if s in states
                  data[:, 1, t_ind, i] = (states .== s)
            else
                  data[:, 1, t_ind, i] .= 1
            end
      end
end

##
sim = read(
      "modelFitting/output/universal_$(charS).jls",
      ModelChains,
)
samples =
      vcat([sim[:, [:s_rates, :ind], :].value[:, :, k] for k = 1:2]...) |> DataFrame
function log_density(sampleRow)
      unnorm_log_post(sim, sampleRow)
end
uD = bridge_sampling(Array(samples), log_density)

##

sz =

sim = read(
      "modelFitting/output/lineage_$(charS).jls",
      ModelChains,
)
samples =
      vcat(
            [
                  sim[1:(size(sim.value)[1]÷1000):end, [:s_rates, :ind], :].value[
                        :,
                        :,
                        k,
                  ] for k = 1:2
            ]...,
      ) |> DataFrame
function log_density(sampleRow)
      unnorm_log_post(sim, sampleRow)
end
lD = bridge_sampling(Array(samples), log_density, verbose=true)

##

try
      mkdir("bsResults")
catch e
end

file = open("bsResults/$(charS).csv", "w")
write(file, join(string.([uD, lD]), ",")*"\n")
close(file)

max_stars_repo_name: erathorn/phylogeneticTypology

max_issues_repo_name: erathorn/phylogeneticTypology

max_forks_repo_name: erathorn/phylogeneticTypology


text:
function part1(input)
    tiles = parse_input(input)
    edges = Dict(n => extract_edges(tile) for (n, tile) in tiles)
    edge_to_tile = Dict{Int, Vector{Int}}()
    for (tile_number, edge_codes) in edges
        for edge_code in edge_codes
            push!(get!(edge_to_tile, edge_code, Int[]), tile_number)
        end
    end
    square, _ = fit_into_square(edges, edge_to_tile)
    return prod(square[[1, end], [1, end]])
end

function part2(input)
    tiles = parse_input(input)
    edges = Dict(n => extract_edges(tile) for (n, tile) in tiles)
    edge_to_tile = Dict{Int, Vector{Int}}()
    for (tile_number, edge_codes) in edges
        for edge_code in edge_codes
            push!(get!(edge_to_tile, edge_code, Int[]), tile_number)
        end
    end
    square, square_edges = fit_into_square(edges, edge_to_tile)
    image = reconstruct_image(square, square_edges, tiles, edges)
    num_sea_monsters = find_sea_monsters.(vcat(rotl90.(Ref(image), 0:3),
                                               rotl90.(Ref(reverse(image, dims = 1)), 0:3)))
    n = only(filter(>(0), num_sea_monsters))
    return sum(image) - 15 * n
end

function parse_input(input)
    return Dict(parse_tile.(split(readchomp(input), "\n\n")))
end

function parse_tile(tile_string)
    lines = split(rstrip(tile_string), "\n")
    tile_number = parse(Int, only(match(r"Tile (\d+):", first(lines)).captures))
    tile = reduce(hcat, (split(line, "") .== "#" for line in lines[2:end]))
    return tile_number => tile
end

function extract_edges(tile)
    return encode_edge.([tile[1, :], reverse(tile[1, :]),
                         tile[:, end], reverse(tile[:, end]),
                         reverse(tile[end, :]), tile[end, :],
                         reverse(tile[:, 1]), tile[:, 1]])
end

function encode_edge(edge)
    return sum(1 .<< (i - 1)
               for (i, e) in enumerate(edge)
               if e)
end

const TOP = 1
const RIGHT = 2
const BOTTOM = 3
const LEFT = 4

function fit_into_square(edges, edge_to_tile)
    side = isqrt(length(edges))
    @assert side^2 == length(edges)
    square = zeros(Int, side, side)
    square_edges = zeros(Int, side, side, 4)
    for (tile, edge_codes) in edges
        used_tiles = Set(tile)
        square[1, 1] = tile
        for i = 1:8
            square_edges[1, 1, :] .= get_tile_permutation(edge_codes, i)
            if fit_rest_of_square!(square, square_edges, used_tiles,
                                   edges, edge_to_tile)
                return square, square_edges
            end
        end
    end
    error("No solution found.")
end

function fit_rest_of_square!(square, square_edges, used_tiles,
                            edges, edge_to_tile)
    side = size(square, 1)
    for (i, j) in Tuple.(CartesianIndices((1:side, 1:side))[2:end])
        if i == 1
            left_edge = square_edges[i, j - 1, RIGHT]
            fitting_tiles = setdiff(edge_to_tile[left_edge], square[i, j - 1])
            isempty(fitting_tiles) && return false
            square[i, j] = only(fitting_tiles)
            square_edges[i, j, :] = get_tile_edges(edges[square[i, j]],
                                                   left_edge, LEFT)
        else
            top_edge = square_edges[i - 1, j, BOTTOM]
            fitting_tiles = setdiff(edge_to_tile[top_edge], square[i - 1, j])
            isempty(fitting_tiles) && return false
            square[i, j] = only(fitting_tiles)
            square_edges[i, j, :] = get_tile_edges(edges[square[i, j]],
                                                   top_edge, TOP)
            if j > 1 && square_edges[i, j - 1, RIGHT] != square_edges[i, j, LEFT]
                return false
            end
        end
    end
    return true
end

function get_tile_edges(edge_codes, given_edge, position)
    @assert given_edge in edge_codes
    @assert position in (TOP, LEFT)
    if position == TOP
        if findfirst(==(given_edge), edge_codes) % 2 == 0
            edge_codes = reverse(edge_codes)
        end
        n = findfirst(==(given_edge), edge_codes)
        edge_codes = edge_codes[mod1.(n .+ (0:7), 8)]
    elseif position == LEFT
        if findfirst(==(given_edge), edge_codes) % 2 == 1
            edge_codes = reverse(edge_codes)
        end
        n = findfirst(==(given_edge), edge_codes)
        edge_codes = edge_codes[mod1.(n .+ (1:8), 8)]
    end
    return edge_codes[[1, 3, 6, 8]]
end

function get_tile_permutation(edge_codes, i)
    if i > 4
        edge_codes = reverse(edge_codes)
        i -= 4
    end
    edge_codes = edge_codes[mod1.(2 * i .+ (-1:6), 8)]
    return edge_codes[[1, 3, 6, 8]]
end

function reconstruct_image(square, square_edges, tiles, edges)
    image = zeros(Int, 8 .* size(square))
    side = size(square, 1)
    for (i, j) in Tuple.(CartesianIndices((1:side, 1:side)))
        tile = extract_tile(square[i, j], square_edges[i, j, :], tiles, edges)
        image[(1 + 8 * (i - 1)):(8 * i), (1 + 8 * (j - 1)):(8 * j)] .= tile
    end
    return image
end

function extract_tile(tile_number, tile_edges, tiles, edges)
    tile = tiles[tile_number]
    n = findfirst(==(tile_edges[1]), edges[tile_number])
    if n % 2 == 1
        tile = rotl90(tile, (n - 1) ÷ 2)
    else
        tile = reverse(tile, dims = 2)
        tile = rotr90(tile, (n - 2) ÷ 2)
    end
    @assert encode_edge(tile[1, :]) == tile_edges[1]
    return tile[2:(end - 1), 2:(end - 1)]
end

function find_sea_monsters(image)
    sea_monster = [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0
                   1 0 0 0 0 1 1 0 0 0 0 1 1 0 0 0 0 1 1 1
                   0 1 0 0 1 0 0 1 0 0 1 0 0 1 0 0 1 0 0 0]
    N = sum(sea_monster)
    num_sea_monsters = 0
    for j = 1:(size(image, 2) - 19)
        for i = 1:(size(image, 1) - 2)
            if sum(image[i:(i + 2), j:(j + 19)] .* sea_monster) == N
                num_sea_monsters += 1
            end
        end
    end
    return num_sea_monsters
end


text:
"""RMP制御器"""

# """
# RMPいろいろ
# """
# module RMP

using LinearAlgebra
using ForwardDiff  # 自動微分パッケージ

using Parameters

# export pullbacked_rmp
# export  get_natural

# export OriginalRMPAttractor
# export OriginalRMPCollisionAvoidance
# export OriginalJointLimitAvoidance
# export RMPfromGDSAttractor
# export RMPfromGDSCollisionAvoidance

"""pullback演算"""
function pullbacked_rmp(f, M, J, dJ, dx)
    pulled_f = J' * (f .- M * dJ * dx)
    pulled_M = J' * M * J
    return pulled_f, pulled_M
end

"""pullback演算"""
function pullbacked_rmp(f, M, J)
    pulled_f = J' * f
    pulled_M = J' * M * J
    return pulled_f, pulled_M
end

#### オリジナルのRMP ###
"""ソフトマックス関数"""
function soft_max(s::T, α::T) where T
    s + 1/α * log(1 + exp(-2 * α * s))
end

"""ソフト正規化関数"""
function soft_normal(v, alpha)
    return v ./ soft_max(norm(v), alpha)
end

"""空間を一方向に伸ばす計量"""
function metric_stretch(v, alpha)
    xi = soft_normal(v, alpha)
    return xi * xi'
end

"""基本の計量"""
function basic_metric_H(f::Vector{T}, alpha::T, beta::T) where T
    dim = length(f)
    return beta .* metric_stretch(f, alpha) + (1 - beta) .* Matrix{T}(I, dim, dim)
end

"""OriginalRMPの目標到達制御器のパラメータ"""
@with_kw struct OriginalRMPAttractor{T}
    max_speed::T
    gain::T
    ddq_damp_r::T
    sigma_W::T
    sigma_H::T
    metric_damp_r::T
end

"""目標加速度 from OirginalRMP"""
function ddz(p::OriginalRMPAttractor{T}, z::Vector{T}, dz::Vector{T}, z0::Vector{T}) where T
    damp = p.gain / p.max_speed
    a = p.gain .* soft_normal(z0.-z, p.metric_damp_r) .- damp*dz
    return a
end

"""目標計量  from OirginalRMP"""
function inertia_matrix(
    p::OriginalRMPAttractor{T}, z::Vector{T}, dz::Vector{T}, z0::Vector{T}, ddq::Vector{T}
) where T
    dis = norm(z0 .- z)
    weight = exp(-dis ./ p.sigma_W)
    beta = 1.0 - exp(-1/2 * (dis / p.sigma_H)^2)
    return weight .* basic_metric_H(ddq, p.ddq_damp_r, beta)
end

"""canonical form []"""
function get_canonical(p::OriginalRMPAttractor{T}, z, dz, z0) where T
    a = ddz(p, z, dz, z0)
    M = inertia_matrix(p, z, dz, z0, a)
    return a, M
end

"""natural form ()"""
function get_natural(p::OriginalRMPAttractor{T}, z, dz, z0) where T
    a, M = get_canonical(p, z, dz, z0)
    f = M * a
    return f, M
end

@with_kw struct OriginalRMPCollisionAvoidance{T}
    scale_rep::T
    scale_damp::T
    ratio::T
    gain::T
    r::T
end

"""障害物回避加速度  from OirginalRMP"""
function ddz(p::OriginalRMPCollisionAvoidance{T}, z, dz, z0) where T

    x = z .- z0
    d = norm(x)
    ∇d = x ./ d  # 勾配

    # 斥力項
    α_rep = p.gain .* exp(-d / p.scale_rep)
    ddq_rep = α_rep .* ∇d

    # ダンピング項
    P_obs = max(0.0, dot(-dz, ∇d)) * ∇d * ∇d' * dz
    damp_gain = p.gain * p.ratio
    α_damp = damp_gain / (d / p.scale_damp + 1e-7)
    ddq_damp = α_damp .* P_obs

    return ddq_rep .+ ddq_damp
end

"""障害物計量 from OriginalRMP"""
function inertia_matrix(p::OriginalRMPCollisionAvoidance{T}, z, dz, z0, ddq) where T
    dim = length(z)
    d = norm(z .- z0)
    weight = (d / p.r)^2 - 2 * d / p.r + 1
    return weight .* Matrix{T}(I, dim, dim)
end

"""canonical form ()"""
function get_canonical(p::OriginalRMPCollisionAvoidance{T}, z, dz, z0) where T
    a = ddz(p, z, dz, z0)
    M = inertia_matrix(p, z, dz, z0, a)
    return a, M
end

"""natural form []"""
function get_natural(p::OriginalRMPCollisionAvoidance{T}, z, dz, z0) where T
    a, M = get_canonical(p, z, dz, z0)
    f = M * a
    return f, M
end

@with_kw struct OriginalJointLimitAvoidance{T}
    gamma_p::T
    gamma_d::T
    lambda::T
end

# """アフィン変換されたシグモイド写像"""
# sigma_L(q, q_min, q_max) = (q_max - q_min) * (1 / (1 + exp.(-q))) + q_min

"""ジョイント制限に関する対角ヤコビ行列"""
function D_sigma(q::Vector{T}, q_min::Vector{T}, q_max::Vector{T}) where T
    diags = (q_max .- q_min) .* (exp.(-q) ./ (1.0 .+ exp.(-q)).^2)
    return diagm(diags)
end

"""ジョイント制限回避加速度 from OriginalRMP"""
function ddz(
    p::OriginalJointLimitAvoidance{T},
    q::Vector{T}, dq::Vector{T}, q_max::Vector{T}, q_min::Vector{T}
) where T
    z = p.gamma_p .* (-q) .- p.gamma_d .* dq
    a = inv(D_sigma(q, q_min, q_max)) * z
    return a
end

"""ジョイント制限回避計量 from OriginalRMP"""
function inertia_matrix(p::OriginalJointLimitAvoidance{T}, q, dq) where T
    n = length(q)
    return p.lambda .* Matrix{T}(I, n, n)
end

"""canonical form ()"""
function get_canonical(p::OriginalJointLimitAvoidance{T}, q, dq, q_max, q_min) where T
    return ddz(p, q, dq, q_max, q_min), inertia_matrix(p, q, dq)
end

"""natural form []"""
function get_natural(p::OriginalJointLimitAvoidance{T}, q, dq, q_max, q_min, q_neutral) where T
    a, M = get_canonical(p, q, dq, q_max, q_min)
    f = M * a
    return f, M
end

### fromGDS ###

"""fromGDSのアトラクター　パラメータ"""
@with_kw struct RMPfromGDSAttractor{T}
    max_speed::T
    gain::T
    f_alpha::T
    sigma_alpha::T
    sigma_gamma::T
    wu::T
    wl::T
    alpha::T
    epsilon::T
end

"""目標吸引ポテンシャル2"""
potential_2 = soft_max

"""ポテンシャルの勾配"""
function ∇potential_2(x, η)
    x_norm = norm(x)
    return (1-exp(-2*η*x_norm)) / (1+exp(-2*η*x_norm)) .* x ./ x_norm
end

"""?"""
α_or_γ(x, σ) = exp(-(norm(x))^2 / (2*σ^2))

"""重み行列（fromGDSのアトラクターで使用）"""
w(x, sigma_gamma, wu, wl) = (wu-wl) * α_or_γ(x, sigma_gamma) + wl

"""fromGDSのアトラクター慣性行列"""
function inertia_matrix(x, p::RMPfromGDSAttractor{T}, x₀) where T
    dim = length(x)
    z = x .- x₀
    ∇pot = ∇potential_2(z, p.alpha)
    α = α_or_γ(z, p.sigma_alpha)
    return w(z, p.sigma_gamma, p.wu, p.wl) .* ((1-α) .* ∇pot * ∇pot' .+ (α + p.epsilon).*Matrix{T}(I, dim, dim))
end

"""力用"""
function xMx(x, p::RMPfromGDSAttractor{T}, dx, x₀) where T
    dx' * inertia_matrix(x, p, x₀) * dx
end

"""曲率項"""
function ξ(p::RMPfromGDSAttractor{T}, x, dx, x₀) where T
    # 第一項を計算
    dim = length(x)  # タスク空間の次元
    #A = Matrix{T}(undef, dim, dim)
    A = zeros(T, dim, dim)
    for i in 1:dim
        _M(x) = inertia_matrix(x, p, x₀)[:, i]
        _jacobian_M = ForwardDiff.jacobian(_M, x)
        #println(_jacobian_M)
        A[:, i] = _jacobian_M * dx
    end
    A *= dx

    # 第二項を計算
    _xMx(x) = xMx(x, p, dx, x₀)
    B = 1/2 .* ForwardDiff.gradient(_xMx, x)  # 便利

    return A .- B
end

"""fromGDSのアトラクター力"""
function f(p::RMPfromGDSAttractor{T}, x, dx, x₀, M) where T
    z = x .- x₀
    damp = p.gain / p.max_speed
    #return M * (-p.gain .* soft_normal(z, p.f_alpha) .- damp .* dx) .- ξ(p, x, dx, x₀)
    return M * (-p.gain .* ∇potential_2(z, p.alpha) .- damp .* dx) .- ξ(p, x, dx, x₀)
end

"""fromGDSのアトラクタの自然形式RMPを取得"""
function get_natural(p::RMPfromGDSAttractor{T}, x, dx, x₀) where T
    M = inertia_matrix(x, p, x₀)
    return f(p, x, dx, x₀, M), M
end

### 障害物回避 ###
"""fromGDSの障害物回避rmpのパラメータ"""
@with_kw struct RMPfromGDSCollisionAvoidance{T}
    rw::T
    sigma::T
    alpha::T
end

#
# w(s) = s^(-4)

# """重み関数の微分"""
# dwds(s) = -4 * s^(-5)

"""距離に関する重み関数"""
function w2(s, rw=1.0)
    if rw - s > 0.0
        return (rw - s)^2 / s
    else
        return 0.0
    end

end

"""重み関数のs微分"""
function dw2(s, rw)
    if rw - s > 0.0
        return (-2*(rw - s) * s + (rw - s)) / s^2
    else
        return 0.0
    end
end

"""速度に関する重み関数"""
function u2(ds, σ)
    if ds < 0.0
        return 1 - exp(-ds^2 / (2*σ^2))
    else
        return 0.0
    end
end

"""重み関数のds微分"""
function du2(ds, σ)
    if ds < 0.0
        return -exp(-ds^2 / (2*σ^2)) * (-ds / σ^2)
    else
        return 0.0
    end
end

function δ(s, ds, σ)
    u2(ds, σ) + 1/2 * ds * du2(ds, σ)
end

"""曲率項"""
function ξ(s, ds, σ, rw)
    1/2 * u2(ds, σ) * dw2(s, rw) * ds^2
end

"""障害物回避ポテンシャル"""
function Φ1(s, α, rw)
    1/2 * α * w2(s, rw)^2
end

"""障害物回避ポテンシャルの勾配"""
function ∇Φ1(s, α, rw)
    α * w2(s, rw) * dw2(s, rw)
end

"""fromGDSの障害物回避力"""
function f(p::RMPfromGDSCollisionAvoidance{T}, s, ds) where T
    return -∇Φ1(s, p.alpha, p.rw) - ξ(s, ds, p.sigma, p.rw)
end

"""fromGDSの障害物回避慣性行列"""
function inertia_matrix(p::RMPfromGDSCollisionAvoidance{T}, s, ds) where T
    return w2(s, p.rw) * δ(s, ds, p.sigma)
end

"""fromGDSの障害物回避rmpの自然形式を取得"""
function get_natural(p::RMPfromGDSCollisionAvoidance{T}, x, dx, x₀, dx₀=zero(x₀)) where T

    s_vec = x .- x₀
    ds_vec = dx .- dx₀
    s = norm(s_vec)
    ds = (1/s .* dot(s_vec, ds_vec))[1]

    m = inertia_matrix(p, s, ds)
    _f = f(p, s, ds)

    J = -s_vec' ./ s
    J̇ = -s^(-2) .* (ds_vec' .- s_vec' .* ds)

    _f, M = pullbacked_rmp(_f, m, J, J̇, dx)

    return _f, M
end

"""RMPfromGDSのジョイント制限回避のパラメータ"""
@with_kw struct RMPfromGDSJointLimitAvoidance{T}
    gamma_p::T
    gamma_d::T
    lambda::T
    sigma::T
end

function α_upper(dq::T, sigma::T) where T
    1.0 - exp(-max(dq, 0)^2 / (2 * sigma^2))
end

function α_lower(dq::T, sigma::T) where T
    1.0 - exp(-min(dq, 0)^2 / (2 * sigma^2))
end

function _s(q::T, q_u::T, q_l::T) where T
    (q - q_l) / (q_u - q_l)
end

function _dsdq(q::T, q_u::T, q_l::T) where T
    1 / (q_u - q_l)
end

function _d(s::T) where T
    4 * s * (1-s)
end

function _dddq(s::T, dsdq::T) where T
    (4.0 - 8 * s) * dsdq
end

function _b(q::T, dq::T, q_u::T, q_l::T, sigma::T) where T
    s = _s(q, q_u, q_l)
    d = _d(s)
    α_u = α_upper(dq, sigma)
    α_l = α_lower(dq, sigma)
    return s*(α_u * d + (1-α_u)) + (1-s)*(α_l * d + (1-α_l))
end

function _dbdq(q::T, dq::T, q_u::T, q_l::T, sigma::T) where T
    s = _s(q, q_u, q_l)
    d = _d(s)
    α_u = α_upper(dq, sigma)
    α_l = α_lower(dq, sigma)
    dsdq = _dsdq(q, q_u, q_l)
    dddq = _dddq(s, dsdq)
    return (dsdq*(α_u * d + (1-α_u)) + s * dddq) + -dsdq*(α_l * d + (1-α_l)) + (1-s) * dddq
end

"""慣性行列の対角要素"""
function _a(q::T, dq::T, q_u::T, q_l::T, sigma::T) where T
    b = _b(q, dq, q_u, q_l, sigma)
    return b^(-2)
end

function _dadq(q::T, dq::T, q_u::T, q_l::T, sigma::T) where T
    -2 * (_b(q, dq, q_u, q_l, sigma))^(-3) * _dbdq(q, dq, q_u, q_l, sigma)
end

"""fromGDSのジョイント制限回避慣性行列"""
function inertia_matrix(
    p::RMPfromGDSJointLimitAvoidance{T}, q::Vector{T}, dq::Vector{T},
    q_max::Vector{T}, q_min::Vector{T}
    ) where T
    dim = length(q)
    A = zeros(T, dim, dim)

    for i in 1:dim
        A[i, i] = _a(q[i], dq[i], q_max[i], q_min[i], p.sigma)
    end
    #println(A)
    return p.lambda * A
end

"""ジョイント制限回避の曲率項"""
function ξ(q::Vector{T}, dq::Vector{T}, q_max::Vector{T}, q_min::Vector{T}, sigma::T) where T

    #z = Vector{T}(undef, 7)
    dim = length(q)
    z = zero(q)

    for i in 1:dim
        z[i] = 1/2 * _dadq(q[i], dq[i], q_max[i], q_min[i], sigma) * dq[i]^2
    end
    return z
end

"""fromGDSのジョイント制限回避慣性力"""
function f(
    p::RMPfromGDSJointLimitAvoidance{T}, q::Vector{T}, dq::Vector{T},
    q_neutral::Vector{T}, A::Matrix{T},
    ) where T
    ξ_A = ξ(q, dq, q_max, q_min, p.sigma)
    #ξ_A = zeros(T, 7)
    return A * (p.gamma_p * (q_neutral .- q) - p.gamma_d * dq) - ξ_A
end

"""fromGDSのジョイント制限回避RMP"""
function get_natural(
    p::RMPfromGDSJointLimitAvoidance{T}, q::Vector{T}, dq::Vector{T},
    q_max::Vector{T}, q_min::Vector{T}, q_neutral::Vector{T}
    ) where T
    A = inertia_matrix(p, q, dq, q_max, q_min)
    _f = f(p, q, dq, q_neutral, A)
    return _f, A
end

#### sice2021 ####
"""インピーダンス制御RMP（自作）

M_d : 所望のインピーダンス慣性行列
D_d : 所望のインピーダンス減衰行列
P_d : 所望のインピーダンス剛性行列
D_e : 環境減衰行列
P_e : 環境剛性行列
eta_d : 一般化位置に関するスケール係数
eta_e : 外力の平衡位置に関するスケール係数
"""
@with_kw struct RMPfromGDSImpedance{T}
    M_d::Matrix{T}
    D_d::Matrix{T}
    P_d::Matrix{T}
    D_e::Matrix{T}
    P_e::Matrix{T}
    a::T  # シグモイドのパラメータ
    eta_d::T
    eta_e::T
    f_alpha::T
    sigma_alpha::T
    sigma_gamma::T
    wu::T
    wl::T
    alpha::T
    epsilon::T
end

"""シグモイド関数"""
sigmoid(x, a) = 1/(1+exp(-a*x))

function s_alpha(x, η)
    x_norm = norm(x)
    (1-exp(-2*η*x_norm)) / (1+exp(-2*η*x_norm))
end

hat(x) = x / norm(x)

# # """インピーダンスポテンシャル

# # ・計算には使いません
# # """
# # function potential_I(e_d, eta_d, e_e, eta_e, a)
# #     alpha = sigmoid(norm(e_d), a)
# #     return (1-alpha)*soft_normal(e_d, eta_d) + alpha*soft_normal(e_e, eta_e)
# # end

"""インピーダンスポテンシャルの勾配"""
function ∇potential_I(p::RMPfromGDSImpedance{T}, y, yd, ye, alpha) where T

    # 古い方

    e_d = y - yd
    e_e = y - ye

    inv_M_d = inv(p.M_d)
    P_tilde_e = inv_M_d * p.P_e
    P_tilde_d = inv_M_d * p.P_d

    t1 = -(1-alpha)*alpha .* hat(e_d) .* soft_max(norm(e_d), p.eta_d)
    t2 = (1-alpha) * s_alpha(e_d, p.eta_d) .* hat(e_d)
    t3 = (1-alpha)*alpha .* hat(e_e) .* soft_max(norm(e_e), p.eta_e)
    t4 = alpha * s_alpha(e_e, p.eta_e) .* hat(e_e)

    #t1 = t3 = zero(t2)

    t1 .+
    P_tilde_d * t2 .+
    t3 .+
    P_tilde_e * t4

    # # 新しい方
    # inv_M_d = inv(p.M_d)
    # P_tilde_e = inv_M_d * p.P_e
    # P_tilde_d = inv_M_d * p.P_d

    # e_d = y - yd
    # e_e = y - ye

    # x2 = norm(yd - ye)

    # y = y - ye

    # a = 0.5
    # b = x2^2 / 4

    # function df(x)
    #     return a*((x - x2/2)^2-b)^2
    # end

    # println(df(norm(y)))
    # return P_tilde_e .* df(norm(y)) * hat(y)

end

"""インピーダンスRMPの慣性行列"""
function inertia_matrix(p::RMPfromGDSImpedance{T}, y, yd, ye, dy, alpha) where T
    e_d = y - yd
    e_e = y - ye
    dim = length(y)
    ∇pot = ∇potential_I(p, y, yd, ye, alpha)
    α = α_or_γ(e_d, p.sigma_alpha)
    return w(e_d, p.sigma_gamma, p.wu, p.wl) .* ((1-α) .* ∇pot * ∇pot' .+ (α + p.epsilon).*Matrix{T}(I, dim, dim))
end

"""力用"""
function xMx(p::RMPfromGDSImpedance{T}, y, yd, ye, dy, alpha) where T
    dy' * inertia_matrix(p, y, yd, ye, dy, alpha) * dy
end

"""曲率項"""
function ξ(p::RMPfromGDSImpedance{T}, y, yd, ye, dy, alpha) where T
    # 第一項を計算
    #A = Matrix{T}(undef, 3, 3)
    e_d = y - yd
    e_e = y - ye

    dim = length(y)
    A = zeros(T, dim, dim)
    for i in 1:length(e_d)
        println("kyokuritu i = ", i)
        _M(y) = inertia_matrix(p, y, yd, ye, dy, alpha)[:, i]
        _jacobian_M = ForwardDiff.jacobian(_M, y)
        #println(_jacobian_M)
        A[:, i] = _jacobian_M * dy
    end
    A *= dy

    # 第二項を計算
    _xMx(y) = xMx(p, y, yd, ye, dy, alpha)
    B = 1/2 .* ForwardDiff.gradient(_xMx, y)  # 便利

    return A .- B
end

"""fromGDSのアトラクター力"""
function f(p::RMPfromGDSImpedance{T}, y, yd, ye, dy, M, alpha) where T
    e_d = y - yd
    e_e = y - ye
    inv_M_d = inv(p.M_d)
    D_tilde_d = inv_M_d * p.D_d
    D_tilde_e = inv_M_d * p.D_e

    # P_tilde_d =  inv_M_d * p.P_d
    # P_tilde_e =  inv_M_d * p.P_e

    #return M * (-p.gain .* soft_normal(z, p.f_alpha) .- damp .* dx) .- ξ(p, x, dx, x₀)

    z = M * (-∇potential_I(p, y, yd, ye, alpha) .- (D_tilde_d .+ alpha .* D_tilde_e) * dy)*1.5 #.-
    #ξ(p, y, yd, ye, dy)

    return z
end

""""""
function get_natural(p::RMPfromGDSImpedance{T}, y, yd, ye, dy, center, r) where T
    #println("呼ばれた！")
    s = r - norm(y-center)
    alpha = sigmoid(s, p.a)
    #println(alpha)
    M = inertia_matrix(p, y, yd, ye, dy, alpha)
    return f(p, y, yd, ye, dy, M, alpha), M
end

#end


text:
# import TrajectoryOptimization.Controllers: RBState

using TrajectoryOptimization: get_trajectory, get_model, AbstractTrajectory

visualize!(vis, solver) = visualize!(vis, get_model(solver), get_trajectory(solver))

"""
    visualize!(vis, model, x)

Visualize a single configuration.
Defaults to using the `position` and `orientation` methods on `model`.

Overload this method to specify the visualization for a specific model.
"""
function visualize!(vis, model::AbstractModel, x::StaticVector, addrobot::Bool=true)
    visualize!(vis, RBState(model, x), addrobot)
end

""" Set state to RBState """
function visualize!(vis, x::RBState{<:Real}, addrobot::Bool=true)
    p = position(x)
    q = orientation(x)
    robot = addrobot ? vis["robot"] : vis
    settransform!(robot, compose(Translation(p), LinearMap(UnitQuaternion(q))))
end

"""
Animate the trajectory of a rigid body in MeshCat
    Assumes the robot geometry is already loaded into `vis["robot"]`
"""
visualize!(vis, model::AbstractModel, Z::AbstractTrajectory) =
    visualize!(vis, model, Z[end].t, states(Z))

"""
    visualize!(vis, probs...; kwargs...)
    visualize!(vis, model, Zs...; kwargs...)
    visualize!(vis, model, tf, Xs...; kwargs...)

Visualize many different trajectories of the same model.

# Arguments
* `probs`: any struct that supports `TO.get_model` and `TO.get_trajectory`
* `Zs`: any `AbstractTrajectory`
* `Xs`: a vector of state vectors
* `colors` (optional): Can be either `nothing` (default), a `Colorant`, or a vector of both, of equal length as the number of trajectories. Changes the color of the models in the order provided.
"""
function visualize!(vis, probs...; kwargs...)
    model = get_model(probs[1])
    visualize!(vis, model, get_trajectory.(probs)...; kwargs...)
end
function visualize!(vis, model::AbstractModel, Zs::Vararg{<:AbstractTrajectory}; kwargs...)
    tf = Zs[1][end].t - Zs[1][1].t
    visualize!(vis, model, tf, states.(Zs)...; kwargs...)
end
function visualize!(vis, model::AbstractModel, tf::Real, Xs...; colors=nothing)
    N = length(Xs[1])
    fps = Int(floor((N-1)/tf))
    anim = MeshCat.Animation(fps)
    num_traj = length(Xs)
    if isnothing(colors)
        colors = fill(nothing, length(Xs))
    elseif colors isa Colorant
        colors = fill(colors, length(Xs)-1)
        colors = [nothing; colors]
    elseif colors isa Vector{<:Union{Nothing, Colorant}}
        @assert length(colors) == length(Xs) "Number of colors must match the number of trajectories"
        set_mesh!(vis, model, color=colors[1])
    end
    for i = 2:num_traj
        TrajOptPlots.set_mesh!(vis["robot_copies/robot$i"], model, color=colors[i])
    end
    for k = 1:N
        atframe(anim, k) do
            for i = 1:num_traj
                if i == 1
                    robot = vis
                else
                    robot = vis["robot_copies"]["robot$i"]
                end
                visualize!(robot, model, Xs[i][k])
            end
        end
    end
    setanimation!(vis, anim)
end

#--- Environment plotting
function plot_cylinder(vis,c1,c2,radius,mat,name="")
    geom = Cylinder(Point3f0(c1),Point3f0(c2),convert(Float32,radius))
    setobject!(vis["cyl"][name],geom,MeshPhongMaterial(color=RGBA(1, 0, 0, 1.0)))
end

function add_cylinders!(vis,x,y,r; height=1.5, idx=0, robot_radius=0.0, color=colorant"red")
    for i in eachindex(x)
        obj = Cylinder(Point3f0(x[i],y[i],0.0), Point3f0(x[i],y[i],height),
            Float32(r[i] - robot_radius))
        mat = MeshPhongMaterial(color=color)
        j = i + idx
        setobject!(vis["obs"]["cyl$j"], obj, mat)
    end
end

function add_cylinders!(vis,solver; kwargs...)
    conSet = TrajectoryOptimization.get_constraints(solver)
    idx = 0
    for con in conSet.constraints
        if con isa CircleConstraint
            add_cylinders!(vis, con.x, con.y, con.radius, idx=idx; kwargs...)
            idx += length(con)
        end
    end
end

function add_point!(vis, x::AbstractVector;
        radius=0.1, color=colorant"green", name="point")
    setobject!(vis[name], HyperSphere(Point3f0(x[1], x[2], x[3]), Float32(radius)),
        MeshPhongMaterial(color=color))
end

#--- Waypoints
const TO = TrajectoryOptimization
waypoints!(vis, solver; kwargs...) =
    waypoints!(vis, TO.get_model(solver), TO.get_trajectory(solver); kwargs...)
function waypoints!(vis, model::AbstractModel, Z::AbstractTrajectory; length=0, inds=Int[],
        color=nothing, color_end=nothing)
    N = size(Z,1)
    if length > 0 && isempty(inds)
        inds = Int.(round.(range(1,N,length=length)))
    elseif !isempty(inds) && length == 0
        length = size(inds,1)
    else
        throw(ArgumentError("Have to pass either length or inds, but not both"))
    end
    if !isnothing(color)
        if isnothing(color_end)
            color_end = color
        end
        colors = range(color, color_end, length=size(inds,1))
        # set_mesh!(vis, model, color=colors[end])
    end

    delete!(vis["waypoints"])
    for (j,i) in enumerate(inds)
        if isnothing(color)
            set_mesh!(vis["waypoints"]["point$i"], model)
        else
            set_mesh!(vis["waypoints"]["point$i"], model, color=colors[j])
        end

        # settransform!(vis["waypoints"]["point$i"]["geom"], compose(Translation(0,0,0.07),LinearMap( RotY(pi/2)*RotZ(-pi/2) )))
        visualize!(vis["waypoints"]["point$i"], model, state(Z[i]))
    end
    visualize!(vis, model, state(Z[end]))
end

clear_waypoints!(vis) = delete!(vis["waypoints"])
clear_copies!(vis) = delete!(vis["robot_copies"])
clear_traj!(vis) = delete!(vis["traj"])

"""
    traj3!(vis, X; [inds, kwargs...])
    traj3!(vis, prob; [inds, kwargs...])

Plot the 3D trajectory in the visualizer. The x,y,z coordinates are given by
`inds[1]`, `inds[2]`, and `inds[3]`, respectively (defaults to 1:3).
Use `kwargs` to modify the properties of the trajectory.
All keyword arguments are passed directly to `MeshCat.LineBasicMaterial`.

The input needs to be the state trajectory, i.e. a vector of vectors. Alternatively,
you can pass in any type for which `TrajectoryOptimization.states` is defined.

# Example
    traj3!(vis, states(solver), inds=1:3, linewidth=3.0, color=colorant"red")
    traj3!(vis, solver, inds=1:3, linewidth=3.0, color=colorant"red")
"""
function traj3!(vis, X::AbstractVector{<:AbstractVector}; inds=SA[1,2,3], kwargs...)
    pts = [Point{3,Float32}(x[inds]) for x in X]
    setobject!(vis[:traj], MeshCat.Line(pts, LineBasicMaterial(;kwargs...)))
end

@inline traj3!(vis, prob; kwargs...) = traj3!(vis, states(prob); kwargs...)


text:
@doc raw"""
    subgradient_method(M, F, ∂F, x)

perform a subgradient method ``x_{k+1} = \mathrm{retr}(x_k, s_k∂F(x_k))``,

where ``\mathrm{retr}`` is a retraction, ``s_k`` can be specified as a function but is
usually set to a constant value. Though the subgradient might be set valued,
the argument `∂F` should always return _one_ element from the subgradient, but
not necessarily deterministic.

# Input
* `M` – a manifold ``\mathcal M``
* `F` – a cost function ``F:\mathcal M→ℝ`` to minimize
* `∂F`– the (sub)gradient ``\partial F: \mathcal M→ T\mathcal M`` of F
  restricted to always only returning one value/element from the subgradient.
  This function can be passed as an allocation function `(M, y) -> X` or
  a mutating function `(M, X, y) -> X`, see `evaluation`.
* `x` – an initial value ``x ∈ \mathcal M``

# Optional
* `evaluation` – ([`AllocatingEvaluation`](@ref)) specify whether the subgradient works by
   allocation (default) form `∂F(M, y)` or [`MutatingEvaluation`](@ref) in place, i.e. is
   of the form `∂F!(M, X, x)`.
* `stepsize` – ([`ConstantStepsize`](@ref)`(1.)`) specify a [`Stepsize`](@ref)
* `retraction` – (`default_retraction_method(M)`) a `retraction(M,x,ξ)` to use.
* `stopping_criterion` – ([`StopAfterIteration`](@ref)`(5000)`)
  a functor, see[`StoppingCriterion`](@ref), indicating when to stop.
* `return_options` – (`false`) – if activated, the extended result, i.e. the
  complete [`Options`](@ref) re returned. This can be used to access recorded values.
  If set to false (default) just the optimal value `x_opt` if returned
...
and the ones that are passed to [`decorate_options`](@ref) for decorators.

# Output
* `x_opt` – the resulting (approximately critical) point of the subgradient method
OR
* `options` - the options returned by the solver (see `return_options`)
"""
function subgradient_method(
    M::AbstractManifold, F::TF, ∂F::TdF, x; kwargs...
) where {TF,TdF}
    x_res = allocate(x)
    copyto!(M, x_res, x)
    return subgradient_method!(M, F, ∂F, x_res; kwargs...)
end
@doc raw"""
    subgradient_method!(M, F, ∂F, x)

perform a subgradient method ``x_{k+1} = \mathrm{retr}(x_k, s_k∂F(x_k))`` in place of `x`

# Input
* `M` – a manifold ``\mathcal M``
* `F` – a cost function ``F:\mathcal M→ℝ`` to minimize
* `∂F`- the (sub)gradient ``\partial F:\mathcal M→ T\mathcal M`` of F
  restricted to always only returning one value/element from the subgradient.
  This function can be passed as an allocation function `(M, y) -> X` or
  a mutating function `(M, X, y) -> X`, see `evaluation`.
* `x` – an initial value ``x ∈ \mathcal M``

for more details and all optional parameters, see [`subgradient_method`](@ref).
"""
function subgradient_method!(
    M::AbstractManifold,
    F::TF,
    ∂F!!::TdF,
    x;
    retraction::TRetr=default_retraction_method(M),
    stepsize::Stepsize=ConstantStepsize(1.0),
    stopping_criterion::StoppingCriterion=StopAfterIteration(5000),
    return_options=false,
    evaluation::AbstractEvaluationType=AllocatingEvaluation(),
    kwargs..., #especially may contain debug
) where {TF,TdF,TRetr}
    p = SubGradientProblem(M, F, ∂F!!; evaluation=evaluation)
    o = SubGradientMethodOptions(M, x, stopping_criterion, stepsize, retraction)
    o = decorate_options(o; kwargs...)
    resultO = solve(p, o)
    if return_options
        return resultO
    else
        return get_solver_result(resultO)
    end
end
function initialize_solver!(p::SubGradientProblem, o::SubGradientMethodOptions)
    o.x_optimal = o.x
    o.∂ = zero_vector(p.M, o.x)
    return o
end
function step_solver!(p::SubGradientProblem, o::SubGradientMethodOptions, iter)
    get_subgradient!(p, o.∂, o.x)
    s = get_stepsize(p, o, iter)
    retract!(p.M, o.x, o.x, -s * o.∂, o.retraction_method)
    (get_cost(p, o.x) < get_cost(p, o.x_optimal)) && (o.x_optimal = o.x)
    return o
end
get_solver_result(o::SubGradientMethodOptions) = o.x_optimal


text:
"""
    Kucoin.get_symbols_list([; market]) -> JSON3.Array{JSON3.Object}

Get a list of available currency pairs for trading. See Kucoin official
[docs](https://docs.kucoin.com/#get-symbols-list) for more details.

# Keywords
- `market::String`: [Optional] a name of the trading market

# Returns
- `JSON3.Array{JSON3.Object}`: a list of symbols for all or a specified markets
```json
[
  ...
  {
    "symbol": "XLM-USDT",
    "name": "XLM-USDT",
    "baseCurrency": "XLM",
    "quoteCurrency": "USDT",
    "feeCurrency": "USDT",
    "market": "USDS",
    "baseMinSize": "0.1",
    "quoteMinSize": "0.01",
    "baseMaxSize": "10000000000",
    "quoteMaxSize": "99999999",
    "baseIncrement": "0.0001",
    "quoteIncrement": "0.000001",
    "priceIncrement": "0.000001",
    "priceLimitRate": "0.1",
    "isMarginEnabled": true,
    "enableTrading": true
  },
  ...
]
```
"""
function get_symbols_list(;
    market::Union{String,Nothing}=nothing
)::JSON3.Array{JSON3.Object}
    return _handle(_kucoin_request(_HTTP_METHOD_GET, "/api/v1/symbols"; market=market))
end

"""
    Kucoin.get_all_tickers() -> JSON3.Array{JSON3.Object}

Get the market information for all the trading pairs in the market. See Kucoin official
[docs](https://docs.kucoin.com/#get-all-tickers) for more details.

# Returns
- `JSON3.Array{JSON3.Object}`
```json
[
  {
    "symbol": "BTC-USDT",  // symbol
    "symbolName":"BTC-USDT",  // Name of trading pairs, it would change after renaming
    "buy": "11328.9",  // bestAsk
    "sell": "11329",  // bestBid
    "changeRate": "-0.0055",  // 24h change rate
    "changePrice": "-63.6",  // 24h change price
    "high": "11610",  // 24h highest price
    "low": "11200",  // 24h lowest price
    "vol": "2282.70993217", // 24h volume，the aggregated trading volume in BTC
    "volValue": "25984946.157790431",  // 24h total, the trading volume in quote currency of last 24 hours
    "last": "11328.9",  // last price
    "averagePrice": "11360.66065903",  // 24h average transaction price yesterday
    "takerFeeRate": "0.001",  // Basic Taker Fee
    "makerFeeRate": "0.001",  // Basic Maker Fee
    "takerCoefficient": "1",  // Taker Fee Coefficient
    "makerCoefficient": "1"  // Maker Fee Coefficient
  }
]
```
"""
function get_all_tickers()::JSON3.Array{JSON3.Object}
    return _handle(_kucoin_request(_HTTP_METHOD_GET, "/api/v1/market/allTickers"))["ticker"]
end


text:
"""
```
threeletter(residue::String)
```

Function to return a three-letter residue code from the one-letter code or residue name. The function is case-insensitive.

### Examples

```julia-repl
julia> threeletter("A")
"ALA"

julia> threeletter("Aspartic acid")
"ASP"

```

"""
function threeletter(residue::String)
  l = length(residue)
  code = uppercase(residue)
  if l != 3
    if l == 1
      ires = findfirst(r->r.one_letter_code == code, natural_aminoacids)
    else
      ires = findfirst(r->uppercase(r.name) == code, natural_aminoacids)
    end
    code = (ires == nothing ? nothing : natural_aminoacids[ires].three_letter_code)
  end
  return code
end


text:
function FieldContinueOnEvent()
    return FieldContinueOnEvent(())
end

function event_occurred(obj::FieldContinueOnEvent, arg0::FieldSpacecraftState, arg1::FieldEventDetector, arg2::jboolean)
    return jcall(obj, "eventOccurred", Action, (FieldSpacecraftState, FieldEventDetector, jboolean), arg0, arg1, arg2)
end

function init(obj::FieldEventHandler, arg0::FieldSpacecraftState, arg1::FieldAbsoluteDate)
    return jcall(obj, "init", void, (FieldSpacecraftState, FieldAbsoluteDate), arg0, arg1)
end

function reset_state(obj::FieldEventHandler, arg0::FieldEventDetector, arg1::FieldSpacecraftState)
    return jcall(obj, "resetState", FieldSpacecraftState, (FieldEventDetector, FieldSpacecraftState), arg0, arg1)
end


text:
"""
The Window would be a fundamental widget.
Decorator is aimed to add different kinds of responsibilities on Window.
The same decorator could be added several times to the same window.
"""

abstract type LCD end

struct Window <: LCD
    size::Float64
end

draw(::Window) = println("Draw basic window view.")

abstract type Decorator <: LCD end

mutable struct Border <: Decorator
    component::LCD
    size::Float64
end

function draw(b::Border)
    println("Draw a $(b.size) width border")
    draw(b.component)
    # do something
end

mutable struct VerticalScrollBar <: Decorator
    component::LCD
end

function draw(vsb::VerticalScrollBar)
    # do something
    draw(vsb.component)
    println("Draw vertical scroll bar.")
end

scroll(vsb::VerticalScrollBar, direction::Symbol) =
    println("Vertical scroll bar scrolls $(direction)")

mutable struct HorizontalScrollBar <: Decorator
    component::LCD
end

function draw(hsb::HorizontalScrollBar)
    # do something
    draw(hsb.component)
    println("Draw horizontal scroll bar.")
end

scroll(hsb::HorizontalScrollBar, direction::Symbol) =
    println("Horizontal scroll bar scrolls $(direction)")


text:
#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl
=#
"""
$(TYPEDEF)

Nonnegative cone of dimension `dim`.

    $(FUNCTIONNAME){T}(dim::Int)
"""
mutable struct Nonnegative{T <: Real} <: Cone{T}
    dim::Int

    point::Vector{T}
    dual_point::Vector{T}
    grad::Vector{T}
    dder3::Vector{T}
    vec1::Vector{T}
    vec2::Vector{T}
    feas_updated::Bool
    grad_updated::Bool
    hess_updated::Bool
    inv_hess_updated::Bool
    is_feas::Bool
    hess::Diagonal{T, Vector{T}}
    inv_hess::Diagonal{T, Vector{T}}

    function Nonnegative{T}(dim::Int) where {T <: Real}
        @assert dim >= 1
        cone = new{T}()
        cone.dim = dim
        return cone
    end
end

use_dual_barrier(::Nonnegative) = false

reset_data(cone::Nonnegative) = (cone.feas_updated = cone.grad_updated =
    cone.hess_updated = cone.inv_hess_updated = false)

use_sqrt_hess_oracles(::Int, cone::Nonnegative) = true

get_nu(cone::Nonnegative) = cone.dim

set_initial_point!(arr::AbstractVector, cone::Nonnegative) = (arr .= 1)

function update_feas(cone::Nonnegative{T}) where T
    @assert !cone.feas_updated
    cone.is_feas = all(>(eps(T)), cone.point)
    cone.feas_updated = true
    return cone.is_feas
end

is_dual_feas(cone::Nonnegative{T}) where T = all(>(eps(T)), cone.dual_point)

function update_grad(cone::Nonnegative)
    @assert cone.is_feas
    @. cone.grad = -inv(cone.point)
    cone.grad_updated = true
    return cone.grad
end

function update_hess(cone::Nonnegative{T}) where T
    cone.grad_updated || update_grad(cone)
    if !isdefined(cone, :hess)
        cone.hess = Diagonal(zeros(T, cone.dim))
    end

    @. cone.hess.diag = abs2(cone.grad)
    cone.hess_updated = true
    return cone.hess
end

function update_inv_hess(cone::Nonnegative{T}) where T
    @assert cone.is_feas
    if !isdefined(cone, :inv_hess)
        cone.inv_hess = Diagonal(zeros(T, cone.dim))
    end

    @. cone.inv_hess.diag = abs2(cone.point)
    cone.inv_hess_updated = true
    return cone.inv_hess
end

function hess_prod!(
    prod::AbstractVecOrMat,
    arr::AbstractVecOrMat,
    cone::Nonnegative,
    )
    @assert cone.is_feas
    @. prod = arr / cone.point / cone.point
    return prod
end

function inv_hess_prod!(
    prod::AbstractVecOrMat,
    arr::AbstractVecOrMat,
    cone::Nonnegative,
    )
    @assert cone.is_feas
    @. prod = arr * cone.point * cone.point
    return prod
end

function sqrt_hess_prod!(
    prod::AbstractVecOrMat,
    arr::AbstractVecOrMat,
    cone::Nonnegative,
    )
    @assert cone.is_feas
    @. prod = arr / cone.point
    return prod
end

function inv_sqrt_hess_prod!(
    prod::AbstractVecOrMat,
    arr::AbstractVecOrMat,
    cone::Nonnegative,
    )
    @assert cone.is_feas
    @. prod = arr * cone.point
    return prod
end

function dder3(cone::Nonnegative, dir::AbstractVector)
    @. cone.dder3 = abs2(dir / cone.point) / cone.point
    return cone.dder3
end

hess_nz_count(cone::Nonnegative) = cone.dim
hess_nz_count_tril(cone::Nonnegative) = cone.dim
inv_hess_nz_count(cone::Nonnegative) = cone.dim
inv_hess_nz_count_tril(cone::Nonnegative) = cone.dim
hess_nz_idxs_col(cone::Nonnegative, j::Int) = [j]
hess_nz_idxs_col_tril(cone::Nonnegative, j::Int) = [j]
inv_hess_nz_idxs_col(cone::Nonnegative, j::Int) = [j]
inv_hess_nz_idxs_col_tril(cone::Nonnegative, j::Int) = [j]

# nonnegative is not primitive, so sum and max proximity measures differ
function get_proximity(
    cone::Nonnegative{T},
    rtmu::T,
    use_sum_prox::Bool, # use sum proximity
    negtol::T = sqrt(eps(T)),
    ) where {T <: Real}
    proxs = (abs(si * zi / rtmu - 1) for (si, zi) in
        zip(cone.point, cone.dual_point))
    if use_sum_prox
        return sum(proxs)
    else
        return maximum(proxs)
    end
end


text:
X, y = @load_crabs

@testset "PCA" begin
    X_array = matrix(X)
    pratio = 0.9999
    # MultivariateStats PCA
    pca_ms = MultivariateStats.fit(
        MultivariateStats.PCA,
        permutedims(X_array),
        pratio=pratio
    )
    # MLJ PCA
    pca_mlj = PCA(pratio=pratio)
    test_composition_model(pca_ms, pca_mlj, X, X_array)
end

@testset "KernelPCA" begin
    X_array = matrix(X)
    # MultivariateStats KernelPCA
    kpca_ms = MultivariateStats.fit(
        MultivariateStats.KernelPCA, permutedims(X_array),
        inverse=true
    )
    # MLJ KernelPCA
    kpca_mlj = KernelPCA(inverse=true)
    test_composition_model(kpca_ms, kpca_mlj, X, X_array)
end

@testset "ICA" begin
    X_array = matrix(X)
    k = 5
    tolerance = 5.0
    # MultivariateStats ICA
    rng = StableRNG(1234) # winit gets randomly initialised
    #Random.seed!(1234) # winit gets randomly initialised
    ica_ms = MultivariateStats.fit(
        MultivariateStats.ICA,
        permutedims(X_array),
        k;
        tol=tolerance,
        winit = randn(rng, eltype(X_array), size(X_array, 2), k)
    )
    # MLJ ICA
    rng = StableRNG(1234) # winit gets randomly initialised
    #Random.seed!(1234) # winit gets randomly initialised
    ica_mlj = ICA(
        k=k,
        tol=tolerance,
        winit=randn(rng, eltype(X_array), size(X_array, 2), k))
    test_composition_model(ica_ms, ica_mlj, X, X_array, test_inverse=false)
end

@testset "PPCA" begin
    X_array = matrix(X)
    tolerance = 5.0
    # MultivariateStats PPCA
    ppca_ms = MultivariateStats.fit(
        MultivariateStats.PPCA,
        permutedims(X_array);
        tol=tolerance
    )
    # MLJ PPCA
    ppca_mlj = PPCA(;tol=tolerance)
    test_composition_model(ppca_ms, ppca_mlj, X, X_array)
end

@testset "FactorAnalysis" begin
    X_array = matrix(X)
    tolerance = 5.0
    eta = 0.01
    # MultivariateStats FactorAnalysis
    factoranalysis_ms = MultivariateStats.fit(
        MultivariateStats.FactorAnalysis,
        permutedims(X_array);
        tol=tolerance,
        η=eta
    )
    factoranalysis_mlj = FactorAnalysis(;tol=tolerance, eta=eta)
    test_composition_model(factoranalysis_ms, factoranalysis_mlj, X, X_array)
end


text:
function calcsortkey(a, key_series, key_values)
    vector_vector_indexes = calcsortindexes(key_series, key_values)
    b = similar(a)
    j = 1
    for vector_indexes in vector_vector_indexes
        for i in vector_indexes
            b[j] = a[i]
            j+=1
        end
    end
    return b
end

function calcsortindexes(key_series, key_values)
    vector_vector_indexes=[Int[] for _ in key_values]

    for i in eachindex(key_series)
        k=indexvalue(key_series[i], key_values)
        push!(vector_vector_indexes[k], i)
    end

    return vector_vector_indexes
end

function indexvalue(v, values::Vector)
    findfirst(x->x == v, values)
end

function calcsort_numzeros(a)
    b = [@view a[:,j] for j in 1:size(a,2)]
    k_s = [length(findall(b[j] .== 0)) for j in 1:length(b)]
    b = calcsortkey(b, k_s, collect(0:size(a,1)))
    return  hcat(b...)
end


text:
"""
    general_tree_test(com::CS.CoM)

Performs some general tests which are nice to test if a new constraint works as intended and implements the necessary methods.
Should be called only if optimize was called already as this works on the final tree structure.
It fails if no backtracking was used as this method should only be called when at least some backtracking steps are needed.
Additionally `keep_logs` needs to be activated to make it work.
What it does:
    - Tests `checkout_from_to!` which basically means jumping from one node in the tree to another.
        - Checks there whether this produces the same search space as before at that stage.
"""
function general_tree_test(com::CS.CoM)
    if length(com.logs) == 0
        @error "Make sure that keep_logs is set to true."
    end
    @assert length(com.logs) > 10
    c_backtrack_idx = com.c_backtrack_idx
    Random.seed!(241)
    n_backtracks = length(com.logs)
    path = Int[c_backtrack_idx]
    path_type = Symbol[]
    all_correct = true
    n_children_tests = 0

    for t in 1:100
        status = :Open
        next_idx = 0
        while status == :Open
            # don't allow root node as that might cause chaos :D
            next_idx = rand(2:n_backtracks)
            status = com.logs[next_idx].status
        end

        CS.checkout_from_to!(com, c_backtrack_idx, next_idx)

        if com.backtrack_vec[next_idx].parent_idx != 0
            c_backtrack_idx = com.backtrack_vec[next_idx].parent_idx

            c_search_space = com.search_space
            # this is a dict =>
            expected_search_space = com.logs[c_backtrack_idx].var_states
            correct = true
            for var in c_search_space
                if sort(CS.values(var)) != sort(expected_search_space[var.idx])
                    println("var.idx: $(var.idx)")
                    println("vals: $(CS.values(var))")
                    println("expected: $(expected_search_space[var.idx])")
                    correct = false
                    all_correct = false
                end
            end
            if !correct
                @error "There were errors on the jump path: $path before pruning the last step"
            end
        end

        c_backtrack_idx = next_idx

        # if it has children
        if length(com.logs[c_backtrack_idx].children) > 0 && n_children_tests < 5
            com.c_backtrack_idx = c_backtrack_idx
            vidx = com.logs[c_backtrack_idx].vidx
            n_children_tests += 1
            # test that pruning produces the same output as before

            for var in com.search_space
                var.changes[c_backtrack_idx] = Vector{Tuple{Symbol,Int,Int,Int}}()
            end
            @assert CS.remove_above!(
                com,
                com.search_space[vidx],
                com.logs[c_backtrack_idx].ub,
            )
            @assert CS.remove_below!(
                com,
                com.search_space[vidx],
                com.logs[c_backtrack_idx].lb,
            )

            @assert CS.prune!(com)
            CS.call_finished_pruning!(com)
            push!(path_type, :prune)
        else
            # prune the last step based on saved information instead
            CS.restore_prune!(com, c_backtrack_idx)
            push!(path_type, :restore_prune)
        end

        push!(path, c_backtrack_idx)
        c_search_space = com.search_space
        # this is a dict =>
        expected_search_space = com.logs[c_backtrack_idx].var_states
        correct = true
        for var in c_search_space
            if sort(CS.values(var)) != sort(expected_search_space[var.idx])
                println("var.idx: $(var.idx)")
                println("vals: $(CS.values(var))")
                println("expected: $(expected_search_space[var.idx])")
                correct = false
                all_correct = false
            end
        end
        if !correct
            @error "There were errors on the jump path: $path"
            @error "More info about the path: $path_type"
        end
        @assert correct
    end
    if n_children_tests != 5
        @error "Make sure that you have more feasible nodes in the search tree to test more."
    end
    @assert n_children_tests == 5

    # back to solved state
    CS.checkout_from_to!(com, c_backtrack_idx, path[1])
    # prune the last step
    CS.restore_prune!(com, path[1])

    return all_correct
end

function is_solved(com::CS.CoM)
    variables = com.search_space
    all_fixed = all(v -> CS.isfixed(v), variables)
    if !all_fixed
        @error "Not all variables are fixed"
        return false
    end
    for constraint in com.constraints
        values = CS.value.(com.search_space[constraint.indices])
        c_solved =
            CS.is_constraint_solved(constraint, constraint.fct, constraint.set, values)
        c_violated =
            CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)
        if !c_solved || c_violated
            @error "Constraint $(constraint.idx) is not solved"
            @error "Info about constraint: $(typeof(constraint)), $(typeof(constraint.fct)), $(typeof(constraint.set))"
            return false
        end
    end
    return true
end


text:
# This file was generated, do not modify it. # hide
X3 = select(X2, [:Lag1, :Lag2])
classif = machine(LogisticClassifier(), X3, y)
fit!(classif, rows=train)
ŷ = predict_mode(classif, rows=test)
accuracy(ŷ, y[test]) |> r3

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
"""
    DensityMatrix(state::AbstractArray{T, 3})
    DensityMatrix(state::AbstractMatrix{T})

Create a `DensityMatrix` with a state represented by array.
"""
YaoBase.DensityMatrix(state::MT) where {T,MT<:AbstractArray{T,3}} =
    DensityMatrix{size(state, 3),T,MT}(state)
YaoBase.DensityMatrix(state::AbstractMatrix) =
    DensityMatrix(reshape(state, size(state)..., 1))

"""
    state(ρ::DensityMatrix)

Return the raw state of density matrix `ρ`.
"""
state(ρ::DensityMatrix) = ρ.state

YaoBase.nqubits(ρ::DensityMatrix) = log2dim1(state(ρ))
YaoBase.nactive(ρ::DensityMatrix) = nqubits(ρ)
YaoBase.nbatch(dm::DensityMatrix{B}) where {B} = B
YaoBase.density_matrix(reg::ArrayReg{1}) = DensityMatrix(reg.state * reg.state')
function YaoBase.density_matrix(reg::ArrayReg{B}) where {B}
    M = size(reg.state, 1)
    s = reshape(reg |> state, M, :, B)
    out = similar(s, M, M, B)
    for b = 1:B
        @inbounds @views out[:, :, b] = s[:, :, b] * s[:, :, b]'
    end
    return DensityMatrix(out)
end

YaoBase.tracedist(dm1::DensityMatrix{B}, dm2::DensityMatrix{B}) where {B} =
    map(b -> trnorm(dm1.state[:, :, b] - dm2.state[:, :, b]), 1:B)

# TODO: use batch_broadcast in the future
"""
    probs(ρ)

Returns the probability distribution from a density matrix `ρ`.
"""
function YaoBase.probs(m::DensityMatrix{B,T}) where {B,T}
    res = zeros(T, size(m.state, 1), B)
    for i = 1:B
        @inbounds res[:, B] = diag(view(m.state, :, :, i))
    end
    return res
end

YaoBase.probs(m::DensityMatrix{1}) = diag(view(m.state, :, :, 1))

function YaoBase.purify(r::DensityMatrix{B}; nbit_env::Int = nactive(r)) where {B}
    Ne = 1 << nbit_env
    Ns = size(r.state, 1)
    state = similar(r.state, Ns, Ne, B)
    for ib = 1:B
        R, U = eigen!(r.state[:, :, ib])
        state[:, :, ib] .= view(U, :, Ns-Ne+1:Ns) .* sqrt.(abs.(view(R, Ns-Ne+1:Ns)'))
    end
    return ArrayReg(state)
end

# obtaining matrix from Yao.DensityMatrix{1}, `1` is the batch size.
LinearAlgebra.Matrix(d::DensityMatrix{1}) = dropdims(d.state, dims = 3)


text:
@testset "target" begin
    @test_throws ArgumentError Target(triple="invalid")
    @test_throws ArgumentError Target(name="invalid")

    host_triple = triple()
    host_t = Target(triple=host_triple)

    host_name = name(host_t)
    description(host_t)

    @test hasjit(host_t)
    @test hastargetmachine(host_t)
    @test hasasmparser(host_t)

    # target iteration
    let ts = targets()
        @test !isempty(ts)
        @test host_t in ts

        @test eltype(ts) == Target

        first(ts)

        for t in ts
            # ...
        end

        @test any(t -> t == host_t, collect(ts))
    end
end


text:
using ArgParse
using Random
using StaticArrays
using NCDatasets
using Test
using DocStringExtensions
using LinearAlgebra
using DelimitedFiles
using Dierckx
using Printf

using ClimateMachine
using ClimateMachine.Atmos
using ClimateMachine.ArtifactWrappers
using ClimateMachine.Orientations
using ClimateMachine.ConfigTypes
using ClimateMachine.Diagnostics
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.GenericCallbacks
using ClimateMachine.ODESolvers
using ClimateMachine.Mesh.Filters
using ClimateMachine.TemperatureProfiles
using ClimateMachine.Thermodynamics
using ClimateMachine.TurbulenceClosures
using ClimateMachine.VariableTemplates
using ClimateMachine.BalanceLaws: vars_state, Prognostic

using CLIMAParameters
using CLIMAParameters.Planet
using CLIMAParameters.Atmos.Microphysics

struct LiquidParameterSet <: AbstractLiquidParameterSet end
struct IceParameterSet <: AbstractIceParameterSet end
struct RainParameterSet <: AbstractRainParameterSet end
struct SnowParameterSet <: AbstractSnowParameterSet end
struct MicropysicsParameterSet{L, I, R, S} <: AbstractMicrophysicsParameterSet
    liq::L
    ice::I
    rai::R
    sno::S
end
struct EarthParameterSet{M} <: AbstractEarthParameterSet
    microphys::M
end

const microphys = MicropysicsParameterSet(
    LiquidParameterSet(),
    IceParameterSet(),
    RainParameterSet(),
    SnowParameterSet(),
)
const param_set = EarthParameterSet(microphys)

"""
  Define initial conditions based on sounding data
"""
function init_squall_line!(problem, bl, state, aux, localgeo, t, args...)

    FT = eltype(state)

    spl_tinit, spl_qinit, spl_uinit, spl_vinit, spl_pinit = args[1]
    # interpolate data
    (x, y, z) = localgeo.coord
    data_t = FT(spl_tinit(z))
    data_q = FT(spl_qinit(z))
    data_u = FT(spl_uinit(z))
    data_v = FT(spl_vinit(z))
    data_p = FT(spl_pinit(z))
    u = data_u
    v = data_v
    w = FT(0)
    if z >= 14000
        data_q = FT(0)
    end
    θ_c = 3.0
    rx = 10000.0
    ry = 1500.0
    rz = 1500.0
    xc = 0.5 * (FT(0) + FT(0))
    yc = 0.5 * (FT(0) + FT(5000))
    zc = 2000.0
    cylinder_flg = 0.0
    r = sqrt(
        (x - xc)^2 / rx^2 +
        cylinder_flg * (y - yc)^2 / ry^2 +
        (z - zc)^2 / rz^2,
    )
    Δθ = 0.0
    if r <= 1.0
        Δθ = θ_c * (cospi(0.5 * r))^2
    end
    θ_liq = data_t + Δθ
    ts = PhaseNonEquil_pθq(param_set, data_p, θ_liq, PhasePartition(FT(data_q)))
    T = air_temperature(ts)
    ρ = air_density(ts)
    e_kin = FT(1 / 2) * FT((u^2 + v^2 + w^2))
    e_pot = gravitational_potential(bl.orientation, aux)
    E = ρ * total_energy(e_kin, e_pot, ts)

    state.ρ = ρ
    state.ρu = SVector(ρ * u, ρ * v, FT(0))
    state.ρe = E

    state.moisture.ρq_tot = ρ * data_q
    if bl.moisture isa NonEquilMoist
        state.moisture.ρq_liq = FT(0)
        state.moisture.ρq_ice = FT(0)
    end

    if bl.precipitation isa RainModel
        state.precipitation.ρq_rai = FT(0)
    end
    if bl.precipitation isa RainSnowModel
        state.precipitation.ρq_rai = FT(0)
        state.precipitation.ρq_sno = FT(0)
    end

    return nothing
end

"""
  Read the original squall sounding
"""
function read_sounding()

    # Artifact creation is not thread-safe
    #      https://github.com/JuliaLang/Pkg.jl/issues/1219
    # To avoid race conditions from multiple jobs running this
    # driver at the same time, we must store artifacts in a
    # separate folder.

    artifact_folder = mktempdir(@__DIR__; prefix = "artifacts_")

    soundings_dataset = ArtifactWrapper(
        joinpath(artifact_folder, "Artifacts.toml"),
        "soundings",
        ArtifactFile[ArtifactFile(
            url = "https://caltech.box.com/shared/static/rjnvt2dlw7etm1c7mmdfrkw5gnfds5lx.nc",
            filename = "sounding_gabersek.nc",
        ),],
    )
    data_folder = get_data_folder(soundings_dataset)
    fsounding = joinpath(data_folder, "sounding_gabersek.nc")
    sounding = Dataset(fsounding, "r")
    height = sounding["height [m]"][:]
    θ = sounding["theta [K]"][:]
    q_vap = sounding["qv [g kg⁻¹]"][:]
    u = sounding["u [m s⁻¹]"][:]
    v = sounding["v [m s⁻¹]"][:]
    p = sounding["pressure [Pa]"][:]
    close(sounding)
    return (height = height, θ = θ, q_vap = q_vap, u = u, v = v, p = p)
end

"""
  Get data from interpolated array onto vectors.
  Accepts data in 6 column format
"""
function spline_int()

    nt = read_sounding()

    # WARNING: Not all sounding data is formatted/scaled the same.
    spl_tinit = Spline1D(nt.height, nt.θ; k = 1)
    spl_qinit = Spline1D(nt.height, nt.q_vap .* 0.001; k = 1)
    spl_uinit = Spline1D(nt.height, nt.u; k = 1)
    spl_vinit = Spline1D(nt.height, nt.v; k = 1)
    spl_pinit = Spline1D(nt.height, nt.p; k = 1)
    return (
        t = spl_tinit,
        q = spl_qinit,
        u = spl_uinit,
        v = spl_vinit,
        p = spl_pinit,
    )
end

function config_squall_line(
    FT,
    N,
    resolution,
    xmax,
    ymax,
    zmax,
    xmin,
    ymin,
    moisture_model = "nonequilibrium",
    precipitation_model = "rainsnow",
)
    # Reference state
    nt = read_sounding()

    zinit = nt.height
    tinit = nt.θ
    qinit = nt.q_vap .* 0.001
    u_init = nt.u
    v_init = nt.v
    p_init = nt.p

    maxz = length(zinit)
    thinit = zeros(maxz)
    piinit = zeros(maxz)
    thinit[1] = tinit[1] / (1 + 0.61 * qinit[1])
    piinit[1] = 1
    for k in 2:maxz
        thinit[k] = tinit[k] / (1 + 0.61 * qinit[k])
        piinit[k] =
            piinit[k - 1] -
            9.81 / (1004 * 0.5 * (tinit[k] + tinit[k - 1])) *
            (zinit[k] - zinit[k - 1])
    end
    T_min = FT(thinit[maxz] * piinit[maxz])
    T_s = FT(thinit[1] * piinit[1])
    Γ_lapse = FT(9.81 / 1004)
    tvmax = T_s * (1 + 0.61 * qinit[1])
    deltatv = -(T_min - tvmax)
    tvmin = T_min * (1 + 0.61 * qinit[maxz])
    htv = 8000.0
    #T = DecayingTemperatureProfile(T_min, T_s, Γ_lapse)
    Tv = DecayingTemperatureProfile{FT}(param_set, tvmax, tvmin, htv)
    rel_hum = FT(0)
    ref_state = HydrostaticState(Tv, rel_hum)

    # Sponge
    c_sponge = FT(0.5)
    # Rayleigh damping
    u_relaxation = SVector(FT(0), FT(0), FT(0))
    zsponge = FT(15000.0)
    rayleigh_sponge =
        RayleighSponge(FT, zmax, zsponge, c_sponge, u_relaxation, 2)

    # Boundary conditions
    # SGS Filter constants
    C_smag = FT(0.18) # 0.21 for stable testing, 0.18 in practice
    C_drag = FT(0.0011)
    LHF = FT(50)
    SHF = FT(10)
    ics = init_squall_line!

    source = (Gravity(), rayleigh_sponge)

    # moisture model and its sources
    if moisture_model == "equilibrium"
        moisture = EquilMoist{FT}(; maxiter = 20, tolerance = FT(1))
    elseif moisture_model == "nonequilibrium"
        source = (source..., CreateClouds()...)
        moisture = NonEquilMoist()
    else
        @warn @sprintf(
            """
%s: unrecognized moisture_model in source terms, using the defaults""",
            moisture_model,
        )
        source = (source..., CreateClouds()...)
        moisture = NonEquilMoist()
    end

    # precipitation model and its sources
    if precipitation_model == "noprecipitation"
        precipitation = NoPrecipitation()
        source = (source..., RemovePrecipitation(true)...)
    elseif precipitation_model == "rain"
        source = (source..., WarmRain_1M()...)
        precipitation = RainModel()
    elseif precipitation_model == "rainsnow"
        source = (source..., RainSnow_1M()...)
        precipitation = RainSnowModel()
    else
        @warn @sprintf(
            """
%s: unrecognized precipitation_model in source terms, using the defaults""",
            precipitation_model,
        )
        source = (source..., RainSnow_1M()...)
        precipitation = RainSnowModel()
    end

    problem = AtmosProblem(
        boundaryconditions = (AtmosBC(), AtmosBC()),
        init_state_prognostic = ics,
    )

    model = AtmosModel{FT}(
        AtmosLESConfigType,
        param_set;
        problem = problem,
        ref_state = ref_state,
        moisture = moisture,
        precipitation = precipitation,
        turbulence = SmagorinskyLilly{FT}(C_smag),
        source = source,
    )

    ode_solver = ClimateMachine.IMEXSolverType()

    config = ClimateMachine.AtmosLESConfiguration(
        "Squall_line",
        N,
        resolution,
        xmax,
        ymax,
        zmax,
        param_set,
        init_squall_line!,
        xmin = xmin,
        ymin = ymin,
        solver_type = ode_solver,
        model = model,
        periodicity = (true, true, false),
        boundary = ((2, 2), (2, 2), (1, 2)),
    )
    return config
end

function config_diagnostics(driver_config, boundaries, resolution)

    interval = "3600ssecs"

    dgngrp_profiles = setup_atmos_default_diagnostics(
        AtmosLESConfigType(),
        interval,
        driver_config.name,
    )

    interpol = ClimateMachine.InterpolationConfiguration(
        driver_config,
        boundaries,
        resolution,
    )
    dgngrp_state = setup_dump_state_diagnostics(
        AtmosLESConfigType(),
        interval,
        driver_config.name,
        interpol = interpol,
    )
    dgngrp_aux = setup_dump_aux_diagnostics(
        AtmosLESConfigType(),
        interval,
        driver_config.name,
        interpol = interpol,
    )

    return ClimateMachine.DiagnosticsConfiguration([
        dgngrp_profiles,
        dgngrp_state,
        dgngrp_aux,
    ])
end

function main()
    # TODO: this will move to the future namelist functionality
    squall_args = ArgParseSettings(autofix_names = true)
    add_arg_group!(squall_args, "SQUALL_LINE")
    @add_arg_table! squall_args begin
        "--moisture-model"
        help = "specify cloud condensate model"
        metavar = "equilibrium|nonequilibrium"
        arg_type = String
        default = "nonequilibrium"
        "--precipitation-model"
        help = "specify precipitation model"
        metavar = "noprecipitation|rain|rainsnow"
        arg_type = String
        default = "rainsnow"
        "--check-asserts"
        help = "should asserts be checked at the end of the simulation"
        metavar = "yes|no"
        arg_type = String
        default = "no"
    end

    cl_args =
        ClimateMachine.init(parse_clargs = true, custom_clargs = squall_args)
    moisture_model = cl_args["moisture_model"]
    precipitation_model = cl_args["precipitation_model"]
    check_asserts = cl_args["check_asserts"]

    FT = Float64

    # DG polynomial order
    N = 4

    # Domain resolution and size
    Δx = FT(250)
    Δy = FT(1000)
    Δv = FT(200)
    resolution = (Δx, Δy, Δv)

    xmax = FT(30000)
    ymax = FT(5000)
    zmax = FT(24000)
    xmin = FT(-30000)
    ymin = FT(0)
    # for diagnostics
    boundaries = [
        xmin ymin FT(0)
        xmax ymax zmax
    ]

    t0 = FT(0)
    timeend = FT(9000)
    spl_tinit, spl_qinit, spl_uinit, spl_vinit, spl_pinit = spline_int()
    Cmax = FT(0.4)

    # driver, solver and diagnostics configs
    driver_config = config_squall_line(
        FT,
        N,
        resolution,
        xmax,
        ymax,
        zmax,
        xmin,
        ymin,
        moisture_model,
        precipitation_model,
    )
    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        (spl_tinit, spl_qinit, spl_uinit, spl_vinit, spl_pinit);
        init_on_cpu = true,
        Courant_number = Cmax,
    )
    dgn_config = config_diagnostics(driver_config, boundaries, resolution)

    filter_vars = ("moisture.ρq_tot",)
    if moisture_model == "nonequilibrium"
        filter_vars = (filter_vars..., "moisture.ρq_liq", "moisture.ρq_ice")
    end
    if precipitation_model == "rain"
        filter_vars = (filter_vars..., "precipitation.ρq_rai")
    end
    if precipitation_model == "rainsnow"
        filter_vars =
            (filter_vars..., "precipitation.ρq_rai", "precipitation.ρq_rai")
    end

    cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(1) do
        Filters.apply!(
            solver_config.Q,
            filter_vars,
            solver_config.dg.grid,
            TMARFilter(),
        )
        nothing
    end

    result = ClimateMachine.invoke!(
        solver_config;
        diagnostics_config = dgn_config,
        user_callbacks = (cbtmarfilter,),
        check_euclidean_distance = true,
    )

    if check_asserts == "yes"

        m = driver_config.bl
        Q = solver_config.Q
        ρ_ind = varsindex(vars_state(m, Prognostic(), FT), :ρ)

        if moisture_model == "equilibrium"
            ρq_tot_ind =
                varsindex(vars_state(m, Prognostic(), FT), :moisture, :ρq_tot)

            min_q_tot = minimum(abs.(
                Array(Q[:, ρq_tot_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_tot = maximum(abs.(
                Array(Q[:, ρq_tot_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))

            @info(min_q_tot, max_q_tot)

            # test that moisture exists and is not NaN
            @test !isnan(max_q_tot)

            # test that there is reasonable amount of moisture
            @test abs(max_q_tot) > FT(1e-2)
        end
        if moisture_model == "nonequilibrium"
            ρq_tot_ind =
                varsindex(vars_state(m, Prognostic(), FT), :moisture, :ρq_tot)
            ρq_liq_ind =
                varsindex(vars_state(m, Prognostic(), FT), :moisture, :ρq_liq)
            ρq_ice_ind =
                varsindex(vars_state(m, Prognostic(), FT), :moisture, :ρq_ice)

            min_q_tot = minimum(abs.(
                Array(Q[:, ρq_tot_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_tot = maximum(abs.(
                Array(Q[:, ρq_tot_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            min_q_liq = minimum(abs.(
                Array(Q[:, ρq_liq_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_liq = maximum(abs.(
                Array(Q[:, ρq_liq_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            min_q_ice = minimum(abs.(
                Array(Q[:, ρq_ice_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_ice = maximum(abs.(
                Array(Q[:, ρq_ice_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            @info(min_q_tot, max_q_tot)
            @info(min_q_liq, max_q_liq)
            @info(min_q_ice, max_q_ice)

            # test that moisture exists and is not NaN
            @test !isnan(max_q_tot)
            @test !isnan(max_q_liq)
            @test !isnan(max_q_ice)

            # test that there is reasonable amount of moisture
            @test abs(max_q_tot) > FT(1e-2)
            @test abs(max_q_liq) > FT(1e-3)
            @test abs(max_q_ice) > FT(1e-3)
        end

        if precipitation_model == "rain"
            ρq_rai_ind = varsindex(
                vars_state(m, Prognostic(), FT),
                :precipitation,
                :ρq_rai,
            )

            min_q_rai = minimum(abs.(
                Array(Q[:, ρq_rai_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_rai = maximum(abs.(
                Array(Q[:, ρq_rai_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))

            @info(min_q_rai, max_q_rai)

            # test that rain variable exists and is not NaN
            @test !isnan(max_q_rai)

            # test that there is reasonable amount of rain water...
            @test abs(max_q_rai) > FT(1e-4)
        end
        if precipitation_model == "rainsnow"
            ρq_rai_ind = varsindex(
                vars_state(m, Prognostic(), FT),
                :precipitation,
                :ρq_rai,
            )
            ρq_sno_ind = varsindex(
                vars_state(m, Prognostic(), FT),
                :precipitation,
                :ρq_sno,
            )
            min_q_rai = minimum(abs.(
                Array(Q[:, ρq_rai_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_rai = maximum(abs.(
                Array(Q[:, ρq_rai_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            min_q_sno = minimum(abs.(
                Array(Q[:, ρq_sno_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_sno = maximum(abs.(
                Array(Q[:, ρq_sno_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))

            @info(min_q_rai, max_q_rai)
            @info(min_q_sno, max_q_sno)

            # test that rain and snow variables exists and are not NaN
            @test !isnan(max_q_rai)
            @test !isnan(max_q_sno)

            # test that there is reasonable amount of precipitation
            @test abs(max_q_rai) > FT(1e-4)
            @test abs(max_q_sno) > FT(1e-6)
        end
    end
end

main()

max_stars_repo_name: jm-c/CLIMA

max_issues_repo_name: jm-c/CLIMA

max_forks_repo_name: jm-c/CLIMA