text:
using SparseArrays
import Base: accumulate
import LinearAlgebra: factorize
export SparseTensor, SparseAssembler,
spdiag, find, spzero, dense_to_sparse, accumulate, assemble, rows, cols,
factorize, solve, trisolve, RawSparseTensor, compress

"""
    SparseTensor

A sparse matrix object. It has two fields

- `o`: internal data structure

- `_diag`: `true` if the sparse matrix is marked as "diagonal".
"""
mutable struct SparseTensor
    o::PyObject
    _diag::Bool
    function SparseTensor(o::PyObject, _diag::Bool=false)
        new(o, _diag)
    end
end

promote_(x::SparseMatrixCSC, y::SparseTensor) =
        (constant(x), y)
promote_(y::SparseTensor, x::SparseMatrixCSC) =
        (y, constant(x))
+(x::SparseMatrixCSC, y::SparseTensor) = +(promote_(x,y)...)
-(x::SparseMatrixCSC, y::SparseTensor) = -(promote_(x,y)...)
*(x::SparseMatrixCSC, y::SparseTensor) = *(promote_(x,y)...)
+(x::SparseTensor, y::SparseMatrixCSC) = +(promote_(x,y)...)
-(x::SparseTensor, y::SparseMatrixCSC) = -(promote_(x,y)...)
*(x::SparseTensor, y::SparseMatrixCSC) = *(promote_(x,y)...)

function Base.:values(o::SparseTensor)
    o.o.values
end

