text:
import Base: show, <, >, ==, findall
import ..CONCEPT: appendSibling
export DataFrameSystem, countAtoms, countResidues, countChains,countSystems, collectAtoms,
    collectResidues, collectChains, collectBonds, findfirst, findall, collectSystems, countBonds,
    viewSystems, viewChains, viewResidues, viewAtoms, eachAtom, eachResidue, eachChain, eachSystem,
    eachBond, sequentialResidues, is_atom_dfr, is_residue_dfr, is_chain_dfr, is_system_dfr,
    is_atom_df, is_residue_df, is_chain_df, is_system_df

const DFR = DataFrameRow

# because a dataframerow from a dataframes/subdataframes has the same column indices, these functions
# work on any AbstractDataFrame

is_atom_dfr(df::DataFrameRow) = begin
    return length(getfield(df, :colindex)) == 18  #atom df has 18 columns
end

is_residue_dfr(df::DataFrameRow) = begin
    return length(getfield(df, :colindex)) == 17  #res df has 17 columns
end

is_chain_dfr(df::DataFrameRow) = begin
    return first(propertynames(df)) == :id_
end

is_system_dfr(df::DataFrameRow) = begin
    return !is_chain_dfr(df) && length(getfield(df, :colindex)) == 13
end

is_atom_df(df::AbstractDataFrame) = begin
    return length(getfield(df, :colindex)) == 18  #atom df has 18 columns
end

#---

is_residue_df(df::AbstractDataFrame) = begin
    return length(getfield(df, :colindex)) == 17  #res df has 17 columns
end

is_chain_df(df::AbstractDataFrame) = begin
    return first(propertynames(df)) == :id_
end

is_system_df(df::AbstractDataFrame) = begin
    return !is_chain_df(df) && length(getfield(df, :colindex)) == 13
end

"""
    countBonds(node::AbstractComposite)
Returns the number of `Bond`s in the subtree rooted in `node`.
"""
countBonds(node::AbstractComposite) = begin
    count::Int = 0
    for at in collectAtoms(node)
        count += length(getBonds(at))
    end
    return count
end

"""
    countAtoms(node::AbstractComposite)
Counts the `Atom`s in the subtree rooted in `node`.
"""
countAtoms(node::AbstractComposite) = begin
    recursive_count(node, Atom)
end

"""
    countAtoms(dfs::DataFrameSystem)
Returns the number of `Atom`s in the `DataFrameSystem`.
"""
countAtoms(dfs::DataFrameSystem) = nrow(dfs.atoms)

"""
    countResidues(node::AbstractComposite)
Counts the `Atom`s in the subtree rooted in `node`.
"""
countResidues(node::AbstractComposite) = begin
    recursive_count(node, Residue)
end

"""
    countResidues(dfs::DataFrameSystem)
Returns the number of `Residue`s in the `DataFrameSystem`.
"""
countResidues(dfs::DataFrameSystem) = nrow(dfs.residues)

"""
    countChains(node::AbstractComposite)
Counts the `Atom`s in the subtree rooted in `node`.
"""
countChains(node::AbstractComposite) = begin
    recursive_count(node, Chain)
end

"""
    countChains(dfs::DataFrameSystem)
Returns the number of `Chain`s in the `DataFrameSystem`.
"""
countChains(dfs::DataFrameSystem) = nrow(dfs.chains)

"""
    countSystems(node::AbstractComposite)
Counts the `Atom`s in the subtree rooted in `node`.
"""
countSystems(node::AbstractComposite) = begin
    recursive_count(node, System)
end

"""
    countSystems(dfs::DataFrameSystem)
Returns the number of `System`s in the `DataFrameSystem`.
"""
countSystems(dfs::DataFrameSystem) = nrow(dfs.systems)

countAtoms(dfr::DataFrameRow) = begin
    return length(eachAtom(dfr))
end

"""
    collectAtoms(dfs::DataFrameSystem) -> Vector{Atom}
Collects (copies of) all the atoms in a `DataFrameSystem`.
"""
collectAtoms(dfs::DataFrameSystem) = begin
    return dfs.atoms
end

eachAtom(df::AbstractDataFrame) = begin
    atoms = DataFrameRow[]

    if is_residue_df(df)
        for row in eachrow(df)
            push!(atoms, getChildren(row)...)
        end
        return atoms
    end

    if is_chain_df(df)
        for chain in eachrow(df)
            for res in getChildren(chain)
                push!(atoms, getChildren(res)...)
            end
        end
        return atoms
    end

    if is_system_df(df)
        for system in eachrow(df)
            for chain in getChildren(system)
                for res in getChildren(chain)
                    push!(atoms, getChildren(res)...)
                end
            end
        end
        return atoms
    end
end

collectAtoms(dfr::DataFrameRow) = begin
    return eachAtom()
end

"""
    countBonds(dfs::DataFrameSystem)
Returns the number of `Bond`s in the `DataFrameSystem`.
"""
countBonds(dfs::DataFrameSystem) = nrow(dfs.bonds)

"""
    eachBond(::DataFrameSystem)
Returns an iterator of all bonds in the `DataFrameSystem`.
"""
eachBond(dfs::DataFrameSystem) = eachrow(dfs.bonds)

"""
    eachBond(::DataFrameSystem)
Returns an iterator of all bonds in the `comp`-tree.
"""
eachBond(comp::AbstractComposite) = collectBonds(comp)

eachBond(dfr::DataFrameRow) = begin
    result_set = Set{DataFrameRow}()
    for at in eachAtom(dfr)
        push!(result_set, getBonds(at)...)
    end
    return result_set
end

countBonds(dfr::DataFrameRow) = length(eachBond(dfr))

eachResidue(df::AbstractDataFrame) = begin
    residues = DataFrameRow[]

    if is_chain_df(df)
        for chain in eachrow(df)
            push!(residues, getChildren(chain)...)
        end
        return residues
    end

    if is_system_df(df)
        for system in eachrow(df)
            for chain in getChildren(system)
                push!(residues, getChildren(chain)...)
            end
        end
        return residues
    end
end

eachChain(df::AbstractDataFrame) = begin
    chains = DataFrameRow[]
    if is_system_df(df)
        for system in eachrow(df)
            push!(chains, getChildren(system)...)
        end
        return chains
    end
end

eachAtom(dfr::DataFrameRow) = begin
    if is_system_dfr(dfr)
        temp_chains = getChildren(dfr)
        temp_residues = DataFrameRow[]
        for chain in temp_chains
            push!(temp_residues, getChildren(chain)...)
        end
        atoms = DataFrameRow[]
        for child in temp_residues
            push!(atoms, getChildren(child)...)
        end

        return atoms
    end

    if is_chain_dfr(dfr)
        temp_residues = getChildren(dfr)
        atoms = DataFrameRow[]
        for child in temp_residues
            push!(atoms, getChildren(child)...)
        end

        return atoms
    end

    # if residue row get all children of row
    if is_residue_dfr(dfr)
        return getChildren(dfr)
    end

    return []
end

"""
    eachAtom(node::AbstractComposite)
Returns an iterator of all `Atom`s of the subtree rooted in `node`. See [`collectAtoms`](@ref)
"""
eachAtom(node::AbstractComposite) = begin
    return collectAtoms(node)
end

"""
    eachAtom(::DataFrameSystem)
Returns an iterator of all atoms in the `DataFrameSystem`.
"""
eachAtom(dfs::DataFrameSystem) = begin
    return eachrow(dfs.atoms)
end

"""
    collectResidues(dfs::DataFrameSystem)
Collects all the residues in a `DataFrameSystem`.
"""
collectResidues(dfs::DataFrameSystem) = begin
    return dfs.residues
end

"""
    collectResidues(node::AbstractComposite, selector::Function)
Collects all the `Residue`s of the subtree rooted in `node`.  Filters the returned list using `selector`.
"""
collectResidues(node::AbstractComposite, selector::Function) = begin
    filter!(selector, collectResidues(node))
end

"""
    eachResidue(node::AbstractComposite)
Collects a the `Residue`s of the subtree rooted in `node`. See [`collectResidues`](@ref)
"""
viewResidues(node::AbstractComposite) = begin
    return collectResidues(node)
end

"""
    eachResidue(::DataFrameSystem)
Returns an iterator of all Residues in the `DataFrameSystem`.
"""
eachResidue(dfs::DataFrameSystem) = begin
    return eachrow(dfs.residues)
end

"""
    collectChains(node::AbstractComposite, selector::Function)
Collects all the `Chain`s of the subtree rooted in `node`.  Filters the returned list using `selector`.
"""
collectChains(node::AbstractComposite, selector::Function) = begin
    filter!(selector, collectChains(node))
end

"""
    collectChains(dfs::DataFrameSystem)
Collects all the chains in a `DataFrameSystem`.
"""
collectChains(dfs::DataFrameSystem) = begin
    return dfs.chains
end

"""
    eachChain(node::AbstractComposite)
Returns an iterator of all `Chain`s of the subtree rooted in `node`. See [`collectAtoms`](@ref)
"""
eachChain(node::AbstractComposite) = begin
    return collectResidues(node)
end

"""
    eachChain(::DataFrameSystem)
Returns an iterator of all chains in the `DataFrameSystem`.
"""
eachChain(dfs::DataFrameSystem) = begin
    return eachrow(dfs.chains)
end

"""
    collectSystems(node::AbstractComposite, selector::Function)
Collects all the `System`s of the subtree rooted in `node`.  Filters the returned list using `selector`.
"""
collectSystems(node::AbstractComposite, selector::Function) = begin
    filter!(selector, collectSystems(node))
end

"""
    collectSystems(dfs::DataFrameSystem)
Collects all the systems in a `DataFrameSystem`.
"""
collectSystems(dfs::DataFrameSystem) = begin
    return dfs.systems
end

"""
    eachSystem(node::AbstractComposite)
Returns an iterator of all `System`s of the subtree rooted in `node`. See [`collectAtoms`](@ref)
"""
eachSystem(node::AbstractComposite) = begin
    return collectSystems(node)
end

"""
    eachSystem(::DataFrameSystem)
Returns an iterator of all systems in the `DataFrameSystem`.
"""
eachSystem(dfs::DataFrameSystem) = begin
    return dfs.system
end

Base.show(io::IO, comp::AbstractComposite) = print(io, typeof(comp)," \"",getName(comp), "\" with ",
countChildren(comp), " child",countChildren(comp) == 1 ? "" : "ren", " containing ", countAtoms(comp)," Atoms")

"""
    getFirstSystem(dfs::DataFrameSystem)
Returns the first System in the `dfs`.
"""
getFirstSystem(dfs::DataFrameSystem) = first(dfs.systems)

"""
    getFirstChain(dfs::DataFrameSystem)
Returns the first Chain in the `dfs`.
"""
getFirstChain(dfs::DataFrameSystem) = first(dfs.chains)

"""
    getFirstResidue(dfs::DataFrameSystem)
Returns the first Residue in the `dfs`.
"""
getFirstResidue(dfs::DataFrameSystem) = first(dfs.residue)

"""
    getFirstAtom(dfs::DataFrameSystem)
Returns the first Atom in the `dfs`.
"""
getFirstAtom(dfs::DataFrameSystem) = first(dfs.atoms)

const comparisonops = [:<,:>]
const kerneltype = [[Atom,AbstractAtom],[Residue,AbstractResidue],[Chain,AbstractChain],
                    [AbstractSystem,System]]
for x in kerneltype
    for y in kerneltype
        for op in comparisonops
            for x_type in x
                for y_type in y
                    @eval Base.$op(::Type{$x_type},::Type{$y_type}) =
                          (Base.$op(findfirst(lambda -> lambda == $x,$kerneltype),
                                   findfirst(lambda -> lambda == $y,$kerneltype)))
                end
            end

        end
    end
end

import Base.findfirst
"""
    findfirst(start_node::AbstractComposite,target_type::Type{T}, chain_attributes=(),
        residue_attributes=(), atom_attributes=()) where T<:AbstractComposite
Finds the first node for which all the attributes in the `Named Tuple`s apply.
Omitting Tuples causes the function to stop early.
"""
findfirst(start_node::AbstractComposite,target_type::Type{T}, chain_attributes=(),
        residue_attributes=(), atom_attributes=()) where T<:AbstractComposite = begin

    target_type > Chain && return nothing
    cur_node::AbstractComposite = start_node

    if typeof(cur_node) > Chain && !isempty(chain_attributes)
        for ch in collectChains(cur_node)
            if all([ch[attr] == value for (attr,value) in pairs(chain_attributes)])
                cur_node = ch
                break
            end
        end
    end

    target_type > Residue && return nothing

    if typeof(cur_node) > Residue && !isempty(residue_attributes)
        for res in collectResidues(cur_node)
            if all([res[attr] == value for (attr,value) in pairs(residue_attributes)])
                cur_node = res
                break
            end
        end
    end

    target_type > Atom && return nothing

    if typeof(cur_node) > Atom && !isempty(atom_attributes)
        for at in collectAtoms(cur_node)
            if all([at[attr] == value for (attr,value) in pairs(atom_attributes)])
                cur_node = at
                break
            end
        end
    end

    cur_node == start_node && return nothing
    return cur_node
