text:
"""
    axpy!(a::Number,X::Union{Matrix,RkMatrix,HMatrix},Y::Union{Matrix,RkMatrix,HMatrix})

Perform `Y <-- a*X + Y` in-place. Note that depending on the types of `X` and
`Y`, this may require converting from/to different formats during intermediate
calculations.

In the case where `Y` is an `RkMatrix`, the call `axpy!(a,X,Y)` should
typically be followed by recompression stage to keep the rank of `Y` under
control.

In the case where `Y` is an `HMatrix`, the call `axpy!(a,X,Y)` sums `X` to the
data in the node `Y` (and not on the leaves). In case `Y` has no `data`, it will
simply be assigned `X`. This means that after the call `axpy(a,X,Y)`, the object
`Y` is in a *dirty* state (see [`isclean`][@ref]) and usually a call to
[`flush_to_leaves!`](@ref) or [`flush_to_children!`](@ref) follows.
"""
function axpy!(a,X::Matrix,Y::RkMatrix)
    axpy!(a,RkMatrix(X),Y)
end

# 1.3
function axpy!(a,X::Matrix,Y::HMatrix)
    if hasdata(Y)
        axpy!(a,X,data(Y))
    else
        setdata!(Y,a*X)
    end
    return Y
end

# 2.1
function axpy!(a,X::RkMatrix,Y::Matrix)
    axpy!(a,Matrix(X),Y)
end

#2.2
function axpy!(a,X::RkMatrix,Y::RkMatrix)
    Y.A   = hcat(a*X.A,Y.A)
    Y.B   = hcat(X.B,Y.B)
    return Y
end

# 2.3
function axpy!(a,X::RkMatrix,Y::HMatrix)
    if hasdata(Y)
        axpy!(a,X,data(Y))
    else
        setdata!(Y,a*X)
    end
    return Y
end

#3.1
function axpy!(a,X::HMatrix,Y::Matrix)
    @debug "calling axpy! with `X` and HMatrix and `Y` a Matrix"
    shift = pivot(X) .- 1
    for block in PreOrderDFS(X)
        irange = rowrange(block) .- shift[1]
        jrange = colrange(block) .- shift[2]
        if hasdata(block)
            axpy!(a,data(block),view(Y,irange,jrange))
        end
    end
    return Y
end

# 3.2
function axpy!(a,X::HMatrix,Y::RkMatrix)
    @debug "calling axpby! with `X` and HMatrix and `Y` an RkMatrix"
    # FIXME: inneficient implementation due to conversion from HMatrix to
    # Matrix. Does it really matter? I don't think this function should be
    # called.
    axpy!(a,Matrix(X;global_index=false),Y)
end

# 3.3
function axpy!(a,X::HMatrix,Y::HMatrix)
    # TODO: assumes X and Y have the same structure. How to reinforce this?
    if hasdata(X)
        if hasdata(Y)
            axpy!(a,data(X),Y)
        else
            setdata!(Y,a*data(X))
        end
    end
    @assert size(children(X)) == size(children(Y)) "adding hierarchical matrices requires identical block structure"
    for (bx,by) in zip(children(X),children(Y))
        axpy!(a,bx,by)
    end
    return Y
end

# add a unifor scaling to an HMatrix return an HMatrix
function axpy!(a,X::UniformScaling,Y::HMatrix)
    @assert isclean(Y)
    if hasdata(Y)
        d = data(Y)
        @assert d isa Matrix
        n = min(size(d)...)
        for i=1:n
            d[i,i] += a*X.λ
        end
    else
        n = min(blocksize(Y)...)
        for i=1:n
            axpy!(a,X,children(Y)[i,i])
        end
    end
    return Y
end

Base.:(+)(X::UniformScaling,Y::HMatrix) = axpy!(true,X,deepcopy(Y))
Base.:(+)(X::HMatrix,Y::UniformScaling) = Y+X

max_stars_repo_name: WaveProp/HMatrices

max_issues_repo_name: WaveProp/HMatrices

max_forks_repo_name: WaveProp/HMatrices


text:
using Flux

struct ReLU
end

(a::ReLU)(x) = relu(x)

der(a::ReLU, x) = ceil(clamp(x, 0, 1))

struct SoftPlus
    beta
end

SoftPlus() = SoftPlus(1.0)

(a::SoftPlus)(x) = softplus(x)

function der(a::SoftPlus, x)
    cx = clamp(x, -20, 20)
    exp_x = exp(convert(eltype(x), a.beta) * cx)
    out = exp_x / (exp_x + convert(eltype(x), 1.0))
    return out
end

struct Tanh
end

(a::Tanh)(x) = tanh(x)

der(a::Tanh, x) = convert(eltype(x), 1.0) - tanh(x)^2

struct Identity
end

(a::Identity)(x) = x

der(a::Identity, x) = clamp(x, 1, 1)

struct Differential{F, G, S, T}
    W::S
    b::T
    g::F
    g_prime::G
end

function Differential(in::Integer, out::Integer, g=Identity();
               initW = Flux.glorot_uniform, initb = Flux.zeros)
    return Differential(initW(out, in), (initb(out)), g, x->der(g, x))
end

@Flux.treelike Differential

function (a::Differential)((x, prev))
    W, b, g, g_prime = a.W, a.b, a.g, a.g_prime
    o = W*x .+ b
    op = g.(o)
    #newd = reshape(g_prime.(o), size(o, 1), 1, size(o, 2)) # hidden x B
    #derW = newd .* W
    derW = reshape(g_prime.(o), size(o, 1), 1, size(o, 2)) .* W
    der = batchmul(derW, prev)
    return op, der
end

struct DifferentialNet
    layers::Chain
    DifferentialNet(xs...) = new(Chain(xs...))
end

Flux.@functor DifferentialNet

function (a::DifferentialNet)(q::AbstractArray)
    xdim, B = size(q)
    qd_dq = repeat(eye(eltype(q), xdim), 1, 1, B)
    qd, qd_dq = a.layers((q, qd_dq))
    return qd, qd_dq
end

struct GeneralizedForces
    forces
end

Flux.@functor GeneralizedForces

function (a::GeneralizedForces)(q::AbstractArray, v::AbstractArray, u::AbstractArray)
    F = sum([f(q, v, u) for f in a.forces])
    return F
end

struct GeneralizedForce
    layers::Chain
end
Flux.@functor GeneralizedForce

function (a::GeneralizedForce)(q, v, u)
    x = vcat(q, v, u)
    return a.layers(x)
end

function GeneralizedForce(qdim, udim, hidden_sizes)
    return GeneralizedForce(Chain(Dense(2*qdim+udim, hidden_sizes[1], tanh),
                                     (Dense(hidden_sizes[i], hidden_sizes[i+1], tanh)
                                      for i in 1:(length(hidden_sizes)-1))...,
                                     Dense(hidden_sizes[end], qdim)))
end

struct ControlAffineForce
    layers::Chain
end

Flux.@functor ControlAffineForce

function (a::ControlAffineForce)(q, v, u)
    bsize = size(q)[end]
    qdim = size(q, 1)
    udim = size(u, 1)
    B = a.layers(q)
    B = reshape(B, (qdim, udim, bsize))
    F = batchmul(B, reshape(u, size(u, 1), 1, size(u, 2)))
    F = dropdims(F, dims=2)
    @assert size(F) == (qdim, bsize)
    return F
end

function ControlAffineForce(qdim, udim, hidden_sizes)
    return ControlAffineForce(Chain(Dense(qdim, hidden_sizes[1], tanh),
                                       (Dense(hidden_sizes[i], hidden_sizes[i+1], tanh)
                                        for i in 1:(length(hidden_sizes)-1))...,
                                       Dense(hidden_sizes[end], qdim * udim)))
end

struct ControlAffineLinearForce
    B
end

Flux.@functor ControlAffineLinearForce

function (a::ControlAffineLinearForce)(q, v, u)
    bsize = size(q)[end]
    qdim = size(q, 1)
    B = reshape(a.B, size(a.B)..., 1)
    F = batchmul(repeat(B, 1, 1, bsize), reshape(u, size(u, 1), 1, size(u, 2)))
    F = dropdims(F, dims=2)
    @assert size(F) == (qdim, bsize)
    return F
end

function ControlAffineLinearForce(qdim, udim)
    return ControlAffineLinearForce(randn(qdim, udim))
end

struct QVForce
    layers::Chain
end

Flux.@functor QVForce

function (a::QVForce)(q, v, u)
    bsize = size(q)[end]
    qdim = size(q, 1)
    F = a.layers(vcat(q, v))
    F = reshape(F, (qdim, bsize))
    @assert size(F) == (qdim, bsize)
    return F
end

function QVForce(qdim, udim, hidden_sizes)
    return QVForce(Chain(Dense(qdim*2, hidden_sizes[1], tanh),
                         (Dense(hidden_sizes[i], hidden_sizes[i+1], tanh)
                          for i in 1:(length(hidden_sizes)-1))...,
                         Dense(hidden_sizes[end], qdim)))
end

function DelanMassMatrix(qdim, hidden_sizes, activation=ReLU())
    return DifferentialNet(Differential(qdim, hidden_sizes[1], activation),
                           (Differential(hidden_sizes[i], hidden_sizes[i+1], activation)
                            for i in 1:(length(hidden_sizes)-1))...,
                           Differential(hidden_sizes[end], round(Int, qdim*(qdim+1)/2)))
end

function DelanPotential(qdim, hidden_sizes, activation=ReLU())
    return DifferentialNet(Differential(qdim, hidden_sizes[1], activation),
                           (Differential(hidden_sizes[i], hidden_sizes[i+1], activation)
                            for i in 1:(length(hidden_sizes)-1))...,
                           Differential(hidden_sizes[end], 1))
end

struct DelanZeroPotential
end

(a::DelanZeroPotential)(q) = (zeros(1, size(q)[end]), zeros(1, size(q, 1), size(q)[end]))

struct DelanRigidBody{M, V, F, I<:Integer, G, H<:AbstractVector{<:Real}} <: AbstractRigidBody
    massmatrix::M
    potential::V
    generalized_forces::F
    qdim::I
    udim::I
    bias::G
    thetamask::H
end

Flux.@functor DelanRigidBody

function DelanRigidBody(qdim::Int, udim::Int, hidden_sizes, thetamask; activation=ReLU(), bias=10.)
    mm = DelanMassMatrix(qdim, hidden_sizes, activation)
    pot = DelanPotential(qdim, hidden_sizes, activation)
    ff = GeneralizedForce(qdim, udim, hidden_sizes)
    return DelanRigidBody(mm, pot, ff, qdim, udim, bias, thetamask)
end

"""
Converts a batch of vectors to a batch of LowerTriangular matrices.
"""
function embedtoLwithbias(x, bias, qdim)
    L_diag = fill_diagonal(x[1:qdim, :] .+ convert(eltype(x), bias))
    L_ltr = fill_lower_triangular(x[qdim+1:end, :])
    L = L_diag .+ L_ltr
    return L
end

function (f::DelanRigidBody)(x, u)
    q = x[1:f.qdim, :]
    v = x[f.qdim+1:2*f.qdim, :]

    embedtoL(x) = embedtoLwithbias(x, 0., f.qdim)

    L_params, dLparamsdq = f.massmatrix(q)

    L = embedtoLwithbias(L_params, f.bias, f.qdim)

    M = batchmul(L, permutedims(L, (2, 1, 3)))

    pot, gradpot = f.potential(q)
    G = dropdims(gradpot, dims=1)

    batched_v = reshape(v, size(v, 1), 1, size(v, 2))

    dLparamsdt = batchmul(dLparamsdq, batched_v)

    dLparamsdt = dropdims(dLparamsdt, dims=2)
    dLdt = embedtoL(dLparamsdt)

    dMdt = batchmul(L, permutedims(dLdt, (2, 1, 3))) + batchmul(dLdt, permutedims(L, (2, 1, 3)))

    dMdtv = batchmul(dMdt, batched_v)
    dMdtv = dropdims(dMdtv, dims=2)

    function dKEdq_i(dLdq_i)
        _mx = batchmul(L, permutedims(dLdq_i, (2, 1, 3))) + batchmul(dLdq_i, permutedims(L, (2, 1, 3)))
        t = batchmul(_mx, batched_v)
        dKEdqi = sum((dropdims(t, dims=2) .* v), dims=1)
        return dKEdqi
    end

    # XXX: Zygote doesn't work with this anymore :/
    #dKEdq = slicemap(x->dKEdq_i(embedtoL(x)), dLparamsdq; dims=(1, 3))
    #dKEdq = dropdims(dKEdq, dims=1)
    dKEdq = map(i->dKEdq_i(embedtoL(dLparamsdq[:, i, :])), 1:f.qdim)
    dKEdq = vcat(dKEdq...)

    Cv = dMdtv - convert(eltype(x), 0.5) * dKEdq

    F = f.generalized_forces(q, v, u)
    totalF = F - Cv - G
    qddot_ls = [(view(M, :, :,i) \ view(totalF, :, i)) for i in 1:size(L, 3)]
    qddot = hcat(qddot_ls...)
    return vcat(v, qddot)
end

"""
Create (a batch of) diagonal matrix from a vector of inputs

$(SIGNATURES)
"""
function fill_diagonal(x)
    m = size(x, 1)
    y = Zygote.Buffer(x, (m, m, size(x)[2:end]...)...)
    for R in CartesianIndices((m, m, size(x)[end]))
        y[R]= convert(eltype(x), 0.)
    end
    for i in 1:size(x)[end]
        vptr = 1
        @inbounds for j in 1:m
            y[j, j, i] = x[vptr, i]
            vptr += 1
        end
    end
    return copy(y)