function rows(o::SparseTensor)
    get(o.o.indices',0)+1
end

function cols(o::SparseTensor)
    get(o.o.indices',1)+1
end

"""
    SparseTensor(I::Union{PyObject,Array{T,1}}, J::Union{PyObject,Array{T,1}}, V::Union{Array{Float64,1}, PyObject}, m::Union{S, PyObject, Nothing}=nothing, n::Union{S, PyObject, Nothing}=nothing) where {T<:Integer, S<:Integer}

Constructs a sparse tensor.
Examples:
```
ii = [1;2;3;4]
jj = [1;2;3;4]
vv = [1.0;1.0;1.0;1.0]
s = SparseTensor(ii, jj, vv, 4, 4)
s = SparseTensor(sprand(10,10,0.3))
```
"""
function SparseTensor(I::Union{PyObject,Array{T,1}}, J::Union{PyObject,Array{T,1}},
      V::Union{Array{Float64,1}, PyObject},
     m::Union{S, PyObject, Nothing}=nothing, n::Union{S, PyObject, Nothing}=nothing; is_diag::Bool=false) where {T<:Integer, S<:Integer}
    if isa(I, PyObject) && size(I,2)==2
        return SparseTensor_(I, J, V)
    end
    I, J, V = convert_to_tensor(I, dtype=Int64), convert_to_tensor(J, dtype=Int64), convert_to_tensor(V)
    m, n = convert_to_tensor(m, dtype=Int64), convert_to_tensor(n, dtype=Int64)
    indices = [I J] .- 1
    value = V
    shape = [m;n]
    sp = tf.SparseTensor(indices, value, shape)
    options.sparse.auto_reorder && (sp = tf.sparse.reorder(sp))
    SparseTensor(sp, is_diag)
end

function dense_to_sparse(o::Union{Array, PyObject})
    if isa(o, Array)
        return SparseTensor(sparse(o))
    else
        idx = tf.where(tf.not_equal(o, 0))
        indices, value = convert_to_tensor([idx, tf.gather_nd(o, idx)], [Int64, Float64])
        return SparseTensor(tf.SparseTensor(indices, value, o.get_shape()), false)
    end
end

"""
    find(s::SparseTensor)

Returns the row, column and values for sparse tensor `s`.
"""
function find(s::SparseTensor)
    ind = s.o.indices
    val = s.o.values
    ii = ind'[1,:]
    jj = ind'[2,:]
    ii+1, jj+1, val
end

function Base.:copy(s::SparseTensor)
    t = SparseTensor(tf.SparseTensor(copy(s.o.indices), copy(s.o.values), s.o.dense_shape), copy(s.is_diag))
end

function Base.:eltype(o::SparseTensor)
    return (eltype(o.o.indices), eltype(o.o.values))
end

function SparseTensor_(indices::Union{PyObject,Array{T,2}}, value::Union{PyObject,Array{Float64,1}},
        shape::Union{PyObject,Array{T,1}}; is_diag::Bool=false) where T<:Integer
    indices = convert_to_tensor(indices, dtype=Int64)
    value = convert_to_tensor(value, dtype=Float64)
    shape = convert_to_tensor(shape, dtype=Int64)
    sp = tf.SparseTensor(indices-1, value, shape)
    options.sparse.auto_reorder && (sp = tf.sparse.reorder(sp))
    SparseTensor(sp, is_diag)
end

"""
    RawSparseTensor(indices::Union{PyObject,Array{T,2}}, value::Union{PyObject,Array{Float64,1}},
        m::Union{PyObject,Int64}, n::Union{PyObject,Int64}; is_diag::Bool=false) where T<:Integer

A convenient wrapper for making custom operators. Here `indices` is 0-based.
"""
function RawSparseTensor(indices::Union{PyObject,Array{T,2}}, value::Union{PyObject,Array{Float64,1}},
    m::Union{PyObject,Int64}, n::Union{PyObject,Int64}; is_diag::Bool=false) where T<:Integer
    indices = convert_to_tensor(indices, dtype=Int64)
    value = convert_to_tensor(value, dtype=Float64)
    m = convert_to_tensor(m, dtype = Int64)
    n = convert_to_tensor(n, dtype = Int64)
    shape = [m;n]
    sp = tf.SparseTensor(indices, value, shape)
    options.sparse.auto_reorder && (sp = tf.sparse.reorder(sp))
    SparseTensor(sp, is_diag)
end

"""
    SparseTensor(A::SparseMatrixCSC)
    SparseTensor(A::Array{Float64, 2})

Creates a `SparseTensor` from numerical arrays.
"""
function SparseTensor(A::SparseMatrixCSC)
    rows = rowvals(A)
    vals = nonzeros(A)
    cols = zeros(eltype(rows), length(rows))
    m, n = size(A)
    k = 1
    for i = 1:n
        for j in nzrange(A, i)
            cols[k] = i
            k += 1
        end
    end
    SparseTensor(rows, cols, vals, m, n; is_diag=isdiag(A))
end

constant(o::SparseMatrixCSC) = SparseTensor(o)
constant(o::SparseTensor) = o
SparseTensor(o::SparseTensor) = o

function SparseTensor(A::Array{Float64, 2})
    SparseTensor(sparse(A))
end

function Base.:show(io::IO, s::SparseTensor)
    shape = size(s)
    s1 = shape[1]===nothing ? "?" : shape[1]
    s2 = shape[2]===nothing ? "?" : shape[2]
    print(io, "SparseTensor($s1, $s2)")
end

function Base.:run(o::PyObject, S::SparseTensor, args...; kwargs...)
    indices, value, shape = run(o, S.o, args...; kwargs...)
    sparse(indices[:,1].+1, indices[:,2].+1, value, shape...)
end

"""
    Array(A::SparseTensor)

Converts a sparse tensor `A` to dense matrix.
"""
function Base.:Array(A::SparseTensor)
    ij = A.o.indices
    vv = values(A)
    m, n = size(A)
    sparse_to_dense_ = load_system_op("sparse_to_dense_ad",multiple=false)
    m_, n_ = convert_to_tensor(Any[m,n], [Int64,Int64])
    out = sparse_to_dense_(ij, vv, m_,n_)
    set_shape(out, (m, n))
end

function Base.:size(s::SparseTensor)
    (get(s.o.shape,0).value,get(s.o.shape,1).value)
end

function Base.:size(s::SparseTensor, i::T) where T<:Integer
    get(s.o.shape, i-1).value
end

function PyCall.:+(s::SparseTensor, o::PyObject)
    if size(s)!=size(o)
        error("size $(size(s)) and $(size(o)) does not match")
    end
    out = tf.sparse_add(s.o, o)
    out
end
PyCall.:+(o::PyObject, s::SparseTensor) = s+o
function Base.:-(s::SparseTensor)
    SparseTensor(s.o.indices+1, -s.o.values, s.o.dense_shape, is_diag=s._diag)
end
PyCall.:-(o::PyObject, s::SparseTensor) = o + (-s)
PyCall.:-(s::SparseTensor, o::PyObject) = s + (-o)

Base.:+(s1::SparseTensor, s2::SparseTensor) = SparseTensor(tf.sparse.add(s1.o,s2.o), s1._diag&&s2._diag)
Base.:-(s1::SparseTensor, s2::SparseTensor) = s1 + (-s2)

function Base.:adjoint(s::SparseTensor)
    indices = [s.o.indices'[2,:] s.o.indices'[1,:]]
    sp = tf.SparseTensor(indices, s.o.values, (size(s,2), size(s,1)))
    options.sparse.auto_reorder && (sp = tf.sparse.reorder(sp))
    SparseTensor(sp, s._diag)
end

function PyCall.:*(s::SparseTensor, o::PyObject)
    flag = false
    if length(size(o))==0
        return o * s
    end
    if length(size(o))==1
        flag = true
        o = reshape(o, length(o), 1)
    end
    out = tf.sparse.sparse_dense_matmul(s.o, o)
    if flag
        out = squeeze(out)
    end
    out
end

function Base.:*(s::SparseTensor, o::Array{Float64})
    s*convert_to_tensor(o)
end

function PyCall.:*(o::PyObject, s::SparseTensor)
    if length(size(o))==0
        SparseTensor(tf.SparseTensor(copy(s.o.indices), o*tf.identity(s.o.values), s.o.dense_shape), s._diag)
    else
        tf.sparse.sparse_dense_matmul(s.o, o, adjoint_a=true, adjoint_b=true)'
    end
end

function Base.:*(o::Array{Float64}, s::SparseTensor)
    convert_to_tensor(o)*s
end

function Base.:*(o::Real, s::SparseTensor)
    o = Float64(o)
    SparseTensor(tf.SparseTensor(copy(s.o.indices), o*tf.identity(s.o.values), s.o.dense_shape), s._diag)
end

Base.:*(s::SparseTensor, o::Real) = o*s
Base.:/(s::SparseTensor, o::Real) = (1/o)*s

function _sparse_concate(A1::SparseTensor, A2::SparseTensor, hcat_::Bool)
    m1,n1 = size(A1)
    m2,n2 = size(A2)
    ii1,jj1,vv1 = find(A1)
    ii2,jj2,vv2 = find(A2)
    sparse_concate_ = load_system_op("sparse_concate"; multiple=true)
    ii1,jj1,vv1,m1_,n1_,ii2,jj2,vv2,m2_,n2_ = convert_to_tensor([ii1,jj1,vv1,m1,n1,ii2,jj2,vv2,m2,n2], [Int64,Int64,Float64,Int32,Int32,Int64,Int64,Float64,Int32,Int32])
    ii,jj,vv = sparse_concate_(ii1,jj1,vv1,m1_,n1_,ii2,jj2,vv2,m2_,n2_,constant(hcat_))
    if hcat_
        SparseTensor(ii,jj,vv, m1, n1+n2)
    else
        SparseTensor(ii,jj,vv,m1+m2,n1)
    end
end

function sparse_concate(args::SparseTensor...; hcat_::Bool)
    reduce((x,y)->_sparse_concate(x,y,hcat_), args)
end

Base.:vcat(args::SparseTensor...) = sparse_concate(args...;hcat_=false)
Base.:hcat(args::SparseTensor...) = sparse_concate(args...;hcat_=true)
function Base.:hvcat(rows::Tuple{Vararg{Int}}, values::SparseTensor...)
    r = Array{SparseTensor}(undef, length(rows))
    k = 1
    for i = 1:length(rows)
        r[i] = hcat(values[k:k+rows[i]-1]...)
        k += rows[i]
    end
    vcat(r...)
end

function Base.:lastindex(o::SparseTensor, i::Int64)
    return size(o,i)
end

function Base.:getindex(s::SparseTensor, i1::Union{Integer, Colon, UnitRange{S}, PyObject,Array{S,1}},
    i2::Union{Integer, Colon, UnitRange{T}, PyObject,Array{T,1}}) where {S<:Real,T<:Real}
    squeeze_dims = Int64[]
    if isa(i1, Integer); i1 = [i1]; push!(squeeze_dims, 1); end
    if isa(i2, Integer); i2 = [i2]; push!(squeeze_dims, 2); end
    if isa(i1, UnitRange) || isa(i1, StepRange); i1 = collect(i1); end
    if isa(i2, UnitRange) || isa(i2, StepRange); i2 = collect(i2); end
    if isa(i1, Colon); i1 = collect(1:lastindex(s,1)); end
    if isa(i2, Colon); i2 = collect(1:lastindex(s,2)); end
    if isa(i1, Array{Bool,1}); i1 = findall(i1); end
    if isa(i2, Array{Bool,1}); i2 = findall(i2); end
    m_, n_ = length(i1), length(i2)
    i1 = convert_to_tensor(i1, dtype=Int64)
    i2 = convert_to_tensor(i2, dtype=Int64)
    ii1, jj1, vv1 = find(s)
    m = tf.convert_to_tensor(get(s.o.shape,0),dtype=tf.int64)
    n = tf.convert_to_tensor(get(s.o.shape,1),dtype=tf.int64)
    ss = load_system_op("sparse_indexing"; multiple=true)
    ii2, jj2, vv2 = ss(ii1,jj1,vv1,m,n,i1,i2)
    ret = SparseTensor(ii2, jj2, vv2, m_, n_)
    if length(squeeze_dims)>0
        ret = squeeze(Array(ret), dims=squeeze_dims)
    end
    ret
end

@doc raw"""
    scatter_update(A::Union{SparseTensor, SparseMatrixCSC{Float64,Int64}},
    i1::Union{Integer, Colon, UnitRange{T}, PyObject,Array{S,1}},
    i2::Union{Integer, Colon, UnitRange{T}, PyObject,Array{T,1}},
    B::Union{SparseTensor, SparseMatrixCSC{Float64,Int64}})  where {S<:Real,T<:Real}

Updates a subblock of a sparse matrix by `B`. Equivalently,
```
A[i1, i2] = B
```
"""
function scatter_update(A::Union{SparseTensor, SparseMatrixCSC{Float64,Int64}},
    i1::Union{Integer, Colon, UnitRange{S}, PyObject,Array{S,1}},
    i2::Union{Integer, Colon, UnitRange{T}, PyObject,Array{T,1}},
    B::Union{SparseTensor, SparseMatrixCSC{Float64,Int64}})  where {S<:Real,T<:Real}
    if isa(i1, Integer); i1 = [i1]; push!(squeeze_dims, 1); end
    if isa(i2, Integer); i2 = [i2]; push!(squeeze_dims, 2); end
    if isa(i1, UnitRange) || isa(i1, StepRange); i1 = collect(i1); end
    if isa(i2, UnitRange) || isa(i2, StepRange); i2 = collect(i2); end
    if isa(i1, Colon); i1 = collect(1:lastindex(A,1)); end
    if isa(i2, Colon); i2 = collect(1:lastindex(A,2)); end
    if isa(i1, Array{Bool,1}); i1 = findall(i1); end
    if isa(i2, Array{Bool,1}); i2 = findall(i2); end
    ii = convert_to_tensor(i1, dtype=Int64)
    jj = convert_to_tensor(i2, dtype=Int64)

    !isa(A, SparseTensor) && (A=SparseTensor(A))
    !isa(B, SparseTensor) && (B=SparseTensor(B))
    ii1, jj1, vv1 = find(A)
    m1_, n1_ = size(A)
    ii2, jj2, vv2 = find(B)
    sparse_scatter_update_ = load_system_op("sparse_scatter_update"; multiple=true)
    ii1,jj1,vv1,m1,n1,ii2,jj2,vv2,ii,jj = convert_to_tensor([ii1,jj1,vv1,m1_,n1_,ii2,jj2,vv2,ii,jj], [Int64,Int64,Float64,Int64,Int64,Int64,Int64,Float64,Int64,Int64])
    ii, jj, vv = sparse_scatter_update_(ii1,jj1,vv1,m1,n1,ii2,jj2,vv2,ii,jj)
    SparseTensor(ii, jj, vv, m1_, n1_)
end

@doc raw"""
    scatter_update(A::Union{SparseTensor, SparseMatrixCSC{Float64,Int64}},
    i1::Union{Integer, Colon, UnitRange{T}, PyObject,Array{S,1}},
    i2::Union{Integer, Colon, UnitRange{T}, PyObject,Array{T,1}},
    B::Union{SparseTensor, SparseMatrixCSC{Float64,Int64}})  where {S<:Real,T<:Real}

Adds `B` to a subblock of a sparse matrix `A`. Equivalently,
```
A[i1, i2] += B
```
"""
function scatter_add(A::Union{SparseTensor, SparseMatrixCSC{Float64,Int64}},
    i1::Union{Integer, Colon, UnitRange{T}, PyObject,Array{S,1}},
    i2::Union{Integer, Colon, UnitRange{T}, PyObject,Array{T,1}},
    B::Union{SparseTensor, SparseMatrixCSC{Float64,Int64}})  where {S<:Real,T<:Real}
    !(isa(A, SparseTensor)) && (A = SparseTensor(A))
    !(isa(B, SparseTensor)) && (B = SparseTensor(B))
    C = A[i1,i2]
    D = B + C
    scatter_update(A, i1, i2, D)
end

function Base.:reshape(s::SparseTensor, shape::T...) where T<:Integer
    SparseTensor(tf.sparse.reshape(s, shape), false)
end

@doc raw"""
    \(A::SparseTensor, o::PyObject, method::String="SparseLU")

Solves the linear equation
$$A x = o$$

# Method
For square matrices $A$, one of the following methods is available
- `auto`: using the solver specified by `ADCME.options.sparse.solver`
- `SparseLU`
- `SparseQR`
- `SimplicialLDLT`
- `SimplicialLLT`

!!! note
    In the case `o` is 2 dimensional, `\` is understood as "batched solve". `o` must have size $n_{b} \times m$, and
    $A$ has a size $m\times n$. It returns the solution matrix of size $n_b \times n$

    $$s_{i,:} = A^{-1} o_{i,:}$$
"""
function PyCall.:\(s::SparseTensor, o::PyObject, method::String="auto")
    local u
    if method=="auto"
        method = options.sparse.solver
    end
    if size(s,1)!=size(s,2)
        _cfun = load_system_op("sparse_least_square")
        ii, jj, vv = find(s)
        ii, jj, vv, o = convert_to_tensor([ii, jj, vv, o], [Int32, Int32, Float64, Float64])
        if length(size(o))==1
            @assert size(s,1)==length(o)
            o = reshape(o, (1, -1))
            u = _cfun(ii, jj, vv, o, constant(size(s, 2), dtype=Int32))
            return u[1]
        end
        @assert size(o,2)==size(s,1)
        u = _cfun(ii, jj, vv, o, constant(size(s, 2), dtype=Int32))
        if size(s,2)!=nothing && size(o,1)!=nothing
            u.set_shape((size(o,1), size(s,2)))
        end
    else
        ss = load_system_op("sparse_solver")
        # in case `indices` has dynamical shape
        ii, jj, vv = find(s)
        ii,jj,vv,o = convert_to_tensor([ii,jj,vv,o], [Int64,Int64,Float64,Float64])
        u = ss(ii,jj,vv,o,method)
        if size(s,2)!=nothing
            u.set_shape((size(s,2),))
        end
    end

    u
end

Base.:\(s::SparseTensor, o::Array{Float64}) = s\constant(o)

"""
    SparseAssembler(handle::Union{PyObject, <:Integer}, n::Union{PyObject, <:Integer}, tol::Union{PyObject, <:Real}=0.0)

Creates a SparseAssembler for accumulating `row`, `col`, `val` for sparse matrices.
- `handle`: an integer handle for creating a sparse matrix. If the handle already exists, `SparseAssembler` return the existing sparse matrix handle. If you are creating different sparse matrices, the handles should be different.
- `n`: Number of rows of the sparse matrix.
- `tol` (optional): Tolerance. `SparseAssembler` will treats any values less than `tol` as zero.

# Example 1
```julia
handle = SparseAssembler(100, 5, 1e-8)
op1 = accumulate(handle, 1, [1;2;3], [1.0;2.0;3.0])
op2 = accumulate(handle, 2, [1;2;3], [1.0;2.0;3.0])
J = assemble(5, 5, [op1;op2])
```
`J` will be a [`SparseTensor`](@ref) object.

# Example 2
```julia
handle = SparseAssembler(0, 5)
op1 = accumulate(handle, 1, [1;2;3], ones(3))
op2 = accumulate(handle, 1, [3], [1.])
op3 = accumulate(handle, 2, [1;3], ones(2))
J = assemble(5, 5, [op1;op2;op3]) # op1, op2, op3 are parallel
Array(run(sess, J))≈[1.0  1.0  2.0  0.0  0.0
                1.0  0.0  1.0  0.0  0.0
                0.0  0.0  0.0  0.0  0.0
                0.0  0.0  0.0  0.0  0.0
                0.0  0.0  0.0  0.0  0.0]
```
"""
function SparseAssembler(handle::Union{PyObject, <:Integer}, n::Union{PyObject, <:Integer}, tol::Union{PyObject, <:Real}=0.0)
    sparse_accumulator = load_system_op("sparse_accumulator", false)
    n = convert_to_tensor(n, dtype=Int32)
    tol = convert_to_tensor(tol, dtype=Float64)
    handle = convert_to_tensor(handle, dtype=Int32)
    sparse_accumulator(tol, n, handle)
end

"""
    accumulate(handle::PyObject, row::Union{PyObject, <:Integer}, cols::Union{PyObject, Array{<:Integer}}, vals::Union{PyObject, Array{<:Real}})

Accumulates `row`-th row. It adds the value to the sparse matrix
```julia
for k = 1:length(cols)
    A[row, cols[k]] += vals[k]
end
```
`handle` is the handle created by [`SparseAssembler`](@ref).

See [`SparseAssembler`](@ref) for an example.

!!! note
    The function `accumulate` returns a `op::PyObject`. Only when `op` is executed, the nonzero values are populated into the sparse matrix.
"""
function accumulate(handle::PyObject, row::Union{PyObject, <:Integer}, cols::Union{PyObject, Array{<:Integer}},
    vals::Union{PyObject, Array{<:Real}})
    sparse_accumulator_add = load_system_op("sparse_accumulator_add", false)
    row = convert_to_tensor(row, dtype=Int32)
    cols = convert_to_tensor(cols, dtype=Int32)
    vals = convert_to_tensor(vals, dtype=Float64)
    return sparse_accumulator_add(handle, row, cols, vals)
end

"""
    assemble(m::Union{PyObject, <:Integer}, n::Union{PyObject, <:Integer}, ops::PyObject)

Assembles the sparse matrix from the `ops` created by [`accumulate`](@ref). `ops` is either a single output from `accumulate`, or concated from several `ops`
```julia
op1 = accumulate(handle, 1, [1;2;3], [1.0;2.0;3.0])
op2 = accumulate(handle, 2, [1;2;3], [1.0;2.0;3.0])
op = [op1;op2] # equivalent to `vcat([op1, op2]...)`
```
`m` and `n` are rows and columns of the sparse matrix.

See [`SparseAssembler`](@ref) for an example.
"""
function assemble(m::Union{PyObject, <:Integer}, n::Union{PyObject, <:Integer}, ops::PyObject)
    sparse_accumulator_copy = load_system_op("sparse_accumulator_copy", false)
    if length(size(ops))==0
        ops = reshape(ops, 1)
    end
    ii, jj, vv = sparse_accumulator_copy(ops)
    return SparseTensor(ii, jj, vv, m, n)
end

"""
    spdiag(n::Int64)

Constructs a sparse identity matrix of size ``n\\times n``, which is equivalent to `spdiag(n, 0=>ones(n))`
"""
function spdiag(n::Int64)
    SparseTensor(sparse(1:n, 1:n, ones(Float64, n)))
end

"""
    spdiag(o::PyObject)

Constructs a sparse diagonal matrix where the diagonal entries are `o`, which is equivalent to `spdiag(length(o), 0=>o)`
"""
function spdiag(o::PyObject)
    if length(size(o))!=1
        error("ADCME: input `o` must be a vector")
    end
    ii = collect(1:length(o))
    SparseTensor(ii, ii, o, length(o), length(o), is_diag=true)
end

spdiag(o::Array{Float64, 1}) = spdiag(constant(o))

"""
    spzero(m::Int64, n::Union{Missing, Int64}=missing)

Constructs a empty sparse matrix of size ``m\\times n``. `n=m` if `n` is `missing`
"""
function spzero(m::Int64, n::Union{Missing, Int64}=missing)
    if ismissing(n)
        n = m
    end
    ii = Int64[]
    jj = Int64[]
    vv = Float64[]
    SparseTensor(ii, jj, vv, m, n, is_diag=true)
end

function Base.:*(s1::SparseTensor, s2::SparseTensor)
    ii1, jj1, vv1 = find(s1)
    ii2, jj2, vv2 = find(s2)
    m, n = size(s1)
    n_, k = size(s2)
    if n!=n_
        error("IGACS: matrix size mismatch: ($m, $n) vs ($n_, $k)")
    end
    mat_mul_fn = load_system_op("sparse_sparse_mat_mul")
    if s1._diag
        mat_mul_fn = load_system_op("diag_sparse_mat_mul")
    elseif s2._diag
        mat_mul_fn = load_system_op("sparse_diag_mat_mul")
    end
    ii3, jj3, vv3 = mat_mul_fn(ii1-1,jj1-1,vv1,ii2-1,jj2-1,vv2, constant(m), constant(n), constant(k))
    SparseTensor(ii3, jj3, vv3, m, k, is_diag=s1._diag&&s2._diag)
end

# missing is treated as zeros
Base.:+(s1::SparseTensor, s2::Missing) = s1
Base.:-(s1::SparseTensor, s2::Missing) = s1
Base.:*(s1::SparseTensor, s2::Missing) = missing
Base.:/(s1::SparseTensor, s2::Missing) = missing
Base.:-(s1::Missing, s2::SparseTensor) = -s2
Base.:+(s1::Missing, s2::SparseTensor) = s2
Base.:*(s1::Missing, s2::SparseTensor) = missing
Base.:/(s1::Missing, s2::SparseTensor) = missing

function Base.:sum(s::SparseTensor; dims::Union{Integer, Missing}=missing)
    if ismissing(dims)
        tf.sparse.reduce_sum(s.o)
    else
        tf.sparse.reduce_sum(s.o, axis=dims-1)
    end
end

@doc raw"""
    spdiag(m::Integer, pair::Pair...)

Constructs a square $m\times m$ [`SparseTensor`](@ref) from pairs of the form
```
offset => array
```
# Example
Suppose we want to construct a $10\times 10$ tridiagonal matrix, where the lower off-diagonals are all -2,
the diagonals are all 2, and the upper off-diagonals are all 3, the corresponding Julia code is
```julia
spdiag(10, -1=>-2*ones(9), 0=>2*ones(10), 1=>3ones(9))
```
"""
function spdiag(n::Integer, pair::Pair...)
    ii = Array{Int64}[]
    jj = Array{Int64}[]
    vv = PyObject[]
    for (k, v) in pair
        @assert -(n-1)<=k<=n-1
        v = convert_to_tensor(v, dtype=Float64)
        if k>=0
            push!(ii, collect(1:n-k))
            push!(jj, collect(k+1:n))
            push!(vv, v)
        else
            push!(ii, collect(-k+1:n))
            push!(jj, collect(1:n+k))
            push!(vv, v)
        end
    end
    ii = vcat(ii...)
    jj = vcat(jj...)
    vv = vcat(vv...)
    indices = [ii jj] .- 1
    sp = tf.SparseTensor(indices, vv, (n, n))
    options.sparse.auto_reorder && (sp = tf.sparse.reorder(sp))
    SparseTensor(sp)
end

@doc raw"""
    factorize(A::Union{SparseTensor, SparseMatrixCSC}, max_cache_size::Int64 = 999999)

Factorizes $A$ for sparse matrix solutions. `max_cache_size` specifies the maximum cache sizes in the C++ kernels,
which determines the maximum number of factorized matrices.
The function returns the factorized matrix, which is basically `Tuple{SparseTensor, PyObject}`.
# Example
```julia
A = sprand(10,10,0.7)
Afac = factorize(A) # factorizing the matrix
run(sess, Afac\rand(10)) # no factorization, solving the equation
run(sess, Afac\rand(10)) # no factorization, solving the equation
```
"""
function factorize(A::Union{SparseTensor, SparseMatrixCSC}, max_cache_size::Int64 = 999999)
    sparse_factorization_ = load_system_op("sparse_factorization"; multiple=false)
    A = constant(A)
    ii, jj, vv = find(A)
    d = size(A, 1)
    ii,jj,vv,d,max_cache_size = convert_to_tensor([ii,jj,vv,d,max_cache_size], [Int64,Int64,Float64,Int64,Int64])
    o = stop_gradient(sparse_factorization_(ii,jj,vv,d,max_cache_size))
    return (A, o)
end

@doc raw"""
    solve(A_factorized::Tuple{SparseTensor, PyObject}, rhs::Union{Array{Float64,1}, PyObject})

Solves the equation `A_factorized * x = rhs` using the factorized sparse matrix. See [`factorize`](@ref).
"""
function solve(A_factorized::Tuple{SparseTensor, PyObject}, rhs::Union{Array{Float64,1}, PyObject})
    A, o = A_factorized
    solve_ = load_system_op("solve"; multiple=false)
    ii, jj, vv = find(constant(A))
    rhs,ii, jj, vv,o = convert_to_tensor([rhs,ii, jj, vv,o], [Float64,Int64, Int64, Float64,Int64])
    out = solve_(rhs,ii, jj, vv,o)
end

@doc raw"""
    Base.:\(A_factorized::Tuple{SparseTensor, PyObject}, rhs::Union{Array{Float64,1}, PyObject})

A convenient overload for [`solve`](@ref). See [`factorize`](@ref).
"""
Base.:\(A_factorized::Tuple{SparseTensor, PyObject}, rhs::Union{Array{Float64,1}, PyObject}) = solve(A_factorized, rhs)

@doc raw"""
    trisolve(a::Union{PyObject, Array{Float64,1}},b::Union{PyObject, Array{Float64,1}},
        c::Union{PyObject, Array{Float64,1}},d::Union{PyObject, Array{Float64,1}})

Solves a tridiagonal matrix linear system. The equation is as follows

$$a_i x_{i-1} + b_i x_i + c_i x_{i+1} = d_i$$

In the matrix format,

```math
\begin{bmatrix}
b_1 & c_1 & &0 \\
a_2 & b_2 & c_2 & \\
   & a_3 & b_3 & &\\
   &     &     & & c_{n-1}\\
0 & & &a_n & b_n
\end{bmatrix}\begin{bmatrix}
x_1\\
x_2\\
\vdots \\
x_n
\end{bmatrix} = \begin{bmatrix}
d_1\\
d_2\\
\vdots\\
d_n\end{bmatrix}
```
"""
function trisolve(a::Union{PyObject, Array{Float64,1}},b::Union{PyObject, Array{Float64,1}},
        c::Union{PyObject, Array{Float64,1}},d::Union{PyObject, Array{Float64,1}})
    n = length(b)
    @assert length(b) == length(d) == length(a)+1 == length(c)+1
    tri_solve_ = load_system_op("tri_solve"; multiple=false)
    a,b,c,d = convert_to_tensor(Any[a,b,c,d], [Float64,Float64,Float64,Float64])
    out = tri_solve_(a,b,c,d)
    set_shape(out, (n,))
end

function Base.:bind(op::SparseTensor, ops...)
    bind(op.o.values, ops...)
end

"""
    compress(A::SparseTensor)

Compresses the duplicated index in `A`.

# Example
```julia
using ADCME
indices = [
    1 1
    1 1
    2 2
    3 3
]
v = [1.0;1.0;1.0;1.0]
A = SparseTensor(indices[:,1], indices[:,2], v, 3, 3)
Ac = compress(A)
sess = Session(); init(sess)

run(sess, A.o.indices) # expected: [0 0;0 0;1 1;2 2]
run(sess, A.o.values) # expected: [1.0;1.0;1.0;1.0]

run(sess, Ac.o.indices) # expected: [0 0;1 1;2 2]
run(sess, Ac.o.values) # expected: [2.0;1.0;1.0]
```

!!! note
    The indices of `A` should be sorted. `compress` does not check the validity of the input arguments.
"""
function compress(A::SparseTensor)
    indices, v = A.o.indices, A.o.values
    sparse_compress_ = load_op_and_grad(libadcme,"sparse_compress", multiple=true)
    ind, vv = sparse_compress_(indices,v)
    RawSparseTensor(ind, vv, size(A)...)
end


text:
# To test a simplified co ax copter.
# Modeling the system as a cube, gimbal and 2 propellers.
# A revolute2 joint between the cube and gimbal
# 1 Rev Joint each between the gimbal and either propeller.

include("../src/plotSol.jl")
include("../src/simulate.jl")
include("../src/OrientationConversion.jl")
##
clearconsole()
j1 = Nothing
j2 = Nothing
j3 = Nothing
j4 = Nothing
## Inertial Properties
mb = 0.755; mg = 0.215; mp = 0.02;

Ib =  [ 0.0009    0.0000   -0.0000;
        0.0000    0.0053    0.0000;
       -0.0000    0.0000    0.0054]#*(10^-3)

Ip =    [ 0.0029    0.0000    0.0000;
          0.0000    0.5824    0.0000;
          0.0000    0.0000    0.5805]*(10^-4)

Ig =    [ 0.1435    0.0000    0.0000;
          0.0000    0.0498    0.0000;
          0.0000    0.0000    0.1381]*(10^-3)

## Initialisation
InFrame = InertialFrameAsRB()
Body = RigidBody(mb,Ib,2)
Gimbal = RigidBody(mg,Ig,3)
Prop1 = RigidBody(mp,Ip,4)
Prop2 = RigidBody(mp,Ip,5)

x0Body = [zeros(3);[1;zeros(3)];zeros(3);zeros(4)]
initialiseRigidBody!(Body,x0Body)

x0Gimbal = [[0.0;0.0;0.143];[1;zeros(3)];zeros(3);zeros(4)]
initialiseRigidBody!(Gimbal,x0Gimbal)

x0Prop1 = [[0.0;0.0;0.234];[1;zeros(3)];zeros(3);zeros(4)]
initialiseRigidBody!(Prop1,x0Prop1)

x0Prop2 = [[0.0;0.0;0.24];[1;zeros(3)];zeros(3);zeros(4)]
initialiseRigidBody!(Prop2,x0Prop2)

## Joint Descriptions
# Joint 1 (Free) between the inertial frame and body.
j1 = Joint(InFrame,Body,zeros(3),zeros(3))

# Joint 2 (Revolute2) between the body and gimbal
axisZ = [0.0 0.0 1.0][:]
rjBody = [0.0 0.0 0.134][:]
rjGimbal1 = [0.0 0.0 -0.009][:]
j2 = Joint(Body,Gimbal,rjBody,rjGimbal1,
     type = "Revolute",axis = axisZ)#, jointTorque = [0.00 0.0 0.00][:])

# Joint 2 (Revolute2) between the body and gimbal
axisZ = [0.0 0.0 1.0][:]
rjGimbal2 = [0.0 0.0 0.091][:]
rjProp1 = [0.0 0.0 0.0][:]
j3 = Joint(Gimbal,Prop1,rjGimbal2,rjProp1,
     type = "Revolute",axis = axisZ, jointTorque = [0.0 0.0 0.001][:])

# Joint 2 (Revolute2) between the body and gimbal
axisZ = [0.0 0.0 1.0][:]
rjGimbal3 = [0.0 0.0 0.097][:]
rjProp2 = [0.0 0.0 0.0][:]
j4 = Joint(Gimbal,Prop2,rjGimbal3,rjProp2,
     type = "Revolute", axis = axisZ, jointTorque = [0.0 0.0 -0.001][:])

## Simulation
tEnd = 1.0
tSpan = 0.01
g = MVector{3}([0.0,0.0,0.0]) # Gravity Vector.
tSim, solFinal = simulate(tEnd,tSpan,j1,j2,j3,j4,g=g)#,extFVec = extFList)

solBody = solFinal[1]
solGimbal = solFinal[2]
solProp1 = solFinal[3]
solProp2 = solFinal[4]

## Check constraints
# Check if joint location has moved
jointLoc2 = Matrix{Float64}(undef,length(tSim),3)
jointLoc3 = Matrix{Float64}(undef,length(tSim),3)
jointLoc4 = Matrix{Float64}(undef,length(tSim),3)
ωBody = Matrix{Float64}(undef,length(tSim),3)
ωGimbal = Matrix{Float64}(undef,length(tSim),3)
ωProp1 = Matrix{Float64}(undef,length(tSim),3)
ωProp2 = Matrix{Float64}(undef,length(tSim),3)
ωGimbalInBody = Matrix{Float64}(undef,length(tSim),3)
ωRel1 = Matrix{Float64}(undef,length(tSim),3)
ωRel2 = Matrix{Float64}(undef,length(tSim),3)
ωRel3 = Matrix{Float64}(undef,length(tSim),3)
for i=1:length(tSim)
    Body.dcm = quat2dcm(solBody.β[i,:])
    Gimbal.dcm = quat2dcm(solGimbal.β[i,:])
    Prop1.dcm = quat2dcm(solProp1.β[i,:])
    Prop2.dcm = quat2dcm(solProp2.β[i,:])
    ωBody[i,:] = angVel(solBody.β[i,:],solBody.βdot[i,:])
    ωGimbal[i,:] = angVel(solGimbal.β[i,:],solGimbal.βdot[i,:])
    ωProp1[i,:] = angVel(solProp1.β[i,:],solProp1.βdot[i,:])
    ωProp2[i,:] = angVel(solProp2.β[i,:],solProp2.βdot[i,:])
    ωGimbalInBody[i,:] = quat2dcm(solBody.β[i,:])*
                         transpose(quat2dcm(solGimbal.β[i,:]))*
                         ωGimbal[i,:]
    ωRel1[i,:] = ωGimbalInBody[i,:] - ωBody[i,:]
    ωRel2[i,:] = quat2dcm(solGimbal.β[i,:])*
                 transpose(quat2dcm(solProp1.β[i,:]))*
                 ωProp1[i,:] - ωGimbal[i,:]
    ωRel3[i,:] = quat2dcm(solGimbal.β[i,:])*
                 transpose(quat2dcm(solProp2.β[i,:]))*
                 ωProp2[i,:] - ωGimbal[i,:]
    jointLoc2[i,:] = solBody.r[i,:] + transpose(Body.dcm)*rjBody -
                     solGimbal.r[i,:] -
                     transpose(Gimbal.dcm)*rjGimbal1
    jointLoc3[i,:] = solGimbal.r[i,:] +
                     transpose(Gimbal.dcm)*rjGimbal2 -
                     solProp1.r[i,:] - transpose(Prop1.dcm)*rjProp1
    jointLoc4[i,:] = solGimbal.r[i,:] +
                     transpose(Gimbal.dcm)*rjGimbal3 -
                     solProp2.r[i,:] - transpose(Prop2.dcm)*rjProp2
end

## Plotting
# Joint Locations
plotPos(tSim,jointLoc4)
plotPos(tSim,solBody.r)
plotPos(tSim,solGimbal.r)

plotErrNorm(tSim,solBody.β)
plotErrNorm(tSim,solGimbal.β)
plotErrNorm(tSim,solProp1.β)
plotErrNorm(tSim,solProp2.β)

plotAngVel(tSim,ωBody)
plotAngVel(tSim,ωGimbal)
plotAngVel(tSim,ωProp1)
plotAngVel(tSim,ωProp2)
plotAngVel(tSim,ωRel1)
plotAngVel(tSim,ωRel2)
plotAngVel(tSim,ωRel3)
# plotAngVel(tSim,ωGimbalInBody)
# plotAngVel(tSim,ωCube - ωProp)

ext: jl

max_stars_repo_name: isrlab/Multibody-Dynamics

max_issues_repo_name: isrlab/Multibody-Dynamics

max_forks_repo_name: isrlab/Multibody-Dynamics


text:
import Base: <

using DataStructures: MutableBinaryMinHeap, push!, pop!, top, _heap_bubble_up!
using DataStructures: DefaultDict

export decode

# Can't come up with a better name
struct Glass
    g_index::Int  # indexing into the `Glass` vector
    indices::Set{Int}
    data::Vector{UInt8}  # data being decoded in progress
end

@inline function heapify!(h::MutableBinaryMinHeap{T}, i::Int) where T
    _heap_bubble_up!(h.comparer, h.nodes, h.node_map, h.node_map[i])
end

function fill_glass(drop::WaterDrop, n_blocks::Int, glass_index::Int)::Glass
    indices = generate_indices(drop.index, n_blocks, Int(drop.degree))
    Glass(glass_index, Set(indices), drop.data)
end

degree(g::Glass) = length(g.indices)
<(g1::Glass, g2::Glass) = degree(g1) < degree(g2)

function percolate!(g::Glass, block_index::Int, data::Vector{UInt8})::Bool
    if block_index ∈ g.indices
        delete!(g.indices, block_index)
        g.data .⊻= data
        true
    else
        false
    end
end

# `n_blocks` is sent to the receiver before the actual sending of data packets.
function decode(water_drops::Vector{WaterDrop}, n_blocks::Int)::Union{Nothing, Vector{UInt8}}
    decoded = Vector{Union{Nothing, Vector{UInt8}}}(nothing, n_blocks)

    glass_heap = MutableBinaryMinHeap{Glass}()

    # Block index -> set of glass mapping
    block_to_glasses_dict = DefaultDict{Int, Set{Glass}}(Set())

    for (i, drop) in enumerate(water_drops)
        glass = fill_glass(drop, n_blocks, i)
        for ii in glass.indices
            push!(block_to_glasses_dict[ii], glass)
        end
        @assert push!(glass_heap, glass) == i
    end

    while !isempty(glass_heap)
        top_degree = degree(top(glass_heap))
        if top_degree > 1
            break
        elseif top_degree == 0
            pop!(glass_heap)
            continue
        end

        glass = pop!(glass_heap)
        block_index = first(glass.indices)
        for glass_neighbor in block_to_glasses_dict[block_index]
            if glass_neighbor != glass
                if percolate!(glass_neighbor, block_index, glass.data)
                    heapify!(glass_heap, glass_neighbor.g_index)
                end
            end
            delete!(block_to_glasses_dict, glass_neighbor)
        end
        decoded[block_index] = glass.data
    end

    if any(decoded .== nothing)
        nothing
    else
        vcat(decoded...)
    end
end

max_stars_repo_name: gyk/TrivialSolutions

max_issues_repo_name: gyk/TrivialSolutions

max_forks_repo_name: gyk/TrivialSolutions


text:
using CUDAapi

using Test

## properties

@test !CUDAapi.gcc_supported(v"5.0", v"5.5")
@test CUDAapi.gcc_supported(v"5.0", v"8.0")
CUDAapi.devices_for_cuda(v"8.0")
CUDAapi.devices_for_llvm(v"5.0")
CUDAapi.isas_for_cuda(v"8.0")
CUDAapi.isas_for_llvm(v"5.0")

## discovery

# generic
find_binary([Sys.iswindows() ? "CHKDSK" : "true"])
find_library([Sys.iswindows() ? "NTDLL" : "c"])

# CUDA

macro test_something(ex...)
    quote
        rv = $(ex...,)
        @test rv != nothing
        rv
    end
end

dirs = find_toolkit()
@test !isempty(dirs)
ver = find_toolkit_version(dirs)

@test_something find_cuda_binary("nvcc", dirs)
@test_something find_cuda_library("cudart", dirs)
@test_something find_libdevice([v"3.0"], dirs)
@test_something find_libcudadevrt(dirs)
@test_something find_host_compiler()
@test_something find_host_compiler(ver)
@test_something find_toolchain(dirs)
@test_something find_toolchain(dirs, ver)


text:
function lab_19(r::Robot)
    steps = 1
    side = Nord
    while ismarker(r)==false
        for _ in 1:2
            special_move!(r,side,steps)
            side=counterclockwise_side(side)
        end
        steps += 1
    end
    println("Marker found!")
end
function special_move!(r::Robot,side::HorizonSide,num::Int)
    for _ in 1:num
        if ismarker(r)
            return nothing
        end
        isborder(r, side) ? go_around(r, side) : move!(r,side)
    end
end

function go_around(r, direct_side)
    move_side = clockwise_side(direct_side)
    left_moves = 0
    right_moves = 0
    while isborder(r, direct_side)
        left_moves += 1
        right_moves = left_moves
        while isborder(r,direct_side) && right_moves > 0
            right_moves -= 1
            move!(r,move_side)
        end
        side=inverse_side(move_side)
    end
    move!(r,direct_side)
    left_moves /= 2
    while left_moves > 0
        left_moves -= 1
        move!(r,move_side)
    end
end


text:
using PDMats
using Test

n = 4
m = 7

for T in [Float64, Float32]
    X = randn(T, n, n)
    Y = randn(T, m, m)
    A = X * X'
    B = Y * Y'
    AkB = kron(A, B)
    PDA = PDMat(A)
    PDB = PDMat(B)
    PDAkB1 = PDMat(AkB)
    PDAkB2 = kron(PDA, PDB)
    @test PDAkB1.dim == PDAkB2.dim
    @test PDAkB1.mat ≈ PDAkB2.mat
    @test PDAkB1.chol.L ≈ PDAkB2.chol.L
    @test PDAkB1.chol.U ≈ PDAkB2.chol.U
    @test Matrix(PDAkB2.chol) ≈ PDAkB2.mat
end


text:
# indicator of an affine set

"""
  IndAffine(A::Array{Real,2}, b::Array{Real,1})

Returns the function `g = ind{x : Ax = b}`.

  IndAffine(A::Array{Real,1}, b::Real)

Returns the function `g = ind{x : dot(a,x) = b}`.
"""

immutable IndAffine{T <: RealOrComplex} <: IndicatorConvex
  A::AbstractArray{T,2}
  b::AbstractArray{T,1}
  R::AbstractArray{T,2}
  function IndAffine(A::AbstractArray{T,2}, b::AbstractArray{T,1})
    if size(A,1) > size(A,2)
      error("A must be full row rank")
    end
    normrows = vec(sqrt.(sum(abs2.(A), 2)))
    A = (1./normrows).*A # normalize rows of A
    b = (1./normrows).*b # and b accordingly
    Q, R = qr(A')
    new(A, b, R)
  end
end

IndAffine{T <: RealOrComplex}(A::AbstractArray{T,2}, b::AbstractArray{T,1}) =
  IndAffine{T}(A, b)

IndAffine{T <: RealOrComplex}(a::AbstractArray{T,1}, b::T) =
  IndAffine{T}(a', [b])

function (f::IndAffine){T <: RealOrComplex}(x::AbstractArray{T,1})
  # the tolerance in the following line should be customizable
  if norm(f.A*x - f.b, Inf) <= 1e-14
    return 0.0
  end
  return +Inf
end

function prox!{T <: RealOrComplex}(f::IndAffine, x::AbstractArray{T,1}, y::AbstractArray{T,1}, gamma::Real=1.0)
  res = f.A*x - f.b
  y[:] = x - f.A'*(f.R\(f.R'\res))
  return 0.0
end

fun_name(f::IndAffine) = "indicator of an affine subspace"
fun_dom(f::IndAffine) = "AbstractArray{Real,1}, AbstractArray{Complex,1}"
fun_expr(f::IndAffine) = "x ↦ 0 if Ax = b, +∞ otherwise"
fun_params(f::IndAffine) =
  string( "A = ", typeof(f.A), " of size ", size(f.A), ", ",
          "b = ", typeof(f.b), " of size ", size(f.b))

function prox_naive{T <: RealOrComplex}(f::IndAffine, x::AbstractArray{T,1}, gamma::Real=1.0)
  res = f.A*x - f.b
  y = x - f.A'*(f.R\(f.R'\res))
  return y, 0.0