end

import Base.getindex

Base.getindex(x::AbstractComposite,sy::Core.Symbol) = Base.getfield(x,sy)
Base.getindex(x::Atom,sy::Core.Symbol) = Base.getfield(x,sy)
Base.getindex(x::Bond, sy::Core.Symbol) = Base.getfield(x,sy)
Base.getindex(x::System, i::Int) = collectChains(x)[i]
Base.getindex(x::System, c::Char) = Base.getindex(x, Int(c- 'A') + 1)
Base.getindex(x::Residue, i::Int) = collectResidues(x)[i]
Base.getindex(x::Chain, i::Int) = collectAtoms(x)[i]
"""
    Base.getindex(vec_chain::Vector{Chain}, id::String)
Idiomatic way to get a `Chain` by name.
"""
Base.getindex(vec_chain::Vector{Chain}, id::String) = Base.getindex(vec_chain, only(id))
Base.getindex(vec_chain::Vector{Chain}, id::Char) =  vec_chain[findfirst(x->x.id_ == id, vec_chain)]

##DEPRECATED
findall(start_node::AbstractComposite,target_type::Type{T}, chain_attributes=(),
        residue_attributes=(), atom_attributes=()) where T<:AbstractComposite = begin

    target_type > Chain && return start_node

    result = target_type[]
    cur_node::AbstractComposite = start_node

    if typeof(cur_node) > Chain && !isempty(chain_attributes)
        for ch in collectChains(cur_node)
            if all([ch[attr] == value for (attr,value) in pairs(chain_attributes)])
                cur_node = ch
                if target_type == Chain
                    push!(result,ch)
                end
                break
            end
        end
    end

    if target_type > Residue
        if cur_node == start_node
            return nothing
        else
            return result
        end
    end

    if typeof(cur_node) > Residue && !isempty(residue_attributes)
        for res in collectResidues(cur_node)
            if all([res[attr] == value for (attr,value) in pairs(residue_attributes)])
                cur_node = res
                if target_type == Residue
                    push!(result,res)
                end
                break
            end
        end
    end

    if target_type > Atom
        if cur_node == start_node
            return nothing
        else
            return result
        end
    end

    if typeof(cur_node) > Atom && !isempty(atom_attributes)
        for at in collectAtoms(cur_node)
            if all([at[attr] == value for (attr,value) in pairs(atom_attributes)])
                cur_node = at
                if target_type == Atom
                    push!(result,at)
                end
                break
            end
        end
    end

    cur_node == start_node && return nothing
    return result

end

"""
    sequentialResidues(res::Union{DFR,AbstractResidue}, res_next::Union{DFR,AbstractResidue})
Returns true if `res` comes before `res_next` in the same Chain and if both are hetero/non-hetero residues.
"""
sequentialResidues(res::Union{DFR,AbstractResidue}, res_next::Union{DFR,AbstractResidue}) = begin
    return getParent(res) == getParent(res_next) &&
            isHetero(res) == isHetero(res_next) &&
            getNext(res) == res_next
end


text:
function print_crayon(io::IO, crayon, color)
    if color
        print(io, crayon)
    end
    nothing
end

"""
    print_sgf([io], sgf; [color = true])

Write the given parsed `sfg` to `io` (defaults to `STDOUT`). If
the keyword parameter `color = true`, then an ANSI based syntax
highlighting will be used.
"""
function print_sgf(A; color = true)
    print_sgf(STDOUT, A; color = color)
end

function print_sgf(io::IO, n::SGFNode; color = true)
    print_crayon(io, Crayon(foreground = 243, bold = true), color)
    print(io, ";")
    print_crayon(io, Crayon(reset = true), color)
    for p in n.properties
        print_sgf(io, p; color = color)
    end
end

function print_sgf(io::IO, p::Pair; color = true)
    print_crayon(io, Crayon(foreground = 67), color)
    print(io, " ", p.first)
    print_crayon(io, Crayon(reset = true), color)
    for v in p.second
        print(io, "[")
        print_crayon(io, Crayon(foreground = 143), color)
        if v isa String
            v2 = replace(v, "]", "\\]")
            v2 = replace(v2, ")", "\\)")
            v2 = replace(v2, ":", "\\:")
            print(io, v2)
        else
            print(io, v)
        end
        print_crayon(io, Crayon(reset = true), color)
        print(io, "]")
    end
end

function print_sgf(io::IO, t::SGFGameTree; color = true)
    print_crayon(io, Crayon(foreground = 173), color)
    print(io, "(")
    print_crayon(io, Crayon(reset = true), color)
    for n in t.sequence
        print_sgf(io, n; color = color)
    end
    for c in t.variations
        println(io)
        print_sgf(io, c; color = color)
    end
    print_crayon(io, Crayon(foreground = 173), color)
    print(io, ")")
    print_crayon(io, Crayon(reset = true), color)
end

function print_sgf(io::IO, col::AbstractVector{SGFGameTree}; color = true)
    for (i,t) in enumerate(col)
        print_sgf(io, t; color = color)
        i < length(col) && println(io)
    end
end

ext: jl


text:
"""
    AP2a(rA, rP)

Compute the semimajor axis given apoapsis and periapsis distances rA and rP.
"""
function AP2a(Q, q)
    a = (Q + q)/2
end

"""
    computeq(rv::Array,p::Array)
computeq
"""
function findrP(rv::Array, p::Array=[0,0,0])
    if size(rv[1]) == (6,) || size(rv[1]) == (3,)
        r = [rv[i][1:3] - p for i = 1:length(rv)]
        rP = minimum(norm,r)
        idx = findall(x -> norm(x) == rP, r)
    else
        rP = norm(rv[1:3] - p)
        idx = 1
    end
    return rP, idx
end

function computeME(rv::Array, μ)
    r_ECI = rv[1:3]
    v_ECI = rv[4:6]
    r = norm(r_ECI)
    v = norm(v_ECI)
    return 0.5*v^2 - μ/r
end