end

"""
Create (a batch of) triangular matrix from a vector of inputs

$(SIGNATURES)
"""
function fill_lower_triangular(x)
    m = size(x)[1]
    n = (sqrt(8m+1)+1)/2
    n = round(Int, n)
    y = Zygote.Buffer(x, (n, n, size(x)[2:end]...)...)

    for R in CartesianIndices((n, n, size(x)[end]))
        y[R]= convert(eltype(x), 0.)
    end
    for i in 1:size(x)[end]
        vptr = 1
        @inbounds for j in 1:n, k in 1:j-1
            y[j, k, i] = x[vptr, i]
            vptr += 1
        end
    end
    return copy(y)
end

ext: jl


text:
module JsonTests

using Test
using Gridap.Helpers
using Gridap.Io
import Gridap.Io: to_dict
import Gridap.Io: from_dict

using JSON

struct Foo <: GridapType
  bar::Int
end

function to_dict(foo::Foo)
  dict = Dict{Symbol,Any}()
  dict[:bar] = foo.bar
  dict
end

function from_dict(::Type{Foo},dict::Dict{Symbol,Any})
  bar = dict[:bar]
  Foo(bar)
end

bar = 3
foo = Foo(bar)

@test foo == from_dict(Foo,to_dict(foo))
@test foo == from_json(Foo,to_json(foo))

d = mktempdir()
f = joinpath(d,"foo.json")
to_json_file(foo,f)
@test foo == from_json_file(Foo,f)
rm(d,recursive=true)

@test foo == from_json(Foo,JSON.json(foo))

end # module


text:
abstract type AbstractDCLineFormulation <: AbstractBranchFormulation end
struct HVDCUnbounded <: AbstractDCLineFormulation end
struct HVDCLossless <: AbstractDCLineFormulation end
struct HVDCDispatch <: AbstractDCLineFormulation end
struct VoltageSourceDC <: AbstractDCLineFormulation end

const UnboundedBranches =
    Union{HVDCUnbounded, StaticLineUnbounded, StaticTransformerUnbounded}

#################################### Branch Variables ##################################################
flow_variables!(
    psi_container::PSIContainer,
    system_formulation::Type{<:PM.AbstractPowerModel},
    devices::IS.FlattenIteratorWrapper{<:PSY.DCBranch},
) = nothing

function flow_variables!(
    psi_container::PSIContainer,
    system_formulation::Type{StandardPTDFModel},
    devices::IS.FlattenIteratorWrapper{B},
) where {B <: PSY.DCBranch}
    time_steps = model_time_steps(psi_container)
    var_name = make_variable_name(FLOW_ACTIVE_POWER, B)
    container = container_spec(
        psi_container.JuMPmodel,
        [PSY.get_name(d) for d in devices],
        time_steps,
    )
    assign_variable!(psi_container, FLOW_ACTIVE_POWER, B, container)
    for (ix, d) in enumerate(devices)
        bus_fr = PSY.get_number(PSY.get_arc(d).from)
        bus_to = PSY.get_number(PSY.get_arc(d).to)
        for t in time_steps
            jvariable = JuMP.@variable(
                psi_container.JuMPmodel,
                base_name = "$(bus_fr), $(bus_to)_{$(PSY.get_name(d)), $(t)}",
            )
            container[PSY.get_name(d), t] = jvariable
            add_to_expression!(
                psi_container.expressions[:nodal_balance_active],
                PSY.get_number(PSY.get_arc(d).from),
                t,
                jvariable,
                -1.0,
            )
            add_to_expression!(
                psi_container.expressions[:nodal_balance_active],
                PSY.get_number(PSY.get_arc(d).to),
                t,
                jvariable,
                1.0,
            )
        end
    end
    return
end