end


text:
ENV["GKSwstype"]="100" #src

#md # # Who am I ?
#md #
#md #  - My name is *Pierre Navaro*
#md #
#md #  - **Fortran 77 + PVM** : during my PhD 1998-2002 (Université du Havre)
#md #
#md #  - **Fortran 90-2003 + OpenMP-MPI** : Engineer in Strasbourg (2003-2015) at IRMA
#md #
#md #  - **Numpy + Cython, R + Rcpp** : Engineer in Rennes (2015-now) at IRMAR
#md #
#md #  - **Julia v1.0** since July 2018
#md #
#md #  ## Instructions to open the notebook and slides
#md #
#md #  - https://github.com/pnavaro/JuliaSMAI2021
#md #  - https://pnavaro.github.io/JuliaSMAI2021
#md #

#md # ---

#md #
#md # # Why Julia?
#md #
#md # - Started in 2009 and first version was released in 2012.
#md # - High-level languages like Python and R let one explore and experiment rapidly, but can run slow.
#md # - Low-level languages like Fortran/C++ tend to take longer to develop, but run fast.
#md # - This is sometimes called the "two language problem" and is something the Julia developers set out to eliminate.
#md # - Julia's promise is to provide a "best of both worlds" experience for programmers who need to develop novel algorithms and bring them into production environments with minimal effort.
#md #