"""
    cart2oe(rv::Array, μ; ang_unit::Symbol=:deg)

Convert cartesian coordinates [km;km/s] in the ECI frame to orbital elements
(a, e, i, Ω, ω, ν, Π, u, l) given gravitational parameter μ.

See also: [`oe2cart`](@ref)

# Examples
```jldoctest
julia> cart2oe([1,0,0,0,1,0], 1, ang_unit=:deg)
1
```

"""
function cart2oe(rv::Array, μ; ang_unit::Symbol=:deg)
    # pull apart input vector
    r_ECI = rv[1:3]
    v_ECI = rv[4:6]
    r = norm(r_ECI)
    v = norm(v_ECI)

    # vectors describing orbit orientation
    h_ECI = cross(r_ECI, v_ECI)
    h = norm(h_ECI)     # [km²/s] specific angular momentum
    n_ECI = cross([0,0,1], h_ECI)
    n = norm(n_ECI)
    e_ECI = (1/μ)*((v^2 - μ/r)*r_ECI - (r_ECI'*v_ECI)*v_ECI)
    e = norm(e_ECI)     # [] eccentricity

    # size of orbit
    ℰ = 0.5*v^2 - μ/r   # [km²/s²] specific mechanical energy
    # ℰ = computeME(rv, μ)

    if e == 1
        a = Inf         # [km] semimajor axis
        p = h^2/μ       # [km] semilatus rectum
    else
        a = -μ/(2*ℰ)    # [km] semimajor axis
        p = a*(1 - e^2) # [km] semilatus rectum
    end

    # orientation of orbit
    i = acosd(h_ECI[3]/h)           # [deg] inclination
    Ω = acosd(n_ECI[1]/n)           # [deg] right ascension of the ascending node
    ω = acosd((n_ECI'*e_ECI)/(e*n)) # [deg] argument of periapsis
    ν = acosd((e_ECI'*r_ECI)/(e*r)) # [deg] true anomaly

    # put angles in correct domain
    if n_ECI[2] < 0;        Ω = 360 - Ω;    end
    if e_ECI[3] < 0;        ω = 360 - ω;    end
    if r_ECI'*v_ECI < 0;    ν = 360 - ν;    end

    # special cases
    Π = NaN;    u = NaN;    l = NaN;
    if i == 0 && e != 0     # elliptical equatorial
        Π = acosd(e_ECI[1]/e)           # [deg] longitude of periapsis (Π = Ω + ω)
        if e_ECI[2] < 0;    Π = 360 - Π;    end
    elseif i != 0 && e == 0 # circular inclined
        u = acosd((n_ECI'*r_ECI)/(n*r)) # [deg] argument of latitude (u = ω + ν)
        if r_ECI[3] < 0;    u = 360 - u;    end
    elseif i == 0 && e == 0 # circular equatorial
        l = acosd(r_ECI[1]/r)           # [deg] true latitude (l = Ω + ω + ν)
        if r_ECI[2] < 0;    l = 360 - l;    end
    end

    if ang_unit == :deg
    elseif ang_unit == :rad
        i = deg2rad(i)
        Ω = deg2rad(Ω)
        ω = deg2rad(ω)
        ν = deg2rad(ν)
        Π = deg2rad(Π)
        u = deg2rad(u)
        l = deg2rad(l)
    else
        error("ang_unit should be :rad or :deg")
    end

    return a, e, i, Ω, ω, ν, Π, u, l, ℰ
end

# function cart2oe(, sys::System, body::Symbol; ang_unit::Symbol=:deg)
#     if body == :prim
#         return cart2oe(rot2inert(rv, θ, μ; origin=:barycenter), sys.μ1, ang_unit=ang_unit)
#     elseif body == :sec
#         return cart2oe(rot2inert(rv, θ, μ; origin=:barycenter), sys.μ2, ang_unit=ang_unit)
#     else
#         error("body should be :prim or :sec")
#     end
# end

"""
    oe2cart(a, e, i, Ω, ω, ν, μ; ang_unit::Symbol=:deg)

Convert orbital elements to cartesian coordinates [km;km/s] in the ECI frame given
gravitational parameter μ.

See also: [`cart2oe`](@ref)

# Examples
```jldoctest
julia> oe2cart(1, 0, 0, 0, 0, 0, ang_unit=:deg)
[1,0,0,0,1,0]
```
"""
function oe2cart(a,e,i,Ω,ω,ν,μ; ang_unit::Symbol=:deg)
    if ang_unit == :deg
    elseif ang_unit == :rad
        i = rad2deg(i)
        Ω = rad2deg(Ω)
        ω = rad2deg(ω)
        ν = rad2deg(ν)
    else
        error("ang_unit should be :rad or :deg")
    end

    p = a*(1 - e^2) # [km] semilatus rectum
    r = p/(1 + e*cosd(ν))   # [km] magnitude of r_ECI
    r_PQW = [r*cosd(ν), r*sind(ν), 0];
    v_PQW = sqrt(μ/p)*[-sind(ν), e + cosd(ν), 0];
    rv_PQW = [r_PQW; v_PQW]

    r_ECI = rotzd(Ω)*rotxd(i)*rotzd(ω)*r_PQW;
    v_ECI = rotzd(Ω)*rotxd(i)*rotzd(ω)*v_PQW;
    rv_ECI = [r_ECI; v_ECI]
    return rv_ECI
end

"""
    nu2E(E, e; ang_unit::Symbol=:deg)

Convert true anomaly to eccentric anomaly given orbit eccentricity, e.

See also: [`E2nu`](@ref), [`M2E`](@ref), [`E2M`](@ref)

# Examples
```jldoctest
julia> nu2E(180, 0.5, ang_unit=:deg)
180
```
"""
function nu2E(ν, e; ang_unit::Symbol=:deg)
    if ang_unit == :deg
    elseif ang_unit == :rad
        ν = rad2deg(ν)
    else
        error("ang_unit should be :rad or :deg")
    end

    E = 2*atand(tand(nu/2)*sqrt((1 - e)/(1 + e)));

    while E >= 360; E -= 360;   end
    while E <= 0;   E += 360;   end

    if ang_unit == :deg
    elseif ang_unit == :rad
        E = deg2rad(E)
    end

    return E
end

"""
    E2nu(E, e; ang_unit::Symbol=:deg)

Convert eccentric anomaly E to true anomaly ν (\nu) given orbit eccentricity, e.

See also: [`nu2E`](@ref), [`M2E`](@ref), [`E2M`](@ref)

# Examples
```jldoctest
julia> E2nu(180, 0.5, ang_unit=:deg)
180
```
"""
function E2nu(E, e; ang_unit::Symbol=:deg)
    if ang_unit == :deg
    elseif ang_unit == :rad
        E = rad2deg(E)
    else
        error("ang_unit should be :rad or :deg")
    end

    ν = 2*atand(sqrt((1 + e)/(1 - e))*tand(E/2));
    while ν >= 360; ν -= 360;   end
    while ν <= 0;   ν += 360;   end

    if ang_unit == :deg
    elseif ang_unit == :rad
        ν = deg2rad(ν)
    end

    return ν
end

"""
    E2M(E, e; ang_unit::Symbol=:deg)

Convert eccentric anomaly to mean anomaly given orbit eccentricity, e.

See also: [`M2E`](@ref)

# Examples
```jldoctest
julia> E2M(180, 0.5, ang_unit=:deg)
180
```
"""
function E2M(E, e; ang_unit::Symbol=:deg)
    if ang_unit == :deg
    elseif ang_unit == :rad
        E = rad2deg(E)
    else
        error("ang_unit should be :rad or :deg")
    end

    M = E - e*sind(E)

    if ang_unit == :deg
    elseif ang_unit == :rad
        M = deg2rad(M)
    end

    return M
end

"""
    M2E(M, e; ang_unit::Symbol=:deg, tol=1e-10, i_max=100)

Convert eccentric anomaly to mean anomaly given orbit eccentricity, e.

See also: [`E2M`](@ref)

# Examples
```jldoctest
julia> M2E(180, 0.5, ang_unit=:deg, tol=1e-10)
180
```
"""
function M2E(E, e; ang_unit::Symbol=:deg, tol=1e-10, i_max=100)
    if ang_unit == :deg
    elseif ang_unit == :rad
        M = rad2deg(M)
    else
        error("ang_unit should be :rad or :deg")
    end

    E = M # initial guess
    δ = -(E - e*sind(E) - M)/(1 - e*cosd(E))
    iter = 0;
    while abs(δ) > abs(tol) && iter < i_max
        E = E + δ
        δ = -(E - e*sind(E) - M)/(1 - e*cosd(E))
        iter += 1
    end

    if ang_unit == :deg
    elseif ang_unit == :rad
        E = deg2rad(E)
    end

    return E
end

"""
    azel2cart(az, el, d=1; ang_unit::Symbol=:deg)

Convert azimuth and elevation angles to cartesian coordinates

See also: [`cart2azel`](@ref)

# Examples
```jldoctest
julia> azel2cart(90, 90, ang_unit=:deg)
[0,0,1]
```
"""
function azel2cart(az, el, d=1; ang_unit::Symbol=:deg)
    if ang_unit == :deg
    elseif ang_unit == :rad
        az = rad2deg(az)
        el = rad2deg(el)
    else
        error("ang_unit should be :rad or :deg")
    end

    # vector in East/North/Up coordinates
    r_ENU = [d*cosd(el)*sind(az);
             d*cosd(el)*cosd(az);
             d*sind(el)];

    return r_ENU
end

"""
    cart2azel(r_ENU; ang_unit::Symbol=:deg)

Convert cartesian coordinates from ENU to azimuth, elevation, and magnitude

See also: [`azel2cart2`](@ref)

# Examples
```jldoctest
julia> cart2azel([1,0,0], ang_unit=:deg)
90
```
"""
function cart2azel(r_ENU; ang_unit::Symbol=:deg)
    d = norm(r_ENU)
    el = asind(r_ENU[3]/d)
    az = atan2d(r_ENU[1], r_ENU[2])

    if ang_unit == :deg
    elseif ang_unit == :rad
        az = deg2rad(az)
        el = deg2rad(el)
    else
        error("ang_unit should be :rad or :deg")
    end

    return az, el, d
end

"""
    delaunay()
    input orbital elements? or rv?
    output delaunay variables
"""
function delaunay()
end


text:
using ApproxFun
using PoincareInvariants: compute_integral, compute_derivative1, compute_derivative2

### Test parameters ###

# nx  = 20
# ny  = 20
# dtol = 1E-12
# itol = 1E-11

nx  = 100
ny  = 100
dtol = 1E-10
itol = 1E-16

scale = 1E-00
# scale = 1E-15

### Set up ApproxFun spaces ###

SC = Chebyshev(0..1)^2
SU = Ultraspherical(1, 0..1)^2
CU = Conversion(SC, SU)

### check derivatives ###

# compute x=(cos(2πτ)cos(2πσ)) on SC
pc = points(SC, nx*ny)
xc = [scale * cos(2π*pc[i][1]) * cos(2π*pc[i][2]) for i in eachindex(pc)]

# compute x=(cos(2πτ)cos(2πσ)) on SU
pu = points(SU, nx*ny)
xu = [scale * cos(2π*pu[i][1]) * cos(2π*pu[i][2]) for i in eachindex(pu)]

# create Fun and compute derivatives
fun_xc = Fun(SC, transform(SC, xc))
fun_xu = Fun(SU, transform(SU, xu))

fx1 = compute_derivative1(fun_xc)
fx2 = compute_derivative2(fun_xc)

# compute τ-derivative of x (-2πsin(2πτ)cos(2πσ))
pu = points(SU, length(fx1))
x1 = [- 2π * scale * sin(2π*pu[i][1]) * cos(2π*pu[i][2]) for i in eachindex(pu)]

# compute σ-derivative of x (-2πcos(2πτ)sin(2πσ))
pu = points(SU, length(fx2))
x2 = [- 2π * scale * cos(2π*pu[i][1]) * sin(2π*pu[i][2]) for i in eachindex(pu)]

# compare Fun derivative and analytical derivative
# println(maximum(abs.(fx1 .- x1)))
# println(maximum(abs.(fx2 .- x2)))

# the following test should fail but for some reason does not
# (with nx=ny=20 success is expected for atol=5E-13 or larger)
# @test fx1 ≈ x1 atol=1E-15
# @test fx2 ≈ x2 atol=1E-15

# this works, but is less compact
@test (fx1 .- x1) ≈ zero(fx1) atol=dtol*scale
@test (fx2 .- x2) ≈ zero(fx2) atol=dtol*scale

### check integral ###

# compute x=(cos(2πτ)cos(2πσ))
pu = points(SU, nx*ny)
x  = zeros(length(pu))

for i in eachindex(x,pu)
    x[i] = scale * cos(2π*pu[i][1]) * cos(2π*pu[i][2])
end

@test compute_integral(x) ≈ 0.0 atol=itol*scale


text:
mutable struct MLP3{T}
    W1::AbstractMatrix{T}
    b1::AbstractVector{T}
    W2::AbstractMatrix{T}
    b2::AbstractVector{T}
    W3::AbstractMatrix{T}
    b3::AbstractVector{T}
end

MLP3{T}(n1, n2, n3, n4) where T = MLP3(randn(T, n2, n1), randn(T, n2),
                                       randn(T, n3, n2), randn(T, n3),
                                       randn(T, n4, n3), randn(T, n4))

MLP3(n1, n2, n3, n4) = MLP3{Float64}(n1, n2, n3, n4)

function Base.show(io::IO, m::MLP3{T}) where T
    n2, n1 = size(m.W1)
    n4, n3 = size(m.W3)
    print(io, "MLP{$T}($n1, $n2, $n3, $n4)")
end

function (m::MLP3{T})(x) where T
    x = relu(m.W1 * x .+ m.b1)
    x = relu(m.W2 * x .+ m.b2)
    x = relu(m.W3 * x .+ m.b3)
    return x
end

loss(m::MLP3, x) = sum(m(x))

function main()
    m = MLP3{Float32}(5, 3, 3, 5)
    x = rand(5, 10)
    v, g = grad(loss, m, x)
end

mutable struct Coupling
    mask  # {0, 1} array for input splitting
    s     # callable struct representing s function (scale)
    t     # callable struct representing t function (translation)
end

function f_func(c::Coupling, x)
    mask, s, t = c.mask, c.s, c.t
    xp = mask .* x
    sv = s(xp) .* (1 .- mask)
    tv = t(xp) .* (1 .- mask)
    z = (1 .- mask) .* (x .- tv) .* exp.(-sv) .+ xp
    return z
end

function g_func(c::Coupling, z)
    mask, s, t = c.mask, c.s, c.t
    zp = mask .* z
    sv = s(zp) .* (1 .- mask)
    tv = t(zp) .* (1 .- mask)
    x = zp .+ (1 .- mask) .* (z .* exp.(sv) .+ tv)
    return x
end


text:
# Computes dE(m)/dm = 2(E(m)-K(m)):

include("ellec_bulirsch.jl")
include("ellk_bulirsch.jl")

using ForwardDiff

dellec_dm = m -> ForwardDiff.derivative(ellec_bulirsch,m)

m = linspace(0.0,1.0,1000)

using PyPlot

dell = zeros(1000)
for i=1:1000
  dell[i] = dellec_dm(m[i])
end
plot(m,dell)
plot(m,.5*(ellec_bulirsch.(m)-ellk_bulirsch.(m))./m)

max_stars_repo_name: tagordon/limbdark

max_issues_repo_name: tagordon/limbdark

max_forks_repo_name: tagordon/limbdark


text:
"""
Blind tip reconstruction from AFM image
"""

function compute_xc_yc(tip)
    tip_xsiz, tip_ysiz = size(tip)
    xc = round(Int, tip_xsiz/2, RoundNearestTiesUp) - 1
    yc = round(Int, tip_ysiz/2, RoundNearestTiesUp) - 1
    return xc, yc
end

function ireflect(surface)
    surface2 = - surface[end:-1:1, end:-1:1]
    return surface2
end

function icmap(image, surface, tip; thresh=0.1)
    xc, yc = compute_xc_yc(tip)
    surf_xsiz, surf_ysiz = size(surface)
    tip_xsiz, tip_ysiz = size(tip)
    r = zeros(Int64, size(surface))
    x = y = 0
    num_contact = 0
    for i = 1:surf_xsiz
        for j = 1:surf_ysiz
            num_contact = 0
            pxmin = max(xc-tip_xsiz, 1-i)
            pymin = max(yc-tip_ysiz, 1-j)
            pxmax = min(xc-1, surf_xsiz-i)
            pymax = min(yc-1, surf_ysiz-j)
            for px = pxmin:pxmax
                for py = pymin:pymax
                    temp1 = image[i, j] - tip[xc-px, yc-py]
                    temp2 = surface[i+px, j+py]
                    if abs(temp1 - temp2) < thresh
                        num_contact += 1
                        x = i+px
                        y = j+py
                    end
                end
            end
            if num_contact == 1
                r[x, y] = 1
            end
        end
    end
    return r
end

function idilation(surface, tip)
    xc, yc = compute_xc_yc(tip)
    surf_xsiz, surf_ysiz = size(surface)
    tip_xsiz, tip_ysiz = size(tip)
    r = zeros(eltype(surface), size(surface))
    for i = 1:surf_xsiz
        for j = 1:surf_ysiz
            pxmin = max(i-surf_xsiz, -xc+1)
            pymin = max(j-surf_ysiz, -yc+1)
            pxmax = min(i-1, -xc+tip_xsiz)
            pymax = min(j-1, -yc+tip_ysiz)
            dil_max = surface[i-pxmin, j-pymin] + tip[xc+pxmin, yc+pymin]
            for px = pxmin:pxmax
                for py = pymin:pymax
                    temp = surface[i-px, j-py] + tip[xc+px, yc+py]
                    dil_max = max(temp, dil_max)
                end
            end
            r[i, j] = dil_max
        end
    end
    return r
end

function idilation_pdiff(surface, tip)
    xc, yc = compute_xc_yc(tip)
    surf_xsiz, surf_ysiz = size(surface)
    tip_xsiz, tip_ysiz = size(tip)
    r = similar(surface)
    r_index = zeros(Int, size(r))
    for i = 1:surf_xsiz
        for j = 1:surf_ysiz
            pxmin = max(i-surf_xsiz, -xc+1)
            pymin = max(j-surf_ysiz, -yc+1)
            pxmax = min(i-1, -xc+tip_xsiz)
            pymax = min(j-1, -yc+tip_ysiz)
            dil_max = surface[i-pxmin, j-pymin] + tip[xc+pxmin, yc+pymin]
            r_index[i, j] = (yc+pymin-1)*size(tip, 1) + (xc+pxmin)
            for px = pxmin:pxmax
                for py = pymin:pymax
                    temp = surface[i-px, j-py] + tip[xc+px, yc+py]
                    if temp > dil_max
                        dil_max = temp
                        r_index[i, j] = (yc+py-1)*size(tip, 1) + (xc+px)
                    end
                end
            end
            r[i, j] = dil_max
        end
    end
    return r, r_index
end

function ChainRulesCore.rrule(::typeof(idilation), surface::AbstractArray, tip::AbstractArray)
    tip_xsiz, tip_ysiz = size(tip)
    xc = round(Int, tip_xsiz/2, RoundNearestTiesUp) - 1
    yc = round(Int, tip_ysiz/2, RoundNearestTiesUp) - 1
    surf_xsiz, surf_ysiz = size(surface)
    image = similar(surface)
    MAX_i = zeros(Int, size(surface))
    MAX_j = zeros(Int, size(surface))
    MAX_u = zeros(Int, size(surface))
    MAX_v = zeros(Int, size(surface))
    for i = 1:surf_xsiz
        for j = 1:surf_ysiz
            pxmin = max(i-surf_xsiz, -xc+1)
            pymin = max(j-surf_ysiz, -yc+1)
            pxmax = min(i-1, -xc+tip_xsiz)
            pymax = min(j-1, -yc+tip_ysiz)
            dil_max = surface[i-pxmin, j-pymin] + tip[xc+pxmin, yc+pymin]
            i_max = i-pxmin
            j_max = j-pymin
            u_max = xc+pxmin
            v_max = yc+pymin
            for px = pxmin:pxmax
                for py = pymin:pymax
                    temp = surface[i-px, j-py] + tip[xc+px, yc+py]
                    if temp > dil_max
                        dil_max = temp
                        i_max = i-px
                        j_max = j-py
                        u_max = xc+px
                        v_max = yc+py
                    end
                end
            end
            MAX_i[i, j] = i_max
            MAX_j[i, j] = j_max
            MAX_u[i, j] = u_max
            MAX_v[i, j] = v_max
            image[i, j] = dil_max
        end
    end

    function idilation_pullback(dI)
        dS = zeros(eltype(surface), size(surface))
        dP = zeros(eltype(tip), size(tip))
        for i = 1:surf_xsiz
            for j = 1:surf_ysiz
                dS[MAX_i[i,j], MAX_j[i,j]] += dI[i,j]
                dP[MAX_u[i,j], MAX_v[i,j]] += dI[i,j]
            end
        end
        #dP .= dP .- dP[xc, yc]
        #dP .= max.(0.0, dP)
        #dP[xc, yc] = 0.0
        return NoTangent(), dS, dP
    end
    return image, idilation_pullback
end

function ierosion(image, tip)
    xc, yc = compute_xc_yc(tip)
    im_xsiz, im_ysiz = size(image)
    tip_xsiz, tip_ysiz = size(tip)
    r = similar(image)
    for i = 1:im_xsiz
        for j = 1:im_ysiz
            pxmin = max(-i+1, -xc+1)
            pymin = max(-j+1, -yc+1)
            pxmax = min(-i+im_xsiz, -xc+tip_xsiz)
            pymax = min(-j+im_ysiz, -yc+tip_ysiz)
            eros_min = image[i+pxmin, j+pymin] - tip[xc+pxmin, yc+pymin]
            for px = pxmin:pxmax
                for py = pymin:pymax
                    temp = image[i+px, j+py] - tip[xc+px, yc+py]
                    eros_min = min(temp, eros_min)
                end
            end
            r[i, j] = eros_min
        end
    end
    return r
end

function ierosion_pdiff(image, tip)
    xc, yc = compute_xc_yc(tip)
    im_xsiz, im_ysiz = size(image)
    tip_xsiz, tip_ysiz = size(tip)
    r = similar(image)
    r_index = zeros(Int, size(r))
    for i = 1:im_xsiz
        for j = 1:im_ysiz
            pxmin = max(-i+1, -xc+1)
            pymin = max(-j+1, -yc+1)
            pxmax = min(-i+im_xsiz, -xc+tip_xsiz)
            pymax = min(-j+im_ysiz, -yc+tip_ysiz)
            eros_min = image[i+pxmin, j+pymin] - tip[xc+pxmin, yc+pymin]
            r_index[i, j] = (yc+pymin-1)*size(tip, 1) + (xc+pxmin)
            for px = pxmin:pxmax
                for py = pymin:pymax
                    temp = image[i+px, j+py] - tip[xc+px, yc+py]
                    if temp < eros_min
                        eros_min = temp
                        r_index[i, j] = (yc+py-1)*size(tip, 1) + (xc+px)
                    end
                end
            end
            r[i, j] = eros_min
        end
    end
    return r, r_index
end

function ChainRulesCore.rrule(::typeof(ierosion), image::AbstractArray, tip::AbstractArray)
    tip_xsiz, tip_ysiz = size(tip)
    xc = round(Int, tip_xsiz/2, RoundNearestTiesUp) - 1
    yc = round(Int, tip_ysiz/2, RoundNearestTiesUp) - 1
    im_xsiz, im_ysiz = size(image)
    surface = similar(image)
    MIN_i = zeros(Int, size(image))
    MIN_j = zeros(Int, size(image))
    MIN_u = zeros(Int, size(image))
    MIN_v = zeros(Int, size(image))
    for i = 1:im_xsiz
        for j = 1:im_ysiz
            pxmin = max(-i+1, -xc+1)
            pymin = max(-j+1, -yc+1)
            pxmax = min(-i+im_xsiz, -xc+tip_xsiz)
            pymax = min(-j+im_ysiz, -yc+tip_ysiz)
            eros_min = image[i+pxmin, j+pymin] - tip[xc+pxmin, yc+pymin]
            i_min = i+pxmin
            j_min = j+pymin
            u_min = xc+pxmin
            v_min = yc+pymin
            for px = pxmin:pxmax
                for py = pymin:pymax
                    temp = image[i+px, j+py] - tip[xc+px, yc+py]
                    if temp < eros_min
                        eros_min = temp
                        i_min = i+px
                        j_min = j+py
                        u_min = xc+px
                        v_min = yc+py
                    end
                end
            end
            MIN_i[i, j] = i_min
            MIN_j[i, j] = j_min
            MIN_u[i, j] = u_min
            MIN_v[i, j] = v_min
            surface[i, j] = eros_min
        end
    end

    function ierosion_pullback(dS)
        dI = zeros(eltype(image), size(image))
        dP = zeros(eltype(tip), size(tip))
        for i = 1:im_xsiz
            for j = 1:im_ysiz
                dI[MIN_i[i,j], MIN_j[i,j]] += dS[i,j]
                dP[MIN_u[i,j], MIN_v[i,j]] -= dS[i,j]
            end
        end
        #dP .= dP .- dP[xc, yc]
        #dP .= max.(0.0, dP)
        #dP[xc, yc] = 0.0
        return NoTangent(), dI, dP
    end
    return surface, ierosion_pullback
end

function iopen(image, tip)
    r = ierosion(image, tip)
    r = idilation(r, tip)
    return r
end

function itip_estimate_point!(tip0, image, ixp, jxp; thresh=0.0)
    xc, yc = compute_xc_yc(tip0)
    im_xsiz, im_ysiz = size(image)
    tip_xsiz, tip_ysiz = size(tip0)
    interior = ((tip_xsiz-1) <= ixp) & (ixp <= (im_xsiz-tip_xsiz)) & ((tip_ysiz-1) <= jxp) & (jxp <= (im_ysiz-tip_ysiz))
    count = 0
    if interior
        for ix = 0:(tip_xsiz-1)
            for jx = 0:(tip_ysiz-1)
                if ((ix+1) == xc) & ((jx+1) == yc)
                    continue
                end
                imagep = image[ixp+1, jxp+1]
                dil = - eltype(image)(Inf)
                for id = 0:(tip_xsiz-1)
                    for jd = 0:(tip_ysiz-1)
                        #if (imagep - image[ixp+xc-id, jxp+yc-jd]) > tip0[id+1, jd+1]
                        #    continue
                        #end
                        #if ((image[ixp+1, jxp+1] - image[ixp+xc-id, jxp+yc-jd]) > 0.0) | (xc == (id+1) & yc == (jd+1))
                        #    continue
                        #end
                        if ((image[ixp+1, jxp+1] - image[ixp+xc-id, jxp+yc-jd]) > 0.0)
                            continue
                        end
                        temp = image[ix+ixp-id+1, jx+jxp-jd+1] + tip0[id+1, jd+1] - imagep + thresh
                        #temp = image[ix+ixp-id+1, jx+jxp-jd+1] - imagep + thresh
                        dil = max(dil, temp)
                    end
                end
                if isinf(-dil)
                    continue
                end
                if dil < tip0[ix+1, jx+1]
                    tip0[ix+1, jx+1] = dil
                    count += 1
                end
            end
        end
    end
    return count
end

function itip_estimate_iter!(tip0, image::Matrix{T}; thresh=0.0) where {T <: Number}
    xc, yc = compute_xc_yc(tip0)
    im_xsiz, im_ysiz = size(image)
    tip_xsiz, tip_ysiz = size(tip0)
    count = 0
    open_image = iopen(image, tip0)
    for ixp = (tip_xsiz-xc):(im_xsiz-xc)
        for jxp = (tip_ysiz-yc):(im_ysiz-yc)
            if (image[ixp+1, jxp+1] - open_image[ixp+1, jxp+1]) > thresh
                c = itip_estimate_point!(tip0, image, ixp, jxp, thresh=thresh)
                if c > 0
                    count += 1
                end
            end
        end
    end
    return count
end

function itip_estimate_single_image!(tip0, image::Matrix{T}; thresh=0.0) where {T <: Number}
    iter = 0
    count = 1
    count_cumsum = 0
    while count > 0
        iter += 1
        count = itip_estimate_iter!(tip0, image, thresh=thresh)
        count_cumsum += count
    end
    return count_cumsum
end

function itip_estimate!(tip0, images::AbstractVector; thresh=0.0)
    nframe = length(images)
    count_cumsum2 = 0
    for iframe = 1:nframe
        count_cumsum = itip_estimate_single_image!(tip0, images[iframe], thresh=thresh)
        count_cumsum2 += count_cumsum
    end
    @printf "Processed %d image\n" nframe
    @printf "%d refinements \n" count_cumsum2
    return
end

function itip_estimate!(tip0, image::AbstractMatrix; thresh=0.0)
    itip_estimate!(tip0, [image], thresh=thresh)
    return
end

function itip_least_squares!(tip0::Matrix{T}, image::Matrix{T}, surface::Matrix{T}; rate=0.01, max_convergence=0.1) where {T <: Number}
    tip_xsiz, tip_ysiz = size(tip0)
    icount = 0
    check_convergence = T(Inf)
    tip0_old = deepcopy(tip0)
    while check_convergence > max_convergence
        xc, yc = MDToolbox.compute_xc_yc(tip0)
        r, r_index = MDToolbox.idilation_pdiff(surface, tip0)
        for ix = 1:tip_xsiz
            for iy = 1:tip_ysiz
                if (ix == xc) & (iy == yc)
                    continue
                end
                id = r_index .== ((iy-1)*size(tip0, 1) + ix)
                if any(id)
                    tip0[ix, iy] -= sum(r[id] .- image[id]) * rate
                end
            end
        end
        icount += 1
        check_convergence = maximum(abs.(tip0_old .- tip0))
        tip0_old .= tip0
        if mod(icount, 1000) == 0
            #@printf "iteration %d : %f\n" icount sum((idilation(surface, tip0) .- image).^2)
            @printf "iteration %d : %f\n" icount check_convergence
        end
    end
end

function itip_least_squares!(tip0, images::Vector{Any}; thresh=0.1, rate=0.1, nstep=100)
    xc, yc = MDToolbox.compute_xc_yc(tip0)
    #tip0 .-= 10^(-8).*rand(Float64, size(tip0))
    #tip0 .= -50.0
    tip0[xc, yc] = 0.0
    nframe = length(images)
    d = zeros(eltype(tip0), size(tip0))
    loss_train_array = []
    for istep = 1:nstep
        d .= eltype(tip0)(0.0)
        for iframe = 1:nframe
            s, e_index = ierosion_pdiff(images[iframe], tip0)
            r, d_index = idilation_pdiff(s, tip0)
            #r .-= minimum(r)
            for ix = 1:size(tip0, 1)
                for iy = 1:size(tip0, 2)
                    if (ix == xc) & (iy == yc)
                        continue
                    end
                    id_d = d_index .== ((iy-1)*size(tip0, 1) + ix)
                    id_e = e_index .== ((iy-1)*size(tip0, 1) + ix)
                    for i = 1:size(images[iframe], 1)
                        for j = 1:size(images[iframe], 2)
                            if id_d[i, j] > 0.5
                                ii = i - (ix - xc)
                                jj = j - (iy - yc)
                                if 1 <= ii <= size(images[iframe], 1)
                                    if 1 <= jj <= size(images[iframe], 2)
                                        if id_e[ii, jj] > 0.5
                                            id_d[i, j] = 0
                                        end
                                    end
                                end
                            end
                        end
                    end
                    d[ix, iy] += - sum((r[id_d] .- images[iframe][id_d]))
                end
            end
        end

        tip0 .+= d .* rate
        #tip0 .= min.(tip0, 0.0)

        if mod(istep, 100) == 0
            loss_train = 0.0
            for iframe = 1:nframe
                s, e_index = ierosion_pdiff(images[iframe], tip0)
                r, d_index = idilation_pdiff(s, tip0)
                loss_train += sum((r .- images[iframe]).^2)
            end
            #loss_train += lambda * sum(abs.(tip0))
            #loss_train /= nframe

            println("step $(istep): loss_train = $(loss_train)")
            push!(loss_train_array, loss_train)
        end
    end
    return loss_train_array
end

function _logsumexp(x; tau=1)
    if tau > 0.0
        c = maximum(x)
    else
        c = minimum(x)
    end
    return c + tau * log(sum(exp.((x .- c) ./ tau)))
end

function _softmax(x; tau=1)
    m = _logsumexp(x; tau=tau)
    if tau > 0.0
        return exp.((x .- m) ./ tau)
    else
        return exp.((x .- m) ./ tau)
    end
end

function itip_least_squares_entropy!(tip0::Matrix{T}, image::Matrix{T}, surface::Matrix{T}; rate=0.01, max_convergence=0.1) where {T <: Number}
    tip_xsiz, tip_ysiz = size(tip0)
    surf_xsiz, surf_ysiz = size(surface)
    xc, yc = MDToolbox.compute_xc_yc(tip0)
    icount = 0
    check_convergence = T(Inf)
    tip0_old = deepcopy(tip0)
    tau = 1.0
    d = zeros(eltype(tip0), size(tip0))
    while (check_convergence > max_convergence) & (icount < 5000)
        r, r_index = MDToolbox.idilation_pdiff(surface, tip0)
        #tip0[ix, iy] -= sum(r[id] .- image[id]) * rate
        d .= 0.0
        for i = 1:surf_xsiz
            for j = 1:surf_ysiz
                pxmin = max(i-surf_xsiz, -xc+1)
                pymin = max(j-surf_ysiz, -yc+1)
                pxmax = min(i-1, -xc+tip_xsiz)
                pymax = min(j-1, -yc+tip_ysiz)
                temp = surface[(i-pxmin):-1:(i-pxmax), (j-pymin):-1:(j-pymax)] .+ tip0[(xc+pxmin):(xc+pxmax), (yc+pymin):(yc+pymax)]
                d[(xc+pxmin):(xc+pxmax), (yc+pymin):(yc+pymax)] .+= (r[i, j] - image[i,j]) .* _softmax(temp; tau=tau)
            end
        end
        d[xc, yc] = 0.0
        tip0 .-= (d .* rate)
        icount += 1
        check_convergence = maximum(abs.(tip0_old .- tip0))
        tip0_old .= tip0
        if mod(icount, 1000) == 0
            #@printf "iteration %d : %f\n" icount sum((idilation(surface, tip0) .- image).^2)
            @printf "iteration %d : %f\n" icount check_convergence
        end
    end
end

function itip_least_squares_entropy!(tip0::Matrix{T}, image::Matrix{T}; tau=1.0, lambda=1.0, rate=0.01, max_convergence=0.1, nstep=1000) where {T <: Number}
    tip_xsiz, tip_ysiz = size(tip0)
    surf_xsiz, surf_ysiz = size(image)
    xc, yc = MDToolbox.compute_xc_yc(tip0)
    #tip0 .-= 10^(-8).*rand(Float64, size(tip0))
    #tip0 .= -200.0
    #tip0 .= tip0 .- maximum(tip0)
    tip0[xc, yc] = 0.0
    icount = 0
    check_convergence = T(Inf)
    tip0_old = deepcopy(tip0)
    d = zeros(eltype(tip0), size(tip0))
    while (check_convergence > max_convergence) & (icount < nstep)
        s, e_index = ierosion_pdiff(image, tip0)
        r, r_index = MDToolbox.idilation_pdiff(s, tip0)
        ##############################################
        #r .= r .- minimum(r)
        ##############################################
        d .= 0.0
        for i = 1:surf_xsiz
            for j = 1:surf_ysiz
                grad1 = zeros(eltype(tip0), size(tip0))
                pxmin = max(i-surf_xsiz, -xc+1)
                pymin = max(j-surf_ysiz, -yc+1)
                pxmax = min(i-1, -xc+tip_xsiz)
                pymax = min(j-1, -yc+tip_ysiz)
                temp = s[(i-pxmin):-1:(i-pxmax), (j-pymin):-1:(j-pymax)] .+ tip0[(xc+pxmin):(xc+pxmax), (yc+pymin):(yc+pymax)]
                grad1[(xc+pxmin):(xc+pxmax), (yc+pymin):(yc+pymax)] .= _softmax(temp; tau=tau)

                grad2 = ones(eltype(tip0), size(tip0))
                t = ones(eltype(tip0), size(tip0))
                for px = pxmin:pxmax
                    for py = pymin:pymax
                        pxmin2 = max(-i+px+1, -xc+1)
                        pymin2 = max(-j+py+1, -yc+1)
                        pxmax2 = min(-i+px+surf_xsiz, -xc+tip_xsiz)
                        pymax2 = min(-j+py+surf_ysiz, -yc+tip_ysiz)
                        if (pxmin2 < pxmax2) & (pymin2 < pymax2)
                            t .= 1.0
                            t[(xc+pxmin2):(xc+pxmax2), (yc+pymin2):(yc+pymax2)] .= _softmax(image[(i-px+pxmin2):(i-px+pxmax2), (j-py+pymin2):(j-py+pymax2)] .- tip0[xc+px, yc+py]; tau=-tau) .+ 1.0
                            grad2[xc+px, yc+py] = t[xc+px, yc+py]
                        end
                    end
                end
                grad2 = ones(eltype(tip0), size(tip0))
                #d .+= (r[i, j] - image[i, j]) .* grad1 .* grad2
                d .+= - grad1 .* grad2
            end
        end
        #temp = 100.0 * (1.0 - icount / nstep)
        #temp = 100000.0
        #d .+= temp .* randn(size(d))
        #d .+= lambda .* (1.0 ./ (tip0 .+ 0.1).^2) .+ temp .* randn(size(d))
        diff = zeros(eltype(d), size(d))
        for i = 2:(tip_xsiz-1)
            for j = 2:(tip_ysiz-1)
                tip[i, j]
            end
        end
        d .+= lambda .* tip0
        #d[xc, yc] = 0.0
        d .-= d[xc, yc]
        #d .= min.(d, 0.0)
        tip0 .-= (d .* rate)
        tip0 .= min.(tip0, 0.0)
        check_convergence = maximum(abs.(tip0_old .- tip0))
        tip0_old .= tip0
        icount += 1
        if mod(icount, 10) == 0
            @printf "iteration %d : %f %f %f\n" icount sum((idilation(ierosion(image, tip0), tip0) .- image).^2) (sum((idilation(ierosion(image, tip0), tip0) .- image).^2) + lambda .* sum(tip0.^2)) check_convergence
        end
    end
end

"""
afmize

simulate afm image from a given structure

example::
pdb = readpbd("1ake.pdb");
decenter!(pdb);
afm_image = afmize(pdb, (150.0, 150.0), (16, 16));
heatmap(afm_image) # using Plots
"""
function afmize_alpha(ta::TrjArray, box, npixel)::Matrix{Float64}
    rpixel = box ./ npixel
    nx = floor.(Int, (ta.xyz[1, 1:3:end] .+ 0.5*box[1]) ./ rpixel[1]);
    ny = floor.(Int, (ta.xyz[1, 2:3:end] .+ 0.5*box[2]) ./ rpixel[2]);
    #@show nx
    #@show ny
    afm_image = zeros(Float64, npixel)
    z_min = minimum(ta.xyz[1, 3:3:end])
    z = ta.xyz[1, 3:3:end] .- z_min
    for iatom in 1:ta.natom
        if (1 <= nx[iatom] <= npixel[1]) && (1 <= ny[iatom] <= npixel[2])
            if z[iatom] > afm_image[nx[iatom], ny[iatom]]
                afm_image[nx[iatom], ny[iatom]] = z[iatom]
            end
        end
    end
    return afm_image
end

function defaultParameters()
    return Dict(
     "H" => 1.20,
    "HE" => 1.40,
     "B" => 1.92,
     "C" => 1.70,
     "CA" => 1.70,
     "CB" => 1.70,
     "CG" => 1.70,
     "CG1" => 1.70,
     "CG2" => 1.70,
     "CG3" => 1.70,
     "CD" => 1.70,
     "CD1" => 1.70,
     "CD2" => 1.70,
     "CD3" => 1.70,
     "CZ" => 1.70,
     "CZ1" => 1.70,
     "CZ2" => 1.70,
     "CZ3" => 1.70,
     "CD" => 1.70,
     "CD1" => 1.70,
     "CD2" => 1.70,
     "CD3" => 1.70,
     "CE" => 1.70,
     "CE1" => 1.70,
     "CE2" => 1.70,
     "CE3" => 1.70,
     "CH" => 1.70,
     "CH1" => 1.70,
     "CH2" => 1.70,
     "CH3" => 1.70,
     "N" => 1.55,
     "NE" => 1.55,
     "NZ" => 1.55,
     "ND1" => 1.55,
     "ND2" => 1.55,
     "NE1" => 1.55,
     "NE2" => 1.55,
     "NH1" => 1.55,
     "NH2" => 1.55,
      "O" => 1.52,
     "OH" => 1.52,
     "OG" => 1.52,
     "OE1" => 1.52,
     "OE2" => 1.52,
     "OG1" => 1.52,
     "OG2" => 1.52,
     "OD1" => 1.52,
     "OD2" => 1.52,
     "OXT" => 1.52,
     "F" => 1.47,
    #"NE" => 1.54,
    "MG" => 1.73,
    "AL" => 1.84,
    "SI" => 2.10,
     "P" => 1.80,
     "S" => 1.80,
     "SD" => 1.80,
     "SG" => 1.80,
    "CL" => 1.75,
    "AR" => 1.88,
     "K" => 2.75,
    "CA" => 2.31,
    "CYS" => 2.75,
    "PHE" => 3.2,
    "LEU" => 3.1,
    "TRP" => 3.4,
    "VAL" => 2.95,
    "ILE" => 3.1,
    "MET" => 3.1,
    "HIS" => 3.05,
    "HSD" => 3.05,
    "TYR" => 3.25,
    "ALA" => 2.5,
    "GLY" => 2.25,
    "PRO" => 2.8,
    "ASN" => 2.85,
    "THR" => 2.8,
    "SER" => 2.6,
    "ARG" => 3.3,
    "GLN" => 3.0,
    "ASP" => 2.8,
    "LYS" => 3.2,
    "GLU" => 2.95,
    )
end

mutable struct Point2D{T<:Real}
    x::T
    y::T
end

mutable struct Point3D{T<:Real}
    x::T
    y::T
    z::T
end

struct Probe
    x::Float64
    y::Float64
    r::Float64
    angle::Float64
end

mutable struct Sphere
    x::Float64
    y::Float64
    z::Float64
    r::Float64
end

mutable struct AfmizeConfig
    probeAngle::Float64  # Radian
    probeRadius::Float64
    range_min::Point2D
    range_max::Point2D
    resolution::Point2D
    atomRadiusDict::Dict{String, Float64}
end

function afmize!(tip::Matrix, config::AfmizeConfig)
    tip_xsiz, tip_ysiz = size(tip)
    xc, yc = compute_xc_yc(tip)
    #xc += 1
    #yc += 1
    @show xc, yc
    for ix = 1:tip_xsiz
        for iy = 1:tip_ysiz
            x = config.resolution.x * abs(ix - xc)
            y = config.resolution.y * abs(iy - yc)
            d = sqrt(x^2 + y^2)
            if d <= config.probeRadius
                z = sqrt(config.probeRadius^2 - d^2)
            else
                #theta = (0.5*pi) - (0.5*config.probeAngle)
                theta = (0.5*pi) - (config.probeAngle)
                z = - tan(theta) * (d - config.probeRadius)
            end
            tip[ix, iy] = z
        end
    end
    tip .= tip .- maximum(tip)
    return 0
end

function surfing(t::TrjArray, config::AfmizeConfig)
    message = checkConfig(t, config)
    if !isnothing(message)
        println(message)
        return zeros(1, 1)
    end

    width = floor(Int, (config.range_max.x - config.range_min.x) / config.resolution.x)
    height = floor(Int, (config.range_max.y - config.range_min.y) / config.resolution.y)
    stage = zeros(height, width)

    #radius_max = maximum(t.radius)
    radius = zeros(Float64, t.natom)
    for iatom = 1:t.natom
        radius[iatom] = config.atomRadiusDict[t.atomname[iatom]]
    end

    x = t.xyz[1, 1:3:end]
    y = t.xyz[1, 2:3:end]
    z = t.xyz[1, 3:3:end]
    z .= z .- minimum(z)
    for w in 1:width
        grid_x = config.range_min.x + (w-0.5) * config.resolution.x
        dx = abs.(grid_x .- x)
        index_x = dx .< radius
        for h in 1:height
            grid_y = config.range_min.y + (h-0.5) * config.resolution.y
            dy = abs.(grid_y .- y)
            index_y = dy .< radius
            index = index_x .& index_y
            if any(index)
                d = sqrt.(dx[index].^2 + dy[index].^2)
                r = radius[index]
                index2 = d .< r
                if any(index2)
                    z_surface = z[index][index2] .+ sqrt.(r[index2].^2 .- d[index2].^2)
                    stage[h, w] = maximum(z_surface)
                end
            end
        end
    end

    return stage
end

function ChainRulesCore.rrule(::typeof(surfing), xyz::AbstractArray, radius::AbstractArray, max_x::Number, min_x::Number, max_y::Number, min_y::Number, resolution_x::Number, resolution_y::Number)
    width = floor(Int, (max_x - min_x) / resolution_x)
    height = floor(Int, (max_y - min_y) / resolution_y)
    S = zeros(height, width)
    natom3 = length(xyz)
    natom = Int(natom3 / 3)

    #radius_max = maximum(t.radius)

    x = xyz[1:3:end]
    y = xyz[2:3:end]
    z = xyz[3:3:end]
    z .= z .- minimum(z)

    S_iatom = zeros(Int, height, width)

    for w in 1:width
        grid_x = min_x + (w-0.5) * resolution_x
        dx = abs.(grid_x .- x)
        index_x = dx .< radius
        for h in 1:height
            grid_y = min_y + (h-0.5) * resolution_y
            dy = abs.(grid_y .- y)
            index_y = dy .< radius
            index = index_x .& index_y
            if any(index)
                d = sqrt.(dx[index].^2 + dy[index].^2)
                r = radius[index]
                index2 = d .< r
                if any(index2)
                    z_surface = z[index][index2] .+ sqrt.(r[index2].^2 .- d[index2].^2)
                    max_s, max_iatom = findmax(z_surface)
                    S[h, w] = max_s
                    S_iatom[h, w] = (1:natom)[index][index2][max_iatom]
                end
            end
        end
    end

    function surfing_pullback(dS)
        dxyz = zeros(eltype(xyz), natom3)
        dradius = zeros(eltype(radius), natom)
        for w in 1:width
            grid_x = min_x + (w-0.5) * resolution_x
            for h in 1:height
                grid_y = min_y + (h-0.5) * resolution_y
                iatom = S_iatom[h, w]
                if iatom != 0
                    ix = 3 * (iatom - 1) + 1
                    dxyz[ix] += - 2.0 * (x[iatom] - grid_x) / sqrt(radius[iatom]^2 - (x[iatom] - grid_x)^2 - (y[iatom] - grid_y)^2) * dS[h, w]
                    iy = 3 * (iatom - 1) + 2
                    dxyz[iy] += - 2.0 * (y[iatom] - grid_y) / sqrt(radius[iatom]^2 - (x[iatom] - grid_x)^2 - (y[iatom] - grid_y)^2) * dS[h, w]
                    iz = 3 * (iatom - 1) + 3
                    dxyz[iz] += dS[h, w]
                    dradius[iatom] += 2.0 * radius[iatom] / sqrt(radius[iatom]^2 - (x[iatom] - grid_x)^2 - (y[iatom] - grid_y)^2) * dS[h, w]
                end
            end
        end
        return NoTangent(), dxyz, dradius, ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent()
    end

    return S, surfing_pullback
end

###################################################
"""
Afmize: calculate pseudo AFM image from a given molecule and a tip shape

the following codes are licensed under the MIT license (Copyright (c) 2018-2021 Tohru Niina), see LICENSE.md
"""

function AfmizeConfig(r::Float64)
    return AfmizeConfig(10.0 * (pi / 180),
                        r,
                        Point2D(-90, -90),
                        Point2D(90, 90),
                        Point2D(15, 15),
                        defaultParameters())
end

function defaultConfig()
    return AfmizeConfig(10.0 * (pi / 180),
                        20.0,
                        Point2D(-90, -90),
                        Point2D(90, 90),
                        Point2D(15, 15),
                        defaultParameters())
end

function checkConfig(tra::TrjArray ,config::AfmizeConfig)::Union{String, Nothing}
    if (config.range_max.x - config.range_min.x) % config.resolution.x != 0
        return "resolution.x must divide range"
    end
    if (config.range_max.y - config.range_min.y) % config.resolution.y != 0
        return "resolution.y must divide range"
    end
    if config.probeAngle == 0
        return "probeAngle must be positive"
    end

    for name in tra.atomname
        if !haskey(config.atomRadiusDict, name)
            return "config doesn't know atom name $(name)"
        end
    end

    return nothing
end

# 底を0にする
function moveBottom(atoms::Array{Sphere})
    bottom = atoms[1].z - atoms[1].r
    for atom in atoms
        bottom = min(bottom, atom.z - atom.r)
    end
    for atom in atoms
        atom.z -= bottom
    end
end

# 針を球として見たとき、どこで衝突するかの計算
function calcCollisionAsSphere(probe::Probe, atom::Sphere)
    distXY = sqrt((probe.x - atom.x)^2 + (probe.y - atom.y)^2)
    dr = probe.r + atom.r
    if 0 < dr^2 - distXY^2
        return atom.z + sqrt(dr^2 - distXY^2) - probe.r
    else
        return 0.0
    end
end

# 針を円形推台として見たとき、どこで衝突するかの計算
function calcCollisionAsCircularThrusters(probe::Probe, atom::Sphere)
    distXY = sqrt((probe.x - atom.x)^2 + (probe.y - atom.y)^2)
    collisionDist = probe.r + atom.r * cos(probe.angle)
    if collisionDist < distXY
        return atom.z + atom.r * sin(probe.angle) - (distXY - collisionDist) / tan(probe.angle) - probe.r
    elseif probe.r < distXY
        return atom.z + sqrt(atom.r^2 - (distXY - probe.r)^2) - probe.r
    else
        return 0.0
    end
end

# 衝突計算をする範囲を返す
function calRectangle(min_point::Point2D, max_point::Point2D, config::AfmizeConfig)
    if max_point.x <= config.range_min.x || max_point.y <= config.range_min.y
        return nothing
    end
    if config.range_max.x <= min_point.x || config.range_max.y <= min_point.y
        return nothing
    end

    min_point = Point2D(max(min_point.x, config.range_min.x),
                        max(min_point.y, config.range_min.y))
    max_point = Point2D(min(max_point.x, config.range_max.x),
                        min(max_point.y, config.range_max.y))
    width  = floor(Int, (config.range_max.x - config.range_min.x) / config.resolution.x)
    height = floor(Int, (config.range_max.y - config.range_min.y) / config.resolution.y)
    return (max(1, floor(Int, (min_point.x - config.range_min.x + config.resolution.x - 1) / config.resolution.x)),
            max(1, floor(Int, (min_point.y - config.range_min.y + config.resolution.y - 1) / config.resolution.y)),
            min(width , floor(Int, (max_point.x - config.range_min.x + config.resolution.x - 1) / config.resolution.x)),
            min(height, floor(Int, (max_point.y - config.range_min.y + config.resolution.y - 1) / config.resolution.y)))
end

function calRectangle_sphere(atom::Sphere, config::AfmizeConfig)
    min_point = Point2D(atom.x - atom.r - config.probeRadius,
                        atom.y - atom.r - config.probeRadius)
    max_point = Point2D(atom.x + atom.r + config.probeRadius,
                        atom.y + atom.r + config.probeRadius)

    return calRectangle(min_point, max_point, config)
end

function calRectangle_circularThrusters(atom::Sphere, config::AfmizeConfig)
    dist_xy = (tan(config.probeAngle) * (atom.z + atom.r * sin(config.probeAngle) - config.probeRadius)
                + config.probeRadius + atom.r * cos(config.probeAngle))
    min_point = Point2D(atom.x - dist_xy,
                        atom.y - dist_xy)
    max_point = Point2D(atom.x + dist_xy,
                        atom.y + dist_xy)

    return calRectangle(min_point, max_point, config)
end

function afmize(tra::TrjArray, config::AfmizeConfig; removeBottom=true)
    message = checkConfig(tra, config)
    if !isnothing(message)
        println(message)
        return zeros(1, 1)
    end

    width = floor(Int, (config.range_max.x - config.range_min.x) / config.resolution.x)
    height = floor(Int, (config.range_max.y - config.range_min.y) / config.resolution.y)
    atoms = [Sphere(tra.xyz[1, 3*(i-1)+1], tra.xyz[1, 3*(i-1)+2], tra.xyz[1, 3*(i-1)+3],
            config.atomRadiusDict[tra.atomname[i]]) for i in 1:tra.natom]
    if removeBottom
        moveBottom(atoms)
    end

    stage = zeros(height, width)
    probes = [Probe(config.range_min.x + (w-0.5) * config.resolution.x,
                    config.range_min.y + (h-0.5) * config.resolution.y,
                    config.probeRadius, config.probeAngle)
             for h in 1:height, w in 1:width]

    # 各原子事に計算する価値のある矩形を求めて、その中で探索を行う
    for atom in atoms
        rectangle = calRectangle_sphere(atom, config)

        if isnothing(rectangle) continue end

        for h in rectangle[2]:rectangle[4], w in rectangle[1]:rectangle[3]
            probe = probes[h, w]

            stage[h, w] = max(stage[h, w], calcCollisionAsSphere(probe, atom))
            #stage[h, w] = calcCollisionAsSphere(probe, atom)
        end
    end

    for atom in atoms
        rectangle = calRectangle_circularThrusters(atom, config)

        if isnothing(rectangle) continue end

        for h in rectangle[2]:rectangle[4], w in rectangle[1]:rectangle[3]
            probe = probes[h, w]

            stage[h, w] = max(stage[h, w], calcCollisionAsCircularThrusters(probe, atom))
            #stage[h, w] = calcCollisionAsCircularThrusters(probe, atom)
        end
    end

    return stage
end

# 高速化前
function afmize_beta(tra::TrjArray, config::AfmizeConfig)
    message = checkConfig(tra, config)
    if !isnothing(message)
        println(message)
        return zeros(1, 1)
    end

    width = floor(Int, (config.range_max.x - config.range_min.x) / config.resolution.x)
    height = floor(Int, (config.range_max.y - config.range_min.y) / config.resolution.y)
    atoms = [Sphere(tra.xyz[1, 3*(i-1)+1], tra.xyz[1, 3*(i-1)+2], tra.xyz[1, 3*(i-1)+3],
            config.atomRadiusDict[tra.atomname[i]]) for i in 1:tra.natom]
    moveBottom(atoms)

    stage = zeros(height, width)

    for h in 1:height, w in 1:width
        probe = Probe(config.range_min.x + (w-0.5) * config.resolution.x,
                      config.range_min.y + (h-0.5) * config.resolution.y,
                      config.probeRadius, config.probeAngle)
        for atom in atoms
            stage[h, w] = max(stage[h, w], calcCollisionAsSphere(probe, atom))
            stage[h, w] = max(stage[h, w], calcCollisionAsCircularThrusters(probe, atom))
        end
    end

    return stage
end

function afmize_gpu(tra::TrjArray, config::AfmizeConfig)
    message = checkConfig(tra, config)
    if !isnothing(message)
        println(message)
        return zeros(1, 1)
    end

    width = floor(Int, (config.range_max.x - config.range_min.x) / config.resolution.x)
    height = floor(Int, (config.range_max.y - config.range_min.y) / config.resolution.y)
    atom_r = tra.mass
    atom_z = tra.z[:]
    atom_z = atom_z .- (minimum(atom_z) - atom_r[argmin(atom_z)])
    #atom_z .= atom_z .- (minimum(atom_z))
    atom_z_max = maximum(atom_z)
    stage = similar(tra.x, 1, height*width)
    stage .= 0.0

    probe_x = similar(tra.x, 1, width)
    probe_y = similar(tra.x, 1, height)
    probe_x[:] .= range(config.range_min.x, step=config.resolution.x, stop=config.range_min.x + (width-1)*config.resolution.x) .+ (0.5*config.resolution.x)
    probe_y[:] .= range(config.range_min.y, step=config.resolution.y, stop=config.range_min.y + (height-1)*config.resolution.y) .+ (0.5*config.resolution.y)
    probe_r = config.probeRadius
    probe_angle = config.probeAngle

    # collision with sphere
    dist_x2 = (tra.x[:] .- probe_x).^2
    dist_y2 = (tra.y[:] .- probe_y).^2
    index_x_to_xy = [j for j=1:width for i=1:height]
    index_y_to_xy = [i for j=1:width for i=1:height]

    dist_xy2 = dist_x2[:, index_x_to_xy] .+ dist_y2[:, index_y_to_xy]
    #stage_xy = similar(dist_xy2)
    #stage_xy .= 0.0
    dr = probe_r .+ atom_r
    dr2 = dr.^2
    s = dr2 .- dist_xy2
    @show typeof(s)
    index_collide = s .> 0.0
    index_not_collide = .!index_collide
    @show typeof(index_collide)
    if any(index_collide)
        s[Array(index_collide)] .= sqrt.(s[index_collide])
    end
    s .= atom_z .+ s .- probe_r
    if any(index_not_collide)
        s[Array(index_not_collide)] .= 0.0
    end
    s_max = maximum(s, dims=1)
    stage .= max.(stage, s_max)

    # collision with circular thruster
    dist_xy2 .= sqrt.(dist_xy2)
    dist_collision = probe_r .+ atom_r .* cos(probe_angle)
    index_side   = dist_collision .< dist_xy2
    index_corner = (probe_r .< dist_xy2) .& .!index_side
    index_not_collide = .!(index_corner .| index_side)
    s[Array(index_side)] .= (atom_r .* sin.(probe_angle) .- (dist_xy2 .- dist_collision) ./ tan.(probe_angle))[index_side]
    s[Array(index_corner)] .= sqrt.( (atom_r.^2 .- (dist_xy2 .- probe_r).^2)[index_corner]  )
    s .= s .+ atom_z .- probe_r
    s[Array(index_not_collide)] .= 0.0
    s_max = maximum(s, dims=1)
    stage .= max.(stage, s_max)

    stage_reshape = reshape(stage, (height, width))
    return stage_reshape
end

###################################################
"""
AFM image analysis: calculate likelihood of many parameters, such as the orientation of molecule from AFM images.
"""

function translateafm(afm, (dx, dy))
    afm_translated = zeros(eltype(afm), size(afm))
    (ny, nx) = size(afm)
    for i in maximum([1, 1-dx]):minimum([nx, nx-dx])
        for j in maximum([1, 1-dy]):minimum([ny, ny-dy])
            afm_translated[j+dy, i+dx] = afm[j, i]
        end
    end
    afm_translated
end

function translateafm_periodic(afm, (dx, dy))
    afm_translated = zeros(eltype(afm), size(afm))
    (nx, ny) = size(afm)
    for x in 1:nx
        for y in 1:ny
            xx = x + dx
            yy = y + dy
            if xx > nx
                xx -= nx
            end
            if yy > ny
                yy -= ny
            end
            afm_translated[xx, yy] = afm[x, y]
        end
    end
    afm_translated
end

function direct_convolution(observed, calculated)
    H, W = size(observed)
    ret = zeros(H, W)
    for y in 1:H, x in 1:W
        for dy in 1:H, dx in 1:W
            ny = y + dy - 1
            nx = x + dx - 1
            if ny > H ny -= H end
            if nx > W nx -= W end
            ret[y, x] += observed[dy, dx] * calculated[ny, nx]
        end
    end
    return ret
end

function direct_convolution_non_periodic(observed, calculated)
    H, W = size(observed)
    ret = zeros(H, W)
    for y in 1:H, x in 1:W
        for dy in 1:(H-y+1), dx in 1:(W-x+1)
            ny = y + dy - 1
            nx = x + dx - 1
            ret[y, x] += observed[dy, dx] * calculated[ny, nx]
        end
    end
    return ret
end

function fft_convolution(observed, calculated)
    #return real.(ifft(fft(observed).*conj.(fft(calculated))))
    return real.(ifftshift(ifft(fft(observed).*conj.(fft(calculated)))))
end

function calcLogProb(observed::AbstractMatrix{T}, calculated::AbstractMatrix{T}, convolution_func=fft_convolution) where {T}
    npix = eltype(observed)(size(observed, 1) * size(observed, 2))
    C_o  = sum(observed)
    C_c  = sum(calculated)
    C_cc = sum(calculated.^2)
    C_oo = sum(observed.^2)
    C_oc = convolution_func(observed, calculated)

    log01 = npix .* (C_cc .* C_oo .- C_oc.^2) .+ eltype(observed)(2) .* C_o .* C_oc .* C_c .- C_cc .* C_o.^2 .- C_oo .* C_c.^2
    #log01[Array(log01 .< log(eps(eltype(observed))))] .= eps(eltype(observed))
    log02 = (npix .- eltype(observed)(2)) .* (npix .* C_cc .- C_c.^2)
    #log02 = log02 <= log(eps(eltype(observed))) ? eps(eltype(observed)) : log02
    logprob = eltype(observed)(0.5) .* (eltype(observed)(3.0) .- npix) .* log.(log01) .+ (eltype(observed)(0.5) .* npix .- eltype(observed)(2.0)) .* log(log02)

    return logprob
end

function getafmposterior(afm::Matrix{Float64}, model_array::TrjArray, q_array::Matrix{Float64}, param_array)
    imax_model = 0
    imax_q = 0
    best_model = model_array[1, :]
    best_param = param_array[1]
    best_translate = [0, 0]
    best_afm = similar(afm)
    best_posterior = -Inf

    observed = afm
    npix = size(observed, 1) * size(observed, 2)
    decenter!(model_array)

    posterior_all = zeros(Float64, size(q_array, 1)*length(param_array)*npix, size(model_array, 1))
    posterior = zeros(Float64, size(model_array, 1))

    ### loop over models(structures)
    Threads.@threads for imodel in 1:size(model_array, 1)
        @show imodel
        icount = 0
        model = model_array[imodel, :]
        ### loop over rotations
        for iq in 1:size(q_array, 1)
            q = q_array[iq, :]
            model_rotated = MDToolbox.rotate(model, q)
            ### loop over afmize parameters
            for iparam in 1:length(param_array)
                param = param_array[iparam]
                calculated = afmize(model_rotated, param)
                logprob = calcLogProb(observed, calculated)
                posterior_all[(icount+1):(icount+npix), imodel] .= logprob[:]
                icount += npix
                maximum_logprob = maximum(logprob)
                if best_posterior < maximum_logprob
                    best_posterior = maximum_logprob
                    imax_model = imodel
                    best_model = model_rotated
                    imax_q = iq
                    best_param = param
                    x_center = ceil(Int64, (size(observed,1)/2)+1.0)
                    y_center = ceil(Int64, (size(observed,2)/2)+1.0)
                    dx_estimated = argmax(logprob)[1] - x_center
                    dy_estimated = argmax(logprob)[2] - y_center
                    best_translate = (dx_estimated, dy_estimated)
                    best_afm = translateafm(calculated, best_translate)
                end
            end
        end
    end

    for imodel in 1:size(model_array, 1)
        posterior[imodel] = logsumexp(posterior_all[:, imodel])
    end
    posterior .= posterior .- maximum(posterior)
    posterior .= exp.(posterior)
    posterior .= posterior ./ sum(posterior)

    return imax_model, imax_q, best_param, best_translate, best_afm, best_posterior, posterior
end

function getafmposterior_gpu(afm::AbstractMatrix{T}, model_array::TrjArray{T, U}, q_array::AbstractMatrix{T}, param_array) where {T, U}
    imax_model = 0
    imax_q = 0
    best_param = param_array[1]
    best_translate = [0, 0]
    best_afm = similar(afm)
    best_posterior = -Inf

    observed = afm
    decenter!(model_array)

    for imodel in 1:size(model_array, 1)
        @show imodel
        model = model_array[imodel, :]
        models_rotated = MDToolbox.rotate(model, q_array)
        for iq in 1:size(q_array, 1)
            model_rotated = models_rotated[iq, :]
            for iparam in 1:length(param_array)
                param = param_array[iparam]
                #@show typeof(model_rotated.x)
                calculated = MDToolbox.afmize_gpu(model_rotated, param)
                #@show typeof(calculated)
                logprob = MDToolbox.calcLogProb(observed, calculated, MDToolbox.fft_convolution)
                maximum_logprob = maximum(logprob)
                if best_posterior < maximum_logprob
                    best_posterior = maximum_logprob
                    imax_model = imodel
                    imax_q = iq
                    best_param = param
                    x_center = ceil(Int32, (size(observed,1)/2)+1.0)
                    y_center = ceil(Int32, (size(observed,2)/2)+1.0)
                    dx_estimated = argmax(logprob)[1] - x_center
                    dy_estimated = argmax(logprob)[2] - y_center
                    best_translate = (dx_estimated, dy_estimated)
                    best_afm = translateafm(calculated, best_translate)
                end
            end
        end
    end

    return imax_model, imax_q, best_param, best_translate, best_afm, best_posterior
end

function logprob_eachmodel(model::TrjArray, afm_array, q_array::Matrix{Float64}, param_array)
    nq = size(q_array, 1)
    nparam = length(param_array)
    nafm = length(afm_array)
    npix = size(afm_array[1], 1) * size(afm_array[1], 2)
    logprob_array = zeros(eltype(afm_array[1]), nafm, nq, nparam)
    dxdy_array = Array{Tuple{Int,Int}}(undef, nafm, nq, nparam)
    dxdy = Array{Tuple{Int,Int}}(undef, nafm)
    x_center = ceil(Int64, (size(afm_array[1],2)/2)+1.0)
    y_center = ceil(Int64, (size(afm_array[2],1)/2)+1.0)

    decenter!(model)
    ### loop over rotations
    Threads.@threads for iq in 1:nq
        @show iq
        q = q_array[iq, :]
        model_rotated = MDToolbox.rotate(model, q)
        ### loop over afmize parameters
        for iparam in 1:nparam
            param = param_array[iparam]
            calculated = afmize(model_rotated, param)
            for iafm in 1:nafm
                observed = afm_array[iafm]
                logprob = calcLogProb(observed, calculated)
                logprob_array[iafm, iq, iparam] = logsumexp(logprob)
                dx = argmax(logprob)[2] - x_center
                dy = argmax(logprob)[1] - y_center
                dxdy_array[iafm, iq, iparam] = (dx, dy)
            end
        end
    end

    for iafm in 1:nafm
        ind = argmax(logprob_array[iafm, :, :])
        dxdy[iafm] = dxdy_array[iafm, :, :][ind]
    end

    return (logprob_array=logprob_array, dxdy=dxdy)
end

function getposterior_parallel(models::TrjArray, afm_array, q_array::Matrix{Float64}, param_array)
    nmodel = models.nframe
    nafm = length(afm_array)
    nq = size(q_array, 1)
    nparam = length(param_array)

    p = @showprogress pmap(x -> logprob_eachmodel(x, afm_array, q_array, param_array), models)

    logprob_all = []
    logprob_model = []
    logprob_q = []
    logprob_param = []
    dxdy_best = []
    for iafm = 1:nafm
        pp = zeros(Float64, nmodel, nq, nparam)
        for imodel = 1:nmodel
            for iq = 1:nq
                for iparam = 1:nparam
                    pp[imodel, iq, iparam] = p[imodel].logprob_array[iafm, iq, iparam]
                end
            end
        end

        push!(logprob_all, pp)

        t = zeros(Float64, nmodel)
        for imodel = 1:nmodel
            t[imodel] = logsumexp(pp[imodel, :, :][:])
        end
        push!(logprob_model, t)

        imax = argmax(t)
        push!(dxdy_best, p[imax].dxdy[iafm])

        t = zeros(Float64, nq)
        for iq = 1:nq
            t[iq] = logsumexp(pp[:, iq, :][:])
        end
        push!(logprob_q, t)

        t = zeros(Float64, nparam)
        for iparam = 1:nparam
            t[iparam] = logsumexp(pp[:, :, iparam][:])
        end
        push!(logprob_param, t)
    end

    return (all=logprob_all,
            model=logprob_model,
            q=logprob_q,
            param=logprob_param,
            dxdy=dxdy_best)
end

mutable struct posteriorResult
    each_quate_id
    each_param_id
    posterior_results
    best_translate
    best_posterior
    best_model_id
    best_quate_id
    best_param_id
    best_afm
end

mutable struct posteriorTmpData
    posteriors
    each_best
end

function outputResults(posteriorResults, fileName)
    open(fileName, "w") do out
        N = size(posteriorResults)[1]
        model_size = size(posteriorResults[1].each_quate_id)[1]
        frame_h, frame_w = size(posteriorResults[1].best_afm)
        println(out, "$(N) $(model_size)")
        for result in posteriorResults
            for i in 1:model_size
                if i != 1 print(out, " ") end
                print(out, "$(result.each_quate_id[i])")
            end
            println(out, "")
            for i in 1:model_size
                if i != 1 print(out, " ") end
                print(out, "$(result.each_param_id[i])")
            end
            println(out, "")
            for i in 1:model_size
                if i != 1 print(out, " ") end
                print(out, "$(result.posterior_results[i])")
            end
            println(out, "")
            println(out, "$(result.best_translate[1]) $(result.best_translate[2])")
            println(out, "$(result.best_posterior)")
            println(out, "$(result.best_model_id)")
            println(out, "$(result.best_quate_id)")
            println(out, "$(result.best_param_id)")
            println(out, "$(frame_h) $(frame_w)")
            for h in 1:frame_h
                for w in 1:frame_w
                    if w != 1 print(out, " ") end
                    print(out, "$(result.best_afm[h, w])")
                end
                println(out, "")
            end
        end
    end
end

function inputResults(fileName)
    ret = []
    open(fileName, "r") do io
        N, model_size = map(x->parse(Int,x),split(readline(io)))
        for result in 1:N
            each_quate_id = map(x->parse(Int,x),split(readline(io)))
            each_param_id = map(x->parse(Int,x),split(readline(io)))
            posterior_results = map(x->parse(Float64,x),split(readline(io)))
            best_translate = Tuple(map(x->parse(Int,x),split(readline(io))))
            best_posterior = parse(Float64, readline(io))
            best_model_id = parse(Int, readline(io))
            best_quate_id = parse(Int, readline(io))
            best_param_id = parse(Int, readline(io))
            frame_h, frame_w = map(x->parse(Int,x),split(readline(io)))
            frame = zeros(Float64, frame_h, frame_w)
            for h in 1:frame_h
                frame[h, :] = map(x->parse(Float64,x),split(readline(io)))
            end
            push!(ret, posteriorResult(each_quate_id,
                                        each_param_id,
                                        posterior_results,
                                        best_translate,
                                        best_posterior,
                                        best_model_id,
                                        best_quate_id,
                                        best_param_id,
                                        frame))
        end
    end
    return ret
end


text:
using HermitianDiagonalization
using Test

max_stars_repo_name: pablosanjose/HermitianDiagonalization

max_issues_repo_name: pablosanjose/HermitianDiagonalization

max_forks_repo_name: pablosanjose/HermitianDiagonalization


text:
include("../imports.jl")

@testset "`Recorder`" begin
    learner = testlearner(coeff = 3, Recorder())
    h = learner.cbstate.history[TrainingPhase()]
    @test h.epochs == 0
    @test h.steps == 0
    fit!(learner, 10)
    h = learner.cbstate.history[TrainingPhase()]
    @test h.epochs == 10
    @test h.stepsepoch == length(learner.data.training)
    @test h.steps == 10 * length(learner.data.training)
end


text:
module TestMultiLens

include("preamble.jl")
using InteractiveUtils
using StaticArrays

function test_positional_set(ml)
    @test set((x=1, y=(z=2,)), ml, ("x", "y.z")) === (x="x", y=(z="y.z",))
    @test set((x=1, y=(z=2,)), ml, (:x, "y.z")) === (x=:x, y=(z="y.z",))
    @test set((y=(z=2,), x=1, a=0), ml, ("x", "y.z")) ===
        (y=(z="y.z",), x="x", a=0)

    for val in Any[
        (10, 20),
        [10, 20],
        [10, 20]',
        Broadcast.broadcasted(x -> 10x, [1, 2]),
        SVector(10, 20),
        SMatrix{1, 2}(10, 20),
    ]
        @test set((x=1, y=(z=2,)), ml, val) === (x=10, y=(z=20,))
    end
end

@testset "Tuple" begin
    ml = MultiLens((
        (@lens _.x),
        (@lens _.y.z),
    ))

    @test get((x=1, y=(z=2,)), ml) === (1, 2)
    @test get((y=(z=2,), x=1, a=0), ml) === (1, 2)

    test_positional_set(ml)
end

function test_namedtuple_set(ml)
    @testset for val in [(a="x", b="y.z")
                         (b="y.z", a="x")]
        @test set((x=1, y=(z=2,)), ml, val) === (x="x", y=(z="y.z",))
        @test set((y=(z=2,), x=1, a=0), ml, val) === (y=(z="y.z",), x="x", a=0)
    end
    @test set((x=1, y=(z=2,)), ml, (a=:x, b="y.z")) === (x=:x, y=(z="y.z",))
end

@testset "NamedTuple" begin
    ml = MultiLens((
        a = (@lens _.x),
        b = (@lens _.y.z),
    ))

    @test get((x=1, y=(z=2,)), ml) === (a=1, b=2)
    @test get((y=(z=2,), x=1, a=0), ml) === (a=1, b=2)

    test_namedtuple_set(ml)
end

@testset "castout = Tuple" begin
    # Input is a `NamedTuple`; Output is a `Tuple`
    ml = MultiLens(
        Tuple,
        (
            a = (@lens _.x),
            b = (@lens _.y.z),
        )
    )

    @test get((x=1, y=(z=2,)), ml) === (1, 2)
    @test get((y=(z=2,), x=1, a=0), ml) === (1, 2)

    test_namedtuple_set(ml)
end

@testset "castout = SVector ∘ Tuple" begin
    # Input is a `NamedTuple`; Output is a `SVector`
    ml = MultiLens(
        SVector ∘ Tuple,
        (
            a = (@lens _.x),
            b = (@lens _.y.z),
        )
    )

    @test get((x=1, y=(z=2,)), ml) === SVector(1, 2)
    @test get((y=(z=2,), x=1, a=0), ml) === SVector(1, 2)

    test_namedtuple_set(ml)
end

function codegen_multilens_tuple()
    obj = (
        a = (b = :x,),
        c = (:y, :z),
    )
    lens = MultiLens((
        (@lens _.a.b),
        # (@lens _.c[1]),
        # (@lens _.c[2]),
        (@lens _.c[$1]),
        (@lens _.c[$2]),
    ))
    x = (1, 2, 3)
    return sum(get(set(obj, lens, x), lens))
end

function codegen_multilens_namedtuple()
    obj = (
        a = (b = :x,),
        c = (:y, :z),
    )
    lens = MultiLens((
        i = (@lens _.a.b),
        j = (@lens _.c[$1]),
        k = (@lens _.c[$2]),
    ))
    x = (k = 1, i = 2, j = 3)
    return sum(Tuple(get(set(obj, lens, x), lens)))
end

function codegen_multilens_svector()
    obj = (
        a = (b = :x,),
        c = (:y, :z),
    )
    lens = MultiLens(
        SVector,
        (
            (@lens _.a.b),
            (@lens _.c[$1]),
            (@lens _.c[$2]),
        )
    )
    x = (1, 2, 3)
    return sum(get(set(obj, lens, x), lens))
end

@testset "Gode gen: $(nameof(f))" for f in [
    codegen_multilens_tuple
    codegen_multilens_namedtuple
    codegen_multilens_svector
]
    @test f() == 6
    llvm = sprint(code_llvm, f, Tuple{})
    @test occursin(r"i(32|64) 6\b", llvm)
end

end  # module


text:
RADWIMPS = stdout

function then(io)
  print(io, "前")
  io
end

function 世(io)
  println(io, "世")
  io
end

function main()
  RADWIMPS |> then |> then |> then |> 世
end

main()

max_stars_repo_name: Clapp-project/RADWIMPS

max_issues_repo_name: Clapp-project/RADWIMPS

max_forks_repo_name: Clapp-project/RADWIMPS


text:
using SetProg
using Polyhedra

using JuMP
const MOI = JuMP.MOI

□ = polyhedron(HalfSpace([1, 0], 1.0) ∩ HalfSpace([-1, 0], 1) ∩ HalfSpace([0, 1], 1) ∩ HalfSpace([0, -1], 1))

mock = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}())
model = JuMP.direct_model(mock);
# Q = [1 0
#      0 1]
# t = √det(Q) = 1                                                                  Q11  Q12  Q22  t
MOI.Utilities.set_mock_optimize!(mock, mock -> MOI.Utilities.mock_optimize!(mock, [1.0, 0.0, 1.0, 1.0]));

@variable(model, ◯, PolySet(degree=4, dimension=2, convex=true))
cref = @constraint(model, ◯ ⊆ □)
@objective(model, Max, nth_root(volume(◯)))

SetProg.optimize!(model)


text:
using DualNumbers
using Random
using StaticArrays
using Test

Random.seed!(0)
@testset "Ring" begin
    T = Int

    randint() = T(rand(-100:100))

    for iter in 1:100
        n = Dual(T(0))
        e = Dual(T(1))
        x = Dual(randint(), randint())
        y = Dual(randint(), randint())
        z = Dual(randint(), randint())
        a = randint()
        b = randint()
        i = rand(1:10)

        @test Dual(a).primal === a
        @test Dual(a).dual === 0
        @test Dual(a, b).primal === a
        @test Dual(a, b).dual === b

        @test +x === x
        @test -(-x) === x

        @test n + x === x
        @test x + n === x

        @test -x === n - x
        @test x - x === n

        @test x + y === y + x

        @test (x + y) + z === x + (y + z)

        @test n * x === n
        @test x * n === n

        @test e * x === x
        @test x * e === x

        @test x * y === y * x
        @test (x * y) * z === x * (y * z)

        @test x^0 == e
        @test x^1 == x
        @test x^i == x * x^(i - 1)

        @test x + a === a + x

        @test Dual(a) == a
        @test (Dual(a) == b) == (a == b)
        @test (Dual(a) < b) == (a < b)
        @test (Dual(a) > b) == (a > b)
    end
end

Random.seed!(0)
@testset "Field" begin
    T = Rational{Int128}

    randint() = T(rand(-100:100))
    randintnz() =
        while true
            i = randint()
            i != 0 && return i
        end
    randrat() = randint() // randintnz()

    for iter in 1:100
        n = Dual(T(0))
        e = Dual(T(1))
        x = Dual(randrat(), randrat())
        y = Dual(randrat(), randrat())
        z = Dual(randrat(), randrat())
        a = randrat()
        b = randrat()

        if x != 0
            @test inv(x) === 1 / x
            @test inv(inv(x)) === x

            @test e / x === inv(x)
            @test x \ y === y / x

            if y != 0
                @test x / y === inv(y / x)
                @test x \ y === inv(y \ x)
            end
        end
    end
end

Random.seed!(0)
@testset "Real" begin
    T = Float64

    randfloat() = rand(T) * 200 - 100

    for iter in 1:100
        n = Dual(T(0))
        e = Dual(T(1))
        x = Dual(randfloat(), randfloat())
        y = Dual(randfloat(), randfloat())
        z = Dual(randfloat(), randfloat())
        a = randfloat()
        b = randfloat()

        @test (abs(x)^a).primal ≈ abs(x).primal^a
        @test cos(x).primal ≈ cos(x.primal)
        @test sin(x).primal ≈ sin(x.primal)
        @test sqrt(abs(x)).primal ≈ sqrt(abs(x).primal)
    end
end

Random.seed!(0)
@testset "Basic Derivatives" begin
    T = Float64

    randfloat() = rand(T)

    for iter in 1:100
        x = randfloat()
        a = randfloat()

        @test derivative(x -> x + a, x) ≈ 1
        @test derivative(x -> a * x, x) ≈ a
        @test derivative(x -> x^2, x) ≈ 2x
        @test derivative(cos, x) ≈ -sin(x)
        @test derivative(sin, x) ≈ cos(x)
        @test derivative(sqrt, x) ≈ 1 / 2sqrt(x)
    end
end

Random.seed!(0)
@testset "Derivative laws" begin
    T = Float64

    randfloat() = rand(T)

    for iter in 1:100
        x = randfloat()

        c0 = randfloat()
        c1 = randfloat()
        fs = [x -> x + c0, x -> c1 * x, x -> x^2, cos, sin, x -> sqrt(x + 1)]
        i = identity
        f = rand(fs)
        g = rand(fs)
        h = rand(fs)

        a = randfloat()

        @test (g ∘ f)(x) ≈ g(f(x))

        @test derivative(x -> f(x) + g(x), x) ≈ derivative(f, x) + derivative(g, x)
        @test derivative(x -> a * f(x), x) ≈ a * derivative(f, x)

        @test derivative(identity ∘ f, x) ≈ derivative(f, x)
        @test derivative(f ∘ identity, x) ≈ derivative(f, x)

        @test derivative(identity, x) ≈ 1
        f′(x) = derivative(f, x)
        g′(x) = derivative(g, x)
        @test derivative(g ∘ f, x) ≈ g′(f(x)) * f′(x)
    end
end

Random.seed!(0)
@testset "Derivatives of multi-valued functions" begin
    T = Float64

    randfloat() = rand(T)

    for iter in 1:100
        xs = SVector(randfloat(), randfloat(), randfloat())

        c0 = randfloat()
        c1 = randfloat()
        fs = [x -> x + c0, x -> c1 * x, x -> x^2, cos, sin, x -> sqrt(x + 1)]
        f1 = rand(fs)
        f2 = rand(fs)
        f3 = rand(fs)
        f(xs) = f1(xs[1]) + f2(xs[2]) * f3(xs[3])
        g(xs) = f1(xs[1]) * f2(xs[2]) * f3(xs[3])

        a = randfloat()

        @test derivative(f, xs, 1) ≈ derivative(f1, xs[1])
        @test derivative(f, xs, 2) ≈ derivative(f2, xs[2]) * f3(xs[3])
        @test derivative(f, xs, 3) ≈ f2(xs[2]) * derivative(f3, xs[3])

        for i in 1:3
            @test derivative(xs -> f(xs) + g(xs), xs, i) ≈ derivative(f, xs, i) + derivative(g, xs, i)
            @test derivative(xs -> a * f(xs), xs, i) ≈ a * derivative(f, xs, i)
        end
    end
end

Random.seed!(0)
@testset "Derivatives of function with complex arguments" begin
    T = Float64

    randfloat() = rand(T)

    for iter in 1:100
        x = Complex(randfloat(), randfloat())
        X = SVector(Complex(randfloat(), randfloat()), Complex(randfloat(), randfloat()))

        c0 = randfloat()
        c1 = randfloat()
        fs = [x -> x + c0, x -> c1 * x, x -> x^2, cos, sin, x -> sqrt(x + 1)]
        f1 = rand(fs)
        f2 = rand(fs)
        f3 = rand(fs)
        f4 = rand(fs)
        g1 = rand(fs)
        g2 = rand(fs)
        g3 = rand(fs)
        g4 = rand(fs)
        f(x) = (f1(real(x)) + im * f2(real(x))) + (f3(imag(x)) + im * f4(imag(x)))
        g(x) = (g1(real(x)) + im * g2(real(x))) * (g3(imag(x)) + im * g4(imag(x)))

        F(xs) = f(xs[1]) + g(xs[2])
        G(xs) = f(xs[2]) * g(xs[1])

        a = randfloat()

        @test derivative(f, x, 1) ≈ derivative(f1, real(x)) + im * derivative(f2, real(x))
        @test derivative(f, x, 2) ≈ derivative(f3, imag(x)) + im * derivative(f4, imag(x))

        for c in 1:2
            @test derivative(x -> f(x) + g(x), x, c) ≈ derivative(f, x, c) + derivative(g, x, c)
            @test derivative(x -> a * f(x), x, c) ≈ a * derivative(f, x, c)
        end

        for i in 1:2, c in 1:2
            @test derivative(X -> F(X) + G(X), X, i, c) ≈ derivative(F, X, i, c) + derivative(G, X, i, c)
            @test derivative(X -> a * F(X), X, i, c) ≈ a * derivative(F, X, i, c)
        end
    end
end