#################################### Flow Variable Bounds ##################################################
#################################### Rate Limits Constraints ##################################################
function branch_rate_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{B},
    model::DeviceModel{B, <:AbstractDCLineFormulation},
    system_formulation::Type{<:PM.AbstractDCPModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {B <: PSY.DCBranch}
    var = get_variable(psi_container, FLOW_ACTIVE_POWER, B)
    time_steps = model_time_steps(psi_container)
    constraint_val =
        JuMPConstraintArray(undef, [PSY.get_name(d) for d in devices], time_steps)
    assign_constraint!(psi_container, FLOW_ACTIVE_POWER, B, constraint_val)
    for t in time_steps, d in devices
        min_rate = max(
            PSY.get_active_power_limits_from(d).min,
            PSY.get_active_power_limits_to(d).min,
        )
        max_rate = min(
            PSY.get_active_power_limits_from(d).max,
            PSY.get_active_power_limits_to(d).max,
        )
        constraint_val[PSY.get_name(d), t] = JuMP.@constraint(
            psi_container.JuMPmodel,
            min_rate <= var[PSY.get_name(d), t] <= max_rate
        )
    end
    return
end

function branch_rate_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{B},
    model::DeviceModel{B, HVDCLossless},
    system_formulation::Type{<:PM.AbstractPowerModel},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {B <: PSY.DCBranch}
    for (var_type, cons_type) in zip(
        (FLOW_ACTIVE_POWER_FROM_TO, FLOW_ACTIVE_POWER_TO_FROM),
        (RATE_LIMIT_FT, RATE_LIMIT_TF),
    )
        var = get_variable(psi_container, var_type, B)
        time_steps = model_time_steps(psi_container)
        constraint_val =
            JuMPConstraintArray(undef, [PSY.get_name(d) for d in devices], time_steps)
        assign_constraint!(psi_container, cons_type, B, constraint_val)
        time_steps = model_time_steps(psi_container)

        for t in time_steps, d in devices
            min_rate = max(
                PSY.get_active_power_limits_from(d).min,
                PSY.get_active_power_limits_to(d).min,
            )
            max_rate = min(
                PSY.get_active_power_limits_from(d).max,
                PSY.get_active_power_limits_to(d).max,
            )
            name = PSY.get_name(d)
            constraint_val[name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                min_rate <= var[name, t] <= max_rate
            )
        end
    end
    return
end

function branch_rate_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{B},
    model::DeviceModel{B, <:AbstractDCLineFormulation},
    system_formulation::Union{
        Type{<:PM.AbstractActivePowerModel},
        Type{<:PM.AbstractPowerModel},
    },
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {B <: PSY.DCBranch}
    time_steps = model_time_steps(psi_container)
    for (var_type, cons_type) in zip(
        (FLOW_ACTIVE_POWER_FROM_TO, FLOW_ACTIVE_POWER_TO_FROM),
        (RATE_LIMIT_FT, RATE_LIMIT_TF),
    )
        var = get_variable(psi_container, var_type, B)
        constraint_val =
            JuMPConstraintArray(undef, [PSY.get_name(d) for d in devices], time_steps)
        assign_constraint!(psi_container, cons_type, B, constraint_val)

        for t in time_steps, d in devices
            min_rate = max(
                PSY.get_active_power_limits_from(d).min,
                PSY.get_active_power_limits_to(d).min,
            )
            max_rate = min(
                PSY.get_active_power_limits_from(d).max,
                PSY.get_active_power_limits_to(d).max,
            )
            constraint_val[PSY.get_name(d), t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                min_rate <= var[PSY.get_name(d), t] <= max_rate
            )
            add_to_expression!(
                psi_container.expressions[:nodal_balance_active],
                PSY.get_number(PSY.get_arc(d).to),
                t,
                var[PSY.get_name(d), t],
                -PSY.get_loss(d).l1,
                -PSY.get_loss(d).l0,
            )
        end
    end
    return
end


text:
@connector function RealInput(;name, nin=1, u_start=nin > 1 ? 0.0 : zeros(nin))
    if nin == 1
        @variables u(t) = u_start
    else
        @variables u[1:nin](t) = u_start
        u = collect(u)
    end
    ODESystem(Equation[], t, [u...], []; name=name)
end
@doc """
    RealInput(;name, nin=1, u_start=nin > 1 ? 0.0 : zeros(nin))

Connector with one input signal of type Real.

# Parameters:
- `nin`: Number of inputs
- `u_start`: Initial value for `u`

# States:
- `u`: Value of of the connector; if nin=1 this is a scalar
""" RealInput

@connector function RealOutput(;name, nout=1, u_start=nout > 1 ? 0.0 : zeros(nout))
    if nout == 1
        @variables u(t) = u_start
    else
        @variables u[1:nout](t) = u_start
        u = collect(u)
    end
    ODESystem(Equation[], t, [u...], []; name=name)
end
@doc """
    RealOutput(;name, nout=1, u_start=nout > 1 ? 0.0 : zeros(nout))

Connector with one output signal of type Real.

# Parameters:
- `nout`: Number of inputs
- `u_start`: Initial value for `u`

# States:
- `u`: Value of of the connector; if nout=1 this is a scalar
""" RealOutput

"""
    SISO(;name, u_start=0.0, y_start=0.0)

Single Input Single Output continuous control block.

# Parameters:
- `u_start`: Initial value for the input
- `y_start`: Initial value for the output
"""
function SISO(;name, u_start=0.0, y_start=0.0)
    @named input = RealInput(u_start=u_start)
    @named output = RealOutput(u_start=y_start)
    @variables begin
        u(t)=u_start
        y(t)=y_start
    end
    eqs = [
        u ~ input.u
        y ~ output.u
    ]
    return ODESystem(eqs, t, [u, y], []; name=name, systems=[input, output])
end

"""
    MIMO(;name, nin=1, nout=1, u_start=zeros(nin), y_start=zeros(nout))

Base class for a multiple Input multiple Output continuous control block.

# Parameters:
- `nin`: Input dimension
- `nout`: Output dimension
- `u_start`: Initial value for the input
- `y_start`: Initial value for the output
"""
function MIMO(;name, nin=1, nout=1, u_start=zeros(nin), y_start=zeros(nout))
    @named input = RealInput(nin=nin, u_start=u_start)
    @named output = RealOutput(nout=nout, u_start=y_start)
    @variables begin
        u[1:nin](t)=u_start
        y[1:nout](t)=y_start
    end
    eqs = [
        [u[i] ~ input.u[i] for i in 1:nin]...,
        [y[i] ~ output.u[i] for i in 1:nout]...,
    ]
    return ODESystem(eqs, t, vcat(u..., y...), []; name=name, systems=[input, output])
end


text:
#=
real symmetric or complex Hermitian sparse positive semidefinite cone (svec format)
sparsity pattern given by lower triangle row and column indices, with all diagonal elements
dual cone is the PSD-completable matrices
=#

mutable struct PosSemidefTriSparse{D <: PrimDual, C <: RealCompF} <: Cache
    oa_s::Vector{AE}
    side::Int
    row_idxs::Vector{Int}
    col_idxs::Vector{Int}
    sep_constr::CR
    PosSemidefTriSparse{D, C}() where {D <: PrimDual, C <: RealCompF} = new{D, C}()
end

function MOIPajarito.Cones.create_cache(
    oa_s::Vector{AE},
    cone::Hypatia.PosSemidefTriSparseCone{<:Hypatia.Cones.PSDSparseImpl, RealF, C},
    opt::Optimizer,
) where {C <: RealCompF}
    D = primal_or_dual(cone.use_dual)
    cache = PosSemidefTriSparse{D, C}()
    cache.oa_s = oa_s
    cache.side = cone.side
    @assert length(oa_s) == MOI.dimension(cone)
    cache.row_idxs = cone.row_idxs
    cache.col_idxs = cone.col_idxs
    if D == Dual
        cache.sep_constr = get_sep_constr(cone, opt)
    end
    return cache
end

# TODO maybe define Hypatia MOI cone equal/hash in Hypatia
function hash_cone(cone::Hypatia.PosSemidefTriSparseCone{I, RealF, C}) where {I, C}
    @assert cone.use_dual
    return hash(I) + hash(C) + hash(cone.side) + hash(cone.row_idxs) + hash(cone.col_idxs)
end

function MOIPajarito.Cones.add_init_cuts(
    cache::PosSemidefTriSparse{<:PrimDual, C},
    opt::Optimizer,
) where {C}
    # add variable bounds for nonnegative diagonal
    # TODO for each offdiag, could add linearization of 2x2 principal matrix condition
    incr = (C == CompF ? 2 : 1)
    i = 1
    for (r, c) in zip(cache.row_idxs, cache.col_idxs)
        if r == c
            JuMP.@constraint(opt.oa_model, cache.oa_s[i] >= 0)
            i += 1
        else
            i += incr
        end
    end
    return
end

# primal cone functions

function MOIPajarito.Cones.get_subp_cuts(
    z::Vector{RealF},
    cache::PosSemidefTriSparse{Prim},
    opt::Optimizer,
)
    cut = dot_expr(z, cache.oa_s, opt)
    return [cut]
end

function MOIPajarito.Cones.get_sep_cuts(
    s::Vector{RealF},
    cache::PosSemidefTriSparse{Prim, C},
    opt::Optimizer,
) where {C}
    # TODO use Arpack for non-dense implementation
    # @warn("no separation oracle implemented for PosSemidefTriSparse", maxlog = 1)
    Λ = svec_to_smat_sparse(s, cache)
    F = eigen!(Hermitian(Λ, :L), -Inf, -opt.tol_feas)
    isempty(F.values) && return AE[]

    cuts = AE[]
    R_i = similar(Λ)
    for i in 1:length(F.values)
        @views r_i = F.vectors[:, i]
        # cuts from eigendecomposition are svec(rᵢ * rᵢ')
        mul!(R_i, r_i, r_i')
        clean_array!(R_i) && continue
        R_vec_i = smat_to_svec_sparse(R_i, cache)
        cut = dot_expr(R_vec_i, cache.oa_s, opt)
        push!(cuts, cut)
    end
    return cuts
end

# dual cone functions

function MOIPajarito.Cones.get_subp_cuts(
    z::Vector{RealF},
    cache::PosSemidefTriSparse{Dual},
    opt::Optimizer,
)
    cut = dot_expr(z, cache.oa_s, opt)
    return [cut]
end

# helpers
# TODO maybe refac with similar Hypatia functions (which should not have the cone as an argument)

function svec_to_smat_sparse(
    vec::AbstractVector{RealF},
    cache::PosSemidefTriSparse{<:PrimDual, RealF},
)
    @assert length(vec) == length(cache.row_idxs)
    mat = zeros(RealF, cache.side, cache.side)
    @inbounds for (idx, (i, j)) in enumerate(zip(cache.row_idxs, cache.col_idxs))
        x = vec[idx]
        if i != j
            x /= rt2
        end
        mat[i, j] = x
    end
    return mat
end

function svec_to_smat_sparse(
    vec::AbstractVector{RealF},
    cache::PosSemidefTriSparse{<:PrimDual, CompF},
)
    mat = zeros(CompF, cache.side, cache.side)
    idx = 1
    @inbounds for (i, j) in zip(cache.row_idxs, cache.col_idxs)
        if i == j
            mat[i, j] = vec[idx]
            idx += 1
        else
            mat[i, j] = Complex(vec[idx], vec[idx + 1]) / rt2
            idx += 2
        end
    end
    @assert idx == length(vec) + 1
    return mat
end

function smat_to_svec_sparse(
    mat::AbstractMatrix{RealF},
    cache::PosSemidefTriSparse{<:PrimDual, RealF},
)
    vec = zeros(length(cache.row_idxs))
    @inbounds for (idx, (i, j)) in enumerate(zip(cache.row_idxs, cache.col_idxs))
        x = mat[i, j]
        if i != j
            x *= rt2
        end
        vec[idx] = x
    end
    return vec
end

function smat_to_svec_sparse(
    mat::AbstractMatrix{CompF},
    cache::PosSemidefTriSparse{<:PrimDual, CompF},
)
    vec = zeros(2 * length(cache.row_idxs) - cache.side)
    idx = 1
    @inbounds for (i, j) in zip(cache.row_idxs, cache.col_idxs)
        x = mat[i, j]
        if i == j
            vec[idx] = real(x)
            idx += 1
        else
            x *= rt2
            vec[idx] = real(x)
            vec[idx + 1] = imag(x)
            idx += 2
        end
    end
    @assert idx == length(vec) + 1
    return vec
end


text:
using ITensors
using ITensorNetworkAD
using ITensorNetworkAD.ITensorNetworks: Square, bonds

@testset "test lattice" begin
  lattice = Square((2, 3))
  bds = bonds(lattice; periodic=false)
  @test length(bds) == 7
end


text:
# functions for a homotopy that goes from one shock sensor to another

"""
  Data for doing homotopy from one shock capturing scheme to another.
  The homotopy function is:

  R_euler (1-lambda)*R_hard + lambda*R_easy

  where R_euler is the physics residual, excluding shock capturing, R_hard
  is the shock sensor that is difficult to converge, and R_easy is the shock
  sensor that is easy to converge.  It is recommended that the initial
  condition satisfy R_euler + R_easy.
"""
mutable struct SCHomotopy
  sensor_easy::AbstractShockSensor
  opts_orig::Dict{String, Bool}

  function SCHomotopy(mesh, sbp, eqn, opts)

    sensor_easy = createShockSensor(mesh, sbp, opts, opts["homotopy_shock_sensor"])
    opts_orig = Dict{String, Bool}

    return new(sensor_easy, opts_orig)
  end
end

const TermNames = ["addVolumeIntegrals", "addBoundaryIntegrals", "addStabilization", "addFaceIntegrals"]

"""
  Disable all terms of the residual evaluation except shock capturing

  **Inputs/Outputs**

   * obj: `HomotopyData`
   * opts

"""
function disableTerms(obj::HomotopyData, opts)

  for key in TermNames
    obj.opts_orig[key] = opts[key]
    opts[key] = false
  end

  return nothing
end

"""
  Reset the options dictionary to its original state

  **Inputs**

   * obj: `HomotopyData`
   * opts: the options dictionary
"""
function enableTerms(obj::HomotopyData, opts)

  for (key, val) in obj.opts_orig
    opts[key] = val
  end
end

"""
  Evaluates the homotopy function R_euler + (1-lambda)*R_hard + lambda*R_easy

  This function assumes parallel communication has already started
"""
function evalHomotopyFunction_sc(obj::HomotopyData, mesh, sbp, eqn, opts, t)
# this probably doesn't make sense if physics_func != evalResidual

  lambda = obj.lambda

  # evaluate R_euler + (1-lambda)R_hard and store it in res_homotopy
  setShockSensorAlpha(eqn, 1-lambda)
  obj.physics_func(mesh, sbp, eqn, opts, t)
  setShockSensorAlpha(eqn, 1.0)  # just in case
  res_homotopy = copy(eqn.res); fill!(eqn.res, 0)

  # evaluate lambda * R_easy and store it in eqn.res
  sensor_orig = getShockSensor(eqn)
  sensor_easy = obj.sensor_easy
  setShockSensor(eqn, obj.sensor_easy)
  setShockSensorAlpha(eqn, lambda)
  disableTerms(obj, opts)
  obj.physics_func(mesh, sbp, eqn, opts, t)

  # reset things
  setShockSensorAlpha(eqn, 1)
  enableTerms(obj, opts)
  setShockSensor(eqn, sensor_orig)

  # combine res_homotopy and eqn.res
  @simd for i=1:length(eqn.res)
    eqn.res[i] += res_homotopy[i]
  end

  return nothing
end

"""
  Evaluates the derivative of the homotopy function with respect to lambda
"""
function calcdHdLambda_sc(obj::HomotopyData, mesh, sbp, eqn, opts, t, res_vec)

  fill!(eqn.res, 0)
  res_homotopy = copy(eqn.res)

  disableTerms(obj, opts)

  # evaluate -1*R_hard
  setShockSensorAlpha(eqn, -1)
  obj.physics_func(mesh, sbp, eqn, opts, t)
  copy!(res_homotopy, eqn.res)
  setShockSensorAlpha(eqn, 1)
  sensor_orig = getShockSensor(eqn)

  # evaluate 1*R_easy
  setShockSensor(eqn, obj.sensor_easy)
  setShockSensorAlpha(eqn, 1)
  obj.physics_func(mesh, sbp, eqn, opts, t)

  @simd for i=1:length(eqn.res)
    eqn.res[i] += res_homotopy[i]
  end

  array3DTo1D(mesh, sbp, eqn, opts, eqn.res, res_vec)

  # reset things
  setShockSensor(eqn, sensor_orig)
  enableTerms(obj, opts)

  return nothing
end

"""
  Computes Jacobian of homotopy function with respect to q
"""
function evalJacobian_homotopy_sc(lo::HomotopyMatLO, mesh, sbp, eqn, opts,
                                  ctx_residual, t)

  hdata = lo.hdata
  lambda = hdata.lambda

  # evaluate jacobian of R_euler + (1-lambda)*R_hard
  setShockSensorAlpha(eqn, 1-lambda)
  calcLinearOperator(lo.lo_inner, mesh, sbp, eqn, opts, ctx_residual, t)
  setShockSensorAlpha(eqn, 1)

  # evaluate Jacobian of lambda*R_easy
  disableTerms(hdata, opts)
  sensor_orig = getShockSensor(eqn)
  setShockSensor(eqn, hdata.sensor_easy)
  setShockSensorAlpha(eqn, lambda)

  A = getBaseLO(lo).A
  physicsJac(mesh, sbp, eqn, opts, A, ctx_residual, t, zero_jac=false)

  # reset things
  setShockSensorAlpha(eqn, 1)
  setShockSensor(eqn, sensor_orig)
  enableTerms(hdata, opts)

  return nothing
end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

## dummy stub for https://github.com/JuliaBinaryWrappers/libblastrampoline_jll.jl

baremodule libblastrampoline_jll
using Base, Libdl
Base.Experimental.@compiler_options compile=min optimize=0 infer=false

const PATH_list = String[]
const LIBPATH_list = String[]

export libblastrampoline

# These get calculated in __init__()
const PATH = Ref("")
const LIBPATH = Ref("")
artifact_dir = ""
libblastrampoline_handle = C_NULL
libblastrampoline_path = ""

const libblastrampoline = if Sys.iswindows()
    "libblastrampoline.dll"
elseif Sys.isapple()
    "@rpath/libblastrampoline.dylib"
else
    "libblastrampoline.so"
end

function __init__()
    global libblastrampoline_handle = dlopen(libblastrampoline)
    global libblastrampoline_path = dlpath(libblastrampoline_handle)
    global artifact_dir = dirname(Sys.BINDIR)
    LIBPATH[] = dirname(libblastrampoline_path)
    push!(LIBPATH_list, LIBPATH[])
end

# JLLWrappers API compatibility shims.  Note that not all of these will really make sense.
# For instance, `find_artifact_dir()` won't actually be the artifact directory, because
# there isn't one.  It instead returns the overall Julia prefix.
is_available() = true
find_artifact_dir() = artifact_dir
dev_jll() = error("stdlib JLLs cannot be dev'ed")
best_wrapper = nothing
get_libblastrampoline_path() = libblastrampoline_path

end  # module libblastrampoline_jll

max_stars_repo_name: jonas-schulze/julia

max_issues_repo_name: jonas-schulze/julia

max_forks_repo_name: jonas-schulze/julia


text:
module Day16
export solvePart1, solvePart2

struct Sample
    before::Array{Int}
    command::Array{Int}
    after::Array{Int}
end

function addr(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    registers[c+1] = registers[a+1] + registers[b+1]
    return registers
end

function addi(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    registers[c+1] = registers[a+1] + b
    return registers
end

function mulr(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    registers[c+1] = registers[a+1] * registers[b+1]
    return registers
end

function muli(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    registers[c+1] = registers[a+1] * b
    return registers
end

function banr(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    registers[c+1] = registers[a+1] & registers[b+1]
    return registers
end

function bani(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    registers[c+1] = registers[a+1] & b
    return registers
end

function borr(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    registers[c+1] = registers[a+1] | registers[b+1]
    return registers
end

function bori(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    registers[c+1] = registers[a+1] | b
    return registers
end

function setr(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    registers[c+1] = registers[a+1]
    return registers
end

function seti(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    registers[c+1] = a
    return registers
end

function gtir(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    if a > registers[b+1]
        registers[c+1] = 1
    else
        registers[c+1] = 0
    end
    return registers
end

function gtri(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    if registers[a+1] > b
        registers[c+1] = 1
    else
        registers[c+1] = 0
    end
    return registers
end

function gtrr(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    if registers[a+1] > registers[b+1]
        registers[c+1] = 1
    else
        registers[c+1] = 0
    end
    return registers
end

function eqir(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    if a == registers[b+1]
        registers[c+1] = 1
    else
        registers[c+1] = 0
    end
    return registers
end

function eqri(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    if registers[a+1] == b
        registers[c+1] = 1
    else
        registers[c+1] = 0
    end
    return registers
end

function eqrr(command::Array{Int}, before::Array{Int})
    a, b, c = command[2:4]
    registers = copy(before)
    if registers[a+1] == registers[b+1]
        registers[c+1] = 1
    else
        registers[c+1] = 0
    end
    return registers
end

function getSamplesInput()
    s = ""
    open("day16input.txt") do f
        s = read(f, String)
    end
    sampleSection, program = split(s, "\n\n\n")
    sampleStrings = split(sampleSection, "\n\n")
    samples = []
    for sample in sampleStrings
        lines = split(sample, '\n')
        before = map(s -> parse(Int, s), split(lines[1][10:end-1], ','))
        command = map(s -> parse(Int, s), split(lines[2]))
        after = map(s -> parse(Int, s), split(lines[3][10:end-1], ','))
        push!(samples, Sample(before, command, after))
    end
    return samples
end

function getProgramInput()
    s = ""
    open("day16input.txt") do f
        s = read(f, String)
    end
    sampleSection, program = split(s, "\n\n\n")
    commands = []
    for line in split(program, '\n')
        command = map(s -> parse(Int, s), split(line))
        if length(command) == 4
            push!(commands, command)
        end
    end
    return commands
end

ops = [addr, addi, mulr, muli, banr, bani, borr, bori, setr, seti, gtir, gtri, gtrr, eqir, eqri, eqrr]

function opMatchesSample(op, sample::Sample)
        result = op(sample.command, sample.before)
        return result == sample.after
end

function solvePart1()
    samples = getSamplesInput()
    ns = []
    for sample in samples
        n = 0
        for op in ops
            if opMatchesSample(op, sample)
                n += 1
            end
        end
        push!(ns, n)
    end
    length(filter(n -> n >= 3, ns))
end

function getMapping(samples)
    possibleMappings = Dict()
    for opcode in 0:15
        possibleMappings[opcode] = Set(ops)
    end

    for sample in samples
        opcode = sample.command[1]
        possibleOps = filter(op -> opMatchesSample(op, sample), ops)
        possibleMappings[opcode] = intersect(possibleMappings[opcode], possibleOps)
    end

    mappings = Dict()
    done = false
    while !done
        certain = filter(kv -> length(kv[2]) == 1, possibleMappings)
        if length(certain) == 0
            done = true
        end
        for kv in certain
            opcode, op = kv
            mappings[opcode] = collect(op)[1]
            ocs = keys(possibleMappings)
            for oc in ocs
                possibleMappings[oc] = setdiff(possibleMappings[oc], op)
            end
            delete!(possibleMappings, opcode)
        end
    end
    return mappings
end

function executeProgram(program, mapping)
    registers = [0, 0, 0, 0]
    for command in program
        op = mapping[command[1]]
        a, b, c = command[2:4]
        registers = op(command, registers)
    end
    return registers
end

function solvePart2()
    samples = getSamplesInput()
    mapping = getMapping(samples)
    program = getProgramInput()
    registers = executeProgram(program, mapping)
end
end


text:
### ============== ### ============== ### ============== ###
##  Numerical simulations of the Discretized version of the
##  inertial spin model Reference: Cavagna...
##  Martin Zumaya Hernandez
##  14 / 11 / 2016
### ============== ### ============== ### ============== ###

using Distributions

### ============== ### ============== ### ============== ###
##                      FLOCK TYPES                       ##
### ============== ### ============== ### ============== ###
"""
    InertialFlock(N, L, v0)
Inertial Flock type
# Constructor Arguments
* N -> number of particles
* L -> size of box
* v0 -> particles speed
# Fields
* pos  -> particles positions
* vel  -> particles velocities
* v_t  -> short-range topological ineraction
* Rij  -> relative distances matrix
* spin -> particles spin
"""
type InertialFlock
    pos ::Array{Array{Float64,1},1}
    vel ::Array{Array{Float64,1},1}
    v_t ::Array{Array{Float64,1},1}
    Rij ::Array{Float64,2}
    spin ::Array{Array{Float64,1},1}

    function InertialFlock(N, L, v0)

        Rij = zeros(Float64, N, N)

        # array of random initial particles' postitions
        pos = [ [2*rand()*L - L, 2*rand()*L - L, 2*rand()*L - L] for i in 1:N ]
        # pos = [ [2*rand()*L - L, 2*rand()*L - L, 0.0] for i in 1:N ]

        # array of particles' velocities
        vel = v0 * [ normalize([2*rand() - 1, 2*rand() - 1, 2*rand() - 1]) for i in 1:N ]
        # vel = v0 * [normalize([1.0, 0.0, 0.0] - [2*δ*rand() - δ, 2*δ*rand() - δ, 2*δ*rand() - δ]) for i in 1:N]

        # short-range topological interactions
        v_t  = [zeros(Float64, 3) for i in 1:N]

        # initialize spins as zero vectors\
        # spin = [zeros(Float64, 3) for i in 1:N]

        # array of  particles' spin (initial random directions)
        spin = [ normalize([2*rand() - 1, 2*rand() - 1, 2*rand() - 1]) for i in 1:N ]
        # # cross product because dot(si, vi) must = 0, then normalize
        map!((x,y) -> normalize(cross(x,y)), spin, spin, vel)

        new(pos, vel, v_t, Rij, spin)
    end
end

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

"""
    InertialExtFlock(N, L, v0)
Inertial Flock type (with non local interactions)
# Constructor Arguments
* N -> number of particles
* L -> size of box
* v0 -> particles speed
# Fields
* pos  -> particles positions
* vel  -> particles velocities
* v_t  -> short-range topological ineraction
* v_nl  -> long-range topological ineraction
* Rij  -> relative distances matrix
* spin -> particles spin
"""
type InertialExtFlock
    pos ::Array{Array{Float64,1},1}
    vel ::Array{Array{Float64,1},1}
    v_t ::Array{Array{Float64,1},1}
    v_nl ::Array{Array{Float64,1},1}
    Rij ::Array{Float64,2}
    spin ::Array{Array{Float64,1},1}

    function InertialExtFlock(N, L, v0)

        Rij = zeros(Float64, N, N)

        # array of random initial particles' postitions
        pos = [ [2*rand()*L - L, 2*rand()*L - L, 2*rand()*L - L] for i in 1:N ]
        # pos = [ [2*rand()*L - L, 2*rand()*L - L, 0.0] for i in 1:N ]

        # array of particles' velocities
        vel = v0 * [ normalize([2*rand() - 1, 2*rand() - 1, 2*rand() - 1]) for i in 1:N ]
        # vel = v0 * [normalize([1.0, 0.0, 0.0] - [2*δ*rand() - δ, 2*δ*rand() - δ, 2*δ*rand() - δ]) for i in 1:N]

        # short-range topological interactions
        v_t  = [zeros(Float64, 3) for i in 1:N]

        # long-range topological interactions
        v_nl = [zeros(Float64, 3) for i in 1:N]

        # initialize spins as zero vectors\
        # spin = [zeros(Float64, 3) for i in 1:N]

        # array of  particles' spin (initial random directions)
        spin = [ normalize([2*rand() - 1, 2*rand() - 1, 2*rand() - 1]) for i in 1:N ]
        # # cross product because dot(si, vi) must = 0, then normalize
        map!((x,y) -> normalize(cross(x,y)), spin, spin, vel)

        new(pos, vel, v_t, v_nl, Rij, spin)
        end
end

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

type SharedInertialExtFlock
    pos  ::SharedArray{Float64}
    vel  ::SharedArray{Float64}
    v_t  ::SharedArray{Float64}
    v_nl ::SharedArray{Float64}
    Rij  ::SharedArray{Float64}
    spin ::SharedArray{Float64}

    function SharedInertialExtFlock(N, L, v0)

        Rij = SharedArray{Float64}(N, N)

        pos  = SharedArray{Float64}(3N) # particles positions
        vel  = SharedArray{Float64}(3N) # array of particles' velocities
        v_t  = SharedArray{Float64}(3N) # short-range topological interactions
        v_nl = SharedArray{Float64}(3N) # long-range topological interactions
        spin = SharedArray{Float64}(3N) # long-range topological interactions

        for i in 1:length(pos)
            pos[i]  = 2*rand()*L - L
            vel[i]  = 2*rand() - 1
            spin[i] = 2*rand() - 1
        end

        for i in 1:3:length(vel)
            norm = sqrt(vel[i]^2 + vel[i+1]^2 + vel[i+2]^2)
            vel[i]   /= norm
            vel[i+1] /= norm
            vel[i+2] /= norm

            vel[i]   *= v0
            vel[i+1] *= v0
            vel[i+2] *= v0

            norm = sqrt(spin[i]^2 + spin[i+1]^2 + spin[i+2]^2)
            spin[i]   /= norm
            spin[i+1] /= norm
            spin[i+2] /= norm

            temp_spin = normalize(cross([spin[i],spin[i+1],spin[i+2]],[vel[i],vel[i+1],vel[i+2]]))

            spin[i]   = temp_spin[1]
            spin[i+1] = temp_spin[2]
            spin[i+2] = temp_spin[3]
        end

        new(pos, vel, v_t, v_nl, Rij, spin)
        end
end

### ============== ### ============== ### ============== ###
##                  SYSTEM'S PARAMETERS                   ##
### ============== ### ============== ### ============== ###
"""
    InertialParameters(χ, J, η, dt, ρ, v0, N, T, n_t, n_nl)
Paramaters for the Inertial Spin Model of Collective Motion
# Fields
* χ  ::Float64 # generalized moment of intertia
* J  ::Float64 # strength of alignment interaction
* η  ::Float64 # friction coefficient
* dt ::Float64  # Integration step
* ρ  ::Float64 # Local initial density
* v0 ::Float64  # Speed
* N  ::Int64  # Number of particles
* T  ::Float64  # generalized temperature
* n_t ::Int64  # number of topological interactions
* n_nl ::Float64 # mean of long-range interactions
"""
immutable InertialParameters
    χ  ::Float64 # generalized moment of intertia
    J  ::Float64 # strength of alignment interaction
    η  ::Float64 # friction coefficient
    dt ::Float64  # Integration step
    ρ  ::Float64 # Local initial density
    v0 ::Float64  # Speed
    N  ::Int64  # Number of particles
    T  ::Float64  # generalized temperature
    d  ::Float64 # spatial dimension
    n_t ::Int64  # number of topological interactions
    κ_dist ::Distributions.Poisson{Float64} # out-degree Poisson Distribution
    # n_nl ::Float64 # mean of non-local interactions

    # default constructor
    InertialParameters() = new(1.25, 0.8, 0.3, 1.0, 1.0, 0.1, 10, 0.5, 3.0, 6, Poisson(1.0))

    # full constructor
    InertialParameters(χ, J, η, dt, ρ, v0, N, T, n_t, n_nl) = new(χ, J, η, dt, ρ, v0, N, T, 3.0, n_t, Poisson(n_nl))
end

### ============== ### ============== ### ============== ###
##      RELATIVE DISTACNCES BETWEEN PARTICLES MATRIX      ##
### ============== ### ============== ### ============== ###
"""
    calc_distance_matrix!(pos, Rij)
Compute the relative distances between particles, store in Rij.
"""
function calc_distance_matrix!(vec, Rij)

    N = size(Rij, 1)

    # compute Rij entries
    for i in 1:N, j in (i+1):N
        Rij[i, j] = norm(vec[i] - vec[j])
        Rij[j, i] = Rij[i, j]
    end
end

### ============== ### ============== ### ============== ###
### COMPUTE RAW RELATIVE (SQUARED) DISTANCES
### ============== ### ============== ### ============== ###

function calc_Rij(R_ij::SharedArray, pos::SharedArray)

    @parallel for i in 1:3:length(pos)

        ri = div(i,3) + 1

        for j in (i+3):3:length(pos)

            rj = div(j,3) + 1

            R_ij[rj,ri] = (pos[i]-pos[j])^2 + (pos[i+1]-pos[j+1])^2 + (pos[i+2]-pos[j+2])^2
        end
    end
end

### ============== ### ============== ### ============== ###
##       COMPUTE LOCAL TOPOLOGICAL INTERACTIONS           ##
### ============== ### ============== ### ============== ###
"""
    calc_topological_interactions!(v_t, Nij, n_t)
Compute and store particles interactions in v_t, given by the adjacency matrix Nij.
#Arguments
* n_t -> number of fixed local topological interactions
"""
function calc_topological_interactions!(vel, v_t, Rij, n_t)

    # compute local topological interaction
    for i in 1:size(Rij,1)
        # v_t[i] = mean( [vel[j] for j in findin(Rij[:,i], sort(Rij[:,i])[2:n_t+1])] )
        v_t[i] = sum( [ vel[j] for j in findin(Rij[:,i], sort(Rij[:,i])[2:n_t+1]) ] )
    end

end

### ============== ### ============== ### ============== ###
##  COMPUTE LOCAL AND NON LOCAL TOPOLOGICAL INTERACTIONS  ##
### ============== ### ============== ### ============== ###
"""
    calc_lr_topological_interactions!(vel, v_t, v_n, Rij, n_t, n_l)
Compute and store particles local topological interactions in v_t and non-local in v_nl, given by the adjacency matrix Nij.
Non-local interactions are chosen at random from the remaining N - (n_t + 1) particles
#Arguments
* n_t -> number of fixed local topological interactions
* n_nl -> number of fixed non-local topological interactions
"""
function calc_lr_topological_interactions!(vel, v_t, v_nl, Rij, n_t, n_nl)

    # compute local topological interaction
    for i in 1:size(Rij,1)
        v_t[i]  = sum( [ vel[j] for j in findin(Rij[:,i], sort(Rij[:,i])[2:n_t+1]) ] )

        # v_nl[i] = sum( [ vel[j] for j in rand(findin(Rij[:,i], sort(Rij[:,i])[n_t+2:end]), n_nl) ] ) # more efficient already knowing n_nl != 0
        n_nl[i] != zero(Int64) ? v_nl[i] = sum( [ vel[j] for j in rand(findin(Rij[:,i], sort(Rij[:,i])[n_t+2:end]), n_nl[i]) ] ) : v_nl[i] = zeros(Float64, length(v_nl[i]))
    end

end

function calc_lr_toplogical_interactions_mean!(vel, v_t, v_nl, Rij, n_t, n_nl)

    # compute local topological interaction
    for i in 1:size(Rij,1)
        v_t[i]  = mean( [ vel[j] for j in findin(Rij[:,i], sort(Rij[:,i])[2:n_t+1]) ] )

        # v_nl[i] = mean( [ vel[j] for j in rand(findin(Rij[:,i], sort(Rij[:,i])[n_t+2:end]), n_nl) ] ) # more efficient already knowing n_nl != 0
        n_nl[i] != zero(Int64) ? v_nl[i] = mean( [ vel[j] for j in rand(findin(Rij[:,i], sort(Rij[:,i])[n_t+2:end]), n_nl[i]) ] ) : v_nl[i] = zeros(Float64, length(v_nl[i]))
    end

end

# function calc_lr_topological_interactions_sh(vel::SharedArray,v_r::SharedArray,v_n::SharedArray,R_ij::SharedArray,N::Int64,k_sh::Int64, κ_dist)
function calc_lr_topological_interactions_sh(vel::SharedArray,v_r::SharedArray,v_n::SharedArray, k_ln::SharedArray, Rij::SharedArray, pars::InertialParameters)

    # f_Rij = Symmetric(R_ij, :L)

    # println(f_Rij)
    println(Rij)

    # for id in first(localindexes(vel)):3:last(localindexes(vel))
    for id in localindexes(k_ln)

        # i = div(id, 3)
        i = id - 1

        v_r[3i+1] = 0.0
        v_r[3i+2] = 0.0
        v_r[3i+3] = 0.0

        v_n[3i+1] = 0.0
        v_n[3i+2] = 0.0
        v_n[3i+3] = 0.0

        # neighbors = sortperm(f_Rij[(i*pars.N)+1:(i+1)*pars.N])
        neighbors = sortperm(view(Rij,(i*pars.N)+1:(i+1)*pars.N))

        # short-range
        for j in neighbors[2:pars.n_t+1]
            v_r[3i+1] += vel[3(j-1)+1]
            v_r[3i+2] += vel[3(j-1)+2]
            v_r[3i+3] += vel[3(j-1)+3]
        end

        # v_r[3i+1] /= k_sh
        # v_r[3i+2] /= k_sh
        # v_r[3i+3] /= k_sh

        # k_ln = rand(pars.κ_dist)

        # possible long range
        if k_ln[id] > 0
            for j in sample(neighbors[pars.n_t+2:end], k_ln[id], replace = true)
                v_n[3i+1] += vel[3(j-1)+1]
                v_n[3i+2] += vel[3(j-1)+2]
                v_n[3i+3] += vel[3(j-1)+3]
            end
            # v_n[3i+1] /= k_ln
            # v_n[3i+2] /= k_ln
            # v_n[3i+3] /= k_ln
        end

    end
end

### ============== ### ============== ### ============== ###
##             DYNAMICAL RULES PER PARTICLE               ##
### ============== ### ============== ### ============== ###
"""
    part_vel_spin_update(vel, v_t, spin, pars, σ)
Dynamical rules for velocity, spin and position update (Updates 1 particle).
# Arguments
* pars -> system's parameters
* σ -> noise related standard deviation
"""
function part_vel_spin_update!(vel, v_t, spin, pars, σ)

    noise = randn(3) * σ

    u_vel = vel + (pars.dt/pars.χ) * cross(spin, vel)

    u_spin =  ( 1.0 - pars.η * pars.dt / pars.χ ) * spin + (pars.J * pars.dt / pars.v0^2) * cross(vel, v_t) + (pars.dt/ pars.v0) * cross(vel, noise)

    spin = u_spin
    vel  = pars.v0  * normalize(u_vel) # codition of constant speed

end

### ============== ### ============== ### ============== ###
##             DYNAMICAL RULES PER PARTICLE               ##
##              WITH NON-LOCAL INTERACTIONS               ##
### ============== ### ============== ### ============== ###
"""
    part_vel_spin_extended_update!(vel, v_t, v_nl, spin, pars, σ)
Dynamical rules for velocity, spin and position update (Updates 1 particle).
# Arguments
* pars -> system's parameters
* σ -> noise related standard deviation
"""
function part_vel_spin_extended_update!(vel, v_t, v_nl, spin, pars, σ)

    noise = randn(3) * σ

    u_vel = vel + (pars.dt/pars.χ) * cross(spin, vel)

    u_spin =  ( 1.0 - pars.η * pars.dt / pars.χ ) * spin[i] + (pars.J * pars.dt / pars.v0^2) * (cross(vel[i], v_t[i]) + cross(vel[i], v_nl[i]) ) + (pars.dt/ pars.v0) * cross(vel[i], noise)

    spin = u_spin
    vel  = pars.v0  * normalize(u_vel) # codition of constant speed

end

### ============== ### ============== ### ============== ###
##             DYNAMICAL RULES (WHOLE SYSTEM)             ##
### ============== ### ============== ### ============== ###
"""
    vel_spin_update!(vel, v_n, spin, pars, σ)
Dynamical rules for velocity, spin and position update.
# Arguments
* pars -> system's parameters
* σ -> noise related standard deviation
"""
function vel_spin_update!(pos, vel, v_t, spin, pars, σ)

    for i in 1:length(vel)
        noise = randn(3) * σ

        u_vel = vel[i] + (pars.dt/pars.χ) * cross(spin[i], vel[i])

        u_spin =  ( 1.0 - pars.η * pars.dt / pars.χ ) * spin[i] + (pars.J * pars.dt / pars.v0^2) * (cross(vel[i], v_t[i])) + (pars.dt/ pars.v0) * cross(vel[i], noise)

        spin[i] = u_spin
        vel[i]  = pars.v0  * normalize(u_vel) # codition of constant speed

        pos[i] += vel[i]*pars.dt
    end
end

### ============== ### ============== ### ============== ###
##             DYNAMICAL RULES (WHOLE SYSTEM)             ##
##              WITH LONG-RANGE INTERACTIONS               ##
### ============== ### ============== ### ============== ###
"""
    vel_spin_nonLocal_update!(vel, v_n, spin, pars, σ)
Dynamical rules for velocity, spin and position update.
# Arguments
* pars -> system's parameters
* σ -> noise related standard deviation
"""
function vel_spin_extended_update!(pos, vel, v_t, v_nl, spin, pars, σ)

    for i in 1:length(vel)
        noise = randn(3) * σ

        u_vel = vel[i] + (pars.dt/pars.χ) * cross(spin[i], vel[i])

        # u_spin =  ( 1.0 - pars.η * pars.dt / pars.χ ) * spin[i] + (pars.J * pars.dt / pars.v0^2) * (cross(vel[i], v_t[i]) + cross(vel[i], v_nl[i]) ) + (pars.dt/ pars.v0) * cross(vel[i], noise)
        u_spin =  ( 1.0 - pars.η * pars.dt / pars.χ ) * spin[i] + (pars.J * pars.dt / pars.v0^2) * (cross(vel[i], v_t[i]) + cross(vel[i], v_nl[i]) ) + (1.0/ pars.v0) * cross(vel[i], noise)

        spin[i] = u_spin
        vel[i]  = pars.v0  * normalize(u_vel) # codition of constant speed

        pos[i] += vel[i]*pars.dt # update positions
    end
end

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

function vel_spin_extended_update_sh(pos::SharedArray, vel::SharedArray, v_t::SharedArray, v_nl::SharedArray, spin::SharedArray, k_ln::SharedArray, pars::InertialParameters, σ::Float64)

    dt_χ = pars.dt/pars.χ

    coeff_1 = 1.0 - (pars.η * (pars.dt / pars.χ))
    coeff_2 = pars.J * pars.dt / pars.v0^2
    coeff_3 = 1.0 / pars.v0

    # println(dt_χ," | " , coeff_1," | ", coeff_2," | ", coeff_3)

    # for id in first(localindexes(vel)):3:last(localindexes(vel))
    for id in localindexes(k_ln)

        # i = div(id, 3)
        i = id - 1

        # u_vel  = zeros(Float64, 3)
        # u_spin = zeros(Float64, 3)

        noise = σ .* randn(3)

        # u_vel = [vel[3i+1],vel[3i+2],vel[3i+3]] + (pars.dt/pars.χ) * cross([spin[3i+1],spin[3i+2],spin[3i+3]], [vel[3i+1],vel[3i+2],vel[3i+3]])

        p_vel  = [vel[3i+1], vel[3i+2], vel[3i+3]]
        p_v_t  = [v_t[3i+1], v_t[3i+2], v_t[3i+3]]
        p_v_nl = [v_nl[3i+1], v_nl[3i+2], v_nl[3i+3]]
        p_spin = [spin[3i+1], spin[3i+2], spin[3i+3]]

        u_vel = p_vel + dt_χ * cross(p_spin, p_vel)
        # u_vel = p_vel + (pars.dt/pars.χ) * cross(p_spin, p_vel)

        u_spin =  coeff_1 * p_spin + coeff_2 * (cross(p_vel, p_v_t) + cross(p_vel, p_v_nl) ) + coeff_3 * cross(p_vel, noise)
        # u_spin =  ( 1.0 - pars.η * pars.dt / pars.χ ) * p_spin + (pars.J * pars.dt / pars.v0^2) * (cross(p_vel, p_v_t) + cross(p_vel, p_v_nl) ) + (1.0/ pars.v0) * cross(p_vel, noise)

        # u_vel[1] = vel[3i+1] + dt_χ * (spin[3i+2]*vel[3i+3]-spin[3i+3]*vel[3i+2])
        # u_vel[2] = vel[3i+2] + dt_χ * (spin[3i+3]*vel[3i+1]-spin[3i+1]*vel[3i+3])
        # u_vel[3] = vel[3i+3] + dt_χ * (spin[3i+1]*vel[3i+2]-spin[3i+2]*vel[3i+1])

        # u_spin =  ( 1.0 - pars.η * pars.dt / pars.χ ) * [spin[3i+1],spin[3i+2],spin[3i+3]] + (pars.J * pars.dt / pars.v0^2) * (cross([vel[3i+1],vel[3i+2],vel[3i+3]], [v_t[3i+1],v_t[3i+2],v_t[3i+3]]) + cross([vel[3i+1],vel[3i+2],vel[3i+3]], [v_nl[3i+1],v_nl[3i+2],v_nl[3i+3]]) ) + (1.0/ pars.v0) * cross([vel[3i+1],vel[3i+2],vel[3i+3]], noise)

        # u_spin[1] = coeff_1 * spin[3i+1] + coeff_2 * ((vel[3i+2]*v_t[3i+3]-vel[3i+3]*v_t[3i+2]) + (vel[3i+2]*v_nl[3i+3]-vel[3i+3]*v_nl[3i+2]) ) + coeff_3 * (vel[3i+2]*noise[3]-vel[3i+3]*noise[2])
        # u_spin[2] = coeff_1 * spin[3i+2] + coeff_2 * ((vel[3i+3]*v_t[3i+1]-vel[3i+1]*v_t[3i+3]) + (vel[3i+3]*v_nl[3i+1]-vel[3i+1]*v_nl[3i+3]) ) + coeff_3 * (vel[3i+3]*noise[1]-vel[3i+1]*noise[3])
        # u_spin[3] = coeff_1 * spin[3i+3] + coeff_2 * ((vel[3i+1]*v_t[3i+2]-vel[3i+2]*v_t[3i+1]) + (vel[3i+1]*v_nl[3i+2]-vel[3i+2]*v_nl[3i+1]) ) + coeff_3 * (vel[3i+1]*noise[2]-vel[3i+2]*noise[1])

        # println(u_spin)

        spin[3i+1] = u_spin[1]
        spin[3i+2] = u_spin[2]
        spin[3i+3] = u_spin[3]

        normalize!(u_vel)  # codition of constant speed

        vel[3i+1] = pars.v0 * u_vel[1]
        vel[3i+2] = pars.v0 * u_vel[2]
        vel[3i+3] = pars.v0 * u_vel[3]

        # println(u_vel)
        # println([vel[3i+1], vel[3i+2], vel[3i+3]])

        pos[3i+1] += vel[3i+1] * pars.dt # update positions
        pos[3i+2] += vel[3i+2] * pars.dt # update positions
        pos[3i+3] += vel[3i+3] * pars.dt # update positions
    end
end

### ============== ### ============== ### ============== ###
##                     SYSTEM EVOLUTION                   ##
### ============== ### ============== ### ============== ###
"""
    evolve_inertial_system(pos, vel, v_t, spin, Rij, pars, σ)
Time evolution of the system
"""
function evolve_system(pos, vel, v_t, spin, Rij, pars, σ)

    ### COMPUTE RELATIVE DISTANCES
    calc_distance_matrix!(pos, Rij)

    ### COMPUTE INTERACTIONS
    calc_topological_interactions!(vel, v_t, Rij, pars.n_t)

    ### SPIN UPDATE
    vel_spin_update!(pos, vel, v_t, spin, pars, σ)

    # ### POSITION UPDATE
    # map!( (p,v) -> p + pars.dt * v, pos, pos, vel )

end

### ============== ### ============== ### ============== ###
##                     SYSTEM EVOLUTION                   ##
##               WITH NON LOCAL INTERACTIONS              ##
### ============== ### ============== ### ============== ###
"""
    evolve_nonLocal_inertial_system(pos, vel, v_t, v_nl, spin, nij, pars, σ)
Time evolution of the system, with non-local interactions
# Arguments
* pos -> particles' positions
* vel -> particles' velocities
* v_r -> local interactions signal
* v_n -> non local interaction signal
* sp_Nij -> non local sparse adjacency matrix
* r0 -> local metric interaction range
* η -> noise intensity
* ω -> interactions relative weight
"""
function evolve_extended_system(pos, vel, v_t, v_nl, spin, Rij, pars, σ)

    ### COMPUTE RELATIVE DISTANCES
    calc_distance_matrix!(pos, Rij)

    n_nl = rand(pars.κ_dist, pars.N)

    ### COMPUTE INTERACTIONS
    calc_lr_topological_interactions!(vel, v_t, v_nl, Rij, pars.n_t, n_nl)

    ### SPIN UPDATE
    vel_spin_extended_update!(pos, vel, v_t, v_nl, spin, pars, σ)

    # ### POSITION UPDATE
    # map!( (p,v) -> p + pars.dt * v, pos, pos, vel )

end

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

function evolve_extended_system_sh(pos::SharedArray, vel::SharedArray, v_t::SharedArray, v_nl::SharedArray, spin::SharedArray, k_nl::SharedArray, Rij::SharedArray, pars::InertialParameters, σ::Float64)

    ### COMPUTE RELATIVE DISTANCES
    calc_Rij(pos, Rij)

    # k_nl = SharedArray(rand(pars.κ_dist, pars.N))
    # f_Rij = SharedArray(Symmetric(fetch(Rij), :L))

    ### COMPUTE INTERACTIONS
    @sync begin
        for p in workers()
            @async remotecall_wait(calc_lr_topological_interactions_sh, p, vel, v_t, v_nl, k_nl, Rij, pars)
        end
    end

    ### PARTICLES UPDATE
    @sync begin
        for p in workers()
            @async remotecall_wait(vel_spin_extended_update_sh, p, pos, vel, v_t, v_nl, spin, k_nl, pars, σ)
        end
    end
end

### ============== ### ============== ### ============== ###
##                 OUTPUT DATA STRUCTURE                  ##
### ============== ### ============== ### ============== ###
"""
    set_output_data_structure(path, N, η, T))
Set up folders for output data
# Arguments
* N -> numer of particles
* η -> friction coefficient
* T -> temperature (noise)
"""
function set_output_data_structure(path, N, η, T)

    parent_folder_path = "$(homedir())/art_DATA/$(path)"
    folder_path        = parent_folder_path * "/DATA/data_N_$(N)"

    reps_path = folder_path * "/eta_$(η)/eta_$(η)_T_$(T)"

    try
        mkdir("$(homedir())/art_DATA")
    catch error
        println("Main data folder already exists")
    end

    try
        mkdir(parent_folder_path)
    catch error
        println("Parent folder already exists")
    end

    try
        mkdir(parent_folder_path * "/DATA")
    catch error
        println("Parent data folder already exists")
    end

    try
        mkdir(folder_path)
    catch error
        println("Folder already exists")
    end

    try
        mkdir(folder_path * "/eta_$(η)")
    catch error
        println("Parent folder already exists")
    end

    try
        mkdir(reps_path)
    catch error
        println("Parameter folder already exists")
    end

    return reps_path
end

### ============== ### ============== ### ============== ###
"""
    set_output_data_structure_inertial_lr(path, N, η, T, n_nl))
Set up folders for output data
# Arguments
* N -> numer of particles
* η -> friction coefficient
* T -> temperature (noise)
* n_nl -> non-local connectivity
"""
function set_output_data_structure_lr(path, N, η, T, n_nl)

    parent_folder_path = "$(homedir())/art_DATA/$(path)"
    folder_path        = parent_folder_path * "/DATA/data_N_$(N)"

    reps_path = folder_path * "/eta_$(η)/eta_$(η)_T_$(T)_nl_$(n_nl)"

    try
        mkdir("$(homedir())/art_DATA")
    catch error
        println("Main data folder already exists")
    end

    try
        mkdir(parent_folder_path)
    catch error
        println("Parent folder already exists")
    end

    try
        mkdir(parent_folder_path * "/DATA")
    catch error
        println("Parent data folder already exists")
    end

    try
        mkdir(folder_path)
    catch error
        println("Folder already exists")
    end

    try
        mkdir(folder_path * "/eta_$(η)")
    catch error
        println("Parent folder already exists")
    end

    try
        mkdir(reps_path)
    catch error
        println("Parameter folder already exists")
    end

    return reps_path
end


text:
@testset "507.perfect-number.jl" begin
    @test perfect_number(28)
    @test perfect_number(7) == false
end


text:
@testset "Elementwise" begin
    @testset "Dense" begin
        @testset "Arithmetic Semiring" begin
            #dense first
            Y = GBMatrix(rand(-10.0:0.05:10.0, 10))
            X = GBMatrix(rand(-10.0:0.05:10.0, 10))
            test_frule(eadd, X, Y)
            test_frule(eadd, X, Y, BinaryOps.PLUS)
            test_rrule(eadd, X, Y)
            test_rrule(eadd, X, Y, BinaryOps.PLUS)
            test_frule(emul, X, Y)
            test_frule(emul, X, Y, BinaryOps.TIMES)
            test_rrule(emul, X, Y)
            test_rrule(emul, X, Y, BinaryOps.TIMES)
        end
    end

    @testset "Sparse" begin
        @testset "Arithmetic Semiring" begin
            Y = GBMatrix(sprand(10, 0.5)) #using matrix for now until I work out transpose(v::GBVector)
            X = GBMatrix(sprand(10, 0.5))
            test_frule(eadd, X, Y)
            test_frule(eadd, X, Y, BinaryOps.PLUS)
            test_rrule(eadd, X, Y)
            test_rrule(eadd, X, Y, BinaryOps.PLUS)
            test_frule(emul, X, Y)
            test_frule(emul, X, Y, BinaryOps.TIMES)
            test_rrule(emul, X, Y)
            test_rrule(emul, X, Y, BinaryOps.TIMES)
        end
    end
end


text:
using LinearAlgebra

import ClimaCore:
    Domains, Fields, Geometry, Meshes, Operators, Spaces, Topologies

import QuadGK
using OrdinaryDiffEq: ODEProblem, solve, SSPRK33

import Logging
import TerminalLoggers
Logging.global_logger(TerminalLoggers.TerminalLogger())

# This example solves the shallow-water equations on a cubed-sphere manifold.
# This file contains five test cases:
# - One, called "steady_state", reproduces Test Case 2 in Williamson et al,
#   "A standard test set for numerical approximations to the shallow water
#   equations in spherical geometry", 1992. This test case gives the steady-state
#   solution to the non-linear shallow water equations. It consists of solid
#   body rotation or zonal flow with the corresponding geostrophic height field.
#   This can be run with an angle α that represents the angle between the north
#   pole and the center of the top cube panel.
# - A second one, called "steady_state_compact", reproduces Test Case 3 in the same
#   reference paper. This test case gives the steady-state solution to the
#   non-linear shallow water equations with nonlinear zonal geostrophic flow
#   with compact support.
# - A third one, called "mountain", reproduces Test Case 5 in the same
#   reference paper. It represents a zonal flow over an isolated mountain,
#   where the governing equations describe a global steady-state nonlinear
#   zonal geostrophic flow, with a corresponding geostrophic height field over
#   a non-uniform reference surface h_s.
# - A fourth one, called "rossby_haurwitz", reproduces Test Case 6 in the same
#   reference paper. It represents the solution of the nonlinear barotropic
#   vorticity equation on the sphere
# - A fifth one, called "barotropic_instability", reproduces the test case in
#   Galewsky et al, "An initial-value problem for testing numerical models of
#   the global shallow-water equations", 2004 (also in Sec. 7.6 of Ullirch et al,
#   "High-order ﬁnite-volume methods for the shallow-water equations on
#   the sphere", 2010). This test case consists of a zonal jet with compact
#   support at a latitude of 45°. A small height disturbance is then added,
#   which causes the jet to become unstable and collapse into a highly vortical
#   structure.

# Physical parameters needed
const R = 6.37122e6
const Ω = 7.292e-5
const g = 9.80616
const D₄ = 1.0e16 # hyperdiffusion coefficient

# Test case specifications
const test_name = get(ARGS, 1, "steady_state") # default test case to run
const test_angle_name = get(ARGS, 2, "alpha0") # default test case to run
const steady_state_test_name = "steady_state"
const steady_state_compact_test_name = "steady_state_compact"
const mountain_test_name = "mountain"
const rossby_haurwitz_test_name = "rossby_haurwitz"
const barotropic_instability_test_name = "barotropic_instability"
const alpha0_test_name = "alpha0"
const alpha30_test_name = "alpha30"
const alpha45_test_name = "alpha45"
const alpha60_test_name = "alpha60"

# Test-specific physical parameters
if test_angle_name == alpha30_test_name
    const α = 30.0
elseif test_angle_name == alpha45_test_name
    const α = 45.0
elseif test_angle_name == alpha60_test_name
    const α = 60.0
else # default test case, α = 0.0
    const α = 0.0
end

if test_name == mountain_test_name
    const u0 = 20.0
    const h0 = 5960
    const a = 20.0 # radius of conical mountain
    const λc = 90.0 # center of mountain long coord, shifted by 180 compared to the paper, because our λ ∈ [-180, 180] (in the paper it was 270, with λ ∈ [0, 360])
    const ϕc = 30.0 # center of mountain lat coord
    const h_s0 = 2e3
elseif test_name == rossby_haurwitz_test_name
    const a = 4.0
    const h0 = 8.0e3
    const ω = 7.848e-6
    const K = 7.848e-6
elseif test_name == steady_state_compact_test_name
    const u0 = 2 * pi * R / (12 * 86400)
    const h0 = 2.94e4 / g
    const ϕᵦ = -30.0
    const ϕₑ = 90.0
    const xₑ = 0.3
elseif test_name == barotropic_instability_test_name
    const u_max = 80.0
    const xₑ = 0.3
    const αₚ = 19.09859
    const βₚ = 3.81971
    const h0 = 10158.18617 # value for initial height from Tempest https://github.com/paullric/tempestmodel/blob/master/test/shallowwater_sphere/BarotropicInstabilityTest.cpp#L86
    const h_hat = 120.0
    const ϕ₀ = 25.71428
    const ϕ₁ = 64.28571
    const ϕ₂ = 45.0
    const eₙ = exp(-4.0 / (deg2rad(ϕ₁) - deg2rad(ϕ₀))^2)
else # default case, steady-state test case
    const u0 = 2 * pi * R / (12 * 86400)
    const h0 = 2.94e4 / g
end

# Plot variables and auxiliary function
ENV["GKSwstype"] = "nul"
using ClimaCorePlots, Plots
Plots.GRBackend()
dirname = "cg_sphere_shallowwater_$(test_name)"
dirname = "$(dirname)_$(test_angle_name)"
path = joinpath(@__DIR__, "output", dirname)
mkpath(path)

function linkfig(figpath, alt = "")
    # Buildkite-agent upload figpath
    # Link figure in logs if we are running on CI
    if get(ENV, "BUILDKITE", "") == "true"
        artifact_url = "artifact://$figpath"
        print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
    end
end

# Set up discretization
ne = 9 # the rossby_haurwitz test case's initial state has a singularity at the pole. We avoid it by using odd number of elements
Nq = 4

domain = Domains.SphereDomain(R)
mesh = Meshes.EquiangularCubedSphere(domain, ne)
grid_topology = Topologies.Topology2D(mesh)
quad = Spaces.Quadratures.GLL{Nq}()
space = Spaces.SpectralElementSpace2D(grid_topology, quad)

coords = Fields.coordinate_field(space)

# Definition of Coriolis parameter
if test_name == rossby_haurwitz_test_name
    f = map(Fields.local_geometry_field(space)) do local_geometry
        coord = local_geometry.coordinates

        ϕ = coord.lat
        λ = coord.long

        f = 2 * Ω * sind(ϕ)

        # Technically this should be a WVector, but since we are only in a 2D space,
        # WVector, Contravariant3Vector, Covariant3Vector are all equivalent.
        # This _won't_ be true in 3D however!
        Geometry.Contravariant3Vector(f)
    end
else # all other test cases share the same Coriolis parameter
    f = map(Fields.local_geometry_field(space)) do local_geometry
        coord = local_geometry.coordinates
        ϕ = coord.lat
        λ = coord.long

        f = 2 * Ω * (-cosd(λ) * cosd(ϕ) * sind(α) + sind(ϕ) * cosd(α))

        # Technically this should be a WVector, but since we are only in a 2D space,
        # WVector, Contravariant3Vector, Covariant3Vector are all equivalent.
        # This _won't_ be true in 3D however!
        Geometry.Contravariant3Vector(f)
    end
end

# Definition of bottom surface topography field
if test_name == mountain_test_name # define the non-uniform reference surface h_s
    h_s = map(Fields.coordinate_field(space)) do coord
        ϕ = coord.lat
        λ = coord.long
        r = sqrt(min(a^2, (λ - λc)^2 + (ϕ - ϕc)^2)) # positive branch
        h_s = h_s0 * (1 - r / a)
    end
else
    h_s = zeros(space)
end

# Set initial condition
if test_name == rossby_haurwitz_test_name
    y0 = map(Fields.local_geometry_field(space)) do local_geometry
        coord = local_geometry.coordinates
        ϕ = coord.lat
        λ = coord.long

        A =
            ω / 2 * (2 * Ω + ω) * cosd(ϕ)^2 +
            1 / 4 *
            K^2 *
            cosd(ϕ)^(2 * a) *
            ((a + 1) * cosd(ϕ)^2 + (2 * a^2 - a - 2) - 2 * a^2 * cosd(ϕ)^-2)
        B =
            2 * (Ω + ω) * K / (a + 1) / (a + 2) *
            cosd(ϕ)^a *
            ((a^2 + 2 * a + 2) - (a + 1)^2 * cosd(ϕ)^2)
        C = 1 / 4 * K^2 * cosd(ϕ)^(2 * a) * ((a + 1) * cosd(ϕ)^2 - (a + 2))

        h =
            h0 +
            (R^2 * A + R^2 * B * cosd(a * λ) + R^2 * C * cosd(2 * a * λ)) / g

        uλ =
            R * ω * cosd(ϕ) +
            R * K * cosd(ϕ)^(a - 1) * (a * sind(ϕ)^2 - cosd(ϕ)^2) * cosd(a * λ)
        uϕ = -R * K * a * cosd(ϕ)^(a - 1) * sind(ϕ) * sind(a * λ)

        u = Geometry.transform(
            Geometry.Covariant12Axis(),
            Geometry.UVVector(uλ, uϕ),
            local_geometry,
        )
        return (h = h, u = u)
    end
elseif test_name == steady_state_compact_test_name
    y0 = map(Fields.local_geometry_field(space)) do local_geometry
        coord = local_geometry.coordinates

        ϕ = coord.lat
        λ = coord.long

        if α == 0.0
            ϕprime = ϕ
            λprime = λ
        else
            ϕprime = asind(sind(ϕ) * cosd(α) - cosd(ϕ) * cosd(λ) * sind(α))
            λprime = asind(sind(λ) * cosd(ϕ) / cosd(ϕprime)) # for alpha45, this experiences numerical precision issues. The test case is designed for either alpha0 or alpha60

            # Temporary angle to ensure λprime is in the right quadrant
            λcond = cosd(α) * cosd(λ) * cosd(ϕ) + sind(α) * sind(ϕ)

            # If λprime is not in the right quadrant, adjust
            if λcond < 0.0
                λprime = -λprime - 180.0 # shifted by 180 compared to the paper, because our λ ∈ [-180, 180]
            end
            if λprime < -180.0
                λprime += 360.0
            end
        end

        # Set auxiliary function needed for initial state of velocity field
        b(x) = x ≤ 0.0 ? 0.0 : exp(-x^(-1))

        x(ϕprime) = xₑ * (ϕprime - ϕᵦ) / (ϕₑ - ϕᵦ)
        uλprime(ϕprime) =
            u0 * b(x(ϕprime)) * b(xₑ - x(ϕprime)) * exp(4.0 / xₑ)
        uϕprime = 0.0

        # Set integral needed for height initial state
        h_int(γ) =
            abs(γ) < 90.0 ?
            (2 * Ω * sind(γ) + uλprime(γ) * tand(γ) / R) * uλprime(γ) : 0.0

        # Set initial state for height field
        h =
            h0 - (R / g) * (pi / 180.0) * QuadGK.quadgk(h_int, -90.0, ϕprime)[1]

        # Set initial state for velocity field
        uϕ = -(uλprime(ϕprime) * sind(α) * sind(λprime)) / cosd(ϕ)
        if abs(cosd(λ)) < 1e-13
            if abs(α) > 1e-13
                if cosd(λ) > 0.0
                    uλ = -uϕ * cosd(ϕ) / tand(α)
                else
                    uλ = uϕ * cosd(ϕ) / tand(α)
                end
            else
                uλ = uλprime(ϕprime)
            end
        else
            uλ =
                (uϕ * sind(ϕ) * sind(λ) + uλprime(ϕprime) * cosd(λprime)) /
                cosd(λ)
        end

        u = Geometry.transform(
            Geometry.Covariant12Axis(),
            Geometry.UVVector(uλ, uϕ),
            local_geometry,
        )

        return (h = h, u = u)
    end
elseif test_name == barotropic_instability_test_name
    y0 = map(Fields.local_geometry_field(space)) do local_geometry
        coord = local_geometry.coordinates

        ϕ = coord.lat
        λ = coord.long

        if α == 0.0
            ϕprime = ϕ
        else
            ϕprime = asind(sind(ϕ) * cosd(α) - cosd(ϕ) * cosd(λ) * sind(α))
        end

        # Set initial state of velocity field
        uλprime(ϕprime) =
            (u_max / eₙ) *
            exp(1.0 / (deg2rad(ϕprime - ϕ₀) * deg2rad(ϕprime - ϕ₁))) *
            (ϕ₀ < ϕprime < ϕ₁)
        uϕprime = 0.0

        # Set integral needed for height initial state
        h_int(γ) =
            abs(γ) < 90.0 ?
            (2 * Ω * sind(γ) + uλprime(γ) * tand(γ) / R) * uλprime(γ) : 0.0

        # Set initial state for height field
        h =
            h0 - (R / g) * (pi / 180.0) * QuadGK.quadgk(h_int, -90.0, ϕprime)[1]

        if λ > 0.0
            λ -= 360.0
        end
        if λ < -360.0 || λ > 0.0
            @info "Invalid longitude value"
        end

        # Add height perturbation
        h += h_hat * cosd(ϕ) * exp(-(λ^2 / αₚ^2) - ((ϕ₂ - ϕ)^2 / βₚ^2))

        uλ = uλprime(ϕprime)
        uϕ = uϕprime

        u = Geometry.transform(
            Geometry.Covariant12Axis(),
            Geometry.UVVector(uλ, uϕ),
            local_geometry,
        )

        return (h = h, u = u)
    end
else # steady-state and mountain test cases share the same form of fields
    y0 = map(Fields.local_geometry_field(space)) do local_geometry
        coord = local_geometry.coordinates

        ϕ = coord.lat
        λ = coord.long
        h =
            h0 -
            (R * Ω * u0 + u0^2 / 2) / g *
            (-cosd(λ) * cosd(ϕ) * sind(α) + sind(ϕ) * cosd(α))^2
        uλ = u0 * (cosd(α) * cosd(ϕ) + sind(α) * cosd(λ) * sind(ϕ))
        uϕ = -u0 * sind(α) * sind(λ)

        u = Geometry.transform(
            Geometry.Covariant12Axis(),
            Geometry.UVVector(uλ, uϕ),
            local_geometry,
        )

        return (h = h, u = u)
    end
end

function rhs!(dydt, y, parameters, t)
    f = parameters.f
    h_s = parameters.h_s

    div = Operators.Divergence()
    wdiv = Operators.WeakDivergence()
    grad = Operators.Gradient()
    wgrad = Operators.WeakGradient()
    curl = Operators.Curl()
    wcurl = Operators.WeakCurl()

    # Compute hyperviscosity first
    @. dydt.h = wdiv(grad(y.h))
    @. dydt.u = wgrad(div(y.u)) - Geometry.Covariant12Vector(wcurl(curl(y.u)))

    Spaces.weighted_dss!(dydt)

    @. dydt.h = -D₄ * wdiv(grad(dydt.h))
    @. dydt.u =
        -D₄ *
        (wgrad(div(dydt.u)) - Geometry.Covariant12Vector(wcurl(curl(dydt.u))))

    # Add in pieces
    @. begin
        dydt.h += -wdiv(y.h * y.u)
        dydt.u +=
            -grad(g * (y.h + h_s) + norm(y.u)^2 / 2) + y.u × (f + curl(y.u))
    end
    Spaces.weighted_dss!(dydt)
    return dydt
end

# Set up RHS function
dydt = similar(y0)
rhs!(dydt, y0, (f = f, h_s = h_s), 0.0)

# Solve the ODE
dt = 9 * 60
T = 86400 * 2

prob = ODEProblem(rhs!, y0, (0.0, T), (f = f, h_s = h_s))

haskey(ENV, "CI_PERF_SKIP_RUN") && exit() # for performance analysis

sol = @timev solve(
    prob,
    SSPRK33(),
    dt = dt,
    saveat = dt,
    progress = true,
    adaptive = false,
    progress_message = (dt, u, p, t) -> t,
)

@info "Test case: $(test_name)"
@info "  with α: $(test_angle_name)"
@info "Solution L₂ norm at time t = 0: ", norm(y0.h)
@info "Solution L₂ norm at time t = $(T): ", norm(sol.u[end].h)
@info "Fluid volume at time t = 0: ", sum(y0.h)
@info "Fluid volume at time t = $(T): ", sum(sol.u[end].h)

if test_name == steady_state_test_name ||
   test_name == steady_state_compact_test_name
    # In these cases, we use the IC as the reference exact solution
    @info "L₁ error at T = $(T): ", norm(sol.u[end].h .- y0.h, 1)
    @info "L₂ error at T = $(T): ", norm(sol.u[end].h .- y0.h)
    @info "L∞ error at T = $(T): ", norm(sol.u[end].h .- y0.h, Inf)
    # Pointwise final L₂ error
    Plots.png(
        Plots.plot(sol.u[end].h .- y0.h),
        joinpath(path, "final_height_L2_error.png"),
    )
    linkfig(
        relpath(
            joinpath(path, "final_height_L2_error.png"),
            joinpath(@__DIR__, "../.."),
        ),
        "Absolute error in height",
    )
    # Height errors over time
    relL1err = Array{Float64}(undef, div(T, dt))
    for t in 1:div(T, dt)
        relL1err[t] = norm(sol.u[t].h .- y0.h, 1) / norm(y0.h, 1)
    end
    Plots.png(
        Plots.plot(
            [1:dt:T],
            relL1err,
            xlabel = "time [s]",
            ylabel = "Relative L₁ err",
            label = "",
        ),
        joinpath(path, "HeightRelL1errorVstime.png"),
    )
    linkfig(
        relpath(
            joinpath(path, "HeightRelL1errorVstime.png"),
            joinpath(@__DIR__, "../.."),
        ),
        "Height relative L1 error over time",
    )

    relL2err = Array{Float64}(undef, div(T, dt))
    for t in 1:div(T, dt)
        relL2err[t] = norm(sol.u[t].h .- y0.h) / norm(y0.h)
    end
    Plots.png(
        Plots.plot(
            [1:dt:T],
            relL2err,
            xlabel = "time [s]",
            ylabel = "Relative L₂ err",
            label = "",
        ),
        joinpath(path, "HeightRelL2errorVstime.png"),
    )
    linkfig(
        relpath(
            joinpath(path, "HeightRelL2errorVstime.png"),
            joinpath(@__DIR__, "../.."),
        ),
        "Height relative L2 error over time",
    )

    RelLInferr = Array{Float64}(undef, div(T, dt))
    for t in 1:div(T, dt)
        RelLInferr[t] = norm(sol.u[t].h .- y0.h, Inf) / norm(y0.h, Inf)
    end
    Plots.png(
        Plots.plot(
            [1:dt:T],
            RelLInferr,
            xlabel = "time [s]",
            ylabel = "Relative L∞ err",
            label = "",
        ),
        joinpath(path, "HeightRelL1InferrorVstime.png"),
    )
    linkfig(
        relpath(
            joinpath(path, "HeightRelLInferrorVstime.png"),
            joinpath(@__DIR__, "../.."),
        ),
        "Height relative L_Inf error over time",
    )
else # In the non steady-state cases, we only plot the latest output of the dynamic problem
    Plots.png(Plots.plot(sol.u[end].h), joinpath(path, "final_height.png"))
    linkfig(
        relpath(
            joinpath(path, "final_height.png"),
            joinpath(@__DIR__, "../.."),
        ),
        "Height field at the final time step",
    )
end


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

mutable struct IoK8sApiCoreV1SecretVolumeSource <: SwaggerModel
    defaultMode::Any # spec type: Union{ Nothing, Int32 } # spec name: defaultMode
    items::Any # spec type: Union{ Nothing, Vector{IoK8sApiCoreV1KeyToPath} } # spec name: items
    optional::Any # spec type: Union{ Nothing, Bool } # spec name: optional
    secretName::Any # spec type: Union{ Nothing, String } # spec name: secretName

    function IoK8sApiCoreV1SecretVolumeSource(;defaultMode=nothing, items=nothing, optional=nothing, secretName=nothing)
        o = new()
        validate_property(IoK8sApiCoreV1SecretVolumeSource, Symbol("defaultMode"), defaultMode)
        setfield!(o, Symbol("defaultMode"), defaultMode)
        validate_property(IoK8sApiCoreV1SecretVolumeSource, Symbol("items"), items)
        setfield!(o, Symbol("items"), items)
        validate_property(IoK8sApiCoreV1SecretVolumeSource, Symbol("optional"), optional)
        setfield!(o, Symbol("optional"), optional)
        validate_property(IoK8sApiCoreV1SecretVolumeSource, Symbol("secretName"), secretName)
        setfield!(o, Symbol("secretName"), secretName)
        o
    end
end # type IoK8sApiCoreV1SecretVolumeSource

const _property_map_IoK8sApiCoreV1SecretVolumeSource = Dict{Symbol,Symbol}(Symbol("defaultMode")=>Symbol("defaultMode"), Symbol("items")=>Symbol("items"), Symbol("optional")=>Symbol("optional"), Symbol("secretName")=>Symbol("secretName"))
const _property_types_IoK8sApiCoreV1SecretVolumeSource = Dict{Symbol,String}(Symbol("defaultMode")=>"Int32", Symbol("items")=>"Vector{IoK8sApiCoreV1KeyToPath}", Symbol("optional")=>"Bool", Symbol("secretName")=>"String")
Base.propertynames(::Type{ IoK8sApiCoreV1SecretVolumeSource }) = collect(keys(_property_map_IoK8sApiCoreV1SecretVolumeSource))
Swagger.property_type(::Type{ IoK8sApiCoreV1SecretVolumeSource }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiCoreV1SecretVolumeSource[name]))}
Swagger.field_name(::Type{ IoK8sApiCoreV1SecretVolumeSource }, property_name::Symbol) =  _property_map_IoK8sApiCoreV1SecretVolumeSource[property_name]

function check_required(o::IoK8sApiCoreV1SecretVolumeSource)
    true
end

function validate_property(::Type{ IoK8sApiCoreV1SecretVolumeSource }, name::Symbol, val)
end


text:
__precompile__()

module PolyChaos

using SpecialFunctions, SparseArrays, Distributions
import LinearAlgebra: I, dot, SymTridiagonal, eigen, issymmetric, pinv
import FFTW: ifft
import Combinatorics: with_replacement_combinations
import Base: show
import AdaptiveRejectionSampling: RejectionSampler, run_sampler!
import Statistics: mean, std, var
import GaussQuadrature: special_eigenproblem!

include("typesAbstract.jl")
include("typesMeasures.jl")
include("typesQuad.jl")
include("typesOrthoPolys.jl")
include("typesTensor.jl")
include("multi_indices.jl")
include("evaluate.jl")
include("quadrature_rules.jl")
include("recurrence_coefficients_monic.jl")
include("stieltjes.jl")
include("densities.jl")
include("auxfuns.jl")
include("scalar_product.jl")
include("tensor.jl")
include("show.jl")
include("polynomial_chaos.jl")

end


text:
############################################################
# Power Layer
############################################################

function forward(backend::GPUBackend, state::PowerLayerState, inputs::Vector{Blob})
  for i = 1:length(inputs)
    input = inputs[i]
    output = state.blobs[i]

    len = length(input)
    data_type = eltype(input)

    # output = input
    copy!(output, input)

    # output *= scale
    if state.layer.scale != 1
      CuBLAS.scal(backend.cublas_ctx, length(output),
          convert(data_type,state.layer.scale), output.ptr, 1)
    end

    if state.layer.shift != 0
      # output += shift
      CuVec.add_scal!(backend, data_type, output.ptr.p, convert(data_type, state.layer.shift), len)
    end

    # output = output ^ power
    if state.layer.power != 1
      if state.layer.power == 2
        CuVec.mul!(backend, data_type, output.ptr.p, output.ptr.p, len)
      else
        CuVec.pow!(backend, data_type, output.ptr.p,
            isinteger(state.layer.power) ? round(Int, state.layer.power) : convert(data_type, state.layer.power),
            len)
      end
    end
  end
end

function backward(backend::GPUBackend, state::PowerLayerState,
    inputs::Vector{Blob}, diffs::Vector{Blob})

  data_type = eltype(inputs[1])
  pow_scale = convert(data_type,state.layer.power * state.layer.scale)
  for i = 1:length(inputs)
    diff = diffs[i]
    if isa(diff, NullBlob)
      continue
    end

    len = length(inputs[i])
    if state.layer.power == 1 || state.layer.scale == 0
      # trivial case, derivative is constant
      fill!(diff, pow_scale)
    else
      input = inputs[i]
      output = state.blobs[i]

      erase!(diff)

      if state.layer.power == 2
        # dO/dI = 2 * scale * (scale * I + shift)
        #       = pow_scale * scale * I + pow_scale * shift
        CuBLAS.axpy(backend.cublas_ctx, length(input), convert(data_type, pow_scale*state.layer.scale),
            input.ptr, 1, diff.ptr, 1)
        if state.layer.shift != 0
          CuVec.add_scal!(backend, data_type, diff.ptr.p, pow_scale * state.layer.shift, len)
        end
      elseif state.layer.shift == 0
        # dO/dI = power * scale * (scale * I) ^ (power - 1)
        #       = power * O / I
        CuBLAS.axpy(backend.cublas_ctx, length(input), convert(data_type,state.layer.power),
            output.ptr, 1, diff.ptr, 1)
        CuVec.div!(backend, data_type, diff.ptr.p, input.ptr.p, len)
      else
        # general case
        # dO/dI = power * scale * (scale * I + shift) ^ (power - 1)
        #       = power * scale * O / (scale * I + shift)
        copy!(diff, input)
        if state.layer.scale != 1
          CuBLAS.scal(backend.cublas_ctx, length(diff),
              convert(data_type,state.layer.scale), diff.ptr, 1)
        end
        CuVec.add_scal!(backend, data_type, diff.ptr.p, state.layer.shift, len)
        CuVec.div2!(backend, data_type, output.ptr.p, diff.ptr.p, len)
        CuBLAS.scal(backend.cublas_ctx, length(diff), pow_scale, diff.ptr, 1)
      end
    end
    CuVec.mul!(backend, data_type, diff.ptr.p, state.blobs_diff[i].ptr.p, len)
  end
end


text:
module IsURL

export isurl, isrelativeurl

# Source: https://github.com/sindresorhus/is-absolute-url (MIT license)
windowsregex = r"^[a-zA-Z]:[\\]"
urlregex = r"^[a-zA-Z][a-zA-Z\d+\-.]*:"

"""
    isurl(str)

Checks if the given string is an absolute URL.

# Examples
```julia-repl
julia> isurl("https://julialang.org")
true

julia> isurl("mailto:someone@example.com")
true

julia> isurl("/foo/bar")
false
```
"""
function isurl(str::AbstractString)
    return !occursin(windowsregex, str) && occursin(urlregex, str)
end

"""
    isrelativeurl(str)

Checks if the given string is a relative URL.

# Examples
```julia-repl
julia> isrelativeurl("../path/to/directory")
true

julia> isrelativeurl("./__file__")
true

julia> isrelativeurl("foo:bar")
false
```
"""
function isrelativeurl(str::AbstractString)
    return !isurl(str)
end

end # module


text:
"""
    map_column_pairs(f, X::Matrix{Float64}[, Y::Matrix{Float64} = X])

Create a matrix by applying function `f` to each pair of columns of input matrices
`X` and `Y`.
"""
function map_column_pairs(f, X::MatF64, Y::MatF64)
    nobs1 = size(X,2)
    nobs2 = size(Y,2)
    D = Array{Float64}(undef, nobs1, nobs2)
    map_column_pairs!(D, f, X, Y)
    return D
end

function map_column_pairs(f, X::MatF64)
    dim, nobsv = size(X)
    D = Array{Float64}(undef, nobsv, nobsv)
    for i in 1:nobsv
        for j in 1:i
            @inbounds D[i,j] = f(X[:,i], X[:,j])
            if i != j; @inbounds D[j,i] = D[i,j]; end;
        end
    end
    return D
end

"""
    map_column_pairs!(D::Matrix{Float64}, f, X::Matrix{Float64}[, Y::Matrix{Float64} = X])

Like [`map_column_pairs`](@ref), but stores the result in `D` rather than a new matrix.
"""
function map_column_pairs!(D::MatF64, f, X::MatF64, Y::MatF64)
    dim, nobs1 = size(X)
    nobs2 = size(Y,2)
    dim == size(Y,1) || throw(ArgumentError("Input observation matrices must have consistent dimensions"))
    size(D,1) == nobs1 || throw(ArgumentError(@sprintf("D has %d rows, while X has %d columns (should be same)",
                                                       size(D,1), nobs1)))
    size(D,2) == nobs2 || throw(ArgumentError(@sprintf("D has %d columns, while Y has %d columns (should be same)",
                                                       size(D,2), nobs2)))
    for i in 1:nobs1, j in 1:nobs2
        @inbounds D[i,j] = f(view(X,:,i), view(Y,:,j))
    end
    return D
end

function map_column_pairs!(D::MatF64, f, X::MatF64)
    dim, nobsv = size(X)
    size(D,1) == nobsv || throw(ArgumentError(@sprintf("D has %d rows, while X has %d columns (should be same)",
                                                       size(D,1), nobsv)))
    size(D,2) == nobsv || throw(ArgumentError(@sprintf("D has %d columns, while X has %d columns (should be same)",
                                                       size(D,2), nobsv)))
    @inbounds for i in 1:nobsv
        for j in 1:i
            D[i,j] = f(view(X,:,i), view(X,:,j))
            if i != j; D[j,i] = D[i,j]; end;
        end
    end
    return D
end


text:
# -*- coding: utf-8 -*-
# +
using Random

using Flux
using Flux.Data:DataLoader
using MLDatasets
using Images
using ProgressMeter
# -

using PyCallChainRules.Torch: TorchModuleWrapper, torch

# +
Base.@kwdef mutable struct Args
    η = 3e-4
    λ = 0
    batchsize=128
    epochs=5
    seed=12345
    use_cuda=false
end

args = Args()
Random.seed!(args.seed)
ENV["DATADEPS_ALWAYS_ACCEPT"] = true
# -

using Flux
using Flux.Data: DataLoader
using MLDatasets
xtrain, ytrain = MLDatasets.MNIST.traindata(Float32)
xtest, ytest = MLDatasets.MNIST.testdata(Float32)
xtrain = Flux.unsqueeze(xtrain, 3) # (28, 28, 60000) -> (28, 28, 1 , 60000)
xtest = Flux.unsqueeze(xtest, 3)   # (28, 28, 10000) -> (28, 28, 1, 10000)
ytrain = Flux.onehotbatch(ytrain, 0:9) # (60000,) -> (10, 60000)
ytest = Flux.onehotbatch(ytest, 0:9)   # (10000,) -> (10, 10000)
train_loader = DataLoader((xtrain, ytrain), batchsize=128, shuffle=true)
test_loader = DataLoader((xtest, ytest),  batchsize=128);

# +
struct LeNet
    cnn_layer
    mlp_layer
    nclasses
end

Flux.@functor LeNet (cnn_layer, mlp_layer) # cnn_layer と mlp_layer が学習パラメータであることを指定する

(net::LeNet)(x) = x |> net.cnn_layer |> Flux.flatten |> net.mlp_layer
# -

function create_model(imsize::Tuple{Int,Int,Int}, nclasses::Int)
    W, H, inC = imsize
    out_conv_size = (W ÷ 4 - 3, H ÷ 4 - 3, 16)
    cnn_layer = Chain(
        TorchModuleWrapper(torch.nn.Conv2d(1, 6, (5,5))),
        x->Flux.relu.(x),
        MaxPool((2,2)),
        TorchModuleWrapper(torch.nn.Conv2d(6, 16, (5,5))),
        x->Flux.relu.(x),
        MaxPool((2,2))
    )
    mlp_layer = Chain(
        TorchModuleWrapper(torch.nn.Linear(prod(out_conv_size), 120)),
        x->Flux.relu.(x),
        TorchModuleWrapper(torch.nn.Linear(120, 84)),
        x->Flux.relu.(x),
        TorchModuleWrapper(torch.nn.Linear(84,nclasses)),
    )
    LeNet(cnn_layer, mlp_layer, nclasses)
end

model = create_model((28, 28, 1), 10) |> f32
ps = Flux.params(model);
opt = Descent(0.05)
loss(ŷ, y) = Flux.Losses.logitcrossentropy(ŷ, y)

first(train_loader)[1] |> size

for e in 1:args.epochs
    @showprogress for (x, y) in train_loader
        gs, = Flux.gradient(model) do m
            ŷ = m(x)
            loss(ŷ, y)
        end
        cnt = 1
        for g in gs
            isnothing(g) && continue
            for lay in g.layers
                isnothing(lay) && continue
                for Δ in lay.params
                    #@show Δ
                    ps[cnt] .-= Flux.Optimise.apply!(opt, ps[cnt], Δ)
                    cnt += 1
                end
            end
        end
    end
    acc = sum(Flux.onecold(model(xtest)) .== Flux.onecold(ytest))
    acc /= size(ytest, 2)
    @info("acc", 100acc, "%")
end

x_test, y_test = first(test_loader);
img = x_test[:, :, :, 1]
label = Flux.onecold(y_test[:, 1]) - 1
println("pred:", Flux.onecold(model(x_test))[1]-1)
println("actual:", label)
Gray.(dropdims(img, dims=3) |> transpose)


text:
using Fermi.Options

@testset "Options" begin
    # Basic check
    @set scf_e_conv 10^-10
    x = @get scf_e_conv
    @test x == 10^-10

    @set basis cc-pvqz
    x = @get basis
    @test x == "cc-pvqz"

    @set basis "6-31g*"
    x = @get basis
    @test x == "6-31g*"

    # Check fetching default
    y = @get cc_max_iter
    @test y == Options.Default["cc_max_iter"]

    # Check error due to invalid key
    @test_throws FermiException @get invalid_options
    @test_throws FermiException @set invalid_option 10

    # Check error due to invalid data type
    @test_throws FermiException @set molstring 13

    # Check attempt to convert data type
    @set scf_max_iter 10.0
    x = @get scf_max_iter
    @test x == 10
    @test_throws FermiException @set cc_max_iter 13.567

    # Check lowercase call (i.e. all options are saved in lowercase)
    @set {
        precision DOUBLE
    }
    x = @get precision
    @test x == "double"

    # Check if reset works
    @reset
    @test Options.Current == Dict()
    @set {
        scf_e_conv 1/10
        cc_max_iter 1
        diis false
    }

    @reset diis
    @test @get diis
    @reset cc_max_iter scf_e_conv
    x = @get scf_e_conv
    y = @get cc_max_iter
    @test x == Options.Default["scf_e_conv"] && y == Options.Default["cc_max_iter"]
    @test_throws FermiException @reset invalid_keyword

    @reset
    x = @capture_out @get
    @test strip(x) == "No user defined keywords found."

    @set {
        basis cc-pvdz
        diis false
        molstring "H 1 2 3\nH 3 4 5"
    }

    # @get (no input)
    x = @capture_out @get
    @test begin
        occursin("Keyword", x)       &&
        occursin("Current Value", x) &&
        occursin("basis", x)         &&
        occursin("diis", x)          &&
        occursin("false", x)         &&
        occursin("molstring", x)     &&
        occursin(r"H\s+?1\.??0*?\s+?2\.??0*?\s+?3\.??0*?", x)
    end

    # Check lookup
    @reset
    x = @capture_out @lookup someinvalidkeyword
    @test strip(x) == "No keywords found containing: someinvalidkeyword"

    @set {
        scf_max_rms 1e-9
        scf_max_iter 50
        molstring "H 1 2 3\nH 3 4 5"
    }

    x = @capture_out @lookup scf_max
    @test begin
        occursin("Keyword", x)       &&
        occursin("Current Value", x) &&
        occursin("scf_max_rms", x)   &&
        occursin("1.00000e-09", x)   &&
        occursin("scf_max_iter", x)  &&
        occursin("50", x)
    end
    x = @capture_out @lookup molstring
    @test begin
        occursin("molstring", x)     &&
        occursin(r"H\s+?1\.??0*?\s+?2\.??0*?\s+?3\.??0*?", x)
    end

    # Test @molecule
    molstring = """
    H 1.0 2.0 3.0
    H 4.0 5.0 6.0"""

    @molecule {
    H 1.0 2.0 3.0
    H 4.0 5.0 6.0
    }
    @test begin
        x = Options.get("molstring")
        occursin(r"H\s+?1\.0*?\s+?2\.0*?\s+?3\.0*?", x) &&
        occursin(r"H\s+?4\.0*?\s+?5\.0*?\s+?6\.0*?", x)
    end
    x = @capture_out showerror(IOContext(stdout, :short => true), Fermi.Options.FermiException("Testing the error!"))
    @test occursin("Testing the error!", x)
end