#md # **Julia: A Fresh Approach to Numerical Computing**
#md #
#md # *Jeff Bezanson, Alan Edelman, Stefan Karpinski, Viral B. Shah*
#md #
#md # SIAM Rev., 59(1), 65–98. (34 pages) 2012

# **SciML Scientific Machine Learning Software** https://sciml.ai/citing/
#
# **Differentialequations.jl--a performant and feature-rich ecosystem for solving differential equations in julia,**
#
# *Christopher Rackauckas and Qing Nie*,
#
# Journal of Open Research Software,, volume 5, number 1, 2017.

#md # ---


text:
module Risk

import ..Clarus

function dv01(;params...)
  return Clarus.api_request("Risk","DV01";params...)
end

function fxdelta(;params...)
  return Clarus.api_request("Risk","FXDelta";params...)
end

function fxvega(;params...)
  return Clarus.api_request("Risk","FXVega";params...)
end

function irdelta(;params...)
  return Clarus.api_request("Risk","IRDelta";params...)
end

function irgamma(;params...)
  return Clarus.api_request("Risk","IRGamma";params...)
end

function irvega(;params...)
  return Clarus.api_request("Risk","IRVega";params...)
end

function stress(;params...)
  return Clarus.api_request("Risk","Stress";params...)
end

function theta(;params...)
  return Clarus.api_request("Risk","Theta";params...)
end

function var(;params...)
  return Clarus.api_request("Risk","VaR";params...)
end

function varattribution(;params...)
  return Clarus.api_request("Risk","VaRAttribution";params...)
end

function varvectors(;params...)
  return Clarus.api_request("Risk","VaRVectors";params...)
end

end #MODULE-END


text:
using KernelEstimator
using Test
using Random
using StatsBase
# write your own tests here
include("testreg.jl")


text:
using Distributions: Uniform
import Distributions.sample

struct ActionSpace{A, N}
    actions::Array{A, N}
end

function ActionSpace(pyspace::PyObject)
    type = pyspace.__class__.__name__
    if type == "Discrete"
        return ActionSpace([pyspace.n])
    elseif type == "MultiDiscrete"
        return ActionSpace(pyspace.nvec)
    elseif type == "Box"
        return ActionSpace(map((l, h) -> (l, h), pyspace.low, pyspace.high))
    else
        error("Unknown python action_space type")
    end
end

Base.size(a::ActionSpace) = size(a.actions)
Base.length(a::ActionSpace) = length(a.actions)
Base.eltype(::Type{ActionSpace{A, N}}) where {A, N} = A
Base.eltype(::Type{ActionSpace{Tuple{A, A}, N}}) where {A, N} = A

sample(a::ActionSpace{A}) where A <: Integer = map(x -> rand(one(A):x), a.actions)
function sample(a::ActionSpace{Tuple{T, T}}) where T <: AbstractFloat
    return map(x -> T(rand(Uniform(x[1], x[2]))), a.actions)
end

struct ObservationSpace{S, N}
    vals::Array{Tuple{S, S}, N}
end

Base.size(o::ObservationSpace) = size(o.vals)
Base.length(o::ObservationSpace) = length(o.vals)
Base.eltype(::Type{ObservationSpace{S, N}}) where {S, N} = S

function ObservationSpace(pyspace::PyObject)
    type = pyspace.__class__.__name__
    if type == "Box"
        return ObservationSpace(map((l, h) -> (l, h), pyspace.low, pyspace.high))
    elseif type == "Tuple"
        return [ObservationSpace(space) for space in pyspace."spaces"]
    else
        error("Unknown python observation_space type")
    end
end


text:
module KelvinletsImage

    include("triangleInterpolator.jl")
    using Images, ProgressMeter, ImageView, LinearAlgebra
    export KelvinletsObject, grab, scale, pinch, grabRectangle, makeVideo, __applyAreaVariation__, __applyCloserAreaVariation__, __applyPerimeterBasedAreaVariation__, __applyEdgeBasedAreaVariation__, __applyPonderedAvgAreaVariation__

    """
        KelvinletsObject(image::AbstractArray{RGB{N0f8}, 2}, ν::Float64, μ::Float64)

    Initializes KelvinletsObject for a given image *image*,
    poisson ratio *ν*
    and elastic shear modulus *μ*

    # Example:
    ```julia-repl
    julia> object = KelvinletsObject(image, 0.4, 1.)
    ```
    """
    struct KelvinletsObject
        a::Float64
        b::Float64
        c::Float64
        sizeX::Int64
        sizeY::Int64
        image::AbstractArray{RGB{N0f8}, 2}
        function KelvinletsObject(image::AbstractArray{RGB{N0f8}, 2},
                                  ν::Float64,
                                  μ::Float64
                )::KelvinletsObject

            a = 1 / (4pi * μ)
            b = a / (4(1 - ν))
            c = 2 / (3a- 2b)

            sizeY, sizeX = size(image)

            new(a, b, c, sizeX, sizeY, image)
        end
    end

    function __applyVariation__(object::KelvinletsObject,
                                variationFunction::Function,
                                retardationFunction::Function,
                                heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        allΔ = zeros(object.sizeY, object.sizeX, 2)
        testImg = fill(RGB(0, 0, 0), object.sizeY, object.sizeX)

        for i=1:object.sizeY
            for j=1:object.sizeX
                Δ = variationFunction([i, j])

                dx1 = j - 1
                dx2 = object.sizeX - j
                dy1 = i - 1
                dy2 = object.sizeY - i

                dx = min(dx1, dx2)
                dy = min(dy1, dy2)

                y = 2(object.sizeY/2 - dy) / object.sizeY
                x = 2(object.sizeX/2 - dx) / object.sizeX

                Δ[1] *= retardationFunction(y)
                Δ[2] *= retardationFunction(x)

                maxnorm = norm([object.sizeY, object.sizeX])
                testImg[i, j] = RGB{N0f8}(norm(Δ)/maxnorm, norm(Δ)/maxnorm, norm(Δ)/maxnorm)

                Δ += [i, j]

                allΔ[i, j, 1] = Δ[1]
                allΔ[i, j, 2] = Δ[2]
            end
        end
        if heatmap
            return testImg
        end
        return __interpolateVariation__(object, allΔ)
    end

    function __applyAreaVariation__(object::KelvinletsObject,
                                    points::Array{Int64, 2},
                                    variationFunction::Function,
                                    retardationFunction::Function,
                                    heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        minX, maxX = points[:, 1]
        minY, maxY = points[:, 2]

        allΔ = zeros(object.sizeY, object.sizeX, 2)

        testImg = fill(RGB(0, 0, 0), object.sizeY, object.sizeX)

        for i=1:object.sizeY
            for j=1:object.sizeX

                if i < minY && j < minX #Top-LEFT corner
                    Δ = variationFunction([i, j], [minY, minX])
                elseif i < minY && j >= minX && j <= maxX # Top
                    Δ = variationFunction([i, j], [minY, j])
                elseif i < minY && j > maxX #Top-RIGHT corner
                    Δ = variationFunction([i, j], [minY, maxX])
                elseif j < minX && i >= minY && i <= maxY #LEFT
                    Δ = variationFunction([i, j], [i, minX])
                elseif j >= minX && j <= maxX && i >= minY && i <= maxY # MIDDLE
                    Δ = variationFunction([i, j], [i, j])
                elseif j > maxX && i >= minY && i <= maxY #RIGHT
                    Δ = variationFunction([i, j], [i, maxX])
                elseif i > maxY && j < minX #Lower-LEFT corner
                    Δ = variationFunction([i, j], [maxY, minX])
                elseif i > maxY && j >= minX && j <= maxX #BASE
                    Δ = variationFunction([i, j], [maxY, j])
                elseif i >  maxY && j > maxX #Lower-RIGHT corner
                    Δ = variationFunction([i, j], [maxY, maxX])
                end

                dx1 = j
                dx2 = object.sizeX - j
                dy1 = i
                dy2 = object.sizeY - i

                dx = min(dx1, dx2)
                dy = min(dy1, dy2)

                y = 2(object.sizeY/2 - dy)/object.sizeY
                x = 2(object.sizeX/2 - dx)/object.sizeX

                Δ[1] *= retardationFunction(y)
                Δ[2] *= retardationFunction(x)

                maxnorm = norm([object.sizeY, object.sizeX])
                testImg[i, j] = RGB{N0f8}(norm(Δ)/maxnorm, norm(Δ)/maxnorm, norm(Δ)/maxnorm)

                Δ += [i, j]

                allΔ[i, j, 1] = Δ[1]
                allΔ[i, j, 2] = Δ[2]
            end
        end
        if heatmap
            return testImg
        end
        return __interpolateVariation__(object, allΔ)
    end

    function __applyCloserAreaVariation__(object::KelvinletsObject,
                                    points::Array{Int64, 2},
                                    variationFunction::Function,
                                    retardationFunction::Function,
                                    heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        minX, maxX = points[:, 1]
        minY, maxY = points[:, 2]

        allΔ = zeros(object.sizeY, object.sizeX, 2)

        testImg = fill(RGB(0, 0, 0), object.sizeY, object.sizeX)
        for i=1:object.sizeY
            for j=1:object.sizeX

                df00 = norm([i, j] - [minY, minX])
                df01 = norm([i, j] - [minY, maxX])
                df10 = norm([i, j] - [maxY, minX])
                df11 = norm([i, j] - [maxY, maxX])

                closer = min(df00, df01, df10, df11)

                if closer == df00
                    Δ = variationFunction([i, j], [minY, minX])
                elseif closer == df01
                    Δ = variationFunction([i, j], [minY, maxX])
                elseif closer == df10
                    Δ = variationFunction([i, j], [maxY, minX])
                elseif closer == df11
                    Δ = variationFunction([i, j], [maxY, maxX])
                end

                dx1 = j
                dx2 = object.sizeX - j
                dy1 = i
                dy2 = object.sizeY - i

                dx = min(dx1, dx2)
                dy = min(dy1, dy2)

                y = 2(object.sizeY/2 - dy)/object.sizeY
                x = 2(object.sizeX/2 - dx)/object.sizeX

                Δ[1] *= retardationFunction(y)
                Δ[2] *= retardationFunction(x)

                maxnorm = norm([object.sizeY, object.sizeX])
                testImg[i, j] = RGB{N0f8}(norm(Δ)/maxnorm, norm(Δ)/maxnorm, norm(Δ)/maxnorm)

                Δ += [i, j]

                allΔ[i, j, 1] = Δ[1]
                allΔ[i, j, 2] = Δ[2]
            end
        end
        if heatmap
            return testImg
        end
        return __interpolateVariation__(object, allΔ)

    end

    function __applyPerimeterBasedAreaVariation__(object::KelvinletsObject,
                                    points::Array{Int64, 2},
                                    variationFunction::Function,
                                    retardationFunction::Function,
                                    heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        minX, maxX = points[:, 1]
        minY, maxY = points[:, 2]

        allΔ = zeros(object.sizeY, object.sizeX, 2)
        numOfPixels = (maxY - minY) * (maxX - minX)

        testImg = fill(RGB(0, 0, 0), object.sizeY, object.sizeX)

        @showprogress for i=1:object.sizeY
            for j=1:object.sizeX

                divided = [0., 0.]
                divisor = 0.

                for k = minX:maxX
                    d = norm([i, j] - [minY, k]) + 1
                    divided += (1/d) * [minY, k]
                    divisor += 1/d
                end

                for k = minX+1:maxX-1
                    d = norm([i, j] - [maxY, k]) + 1
                    divided += (1/d) * [maxY, k]
                    divisor += 1/d
                end

                for k = minY+1:maxY
                    d = norm([i, j] - [k, minX]) + 1
                    divided += (1/d) * [k, minX]
                    divisor += 1/d
                end

                for k = minY+1:maxX
                    d = norm([i, j] - [k, maxX]) + 1
                    divided += (1/d) * [k, maxX]
                    divisor += 1/d
                end

                newReference = divided/divisor
                Δ = variationFunction([i, j], round.(Int64, newReference))

                dx1 = j
                dx2 = object.sizeX - j
                dy1 = i
                dy2 = object.sizeY - i

                dx = min(dx1, dx2)
                dy = min(dy1, dy2)

                y = 2(object.sizeY/2 - (dy - 1))/object.sizeY
                x = 2(object.sizeX/2 - (dx - 1))/object.sizeX

                Δ[1] *= retardationFunction(y)
                Δ[2] *= retardationFunction(x)

                maxnorm = norm([object.sizeY, object.sizeX])
                testImg[i, j] = RGB{N0f8}(norm(Δ)/maxnorm, norm(Δ)/maxnorm, norm(Δ)/maxnorm)

                Δ += [i, j]
                allΔ[i, j, 1] = Δ[1]
                allΔ[i, j, 2] = Δ[2]
            end
        end
        if heatmap
            return testImg
        end
        return __interpolateVariation__(object, allΔ)
    end

    function __applyEdgeBasedAreaVariation__(object::KelvinletsObject,
                                    points::Array{Int64, 2},
                                    variationFunction::Function,
                                    retardationFunction::Function,
                                    heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        minX, maxX = points[:, 1]
        minY, maxY = points[:, 2]

        allΔ = zeros(object.sizeY, object.sizeX, 2)
        numOfPixels = (maxY - minY) * (maxX - minX)

        testImg = fill(RGB(0, 0, 0), object.sizeY, object.sizeX)

        @showprogress for i=1:object.sizeY
            for j=1:object.sizeX

                d1 = norm([i, j] - [minY, minX]) + 1
                d2 = norm([i, j] - [minY, maxX]) + 1
                d3 = norm([i, j] - [maxY, minX]) + 1
                d4 = norm([i, j] - [maxY, maxX]) + 1

                newReference = ((1/d1) * [minY, minX] + (1/d2) * [minY, maxX] + (1/d3) * [maxY, minX] + (1/d4) * [maxY, maxX]) / ((1/d1) + (1/d2) + (1/d3) + (1/d4))
                Δ = variationFunction([i, j], round.(Int64, newReference))

                dx1 = j
                dx2 = object.sizeX - j
                dy1 = i
                dy2 = object.sizeY - i

                dx = min(dx1, dx2)
                dy = min(dy1, dy2)

                y = 2(object.sizeY/2 - (dy - 1))/object.sizeY
                x = 2(object.sizeX/2 - (dx - 1))/object.sizeX

                Δ[1] *= retardationFunction(y)
                Δ[2] *= retardationFunction(x)

                maxnorm = norm([object.sizeY, object.sizeX])

                testImg[i, j] = RGB{N0f8}(norm(Δ)/maxnorm, norm(Δ)/maxnorm, norm(Δ)/maxnorm)

                Δ += [i, j]
                allΔ[i, j, 1] = Δ[1]
                allΔ[i, j, 2] = Δ[2]
            end
        end
        if heatmap
            return testImg
        end
        return __interpolateVariation__(object, allΔ)
    end

    function __applyPonderedAvgAreaVariation__(object::KelvinletsObject,
                                    points::Array{Int64, 2},
                                    variationFunction::Function,
                                    retardationFunction::Function,
                                    heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        minX, maxX = points[:, 1]
        minY, maxY = points[:, 2]

        allΔ = zeros(object.sizeY, object.sizeX, 2)
        numOfPixels = (maxY - minY) * (maxX - minX)

        testImg = fill(RGB(0, 0, 0), object.sizeY, object.sizeX)

        @showprogress for i=1:object.sizeY
            for j=1:object.sizeX

                d1 = norm([i, j] - [minY, minX]) + 1
                d2 = norm([i, j] - [minY, maxX]) + 1
                d3 = norm([i, j] - [maxY, minX]) + 1
                d4 = norm([i, j] - [maxY, maxX]) + 1

                deform1 = variationFunction([i,j], [minY, minX])
                deform2 = variationFunction([i,j], [minY, maxX])
                deform3 = variationFunction([i,j], [maxY, minX])
                deform4 = variationFunction([i,j], [maxY, maxX])

                Δ = ((1/d1) * deform1 + (1/d2) * deform1 + (1/d3) * deform1 + (1/d4) * deform1) /
                     ((1/d1) + (1/d2) + (1/d3) + (1/d4))

                dx1 = j
                dx2 = object.sizeX - j
                dy1 = i
                dy2 = object.sizeY - i

                dx = min(dx1, dx2)
                dy = min(dy1, dy2)

                y = 2(object.sizeY/2 - (dy - 1))/object.sizeY
                x = 2(object.sizeX/2 - (dx - 1))/object.sizeX

                Δ[1] *= retardationFunction(y)
                Δ[2] *= retardationFunction(x)

                maxnorm = norm([object.sizeY, object.sizeX])
                testImg[i, j] = RGB{N0f8}(norm(Δ)/maxnorm, norm(Δ)/maxnorm, norm(Δ)/maxnorm)

                Δ += [i, j]
                allΔ[i, j, 1] = Δ[1]
                allΔ[i, j, 2] = Δ[2]
            end
        end
        if heatmap
            return testImg
        end
        return __interpolateVariation__(object, allΔ)
    end

    function __interpolateVariation__(object::KelvinletsObject,
                                      allΔ::Array{Float64, 3}
            )::Array{RGB{N0f8}, 2}

        interpImg = fill(RGB(0, 0, 0), object.sizeY, object.sizeX)

        rasterize = function(A, B, C)
            colorA = object.image[A[1], A[2]]
            colorB = object.image[B[1], B[2]]
            colorC = object.image[C[1], C[2]]
            triangleInterpolator.rasterizationBBOX(interpImg,
                            [allΔ[A[1], A[2], 1], allΔ[A[1], A[2], 2]],
                            [allΔ[B[1], B[2], 1], allΔ[B[1], B[2], 2]],
                            [allΔ[C[1], C[2], 1], allΔ[C[1], C[2], 2]],
                            colorA, colorB, colorC
            )
        end

        for i=1:object.sizeY-1
          for j=1:object.sizeX-1
            rasterize([i, j], [i, j+1], [i+1, j+1])
            rasterize([i, j], [i+1, j], [i+1, j+1])
          end
        end
        return interpImg
    end

    """
        grab(object::KelvinletsObject, x0::Array{Int64}, force::Array{Float64}, ϵ::Float64, heatmap::Bool)

    grabs a point in an image given a KelvinletsObject *obj*,
    a pressure point *x0*,
    a force vector *force*,
    a brush size *ϵ*
    and whether you would like to generate the heatmap of the deformation
    # Example:
    ```julia-repl
        julia> newImage = grab(obj, [100, 100], [100., 0.], 50.)
    ```
    """
    function grab(object::KelvinletsObject,
                  x0::Array{Int64},
                  force::Array{Float64},
                  ϵ::Float64,
                  heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        grabFunc = function(x::Array{Int64})

            r = x - x0
            rLength = norm(r)
            rϵ = sqrt(rLength^2 + ϵ^2)
            kelvinState = (((object.a - object.b)/rϵ) * I +
                            (object.b / rϵ^3) * r * r' +
                            (object.a / 2) * (ϵ^2 / rϵ^3) * I)
            object.c * ϵ * kelvinState * force
        end

        retardationFunc = α -> (cos(π * α) + 1) / 2
        return __applyVariation__(object, grabFunc, retardationFunc, heatmap)
    end

    """
        scale(object::KelvinletsObject, x0::Array{Int64}, scale::Float64, ϵ::Float64, heatmap::Bool)

    Scales a point on an image given a KelvinletsObject *obj*,
    a pressure point *x0*,
    a scale alpha *scale*,
    a brush size *ϵ*,
    and whether you would like to generate the heatmap of the deformation
    # Example:
    ```julia-repl
        julia> newImage = scale(obj, [100, 100], 50., 50.)
    ```
    """
    function scale(object::KelvinletsObject,
                   x0::Array{Int64},
                   force::Float64,
                   ϵ::Float64,
                   heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        scaleFunc = function(x::Array{Int64})

            r = x - x0
            rLength = norm(r)
            rϵ = sqrt(rLength^2 + ϵ^2)

            return (2 * object.b - object.a) *
                   ( (1 / rϵ^2) +
                   ((ϵ^2)) / (2 * (rϵ^4))) *
                   (force * r)
        end

        retardationFunc = α -> (cos(π * α) + 1) / 2
        return __applyVariation__(object, scaleFunc, retardationFunc, heatmap)
    end

    """
         pinch(object::KelvinletsObject, x0::Array{Int64}, force::Array{Float64, 2}, ϵ::Float64, heatmap::Bool)

    Pinches the image given a KelvinletsObject *obj*,
    a pressure point *x0*,
    a force matrix *force*,
    a brush size *ϵ*,
    and whether you would like to generate the heatmap of the deformation
    # Example:
    ```julia-repl
        julia> newImage = pinch(obj, [100, 100], [10000. 0. ; 10000. 0.], 50.)
    ```
    """
    function pinch(object::KelvinletsObject,
                   x0::Array{Int64},
                   force::Array{Float64, 2},
                   ϵ::Float64,
                   heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        pinchFunc = function(x::Array{Int64})

            r = x - x0
            rLength = norm(r)
            rϵ = sqrt(rLength^2 + ϵ^2)
            return  -2 * object.a * ((1 / rϵ^2) +
                    (ϵ^2 / rϵ^4)) * force * r +
                    4 * object.b * ((1 / rϵ^2) * force -
                    (1 / rϵ^4) * (r' * force * r) * I) * r
        end

        retardationFunc = α -> (cos(π * α) + 1) / 2
        return __applyVariation__(object, pinchFunc, retardationFunc, heatmap)
    end

    """
        grabRectangle(object::KelvinletsObject, points::Array{Int64, 2}, force::Array{Float64}, ϵ::Float64, heatmap::Bool)

    Grabs the image using a square brush given a KelvinletsObject *obj*,
    a given matrix representation of a rectangle *points* in the following format --> [minY, minX ; maxY, maxX],
    a force vector *force*,
    a given brush size *ϵ* (to calculate te variation of outside pixels),
    and whether you would like to generate the heatmap of the deformation
    # Example:
    ```julia-repl
        julia> newImage = grabRectangle(obj, [100 100 ; 250 250], [100., 0.], 50.)
    ```
    """
    function grabRectangle(object::KelvinletsObject,
                           points::Array{Int64, 2},
                           force::Array{Float64},
                           ϵ::Float64,
                           areaVariation::Function,
                           heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        grabFunc = function(x::Array{Int64},
                            x0::Array{Int64}
                    )

            r = x - x0
            rLength = norm(r)
            rϵ = sqrt(rLength^2 + ϵ^2)
            kelvinState = (((object.a - object.b)/rϵ) * I +
                            (object.b / rϵ^3) * r * r' +
                            (object.a / 2) * (ϵ^2 / rϵ^3) * I)
            object.c * ϵ * kelvinState * force
        end

        retardationFunc = α -> (cos(π * α) + 1) / 2
        return areaVariation(object, points, grabFunc, retardationFunc, heatmap)

    end

    function __r__(a, b, c, d, x)
        ax = a - x
        bx = b - x
        cx = c - x
        dx = d - x

        mainDiag = a - d

        Aa = abs(ax[1]) * abs(ax[2])
        Ab = abs(bx[1]) * abs(bx[2])
        Ac = abs(cx[1]) * abs(cx[2])
        Ad = abs(dx[1]) * abs(dx[2])

        At = abs(mainDiag[1]) * abs(mainDiag[2])

        da = norm(ax) + 1
        db = norm(bx) + 1
        dc = norm(cx) + 1
        dd = norm(dx) + 1

        return 1/512 * ((At/(Aa + Ab + Ac + Ad)) - 1) *
                   (x -
                    ((1/da) * a + (1/db) * b + (1/dc) * c + (1/dd) * d)) /
                    ((1/da) + (1/db) + (1/dc) + (1/dd)
                   )
    end

    function grabRectangle__new(object::KelvinletsObject,
                           points::Array{Int64, 2},
                           force::Array{Float64},
                           ϵ::Float64,
                           heatmap::Bool
            )::Array{RGB{N0f8}, 2}

        minX, maxX = points[:, 1]
        minY, maxY = points[:, 2]

        a = [minY, minX]
        b = [minY, maxX]
        c = [maxY, minX]
        d = [maxY, maxX]

        grabFunc = function(x::Array{Int64})

            r = __r__(a, b, c, d, x)
            rLength = norm(r)
            rϵ = sqrt(rLength^2 + ϵ^2)
            kelvinState = (((object.a - object.b)/rϵ) * I +
                            (object.b / rϵ^3) * r * r' +
                            (object.a / 2) * (ϵ^2 / rϵ^3) * I)
            object.c * ϵ * kelvinState * force
        end

        retardationFunc = α -> (cos(π * α) + 1) / 2
        return __applyVariation__(object, grabFunc, retardationFunc, heatmap)

    end

    """
        makeVideo(object::KelvinletsObject, kelvinletsFunction::Function, x0::Array{Int64}, force, ϵ::Float64, frames::Int64)

    Computes a video for a given KelvinletsImage deformation function *KelvinletsFunction*,
    using a *KelvinletsObject* as reference,
    a *x0* point of force application,
    a *force* matrix/vector (depending on the function),
    a *ϵ* brush size
    and a number of frames *frames*
    # Example
    ```julia-repl
        julia> video = (object, grab, [100, 100], [200., 0.], 70., 20)
    ```
    """
    function makeVideo(object::KelvinletsObject,
                       kelvinletsFunction::Function,
                       x0::Array{Int64},
                       force,
                       ϵ::Float64,
                       frames::Int64
        )::Array{RGB{N0f8},3}

        if typeof(force) == Float64
            var = range(0, stop=force, length=frames)
        else
            var = range(fill(0, size(force)), stop=force, length=frames)
        end

        video = Array{RGB{N0f8}}(undef, object.sizeY, object.sizeX, frames)
        @showprogress for i=1:frames
            video[:,:,i] = kelvinletsFunction(object, x0, var[i], ϵ, false)
        end
        imshow(video)
        return video
    end
end


text:
"""
"""
function run_YORP(shape, orbit, spin, params_thermo)
    @unpack P, Δt, t_bgn, t_end, Nt, Nz = params_thermo

    init_temperature!(shape, orbit, spin, params_thermo)

    τ̄ = zeros(3)  # Net YORP torque

    for t in (t_bgn:Δt:t_end)*P
        spin_phase = spin.ω * t
        F☉, r̂☉ = getSolarCondition(orbit, spin, t)

        update_flux_sun!(shape, F☉, r̂☉)  # あってる
        update_flux_scat_single!(shape, params_thermo)  # ほんのちょっと違う。
        update_flux_rad_single!(shape, params_thermo)  # ほんのちょっと違う。

        update_force!(shape, params_thermo)  # ここからけっこう違う！
        # update_force_Rubincam!(shape, params_thermo)

        τ̄ .+= body_to_orbit(SVector{3}(shape.torque), spin.γ, spin.ε, spin_phase)

        update_temperature!(shape, params_thermo)

        @unpack Tz, flux, df = shape.smeshes[1]
        println(Tz[begin], ", ", flux.sun, ", ", flux.scat, ", ", flux.rad, ", ", df)
    end
    τ̄ /= Nt
end

"""
"""
function run_Yarkovsky(shape, orbit, spin, params_thermo)
    @unpack P, Δt, t_bgn, t_end, Nt, Nz = params_thermo

    init_temperature!(shape, orbit, spin, params_thermo)
    shape.smeshes

    Fs = typeof(shape.F)[]
    # τ̄ = zeros(3)  # Net YORP torque

    for t in (t_bgn:Δt:t_end)*P
        spin_phase = spin.ω * t
        F☉, r̂☉ = getSolarCondition(orbit, spin, t)

        update_flux_sun!(shape, F☉, r̂☉)
        update_flux_scat_single!(shape, params_thermo)
        update_flux_rad_single!(shape, params_thermo)

        update_force!(shape, params_thermo)
        # update_force_Rubincam!(shape, params_thermo)

        F = SVector{3}(shape.F)
        # τ = SVector{3}(shape.τ)

        F = body_to_orbit(F, spin.γ, spin.ε, spin_phase)
        # τ̄ .+= body_to_orbit(τ, spin.γ, spin.ε, spin_phase)
        push!(Fs, F)

        update_temperature!(shape, params_thermo)
    end
    Fs
    # τ̄ /= Nt
end

function init_temperature!(shape, orbit, spin, params_thermo)
    @unpack P, Δt, t_bgn, t_end, Nt, Nz = params_thermo

    for facet in shape.facets
        length(facet.Tz) == 0 ? append!(facet.Tz, zeros(Nz)) : facet.Tz .= 0
    end
    length(shape.Tz⁺) == 0 ? append!(shape.Tz⁺, zeros(Nz)) : shape.Tz⁺ .= 0
end

# ****************************************************************
#        Energy flux of sunlight, scattering, and radiation
# ****************************************************************

"""
    update_flux_sun!(shape, F☉, r̂☉)

Update illumination
"""
function update_flux_sun!(shape, F☉, r̂☉)
    for facet in shape.facets
        if isIlluminated(facet, r̂☉, shape)
            facet.flux.sun = F☉ * (facet.normal ⋅ r̂☉)
        else
            facet.flux.sun = 0.
        end
    end
end

"""
    update_flux_scat_single!(shape, params_thermo)

Single scattering of sunlight is considered.
"""
function update_flux_scat_single!(shape, params_thermo)
    @unpack A_B = params_thermo

    for facet in shape.facets
        facet.flux.scat = 0.
        for (id, f) in zip(facet.visiblefacets.id, facet.visiblefacets.f)
            facet.flux.scat += f * shape.facets[id].flux.sun
        end
        facet.flux.scat *= A_B
    end
end

"""
    update_flux_scat_mult!(shape, params_thermo)

Multiple scattering of sunlight is considered.
"""
function update_flux_scat_mult!(shape, params_thermo)
    @unpack A_B = params_thermo

    # for m in shape.smeshes
    #     m.flux.scat = 0.
    #     for (id, f) in zip(m.visiblefaces.id, m.visiblefaces.f)
    #         m.flux.scat += f * shape.smeshes[id].flux.sun
    #     end
    #     m.flux.scat *= A_B
    # end
end

"""
    update_flux_rad_single!(shape, params_thermo)

Single radiation-reabsorption is considered,
assuming albedo is close to zero at thermal infrared wavelength.
"""
function update_flux_rad_single!(shape, params_thermo)
    @unpack ϵ, A_TH = params_thermo

    for facet in shape.facets
        facet.flux.rad = 0.
        for (id, f) in zip(facet.visiblefacets.id, facet.visiblefacets.f)
            T = shape.facets[id].Tz[begin]
            facet.flux.rad += f * T^4
        end
        facet.flux.rad *= ϵ * σ_SB * (1 - A_TH)
    end
end

# ****************************************************************
#                  Photon recoil force / torque
# ****************************************************************

"""
    update_force!(shape, params_thermo)

Update photon recoil force on every facet (df)
"""
function update_force!(shape, params_thermo)  # ここが何か変わってしまったはず！
    @unpack A_B, ϵ = params_thermo

    for facet in shape.facets
        E = A_B * facet.flux.scat + ϵ * σ_SB * facet.Tz[begin]^4

        @. facet.force = facet.normal
        for vf in facet.visiblefacets
            @. facet.force -= 1.5 * vf.f * vf.d̂
        end
        @. facet.force *= - 2*E*facet.area / (3*c₀)
    end

    shape.force  .= 0
    shape.torque .= 0
    for facet in shape.facets
        r  = SVector{3}(facet.center)
        r̂  = normalize(r)
        df = SVector{3}(facet.force)

        shape.force  .+= (r̂ ⋅ df) * r̂  # Photon recoil force
        shape.torque .+= r × df        # Photon recoil torque
    end
end

"""
    update_force_Rubincam!(shape, params_thermo)

Update photon recoil force on every facet (df) based on Rubincam (2000) approximation
"""
function update_force_Rubincam!(shape, params_thermo)
    shape.torque .= 0.
    for facet in shape.facets
        facet.force .= - 2 * facet.flux.sun * facet.area / (3*c₀) .* facet.normal
        shape.torque .+= facet.center × SVector{3}(facet.force)
    end
end

# ****************************************************************
#
# ****************************************************************

"""
    getSolarIrradiation(distance) -> solar_irrad

Calculate the solar irradiation on the body

# Parameter
- `rₕ` : heliocentric distance of the body [m]

# Return
- `F☉` : solar irradiation [W/m^2]
"""
getSolarIrradiation(rₕ) = SOLAR_CONST / (rₕ / AU)^2

"""
    getSolarCondition(orbit, spin, time) -> Φ, r_sun

Get the solar irradition and the direction of the Sun

# Parameters
- `orbit` :
- `spin`  :
- `time`  : [sec]

# Returns
- `F☉` : solar irradiation [W/m^2]
- `r̂☉` : solar direction in the body-fixed frame
"""
function getSolarCondition(orbit, spin, time)
    u = solveKeplerEquation2(orbit, time)
    r = get_r(orbit, u)
    F☉ = getSolarIrradiation(norm(r))

    r̂☉ = normalize(r) * -1  # Shift the origin from the sun to the body

    spin_phase = spin.ω * time
    r̂☉ = orbit_to_body(r̂☉, spin.γ, spin.ε, spin_phase)

    F☉, r̂☉
end

"""

Integrate YORP torque over the entire surface of the shape [Rubincam2000]

# Parameters
- `shape` :
- `r_sun` : direction of the sun in the body-fixed frame (normalized)
- `Φ`     : solar irradiation [W/m^2]

Return
- `τ`

In the future, we need to improve:
- thermal pressure model: Rubincam (2000) model currently implemented.
- local shadow detection
- Tangential YORP effect
"""
function sumTorqueOverSurface(shape, F☉, r̂☉)
    τ = MVector(0., 0., 0.)  # YORP torque

    for mesh in shape.smeshes
        Ψ = mesh.normal ⋅ r̂☉  # cosine of the Sun illumination angle
        if Ψ > 0  # daytime hemisphere of the body
            df = Ψ * mesh.area * mesh.normal  # force on each facet
            dτ = mesh.center × df             # torque on each facet
            τ .+= dτ
        end
    end
    τ *= - 2/3 * F☉ / c₀
    return SVector(τ)
end

function sumTorqueOverSurface_shadowing(shape, F☉, r̂☉)
    τ = MVector(0., 0., 0.)  # YORP torque

    for mesh in shape.smeshes
        Ψ = mesh.normal ⋅ r̂☉  # cosine of the Sun illumination angle
        if Ψ > 0  # daytime hemisphere of the body
            if isIlluminated(mesh, r̂☉, shape.smeshes)
                df = Ψ * mesh.area * mesh.normal  # force on each facet
                dτ = mesh.center × df             # torque on each facet
                τ .+= dτ
            end
        end
    end
    τ *= - 2/3 * F☉ / c₀
    return SVector(τ)
end

"""
    getNetTorque(shape, orbit, spin, times) -> τ̄

Average YORP torque over given time steps
"""
function getNetTorque(shape, orbit, spin, times)

    τ̄ = MVector(0., 0., 0.)  # net YORP torque

    for time in times
        spin_phase = spin.ω * time
        F☉, r̂☉ = getSolarCondition(orbit, spin, time)
        τ = sumTorqueOverSurface(shape, F☉, r̂☉)
        τ = body_to_orbit(τ, spin.γ, spin.ε, spin_phase)

        τ̄ .+= τ
    end
    τ̄ /= length(times)
end

function getNetTorque_shadowing(shape, orbit, spin, times)

    τ̄ = MVector(0., 0., 0.)  # net YORP torque

    for time in times
        spin_phase = spin.ω * time
        F☉, r̂☉ = getSolarCondition(orbit, spin, time)
        τ = sumTorqueOverSurface_shadowing(shape, F☉, r̂☉)
        τ = body_to_orbit(τ, spin.γ, spin.ε, spin_phase)

        τ̄ .+= τ
    end
    τ̄ /= length(times)
end

# ****************************************************************
#
# ****************************************************************

"""
Integrate torque over a whole surface: Pseudo-convex model
"""
function getTotalTorque_convex(shape, F☉, r̂☉)
    τ = MVector(0., 0., 0.)
    for mesh in shape.smeshes
        Ψ = mesh.normal ⋅ r̂☉  # cosine of the Sun illumination angle
        Ψ > 0 && (τ .+= gettorque(mesh, F☉ * Ψ))
    end
    return SVector(τ)
end

"""
Integrate torque over a whole surface: Self-shadowing model
"""
function getTotalTorque_shadow(shape, F☉, r̂☉)
    τ = MVector(0., 0., 0.)
    for mesh in shape.smeshes
        Ψ = mesh.normal ⋅ r̂☉  # cosine of the Sun illumination angle
        Ψ > 0 && isIlluminated(mesh, r̂☉, shape) && (τ .+= gettorque(mesh, F☉ * Ψ))
    end
    return SVector(τ)
end

#####################

"""
Net torque during a revolutional cycle
"""
function getAnnualNetTorque(shape, orbit, spin, times; shadow)
    if shadow
        getTotalTorque = getTotalTorque_shadow
    else
        getTotalTorque = getTotalTorque_convex
    end

    getAnnualNetTorque(shape, orbit, spin, times, getTotalTorque)
end

function getAnnualNetTorque(shape, orbit, spin, times, getTotalTorque)
    τ̄ = MVector(0., 0., 0.)

    for time in times
        spin_phase = spin.ω * time
        F☉, r̂☉ = getSolarCondition(orbit, spin, time)
        τ = getTotalTorque(shape, F☉, r̂☉)
        τ = body_to_orbit(τ, spin.γ, spin.ε, spin_phase)

        τ̄ .+= τ
    end
    τ̄ /= length(times)
end

# ****************************************************************
#                            Analysis
# ****************************************************************

getτω(τ, spin) = τ ⋅ spin.ŝ
getτε(τ, spin) = τ ⋅ getSpinUnit1(spin)
getτψ(τ, spin) = τ ⋅ getSpinUnit2(spin)

"""
    torque2rate(τ, spin, C) -> ω̇, ωε̇, ωψ̇

# Parameters
- `τ`    : torque in a body's orbital plane frame
- `spin` :
- `C`    : moment of inertia
"""
function torque2rate(τ, spin, C)
    τ_ω = getτω(τ, spin)
    τ_ε = getτε(τ, spin)
    τ_ψ = getτψ(τ, spin)

    ## [rad/sec/sec]
    ω̇  = τ_ω / C
    ωε̇ = τ_ε / C
    ωψ̇ = τ_ψ / C

    ## [deg/day/day]
    ω̇  = rad2deg(ω̇)  * (3600*24)^2
    ωε̇ = rad2deg(ωε̇) * (3600*24)^2
    ωψ̇ = rad2deg(ωψ̇) * (3600*24)^2

    return ω̇, ωε̇, ωψ̇
end

function torque2rate(τ, spin)
    τ_ω = getτω(τ, spin)
    τ_ε = getτε(τ, spin)
    τ_ψ = getτψ(τ, spin)

    ## [rad/sec/sec]
    ω̇  = τ_ω / shape.I[3, 3]
    ωε̇ = τ_ε / shape.I[3, 3]
    ωψ̇ = τ_ψ / shape.I[3, 3]

    ## [deg/day/day]
    ω̇  = rad2deg(ω̇)  * (3600*24)^2
    ωε̇ = rad2deg(ωε̇) * (3600*24)^2
    ωψ̇ = rad2deg(ωψ̇) * (3600*24)^2

    return ω̇, ωε̇, ωψ̇
end

"""
Caluculate the YORP time scale given a constant acceleration/deceleration

Parameters
- `T_start` [hour]
- `T_end`   [hour]
- `ω̇`       [deg/day/day]
"""
function getTimeScale(T_start, T_end, ω̇)
    T_start *= 3600  # [sec]
    T_end   *= 3600  # [sec]

    ω_start = 2π / T_start  # initial spin rate [rad/sec]
    ω_end   = 2π / T_end    # final spin rate [rad/sec]

    Δω = ω_end - ω_start  # [rad/sec]
    Δω = rad2deg(Δω)      # [deg/sec]
    Δω *= 3600*24         # [deg/day]

    timescale = Δω / ω̇   # [day]
    timescale /= 365e6   # [Myr]
end

# ****************************************************************
#                   Time variance of torque
# ****************************************************************

"""
    getTorqueVariation(shape, orbit, spin, times) -> τs

Get time variation of YORP torque at each time step

"""
function getTorqueVsTime(shape, orbit, spin, times)

    τs = SVector{3,Float64}[]

    for time in times
        spin_phase = spin.ω * time
        F_sun, r̂_sun = getSolarCondition(orbit, spin, time)
        τ = sumTorqueOverSurface(shape, F_sun, r̂_sun)
        τ = body_to_orbit(τ, spin.γ, spin.ε, spin_phase)

        push!(τs, τ)
    end
    τs
end

function getTorqueVsTime_shadowing(shape, orbit, spin, times)

    τs = SVector{3,Float64}[]

    for time in times
        spin_phase = spin.ω * time
        F_sun, r̂_sun = getSolarCondition(orbit, spin, time)
        τ = sumTorqueOverSurface_shadowing(shape, F_sun, r̂_sun)
        τ = body_to_orbit(τ, spin.γ, spin.ε, spin_phase)

        push!(τs, τ)
    end
    τs
end

# ****************************************************************
#              Spatial distribution of net torque
# ****************************************************************

"""
"""
function addTorqueDistribution!(dτ̄s, shape, F_sun, r̂_sun, spin_phase)
    # c₀ = 299792458.0  # speed of light [m/s]

    for i in eachindex(shape.smeshes)
        mesh = shape.smeshes[i]
        Ψ = mesh.normal ⋅ r̂_sun  # cosine of the Sun illumination angle
        if Ψ > 0  # daytime hemisphere of the body
            df = - 2/3 * F_sun * Ψ / c₀ * mesh.area * mesh.normal
            dτ = mesh.center × df
            dτ = rotateZ(dτ, -spin_phase)

            dτ̄s[i] .+= dτ
        end
    end
end

function addTorqueDistribution_shadowing!(dτ̄s, shape, F_sun, r̂_sun, spin_phase)
    # c₀ = 299792458.0  # speed of light [m/s]

    for i in eachindex(shape.smeshes)
        mesh = shape.smeshes[i]
        Ψ = mesh.normal ⋅ r̂_sun  # cosine of the Sun illumination angle
        if Ψ > 0  # daytime hemisphere of the body
            if isIlluminated(mesh, r̂_sun, shape.smeshes)
                df = - 2/3 * F_sun * Ψ / c₀ * mesh.area * mesh.normal
                dτ = mesh.center × df
                dτ = rotateZ(dτ, -spin_phase)  # body's equatorial coordinate

                dτ̄s[i] .+= dτ
            end
        end
    end
end

"""
"""
function getNetTorqueDistribution(shape, orbit, spin, times)

    ## Spatial distribution of net torque
    dτ̄s = [MVector(0.,0.,0.) for i in eachindex(shape.smeshes)]

    for time in times
        spin_phase = spin.ω * time
        F_sun, r̂_sun = getSolarCondition(orbit, spin, time)
        addTorqueDistribution!(dτ̄s, shape, F_sun, r̂_sun, spin_phase)
    end
    dτ̄s ./= length(times)

    for dτ̄ in dτ̄s
        rotateX!(dτ̄, -spin.ε)  # body's ecliptic coordinate
        rotateZ!(dτ̄, -spin.γ)  # orbital plane frame
    end
    return dτ̄s
end

function getNetTorqueDistribution_shadowing(shape, orbit, spin, times)

    ## Spatial distribution of net torque
    dτ̄s = [MVector(0.,0.,0.) for i in eachindex(shape.smeshes)]

    for time in times
        spin_phase = spin.ω * time
        F_sun, r̂_sun = getSolarCondition(orbit, spin, time)
        addTorqueDistribution_shadowing!(dτ̄s, shape, F_sun, r̂_sun, spin_phase)
    end
    dτ̄s ./= length(times)

    for dτ̄ in dτ̄s
        rotateX!(dτ̄, -spin.ε)  # body's ecliptic coordinate
        rotateZ!(dτ̄, -spin.γ)  # orbital plane frame
    end
    return dτ̄s
end

max_stars_repo_name: MasanoriKanamaru/Astroshaper

max_issues_repo_name: MasanoriKanamaru/Astroshaper

max_forks_repo_name: MasanoriKanamaru/Astroshaper


text:
function predict_output_type(x...)::Type
    zero_out_ = predict_output_type_zero(x...)

    return typeof(zero_out_)
end

function predict_output_type_zero_(x...)::Number
    zero_out_ = sum(y -> predict_output_type_zero(y), x)

    return zero_out_
end

function predict_output_type_zero(x::abstractArray)::Number where {abstractArray <: AbstractArray{T}} where {T}
    zero_out_ = sum(y -> predict_output_type_zero(y), x)

    return zero_out_
end

function predict_output_type_zero(x::num)::num where {num <: Number}
    return zero(x)
end

function predict_output_type_zero(und::AbstractUnderlying)::Number
    return predict_output_type_zero(und.S0) + predict_output_type_zero(und.d)
end

function predict_output_type_zero(und::Curve)::Number
    return zero(eltype(keys_(und))) + zero(eltype(values_(und)))
end

function predict_output_type_zero(rfCurve::AbstractZeroRateCurve)::Number
    return predict_output_type_zero(rfCurve.r)
end

function predict_output_type_zero(::Any)::Int8
    zero(Int8)
end

function predict_output_type_zero(::Spot)::Int8
    return zero(Int8)
end

function predict_output_type_zero(::basePayoff) where {basePayoff <: AbstractPayoff{num}} where {num <: Number}
    return zero(num)
end

function predict_output_type_zero(x::baseProcess) where {baseProcess <: ScalarMonteCarloProcess{num}} where {num <: Number}
    return zero(num) + predict_output_type_zero(x.underlying)
end

function predict_output_type_zero(::baseProcess) where {baseProcess <: VectorialMonteCarloProcess{num}} where {num <: Number}
    return zero(num)
end

##Vectorization "utils"
import Base.iterate
iterate(x::Union{BaseProcess, AbstractZeroRateCurve, AbstractPayoff, AbstractMonteCarloConfiguration}) = (x, nothing)
iterate(::Union{BaseProcess, AbstractZeroRateCurve, AbstractPayoff, AbstractMonteCarloConfiguration}, p::Nothing) = nothing

import Base.length
length(x::Union{BaseProcess, AbstractZeroRateCurve, AbstractPayoff, AbstractMonteCarloConfiguration}) = 1


text:
"""
    get_neighbor_indices(model::Ising, (i, j)::NTuple{2, Int})

Returns the indices of the nearest neighbors for a given lattice site (i, j).

Arguments
* `model`  : AbstractMagnetismModel type
* `(i, j)` : Two-tuple containing the row and column indices of a lattice site
Returns
* Four-tuple of two-tuples representing indices of the neighbors of site (i, j)
"""
function get_neighbor_indices(model::AbstractMagnetismModel, (i, j)::NTuple{2,Int})
    return (
        (plus(i, model.params.L), j),
        (minus(i, model.params.L), j),
        (i, plus(j, model.params.L)),
        (i, minus(j, model.params.L)),
    )
end

"""
    sum_of_neighbors(model::Union{Ising, XY}, (i, j)::NTuple{2, Int})

Sums the values of the nearest neighbors for a given lattice site (i, j).

Arguments
* `model`  : Ising or XY type
* `(i, j)` : Two-tuple containing the row and column indices of a lattice site
Returns
* Sum of the spin values of the nearest-neighbors of site (i, j)
"""
function sum_of_neighbors(model::Union{Ising,XY}, (i, j)::NTuple{2,Int})
    return @. (
        model.σ[plus(i, model.params.L), j] +
        model.σ[minus(i, model.params.L), j] +
        model.σ[i, plus(j, model.params.L)] +
        model.σ[i, minus(j, model.params.L)]
    )
end

"""
    dot(a::NTuple{2, Float64}, b::NTuple{2, Float64})

Inner product of two 2D vectors.

Arguments
* `a`: Two-tuple representing a 2D vector
* `b`: Two-tuple representing a 2D vector
Returns
* Dot product between `a` and `b`
"""
function dot(a::NTuple{2,Float64}, b::NTuple{2,Float64})
    return a[1] * b[1] + a[2] * b[2]
end

"""
    compute_counts(model::Potts, (i, j)::NTuple{2, Int})

Computes the counts/frequency of the different spin vaues neighboring `σ[i, j]`.

Arguments
* `model`  : Potts type
* `(i, j)` : Two-tuple containing the row and column indices of a lattice site
Returns
* `counts` : Array where `counts[i]` is the number of neighbors of (i, j) with spin value i
"""
function compute_counts(model::Potts, (i, j)::NTuple{2,Int})
    counts = zeros(Int, model.params.q)

    counts[model.σ[plus(i, model.params.L), j]] += 1
    counts[model.σ[minus(i, model.params.L), j]] += 1
    counts[model.σ[i, plus(j, model.params.L)]] += 1
    counts[model.σ[i, minus(j, model.params.L)]] += 1

    return counts
end

"""
    flip_spin(model::XY, (i, j)::NTuple{2, Int}, r::NTuple{2, Float64})

Flips an XY model spin w.r.t. a random vector `r`.

Arguments
* `model`  : XY type
* `(i, j)` : Two-tuple containing the row and column indices of a lattice site
* `r`      : The vector to flip the spin vector at (i, j) over
Returns
* The spin vector at site (i, j) flipped over the `r` vector
"""
function flip_spin(model::XY, (i, j)::NTuple{2,Int}, r::NTuple{2,Float64})
    return model.σ[i, j] .- 2 * dot(model.σ[i, j], r) .* r
end

"""
    flip_spin!(model::XY, (i, j)::NTuple{2, Int}, r::NTuple{2, Float64})

Flips an XY model spin w.r.t. a random vector `r`.

Arguments
* `model`  : XY type
* `(i, j)` : Two-tuple containing the row and column indices of a lattice site
* `r`      : The vector to flip the spin vector at (i, j) over
"""
function flip_spin!(model::XY, (i, j)::NTuple{2,Int}, r::NTuple{2,Float64})
    model.σ[i, j] = flip_spin(model, (i, j), r)
end

"""
    minus(i::Int, L::Int)

Returns the "minus" nearest neighbor index using periodic boundary conditions.

Arguments
* `i`: Linear index
* `L`: Maximum linear index
Returns
* The "minus" nearest neighbor index using periodic boundary conditions.
"""
minus(i::Int, L::Int) = i ≠ 1 ? i - 1 : L

"""
    plus(i::Int, L::Int)

Returns the "plus" nearest neighbor index using periodic boundary conditions.

Arguments
* `i`: Linear index
* `L`: Maximum linear index
Returns
* The "plus" nearest neighbor index using periodic boundary conditions.
"""
plus(i::Int, L::Int) = i ≠ L ? i + 1 : 1


text:
function binarySearch(f, key, low, high)
    #given monotonically increasing function f, find lowest index i such that
    #f(i) >= key. For monotonically decreasing simply negate f and key, and this
    #will return lowest i such that f(i) <= key
    low > high && error("binarySearch: low > high")
    f(high) < key && error("binarySearch: search area does not include value")
    val = 0
    mid = 0
    while low != high
        #compute function
        mid = div(low + high, 2);
        val = f(mid)
        if val <= key
            low = mid + 1
        else
            high = mid
        end
    end
    if low > mid
        return low,-f(low) #have to recompute in case we're done after "low=mid+1"
    else
        return low,-val
    end
end

function pLeftOfLine(p, l1, l2)
    #tests if a point p is left (>0), on (0) or right (<0) of line between l1
    #and l2
    return ((l2[1] - l1[1])*(p[2] - l1[2])
            - (p[1] - l1[1])*(l2[2] - l1[2]))
end

function pInPolygon(p, ft)
    #use winding number algorithm to test if a point p is in polygon a.
    #returns 1 if it is in, -1 if out, 0 if on.
    #TODO: Currently returns false if it is on a corner, fix this...
    wn = Int(0)
    n = size(ft,1)
    for i = 1:n
        i_1 = i == n ? 1 : i + 1
        if (ft[i,2] <= p[2])
            #if point is between points of line, and to the left of it, +1
            if (ft[i_1,2] > p[2])
                ret = pLeftOfLine(p, ft[i,:], ft[i_1,:])
                (ret > 0) && (wn += 1)
            end
        else
            #if point is between points of line, and to the right of it, -1
            if (ft[i_1,2] <= p[2])
                ret = pLeftOfLine(p, ft[i,:], ft[i_1,:])
                (ret < 0) && (wn -= 1)
            end
        end
    end
    wn != 0
end

"""
    uniqueind(v::Vector{T}) where T <: Number -> inds,u

Given a vector of numbers `v` of length `n`, returns the unique subset `u` as
well as a vector of indices `inds` of length `n` such that `v == u[inds]`.
"""
function uniqueind(v::Vector{T}) where T <: Number
    #returns inds,vals such that vals[inds[i]] == v[i]
    inds = Array{Int}(undef, 0)
    u = Array{T}(undef, 0)
    k = 0
    for val in v
        ind = findfirst(isequal(val), u)
        if ind == nothing
            k += 1
            push!(u, val)
            push!(inds, k)
        else
            push!(inds, ind)
        end
    end
    inds,u
end

"""
    find_border(sp::ScatteringProblem) -> [x_min; x_max; y_min; y_max]

Returns bounding box that contains all of the shapes in `sp`.
"""
function find_border(sp::ScatteringProblem)
    Rmax = maximum(s.R for s in sp.shapes)
    x_max,y_max = maximum(sp.centers, dims=1) .+ 2*Rmax
    x_min,y_min = minimum(sp.centers, dims=1) .- 2*Rmax
    border = [x_min; x_max; y_min; y_max]
end

"""
    find_border(sp::ScatteringProblem, points::Array{Float64,2}) -> [x_min; x_max; y_min; y_max]

Returns bounding box that contains all of the shapes in `sp` as well as specified
`points`.
"""
function find_border(sp::ScatteringProblem, points::Array{Float64,2})
    Rmax = maximum(s.R for s in sp.shapes)
    x_max,y_max = maximum([sp.centers;points], dims=1) .+ 2*Rmax
    x_min,y_min = minimum([sp.centers;points], dims=1) .- 2*Rmax
    border = [x_min; x_max; y_min; y_max]
end

function cartesianrotation(φ)
    [cos(φ) -sin(φ);
     sin(φ) cos(φ)]
end


text:
using Test
import ArchGDAL as AG

@testset "test_drivers.jl" begin
    @testset "Testing ConfigOptions" begin
        @test AG.getconfigoption("GDAL_CACHEMAX") == ""
        AG.setconfigoption("GDAL_CACHEMAX", "64")
        @test AG.getconfigoption("GDAL_CACHEMAX") == "64"
        AG.clearconfigoption("GDAL_CACHEMAX")
        @test AG.getconfigoption("GDAL_CACHEMAX", "128") == "128"

        @test AG.getthreadconfigoption("GDAL_CACHEMAX") == ""
        AG.setthreadconfigoption("GDAL_CACHEMAX", "32")
        @test AG.getthreadconfigoption("GDAL_CACHEMAX") == "32"
        AG.clearthreadconfigoption("GDAL_CACHEMAX")
        @test AG.getthreadconfigoption("GDAL_CACHEMAX", "128") == "128"

        @test AG.getconfigoption("GDAL_CACHEMAX") == ""
        @test AG.getconfigoption("CPL_LOG_ERRORS") == ""
        @test AG.getthreadconfigoption("GDAL_CACHEMAX") == ""
        @test AG.getthreadconfigoption("CPL_LOG_ERRORS") == ""

        AG.environment(
            globalconfig = [("GDAL_CACHEMAX", "64"), ("CPL_LOG_ERRORS", "ON")],
            threadconfig = [("GDAL_CACHEMAX", "32"), ("CPL_LOG_ERRORS", "OFF")],
        ) do
            # it seems that thread settings overwrites global settings?
            @test AG.getconfigoption("GDAL_CACHEMAX") == "32"
            @test AG.getconfigoption("CPL_LOG_ERRORS") == "OFF"
            @test AG.getthreadconfigoption("GDAL_CACHEMAX") == "32"
            @test AG.getthreadconfigoption("CPL_LOG_ERRORS") == "OFF"
        end

        AG.environment(
            globalconfig = [("GDAL_CACHEMAX", "64"), ("CPL_LOG_ERRORS", "ON")],
        ) do
            # everything normal here
            @test AG.getconfigoption("GDAL_CACHEMAX") == "64"
            @test AG.getconfigoption("CPL_LOG_ERRORS") == "ON"
            @test AG.getthreadconfigoption("GDAL_CACHEMAX") == ""
            @test AG.getthreadconfigoption("CPL_LOG_ERRORS") == ""
        end
    end

    @testset "Test Driver Capabilities" begin
        drivers = AG.listdrivers()
        @test drivers["GTiff"] == "GeoTIFF"
        @test length(AG.driveroptions("GTiff")) > 100
        @test sprint(print, AG.identifydriver("data/point.geojson")) ==
              "Driver: GeoJSON/GeoJSON"
        @test sprint(print, AG.identifydriver("data/utmsmall.tif")) ==
              "Driver: GTiff/GeoTIFF"

        driver = AG.getdriver("GTiff")
        @test isnothing(AG.deregister(driver))
        @test isnothing(AG.register(driver))
        @test AG.validate(driver, ["COMPRESS=LZW", "INTERLEAVE=PIXEL"]) == true
        @test AG.validate(driver, ["COMPRESS=LZW"]) == true
        @test AG.validate(driver, ["INTERLEAVE=PIXEL"]) == true

        AG.read("data/point.geojson") do dataset
            @test AG.listcapability(dataset) == Dict(
                "CreateLayer" => false,
                "DeleteLayer" => false,
                "CreateGeomFieldAfterCreateLayer" => false,
                "CurveGeometries" => false,
                "Transactions" => false,
                "EmulatedTransactions" => false,
            )
            @test AG.listcapability(AG.getlayer(dataset, 0)) == Dict(
                "SequentialWrite" => false,
                "DeleteField" => false,
                "IgnoreFields" => false,
                "FastSpatialFilter" => false,
                "DeleteFeature" => false,
                "FastFeatureCount" => true,
                "StringsAsUTF8" => true,
                "CreateGeomField" => false,
                "ReorderFields" => false,
                "MeasuredGeometries" => true,
                "FastSetNextByIndex" => true,
                "CreateField" => false,
                "RandomWrite" => false,
                "RandomRead" => true,
                "CurveGeometries" => false,
                "FastGetExtent" => false,
                "Transactions" => false,
                "AlterFieldDefn" => false,
            )
        end
    end

    @testset "Test extensions list" begin
        exts = AG.extensions()
        @test exts[".tif"] == "GTiff"
        @test exts[".grb"] == "GRIB"
        @test exts[".geojson"] == "GeoJSON"
    end

    @testset "Test getting extensiondriver" begin
        @test AG.extensiondriver("filename.tif") == "GTiff"
        @test AG.extensiondriver(".tif") == "GTiff"
        @test AG.extensiondriver("filename.asc") == "AAIGrid"
        @test AG.extensiondriver(".asc") == "AAIGrid"
        @test_throws ArgumentError AG.extensiondriver(".not_an_extension")
    end
end


text:
## for symbolic functions (dsolve.jl)
(u::SymFunction)(x::Base.Dict) = throw(ArgumentError("IVPsolutions can only be called with symbolic objects"))
(u::SymFunction)(x::Base.Pair) = throw(ArgumentError("IVPsolutions can only be called with symbolic objects"))
function (u::SymFunction)(x)
    if u.n == 0
        u.u(SymPy.project(x))
    else
        __x = Sym("__x")
        diff(u.u(__x.x), __x, u.n)(__x => x)
    end
end

(u::SymFunction)(x, y...) = u.n== 0 ? u.u(map(SymPy.project, vcat(x, y...))...) : error("Need to implement derivatives of symbolic functions of two or more variables")

## Call symbolic object with natural syntax
## ex(x=>val)
## how to do from any symbolic object?
function (ex::Sym)(args...)
    xs = free_symbols(ex)
    subs(ex, collect(zip(xs, args))...)
end
(ex::Sym)(x::Dict) = subs(ex, x)
(ex::Sym)(x::Pair...) = subs(ex, x...)


text:
export benchmark_dopri

function benchmark_dopri(times::Int)
    y = [8.59072560e+02, -4.13720368e+03, 5.29556871e+03, 7.37289205e+00, 2.08223573e+00, 4.39999794e-01]
    y0 = copy(y)
    mu = 3.986004418e5
    el = elements(y[1:3], y[4:6], mu)
    rpar = [mu]
    x = 0.0
    xend = period(el[1], mu)
    worst = -Inf
    best = Inf
    total = 0.0
    dopri!(x, y, xend, rpar)
    for i = 1:times
        gc_enable(false)
        t = @elapsed dopri!(x, y, xend, rpar)
        gc_enable(true)
        y = copy(y0)
        x = 0.0
        if t > worst
            worst = t
        end
        if t < best
            best = t
        end
        total += t
    end
    println("[$(total/times),$best,$worst]")
end

function gravity!(_n::Ptr{Cint}, _x::Ptr{Cdouble}, _y::Ptr{Cdouble}, _f::Ptr{Cdouble},
    _rpar::Ptr{Cdouble}, _ipar::Ptr{Cint})
    n = unsafe_load(_n, 1)
    t = unsafe_load(_x, 1)
    rpar = unsafe_load(_rpar, 1)
    y = unsafe_wrap(Array, _y, n, false)
    f = unsafe_wrap(Array, _f, n, false)

    r = sqrt(y[1]*y[1]+y[2]*y[2]+y[3]*y[3])
    r3 = r*r*r
    f[1] = y[4]
    f[2] = y[5]
    f[3] = y[6]
    f[4] = -rpar[1] * y[1] / r3
    f[5] = -rpar[1] * y[2] / r3
    f[6] = -rpar[1] * y[3] / r3
    return nothing
end

function _solout(_nr::Ptr{Cint}, _xold::Ptr{Cdouble}, _x::Ptr{Cdouble},
    _y::Ptr{Cdouble}, _n::Ptr{Cint}, _con::Ptr{Cdouble}, _icomp::Ptr{Cint},
    _nd::Ptr{Cint}, _rpar::Ptr{Cdouble}, _ipar::Ptr{Cint}, _irtrn::Ptr{Cint},
    _xout::Ptr{Cdouble})
    return nothing
end

cfcn = cfunction(gravity!, Void, (Ptr{Cint}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cint}))
csolout = cfunction(_solout, Void, (Ptr{Cint}, Ptr{Cdouble}, Ptr{Cdouble},
    Ptr{Cdouble}, Ptr{Cint}, Ptr{Cdouble}, Ptr{Cint},
    Ptr{Cint}, Ptr{Cdouble}, Ptr{Cint}, Ptr{Cint},
    Ptr{Cdouble}))

function dopri!(x::Float64, y::Vector{Float64}, xend::Float64,
    rpar::Vector{Float64}=Float64[], reltol::Float64=1e-6, abstol::Float64=1e-8)
    n = length(y)
    lwork = 11*n + 8*n + 21
    liwork = n + 21
    work = zeros(Cdouble, lwork)
    iwork = zeros(Cint, liwork)
    ipar = Cint[]
    iout = 0
    idid = 0
    itol = 0
    rtol = collect(reltol)
    atol = collect(abstol)
    ccall((:c_dop853, :libdopri), Void, (Ptr{Cint}, Ptr{Void}, Ptr{Cdouble},
    Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
    Ptr{Cint}, Ptr{Void}, Ptr{Cint}, Ptr{Cdouble},
    Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cdouble},
    Ptr{Cint}, Ptr{Cint}),
    &n, cfcn, &x, y, &xend, rtol, atol, &itol, csolout, &iout,
    work, &lwork, iwork, &liwork, rpar, ipar, &idid)
end

max_issues_repo_name: OpenAstrodynamics/benchmarks

max_forks_repo_name: OpenAstrodynamics/benchmarks


text:
import JSON
using Test
using Distributed
using Statistics

arg_mode = "std"
if length(ARGS)>1 && ARGS[1]=="-mode"
    if ARGS[2] ∈ ["matrix", "parallel"]
        addprocs()
        arg_mode = ARGS[2]
    elseif ARGS[2]≠"std"
        error("-mode: invalid choice: $(ARGS[2]) (choose from 'std', 'matrix', 'parallel')")
    end
else
    error("use -mode MODE and choose MODE from 'std', 'matrix', 'parallel'")
end

data = JSON.parsefile("data.json")

"Generate GBM price paths"
function paths(S, tau, r, q, v, M::Int, N::Int)
	dt = tau/M
    S*exp.(cumsum((r-q-v/2)dt .+ √(v*dt)randn(Float32, M, N), dims=1))
end

"Price a barrier option (using path matrix)"
function barrier_mat(S, K, B, tau, r, q, v, M::Int, N::Int)
	S = paths(S, tau, r, q, v, M, N)
	l = mapslices(x->all(x.>B), S, dims=1)
    payoffs = l.*max.(S[end, :].-K, 0)'
	exp(-r*tau)*mean(payoffs)
end

"Price a barrier option"
function barrier(S, K, B, tau, r, q, v, M::Int, N::Int)
    dt = tau/M
    g1 = (r-q-v/2)*dt
    g2 = sqrt(v*dt)

    payoffs = @distributed (+) for i in 1:N
        cumsum = log(S)
        min_price = Inf32
        for j in 1:M
            cumsum += g1+g2*randn(Float32)
            price = exp(cumsum)
            if price < min_price
                min_price = price
            end
        end
        min_price>B ? max(exp(cumsum)-K, 0) : 0
    end
    exp(-r*tau)*payoffs/N
end

"Benchmark barrier option pricing"
function benchmark(mode::String)
    fun = (mode=="matrix") ? barrier_mat : barrier
    call_price() = fun(data["price"], data["strike"],
              data["barrier"], data["tau"], data["rate"], data["dy"],
              data["vol"], data["time_steps"], data["n_rep"])
    @test call_price() ≈ data["val"] atol=data["tol"]  #warm-up and validate output
    @elapsed call_price()
end

println(JSON.json(Dict("julia-$arg_mode" => benchmark(arg_mode))))

max_stars_repo_name: AndreyKolev/MLLibMark

max_issues_repo_name: AndreyKolev/MLLibMark

max_forks_repo_name: AndreyKolev/MLLibMark


text:
module HighOrderFunctionKoans

function round_list(list)
end

function apply_fn(f, x)
end

function add_n(list, n)
end

function remove_multiples(list, n)
end

end