text:
using LLVM
using Libdl
import GPUCompiler: IRError, InvalidIRError

function check_ir(job, args...)
    errors = check_ir!(job, IRError[], args...)
    unique!(errors)
    if !isempty(errors)
        throw(InvalidIRError(job, errors))
    end

    return
end

function check_ir!(job, errors, mod::LLVM.Module)
    for f in functions(mod)
        check_ir!(job, errors, f)
    end

    return errors
end

function check_ir!(job, errors, f::LLVM.Function)
    for bb in blocks(f), inst in instructions(bb)
        if isa(inst, LLVM.CallInst)
            check_ir!(job, errors, inst)
        end
    end

    return errors
end

const libjulia = Ref{Ptr{Cvoid}}(C_NULL)

import GPUCompiler: DYNAMIC_CALL, DELAYED_BINDING, RUNTIME_FUNCTION, UNKNOWN_FUNCTION, POINTER_FUNCTION
import GPUCompiler: backtrace, isintrinsic
function check_ir!(job, errors, inst::LLVM.CallInst)
    bt = backtrace(inst)
    dest = called_value(inst)
    if isa(dest, LLVM.Function)
        fn = LLVM.name(dest)

        # some special handling for runtime functions that we don't implement
        if fn == "jl_get_binding_or_error"
            try
                m, sym, _ = operands(inst)
                sym = first(operands(sym::ConstantExpr))::ConstantInt
                sym = convert(Int, sym)
                sym = Ptr{Cvoid}(sym)
                sym = Base.unsafe_pointer_to_objref(sym)
                push!(errors, (DELAYED_BINDING, bt, sym))
            catch e
                isa(e,TypeError) || rethrow()
                @debug "Decoding arguments to jl_get_binding_or_error failed" inst bb=LLVM.parent(inst)
                push!(errors, (DELAYED_BINDING, bt, nothing))
            end
        elseif fn == "jl_invoke"
            try
                if VERSION < v"1.3.0-DEV.244"
                    meth, args, nargs, _ = operands(inst)
                else
                    f, args, nargs, meth = operands(inst)
                end
                meth = first(operands(meth::ConstantExpr))::ConstantExpr
                meth = first(operands(meth))::ConstantInt
                meth = convert(Int, meth)
                meth = Ptr{Cvoid}(meth)
                meth = Base.unsafe_pointer_to_objref(meth)::Core.MethodInstance
                push!(errors, (DYNAMIC_CALL, bt, meth.def))
            catch e
                isa(e,TypeError) || rethrow()
                @debug "Decoding arguments to jl_invoke failed" inst bb=LLVM.parent(inst)
                push!(errors, (DYNAMIC_CALL, bt, nothing))
            end
        elseif fn == "jl_apply_generic"
            try
                if VERSION < v"1.3.0-DEV.244"
                    args, nargs, _ = operands(inst)
                    ## args is a buffer where arguments are stored in
                    f, args = user.(uses(args))
                    ## first store into the args buffer is a direct store
                    f = first(operands(f::LLVM.StoreInst))::ConstantExpr
                else
                    f, args, nargs, _ = operands(inst)
                end

                f = first(operands(f))::ConstantExpr # get rid of addrspacecast
                f = first(operands(f))::ConstantInt # get rid of inttoptr
                f = convert(Int, f)
                f = Ptr{Cvoid}(f)
                f = Base.unsafe_pointer_to_objref(f)
                push!(errors, (DYNAMIC_CALL, bt, f))
            catch e
                isa(e,TypeError) || rethrow()
                @debug "Decoding arguments to jl_apply_generic failed" inst bb=LLVM.parent(inst)
                push!(errors, (DYNAMIC_CALL, bt, nothing))
            end

        # detect calls to undefined functions
        elseif isdeclaration(dest) && intrinsic_id(dest) == 0 && !isintrinsic(job, fn)
            # figure out if the function lives in the Julia runtime library
            if libjulia[] == C_NULL
                paths = filter(Libdl.dllist()) do path
                    name = splitdir(path)[2]
                    startswith(name, "libjulia")
                end
                libjulia[] = Libdl.dlopen(first(paths))
            end

            if Libdl.dlsym_e(libjulia[], fn) != C_NULL
                push!(errors, (RUNTIME_FUNCTION, bt, LLVM.name(dest)))
            else
                push!(errors, (UNKNOWN_FUNCTION, bt, LLVM.name(dest)))
            end
        end

    elseif isa(dest, InlineAsm)
        # let's assume it's valid ASM

    elseif isa(dest, ConstantExpr)
        # Enzyme should be able to handle these
    #     # detect calls to literal pointers
    #     if occursin("inttoptr", string(dest))
    #         # extract the literal pointer
    #         ptr_arg = first(operands(dest))
    #         GPUCompiler.@compiler_assert isa(ptr_arg, ConstantInt) job
    #         ptr_val = convert(Int, ptr_arg)
    #         ptr = Ptr{Cvoid}(ptr_val)

    #         # look it up in the Julia JIT cache
    #         frames = ccall(:jl_lookup_code_address, Any, (Ptr{Cvoid}, Cint,), ptr, 0)
    #         if length(frames) >= 1
    #             GPUCompiler.@compiler_assert length(frames) == 1 job frames=frames
    #             if VERSION >= v"1.4.0-DEV.123"
    #                 fn, file, line, linfo, fromC, inlined = last(frames)
    #             else
    #                 fn, file, line, linfo, fromC, inlined, ip = last(frames)
    #             end
    #             push!(errors, (POINTER_FUNCTION, bt, fn))
    #         else
    #             push!(errors, (POINTER_FUNCTION, bt, nothing))
    #         end
    #     end
    end

    return errors
end


text:
struct LinearCombination{T, As} <: LinearMap{T}
    maps::As
    function LinearCombination{T, As}(maps::As) where {T, As}
        N = length(maps)
        sz = size(maps[1])
        for n in 1:N
            size(maps[n]) == sz || throw(DimensionMismatch("LinearCombination"))
            promote_type(T, eltype(maps[n])) == T || throw(InexactError())
        end
        new{T, As}(maps)
    end
end

LinearCombination{T}(maps::As) where {T, As} = LinearCombination{T, As}(maps)

function Base.show(io::IO, A::LinearCombination)
    write(io, "LinearMaps.LinearCombination(")
    n = length(A.maps)
    if get(io, :limit, true) && n > 10
        s = 1:5
        e = n-5:n
        if e[1]-s[end] > 1
            write(io, "(")
            for i in s
                show(io, A.maps[i])
                write(io, ", ")
            end
            write(io, "…")
            for i in e
                write(io, ", ")
                show(io, A.maps[i])
            end
            write(io, ")")
        else
            show(io, A.maps)
        end
    else
        show(io, A.maps)
    end
    write(io, ")")
end

MulStyle(A::LinearCombination) = MulStyle(A.maps...)

# basic methods
Base.size(A::LinearCombination) = size(A.maps[1])
LinearAlgebra.issymmetric(A::LinearCombination) = all(issymmetric, A.maps) # sufficient but not necessary
LinearAlgebra.ishermitian(A::LinearCombination) = all(ishermitian, A.maps) # sufficient but not necessary
LinearAlgebra.isposdef(A::LinearCombination) = all(isposdef, A.maps) # sufficient but not necessary

# adding linear maps
"""
    A::LinearMap + B::LinearMap

Construct a `LinearCombination <: LinearMap`, a (lazy) representation of the sum
of the two operators. Sums of `LinearMap`/`LinearCombination` objects and
`LinearMap`/`LinearCombination` objects are reduced to a single `LinearCombination`.
In sums of `LinearMap`s and `AbstractMatrix`/`UniformScaling` objects, the latter
get promoted to `LinearMap`s automatically.

# Examples
```jldoctest; setup=(using LinearAlgebra, LinearMaps)
julia> CS = LinearMap{Int}(cumsum, 3)::LinearMaps.FunctionMap;

julia> LinearMap(ones(Int, 3, 3)) + CS + I + rand(3, 3);
```
"""
function Base.:(+)(A₁::LinearMap, A₂::LinearMap)
    size(A₁) == size(A₂) || throw(DimensionMismatch("+"))
    T = promote_type(eltype(A₁), eltype(A₂))
    return LinearCombination{T}(tuple(A₁, A₂))
end
function Base.:(+)(A₁::LinearMap, A₂::LinearCombination)
    size(A₁) == size(A₂) || throw(DimensionMismatch("+"))
    T = promote_type(eltype(A₁), eltype(A₂))
    return LinearCombination{T}(tuple(A₁, A₂.maps...))
end
Base.:(+)(A₁::LinearCombination, A₂::LinearMap) = +(A₂, A₁)
function Base.:(+)(A₁::LinearCombination, A₂::LinearCombination)
    size(A₁) == size(A₂) || throw(DimensionMismatch("+"))
    T = promote_type(eltype(A₁), eltype(A₂))
    return LinearCombination{T}(tuple(A₁.maps..., A₂.maps...))
end
Base.:(-)(A₁::LinearMap, A₂::LinearMap) = +(A₁, -A₂)

# comparison of LinearCombination objects, sufficient but not necessary
Base.:(==)(A::LinearCombination, B::LinearCombination) = (eltype(A) == eltype(B) && A.maps == B.maps)

# special transposition behavior
LinearAlgebra.transpose(A::LinearCombination) = LinearCombination{eltype(A)}(map(transpose, A.maps))
LinearAlgebra.adjoint(A::LinearCombination)   = LinearCombination{eltype(A)}(map(adjoint, A.maps))

# multiplication with vectors & matrices
for Atype in (AbstractVector, AbstractMatrix)
    @eval Base.@propagate_inbounds function LinearAlgebra.mul!(y::$Atype, A::LinearCombination, x::$Atype,
                             α::Number=true, β::Number=false)
        @boundscheck check_dim_mul(y, A, x)
        if iszero(α) # trivial cases
            iszero(β) && (fill!(y, zero(eltype(y))); return y)
            isone(β) && return y
            # β != 0, 1
            rmul!(y, β)
            return y
        else
            mul!(y, first(A.maps), x, α, β)
            return _mul!(MulStyle(A), y, A, x, α)
        end
    end
end

# For tuple-like storage of the maps (default), we recurse on the tail
# of the tuple.

@inline _mul!(::FiveArg, y, A::LinearCombination, x, α::Number) =
    __mul!(y, Base.tail(A.maps), x, α, nothing)
@inline function _mul!(::ThreeArg, y, A::LinearCombination, x, α::Number)
    z = similar(y)
    __mul!(y, Base.tail(A.maps), x, α, z)
end

@inline __mul!(y, As::Tuple{Vararg{LinearMap}}, x, α, z) =
    __mul!(__mul!(y, first(As), x, α, z), Base.tail(As), x, α, z)
@inline __mul!(y, A::Tuple{LinearMap}, x, α, z) = __mul!(y, first(A), x, α, z)
@inline __mul!(y, A::LinearMap, x, α, z) = muladd!(MulStyle(A), y, A, x, α, z)

# For vector-like storage of the maps, we simply loop over the maps.

@inline function _mul!(::FiveArg, y, A::LinearCombination{<:Any,<:AbstractVector}, x, α::Number)
    for i = 2:length(A.maps)
        mul!(y, A.maps[i], x, α, true)
    end
    y
end
@inline function _mul!(mulstyle::ThreeArg, y, A::LinearCombination{<:Any,<:AbstractVector}, x, α::Number)
    z = similar(y)
    for i = 2:length(A.maps)
        muladd!(mulstyle, y, A.maps[i], x, α, z)
    end
    y
end

@inline muladd!(::FiveArg, y, A, x, α, _) = mul!(y, A, x, α, true)
@inline function muladd!(::ThreeArg, y, A, x, α, z)
    mul!(z, A, x)
    if isone(α)
        y .+= z
    else
        y .+= z .* α
    end
    return y
end

A_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector) = mul!(y, A, x)

At_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector) = mul!(y, transpose(A), x)

Ac_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector) = mul!(y, adjoint(A), x)

ext: jl


text:
#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   General functions to compute the gravity field potential and gravity
#   acceleration.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#   [2] http://icgem.gfz-potsdam.de/home
#   [3] Delgado, M. R (2008). Aspherical Gravitational Field I: Modeling the
#       Space Environment. Universidad Poletécnica de Madrid, Madrid, Spain.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export create_gravity_model_coefs, compute_dU, compute_U, compute_g

################################################################################
#                               Public Functions
################################################################################

"""
    function create_gravity_model_coefs(icgem::ICGEM)

Return an instance of the structure `GravityModel_Coefs` based on the
information obtained from an ICGEM file in `icgem` (see `parse_icgem`).

"""
function create_gravity_model_coefs(icgem::ICGEM)

    legendre_norm = (icgem.norm == :fully_normalized) ? :full : :conv

    # Create and return the gravity model coefficients.
    gm_coefs::GravityModel_Coefs{Float64} =
        GravityModel_Coefs(name          = icgem.modelname,
                           μ             = icgem.gravity_constant,
                           R0            = icgem.radius,
                           n_max         = icgem.max_degree,
                           C             = copy(icgem.Clm),
                           S             = copy(icgem.Slm),
                           legendre_norm = legendre_norm)
end

"""
    function compute_dU(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector, n_max::Number = -1, m_max::Number = -1) where T<:Number

Compute the derivatives w.r.t. the spherical coordinates of the gravitational
field (`∂U/∂r`, `∂U/∂ϕ`, `∂U/∂λ`) defined by the coefficients `gm_coefs` (see
`GravityModel_Coefs`) at the position `r` [m] in ITRF. The maximum degree that
will be used while computing the spherical harmonics will be `n_max` and the
maximum order is `m_max`.

If `n_max` is negative, then the maximum available degree will be used. If
`n_max` is omitted, then it defaults to 0.

If `m_max` is negative or if it is greater than `n_max`, then it will be set to
`n_max`. If `m_max` is omitted, then it defaults to 0.

!!! info

    By convention, the result with `n_max` 0 and 1 will be the same.

# Returns

* The derivative of the gravitational field w.r.t. the radius (`∂U/∂r`).
* The derivative of the gravitational field w.r.t. the latitude (`∂U/∂ϕ`).
* The derivative of the gravitational field w.r.t. the longitude (`∂U/∂λ`).

# Remarks

In this case, `ϕ` is the geocentric latitude and `λ` is the geocentric
longitude.

"""
function compute_dU(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector,
                    n_max::Number = -1, m_max::Number = -1) where T<:Number

    # Unpack gravity model coefficients.
    #
    # We must not unpack all the variables because `n_max` will be overwritten.
    @unpack μ, R0, C, S, legendre_norm = gm_coefs

    # Check n_max value.
    if (n_max < 0) || (gm_coefs.n_max < n_max)
        n_max = gm_coefs.n_max
    end

    # Check m_max value.
    if (m_max < 0) || (m_max > n_max)
        m_max = n_max
    end

    # Get the geocentric latitude and longitude
    # =========================================
    r_gc = norm(r)
    ρ_gc = sqrt(r[1]^2 + r[2]^2)
    ϕ_gc = atan(r[3], ρ_gc)
    λ_gc = atan(r[2], r[1])

    # Auxiliary variables
    # ===================

    # Sine and cosine of the geocentric longitude.
    #
    # This values were be used in the algorithm to decrease the computational
    # burden.
    sin_λ,  cos_λ  = sincos(1λ_gc)
    sin_2λ, cos_2λ = sincos(2λ_gc)

    # First derivative of the non-spherical portion of the grav. field
    # ================================================================
    ∂Ur = T(1)  # Derivative w.r.t. the radius.
    ∂Uϕ = T(0)  # Derivative w.r.t. the geocentric latitude.
    ∂Uλ = T(0)  # Derivative w.r.t. the geocentric longitude.

    # Compute the associated Legendre functions `P_n,m[sin(ϕ_gc)]` with the
    # required normalization and its time derivative.
    #
    # Notice that cos(ϕ_gc-pi/2) = sin(ϕ_gc).
    dP, P = dlegendre(Val{legendre_norm}, ϕ_gc-pi/2, n_max, n_max, false)

    # Auxiliary variables.
    rn_fact = R0/r_gc
    rn      = rn_fact

    # Compute the derivatives.
    @inbounds for n = 2:n_max
        rn *= rn_fact # -> rn = (R0/r_gc)^n

        aux_∂Ur = T(0)
        aux_∂Uϕ = T(0)
        aux_∂Uλ = T(0)

        # Sine and cosine with m = 0.
        #
        # This values will be used to update recursively `sin(m*λ_gc)` and
        # `cos(m*λ_gc`), reducing the computational burden.
        sin_mλ   = T(0)      # sin( 0*λ_gc)
        sin_m_1λ = -sin_λ    # sin(-1*λ_gc)
        sin_m_2λ = -sin_2λ   # sin(-2*λ_gc)
        cos_mλ   = T(1)      # cos( 0*λ_gc)
        cos_m_1λ = cos_λ     # cos(-1*λ_gc)
        cos_m_2λ = cos_2λ    # cos(-2*λ_gc)

        for m = 0:n
            # Compute recursively `sin(m*λ_gc)` and `cos(m*λ_gc)`.
            sin_mλ = 2cos_λ*sin_m_1λ-sin_m_2λ
            cos_mλ = 2cos_λ*cos_m_1λ-cos_m_2λ

            CcSs_nm = C[n+1,m+1]*cos_mλ + S[n+1,m+1]*sin_mλ

            aux_∂Ur +=   P[n+1,m+1]*CcSs_nm
            aux_∂Uϕ +=  dP[n+1,m+1]*CcSs_nm
            aux_∂Uλ += m*P[n+1,m+1]*(S[n+1,m+1]*cos_mλ - C[n+1,m+1]*sin_mλ)

            # Check if we reached the maximum desired order.
            m >= m_max && break

            # Update the values for the next step.
            sin_m_2λ = sin_m_1λ
            sin_m_1λ = sin_mλ
            cos_m_2λ = cos_m_1λ
            cos_m_1λ = cos_mλ
        end

        ∂Ur += rn*(n+1)*aux_∂Ur
        ∂Uϕ += rn*aux_∂Uϕ
        ∂Uλ += rn*aux_∂Uλ
    end

    ∂Ur *= -μ/r_gc^2
    ∂Uϕ *= +μ/r_gc
    ∂Uλ *= +μ/r_gc

    ∂Ur, ∂Uϕ, ∂Uλ
end

"""
    function compute_g(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector, n_max::Number = -1, m_max::Number = -1) where T<:Number

Compute the gravitational acceleration (ITRF) \\[m/s²] at position `r` \\[m]
(ITRF) using the coefficients `gm_coefs` (see `GravityModel_Coefs`). The maximum
degree that will be used while computing the spherical harmonics will be
`n_max` and the maximum order it `m_max`.

If `n_max` is negative, then the maximum available degree will be used. If
`n_max` is omitted, then it defaults to 0.

If `m_max` is negative or if it is greater than `n_max`, then it will be set to
`n_max`. If `m_max` is omitted, then it defaults to 0.

!!! info

    By convention, the result with `n_max` 0 and 1 will be the same.

# Remarks

Notice that this function computes the **gravitational acceleration**. Hence,
the acceleration due to Earth rotation rate **is not** included.

"""
function compute_g(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector,
                   n_max::Number = -1, m_max::Number = -1) where T<:Number

    # Compute the partial derivatives of the gravitational field w.r.t. the
    # spherical coordinates.
    ∂Ur, ∂Uϕ, ∂Uλ = compute_dU(gm_coefs, r, n_max, m_max)

    # Auxiliary variables.
    r_gc = norm(r)
    ρ_gc = sqrt(r[1]^2 + r[2]^2)
    ϕ_gc = atan(r[3], ρ_gc)
    λ_gc = atan(r[2], r[1])

    # Compute the acceleration represented in the ITRF
    # ================================================

    # Compute the partial derivatives in spherical coordinate systems:
    #
    #     ∂U        1      ∂U     1   ∂U
    #    ---- , ---------.---- , ---.----
    #     ∂r     r.cos ϕ   ∂λ     r   ∂ϕ
    #
    # Notice that the singularity is not a problem here. When computing
    # `cos(pi/2)` a very small number will be returned and `∂U/∂λ` is 0. Hence,
    # the 2nd component will be 0.

    a_l = SVector{3,T}(∂Ur,
                       ∂Uλ/(r_gc*cos(ϕ_gc)),
                       ∂Uϕ/r_gc)

    # The vector `a_l` is represented in the local UEN (Up-Earth-North)
    # reference frame. Hence, we need to describe the unitary vectors of this
    # frame in the ECEF reference frame. This can be accomplished by the
    # following rotations matrix.

    Del = angle_to_dcm(0, ϕ_gc, -λ_gc, :XYZ)
    a_ITRF = Del*a_l

    a_ITRF
end

"""
    function compute_U(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector, n_max::Number = -1, m_max::Number = -1) where T<:Number

Compute the gravitational potential [J/kg] at `r` (ITRF) \\[m] using the
coefficients `gm_coefs` (see `GravityModel_Coefs`). The maximum degree that will
be used while computing the spherical harmonics will be `n_max` and the maximum
order is `m_max`.

If `n_max` is negative, then the maximum available degree will be used. If
`n_max` is omitted, then it defaults to 0.

If `m_max` is negative or if it is greater than `n_max`, then it will be set to
`n_max`. If `m_max` is omitted, then it defaults to 0.

!!! info

    By convention, the result with `n_max` 0 and 1 will be the same.

"""
function compute_U(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector,
                   n_max::Number = -1, m_max::Number = -1) where T<:Number

    # Unpack gravity model coefficients.
    #
    # We must not unpack all the variables because `n_max` will be overwritten.
    @unpack μ, R0, C, S, legendre_norm = gm_coefs

    # Check n_max value.
    if (n_max < 0) || (gm_coefs.n_max < n_max)
        n_max = gm_coefs.n_max
    end

    # Check m_max value.
    if (m_max < 0) || (m_max > n_max)
        m_max = n_max
    end

    # Auxiliary variables.
    r_gc = norm(r)
    ρ_gc = sqrt(r[1]^2 + r[2]^2)
    ϕ_gc = atan(r[3], ρ_gc)
    λ_gc = atan(r[2], r[1])

    # Sine and cosine of the geocentric longitude.
    #
    # This values were be used in the algorithm to decrease the computational
    # burden.
    sin_λ,  cos_λ  = sincos(1λ_gc)
    sin_2λ, cos_2λ = sincos(2λ_gc)

    # Consider the zero degree term.
    U = T(1)

    # Compute the associated Legendre functions `P_n,m[sin(ϕ_gc)]` with the
    # required normalization.
    #
    # Notice that cos(ϕ_gc-pi/2) = sin(ϕ_gc).
    P = legendre(Val{legendre_norm}, ϕ_gc-pi/2, n_max, n_max, false)

    # Auxiliary variables.
    rn_fact = R0/r_gc
    rn      = rn_fact

    @inbounds for n = 2:n_max
        aux_U = T(0)

        # Sine and cosine with m = 0.
        #
        # This values will be used to update recursively `sin(m*λ_gc)` and
        # `cos(m*λ_gc`), reducing the computational burden.
        sin_mλ   = T(0)      # sin( 0*λ_gc)
        sin_m_1λ = -sin_λ    # sin(-1*λ_gc)
        sin_m_2λ = -sin_2λ   # sin(-2*λ_gc)
        cos_mλ   = T(1)      # cos( 0*λ_gc)
        cos_m_1λ = cos_λ     # cos(-1*λ_gc)
        cos_m_2λ = cos_2λ    # cos(-2*λ_gc)

        for m = 0:n
            # Compute recursively `sin(m*λ_gc)` and `cos(m*λ_gc)`.
            sin_mλ = 2cos_λ*sin_m_1λ-sin_m_2λ
            cos_mλ = 2cos_λ*cos_m_1λ-cos_m_2λ

            aux_U += P[n+1,m+1]*(C[n+1,m+1]*cos_mλ + S[n+1,m+1]*sin_mλ)

            # Check if we reached the maximum desired order.
            m >= m_max && break

            # Update the values for the next step.
            sin_m_2λ = sin_m_1λ
            sin_m_1λ = sin_mλ
            cos_m_2λ = cos_m_1λ
            cos_m_1λ = cos_mλ
        end

        rn *= rn_fact # -> rn = (R0/r_gc)^n
        U  += rn*aux_U
    end

    U *= μ/r_gc
end


text:
romeo = unwrap # access unwrap function via alias romeo
romeo! = unwrap!

function mask_from_voxelquality(qmap, threshold)
    qmap_bin = qmap .> threshold # NaN defaults to false (0)
    max_hole_size = length(qmap) / 20
    qmap_bin = .!imfill(.!qmap_bin, (1, max_hole_size)) # fills all holes up to max_hole_size (uses 6 connectivity as default)
    return gaussiansmooth3d(qmap_bin; padding=true) .> 0.8 # hardcoded final threshold
end

function ROMEO.calculateweights(phase::AbstractArray{T,4}; TEs, template=2, p2ref=1, keyargs...) where T
    args = Dict{Symbol, Any}(keyargs)
    args[:phase2] = phase[:,:,:,p2ref]
    args[:TEs] = TEs[[template, p2ref]]
    if haskey(args, :mag)
        args[:mag] = args[:mag][:,:,:,template]
    end
    return ROMEO.calculateweights(view(phase,:,:,:,template); args...)
end

# calculates B0 in [Hz]
function calculateB0_unwrapped(unwrapped_phase, mag, TEs)
    dims = 4
    TEs = to_dim(TEs, 4)
    return (1000 / 2π) * sum(unwrapped_phase .* mag .* mag .* TEs; dims) ./ sum(mag .* mag .* TEs.^2; dims) |> I -> dropdims(I; dims)
end

"""
Calculates a quality for each voxel. Takes the same inputs as ROMEO
"""
function romeovoxelquality(phase; keyargs...)
    weights = ROMEO.calculateweights(phase; type=Float32, rescale=x->x, keyargs...)
    return dropdims(sum(weights; dims=1); dims=1) ./ 3 # [0;1]
end


text:
#
# Optimizations:
#  - replace leftmap, rightmap with sorted vectors
#

export SiteOp,
       MPOTerm,
       AutoMPO,
       terms,
       add!,
       toMPO,
       MPOTerm,
       MatElem,
       SiteOp

import LinearAlgebra.svd

###########################
# SiteOp                  #
###########################

struct SiteOp
  name::String
  site::Int
end
name(s::SiteOp) = s.name
site(s::SiteOp) = s.site
Base.show(io::IO,s::SiteOp) = print(io,"\"$(name(s))\"($(site(s)))")

function Base.isless(s1::SiteOp,s2::SiteOp)::Bool
  if site(s1) < site(s2)
    return true
  end
  return name(s1) < name(s2)
end

###########################
# OpTerm                  #
###########################

const OpTerm = Vector{SiteOp}
mult(t1::OpTerm,t2::OpTerm) = isempty(t2) ? t1 : vcat(t1,t2)

###########################
# MPOTerm                 #
###########################

struct MPOTerm
  coef::ComplexF64
  ops::OpTerm
end
coef(op::MPOTerm) = op.coef
ops(op::MPOTerm) = op.ops
length(op::MPOTerm) = length(op.ops)

function ==(t1::MPOTerm,t2::MPOTerm)
  return (t1.ops==t2.ops && isapprox(t1.coef,t2.coef))
end

function Base.isless(t1::MPOTerm,t2::MPOTerm)::Bool
  if !isapprox(coef(t1),coef(t2))
    return coef(t1) < coef(t2)
  end
  return ops(t1) < ops(t2)
end

function MPOTerm(c::Number,
                 op1::String,
                 i1::Int)
  return MPOTerm(convert(ComplexF64,c),[SiteOp(op1,i1)])
end

function MPOTerm(c::Number,
                 op1::String,i1::Int,
                 op2::String,i2::Int)
  return MPOTerm(convert(ComplexF64,c),[SiteOp(op1,i1),SiteOp(op2,i2)])
end

function MPOTerm(c::Number,
                 op1::String,i1::Int,
                 op2::String,i2::Int,
                 ops...)
  vop = OpTerm(undef,2+div(length(ops),2))
  vop[1] = SiteOp(op1,i1)
  vop[2] = SiteOp(op2,i2)
  for n=1:2:length(ops)
    vop[2+n] = SiteOp(ops[n],ops[n+1])
  end
  return MPOTerm(convert(ComplexF64,c),vop)
end

function MPOTerm(c::Number,
                 ops::OpTerm)
  return MPOTerm(convert(ComplexF64,c),ops)
end

function Base.show(io::IO,
              op::MPOTerm)
  c = coef(op)
  if c != 1.0+0.0im
    if imag(c) == 0.0
      print(io,"$(real(c)) ")
    elseif real(c) == 0.0
      print(io,"$(imag(c))im ")
    else
      print(io,"($c) ")
    end
  end
  for o in ops(op)
    print(io,"\"$(name(o))\"($(site(o)))")
  end
end

############################
## AutoMPO                 #
############################

struct AutoMPO
  sites::SiteSet
  terms::Vector{MPOTerm}
  AutoMPO(s::SiteSet) = new(s,Vector{MPOTerm}())
end
sites(ampo::AutoMPO) = ampo.sites
terms(ampo::AutoMPO) = ampo.terms

function add!(ampo::AutoMPO,
              op::String, i::Int)
  push!(terms(ampo),MPOTerm(1.0,op,i))
end

function add!(ampo::AutoMPO,
              coef::Number,
              op::String, i::Int)
  push!(terms(ampo),MPOTerm(coef,op,i))
end

function add!(ampo::AutoMPO,
              op1::String, i1::Int,
              op2::String, i2::Int)
  push!(terms(ampo),MPOTerm(1.0,op1,i1,op2,i2))
end

function add!(ampo::AutoMPO,
              coef::Number,
              op1::String, i1::Int,
              op2::String, i2::Int)
  push!(terms(ampo),MPOTerm(coef,op1,i1,op2,i2))
end

function add!(ampo::AutoMPO,
              op1::String, i1::Int,
              op2::String, i2::Int,
              ops...)
  push!(terms(ampo),MPOTerm(1.0,op1,i1,op2,i2,ops...))
end

function add!(ampo::AutoMPO,
              coef::Number,
              op1::String, i1::Int,
              op2::String, i2::Int,
              ops...)
  push!(terms(ampo),MPOTerm(coef,op1,i1,op2,i2,ops...))
end

function show(io::IO,
              ampo::AutoMPO)
  println(io,"AutoMPO:")
  for term in terms(ampo)
    println(io,"  $term")
  end
end

##################################
# MatElem (simple sparse matrix) #
##################################

struct MatElem{T}
  row::Int
  col::Int
  val::T
end

function Base.show(io::IO,m::MatElem)
  print(io,"($(m.row),$(m.col),$(m.val))")
end

function toMatrix(els::Vector{MatElem{T}})::Matrix{T} where {T}
  nr = 0
  nc = 0
  for el in els
    nr = max(nr,el.row)
    nc = max(nr,el.col)
  end
  M = zeros(T,nr,nc)
  for el in els
    M[el.row,el.col] = el.val
  end
  return M
end

function ==(m1::MatElem{T},m2::MatElem{T})::Bool where {T}
  return (m1.row==m2.row && m1.col==m2.col && m1.val==m2.val)
end

function Base.isless(m1::MatElem{T},m2::MatElem{T})::Bool where {T}
  if m1.row != m2.row
    return m1.row < m2.row
  elseif m1.col != m2.col
    return m1.col < m2.col
  end
  return m1.val < m2.val
end

function posInLink!(linkmap::Vector{OpTerm},
                    op::OpTerm)::Int
  isempty(op) && return -1
  for n=1:length(linkmap)
    (linkmap[n]==op) && return n
  end
  push!(linkmap,op)
  return length(linkmap)
end

function determineValType(terms::Vector{MPOTerm})
  for t in terms
    (!isreal(coef(t))) && return ComplexF64
  end
  return Float64
end

function computeSiteProd(sites::SiteSet,
                         ops::OpTerm)::ITensor
  i = ops[1].site
  T = op(sites,ops[1].name,i)
  for j=2:length(ops)
    (ops[j].site != i) && error("Mismatch of site number in computeSiteProd")
    opj = op(sites,ops[j].name,i)
    T = multSiteOps(T,opj)
  end
  return T
end

function remove_dups!(v::Vector{T}) where {T}
  N = length(v)
  (N==0) && return
  sort!(v)
  n = 1
  u = 2
  while u <= N
    while u < N && v[u]==v[n]
      u += 1
    end
    if v[u] != v[n]
      v[n+1] = v[u]
      n += 1
    end
    u += 1
  end
  resize!(v,n)
  return
end

function svdMPO(ampo::AutoMPO;
                kwargs...)::MPO

  mindim::Int = get(kwargs,:mindim,1)
  maxdim::Int = get(kwargs,:maxdim,10000)
  cutoff::Float64 = get(kwargs,:cutoff,1E-13)

  N = length(sites(ampo))

  ValType = determineValType(terms(ampo))

  Vs = [Matrix{ValType}(undef,1,1) for n=1:N]
  tempMPO = [MatElem{MPOTerm}[] for n=1:N]

  crosses_bond(t::MPOTerm,n::Int) = (ops(t)[1].site <= n <= ops(t)[end].site)

  rightmap = OpTerm[]
  next_rightmap = OpTerm[]

  for n=1:N

    leftbond_coefs = MatElem{ValType}[]

    leftmap = OpTerm[]
    for term in terms(ampo)
      crosses_bond(term,n) || continue

      left::OpTerm   = filter(t->(t.site < n),ops(term))
      onsite::OpTerm = filter(t->(t.site == n),ops(term))
      right::OpTerm  = filter(t->(t.site > n),ops(term))

      bond_row = -1
      bond_col = -1
      if !isempty(left)
        bond_row = posInLink!(leftmap,left)
        bond_col = posInLink!(rightmap,mult(onsite,right))
        bond_coef = convert(ValType,coef(term))
        push!(leftbond_coefs,MatElem(bond_row,bond_col,bond_coef))
      end

      A_row = bond_col
      A_col = posInLink!(next_rightmap,right)
      site_coef = 1.0+0.0im
      if A_row == -1
        site_coef = coef(term)
      end
      isempty(onsite) && push!(onsite,SiteOp("Id",n))
      el = MatElem(A_row,A_col,MPOTerm(site_coef,onsite))
      push!(tempMPO[n],el)
    end
    rightmap = next_rightmap
    next_rightmap = OpTerm[]

    remove_dups!(tempMPO[n])

    if n > 1 && !isempty(leftbond_coefs)
      M = toMatrix(leftbond_coefs)
      U,S,V = svd(M)
      P = S.^2
      truncate!(P;maxdim=maxdim,cutoff=cutoff,mindim=mindim)
      tdim = length(P)
      nc = size(M,2)
      Vs[n-1] = Matrix{ValType}(V[1:nc,1:tdim])
    end

  end

  llinks = [Index() for n=1:N+1]
  llinks[1] = Index(2,"Link,n=0")

  H = MPO(sites(ampo))

  # Constants which define MPO start/end scheme
  rowShift = 2
  colShift = 2
  startState = 2
  endState = 1

  for n=1:N
    VL = Matrix{ValType}(undef,1,1)
    if n > 1
      VL = Vs[n-1]
    end
    VR = Vs[n]
    tdim = size(VR,2)

    llinks[n+1] = Index(2+tdim,"Link,n=$n")

    finalMPO = Dict{OpTerm,Matrix{ValType}}()

    ll = llinks[n]
    rl = llinks[n+1]

    idTerm = [SiteOp("Id",n)]
    finalMPO[idTerm] = zeros(ValType,dim(ll),dim(rl))
    idM = finalMPO[idTerm]
    idM[1,1] = 1.0
    idM[2,2] = 1.0

    defaultMat() = zeros(ValType,dim(ll),dim(rl))

    for el in tempMPO[n]
      A_row = el.row
      A_col = el.col
      t = el.val
      (abs(coef(t)) > eps()) || continue

      M = get!(finalMPO,ops(t),defaultMat())

      ct = convert(ValType,coef(t))
      if A_row==-1 && A_col==-1 #onsite term
        M[startState,endState] += ct
      elseif A_row==-1 #term starting on site n
        for c=1:size(VR,2)
          z = ct*VR[A_col,c]
          M[startState,colShift+c] += z
        end
      elseif A_col==-1 #term ending on site n
        for r=1:size(VL,2)
          z = ct*conj(VL[A_row,r])
          M[rowShift+r,endState] += z
        end
      else
        for r=1:size(VL,2),c=1:size(VR,2)
          z = ct*conj(VL[A_row,r])*VR[A_col,c]
          M[rowShift+r,colShift+c] += z
        end
      end
    end

    s = sites(ampo)[n]
    H[n] = ITensor(dag(s),s',ll,rl)
    for (op,M) in finalMPO
      T = ITensor(M,ll,rl)
      H[n] += T*computeSiteProd(sites(ampo),op)
    end

  end

  L = ITensor(llinks[1])
  L[startState] = 1.0

  R = ITensor(llinks[N+1])
  R[endState] = 1.0

  H[1] *= L
  H[N] *= R

  return H
end

function sortEachTerm!(ampo::AutoMPO)
  isless_site(o1::SiteOp,o2::SiteOp) = (site(o1)<site(o2))
  for t in terms(ampo)
    sort!(ops(t),alg=InsertionSort,lt=isless_site)
  end
end

function toMPO(ampo::AutoMPO;
               kwargs...)::MPO
  sortEachTerm!(ampo)
  return svdMPO(ampo;kwargs...)
end


text:
# Содержание:
# Функция вычисления          неоднородности в сопряженной задаче               |heterogeneity|
#
# Функция вычисления          правой части сопряженной задачи                   |adjointRP|
# Функция вычисления якобиана правой части через автодифференцирование          |∂adjointRP_∂y|
#
# Функция вычисления якобиана правой части в явном виде, возвращает 3 диагонали |ARP_y|
# Обертка            якобиана в явном виде, возвращает трехдиагональную марицу  |∂ARP_∂y|
#
# Функция нахождения решения сопряженной задачи                                 |solve_adjoint|
#

@doc raw"""
    heterogeneity(n::Int, m::Int,
                 Xₙ::Vector, N::Int,
                 Tₘ::Vector, M::Int,
                 Uₙₘ::Matrix,
                 f1::Vector, f2::Vector,
                 w::Real)

Неоднородность выражающая невязку текущего решения с искомым.
`` 2 \delta( x - f_1(t) ) ( u^s(x,t) - f_2(t) ) ``.
"""
function heterogeneity(n::Int, m::Int,
                      Xₙ::Vector, N::Int,
                      Tₘ::Vector, M::Int,
                      Uₙₘ::Matrix,
                      f1::Vector, f2::Vector,
                      w::Real)

    @assert length(Xₙ) == N-1

    @assert length(Tₘ) == M+1
    @assert length(f1) == M+1
    @assert length(f2) == M+1

    @assert size(Uₙₘ) == (N-1, M+1)

    return 2 * deltaw( n, f1[m], Xₙ, N, w) * ( Uₙₘ[n, m] - f2[m] )
end

@doc raw"""
    adjointRP(y::Vector, m::Int,
              Xₙ::Vector, N::Int,
              Tₘ::Vector, M::Int,
              ε::Real,
              ulₘ::Vector, urₘ::Vector,
              qₙ::Vector,
              Uₙₘ::Matrix, f1::Vector, f2::Vector)

Вычисляет вектор правой части с помощью конечно-разностной
аппроксимации пространственных производных.

# Arguments
- `y::Vector`:          Вектор размера `N-1`, решение в текущий момент без ГТ.
- `m::Int`:             Номер шага в сетке по времени.
- `Xₙ::Vector`:         Вектор размера `N+1` сетки по ``X``, вместе с ГТ.
- `N::Int`:             Число **интервалов** в полной сетке по ``X``.
- `Tₘ::Vector`:         Вектор размера `M+1` сетки по ``T``.
- `M::Int`:             Число **интервалов** в полной сетке по ``T``.
- `ε::Real`:            Малый параметр при старшей производной.
- `ulₘ::Vector`:        Вектор сеточных значений левого ГУ.
- `urₘ::Vector`:        Вектор сеточных значений правого ГУ.
- `qₙ::Vector`:         Вектор размера `N-1` сеточные значений  неоднородности без ГТ.
- `Uₙₘ::Matrix`:        Матрица размера `N-1, M+1` решения прямой задачи при данном `qₙ`.
- `f1::Vector`:         Эспериментальные данные — положение внутреннего слоя, размера `M+1`.
- `f2::Vector`:         Эспериментальные данные — значение функции на внутреннеем слое, размера `M+1`.
- `w::Real`:            Априорный параметр в аппроксимации дельта-функции (см. [`heterogeneity`](@ref), [`δw`](@ref)).

!!! warn
    Функция не входит в публичный API, поэтому размерность
    входных векторов отличается от аналогичных в [`solve_adjoint`](@ref).
     - `qₙ`, `y`    размера `N-1`,
     - `Xₙ`         размера `N+1`,
     - `Uₙₘ`        размера `N-1, M+1`
"""
function adjointRP(y::Vector, m::Int,
                   Xₙ::Vector, N::Int,
                   Tₘ::Vector, M::Int,
                   ε::Real,
                   ulₘ::Vector, urₘ::Vector,
                   qₙ::Vector,
                   Uₙₘ::Matrix, f1::Vector, f2::Vector,
                   w::Real)

    @assert length(Xₙ) == N+1
    @assert length(qₙ) == N-1
    @assert length(y)  == N-1

    @assert (Tₘ[end] - Tₘ[1]) < 0
    @assert length(Tₘ) == M+1
    @assert length(ulₘ) == M+1
    @assert length(urₘ) == M+1
    @assert length(f1) == M+1
    @assert length(f2) == M+1

    @assert size(Uₙₘ) == (N-1, M+1)

    @assert m < length(ulₘ)

    RP = zero(y)                # Создаем нулевой вектор того же типа и размера
    X = Xₙ[2:N];                # Передается внутрь **только** `heterogeneity`

    # Здесь нужно применить сдвиг индексов в `Xₙ` на +1!
    # Xₙ[1] соответствует первому узлу нашей сетки — ``x_0``
    # Xₙ[2] — ``x_1``
    # Xₙ[3] — ``x_2``
    #
    # Для удобства, формулы записаны в следующем виде
    # d[1] = ... Xₙ[1+1+1] ...,
    # первая единица соответствует единице в d[1],
    # вторая — берем следующий элемент,
    # третья — совершаем сдвиг индексов для Xₙ
    RP[1]     = - 2ε / (Xₙ[1+1+1] - Xₙ[1-1+1]) * ( (y[2] - y[1]) /  (Xₙ[1+1+1] - Xₙ[1+1]) - (y[1] - ulₘ[m])/(Xₙ[1+1] - Xₙ[1-1+1])) +
        Uₙₘ[1,m]   * (y[2]   - ulₘ[m])/(Xₙ[1+1+1] - Xₙ[1-1+1]) + y[1] * qₙ[1] - heterogeneity(1, m, X, N, Tₘ, M, Uₙₘ, f1, f2, w)

    # Xₙ[n  ] — соответствует ``x_{n-1}``
    # Xₙ[n+1] — соответствует ``x_{n  }``
    # Xₙ[n+2] — соответствует ``x_{n+1}``
    #
    # Для удобства, формулы записаны в следующем виде
    # d[n] = ... Xₙ[n-1+1],
    # первое слагаемое соответствует d[n],
    # второе — берем предыдущий элемент,
    # третье — совершаем сдвиг индексов для Xₙ
    for n in 2:N-2
        RP[n] = - 2ε / (Xₙ[n+1+1] - Xₙ[n-1+1]) * ( (y[n+1] - y[n]) /  (Xₙ[n+1+1] - Xₙ[n+1]) - (y[n] - y[n-1])/(Xₙ[n+1] - Xₙ[n-1+1])) +
        + Uₙₘ[n, m]   * (y[n+1] - y[n-1])/(Xₙ[n+1+1] - Xₙ[n-1+1]) + y[n]   * qₙ[n] - heterogeneity(n, m, X, N, Tₘ, M, Uₙₘ, f1, f2, w)
    end

    RP[N-1]   =  - 2ε / (Xₙ[N+1] - Xₙ[N-2+1]) * ( (urₘ[m] - y[N-1]) /  (Xₙ[N+1] - Xₙ[N-1+1]) - (y[N-1] - y[N-2])/(Xₙ[N-1+1] - Xₙ[N-2+1])) +
        Uₙₘ[N-1, m] * (urₘ[m] - y[N-2])/(Xₙ[N+1] - Xₙ[N-2+1]) + y[N-1] * qₙ[N-1] - heterogeneity(N-1, m, X, N, Tₘ, M, Uₙₘ, f1, f2, w)

    return RP;
end

@doc raw""""
    ARP_y(y::Vector, m::Int,
          Xₙ::Vector, N::Int,
          Tₘ::Vector, M::Int,
          ε::Real,
          ulₘ::Vector, urₘ::Vector,
          qₙ::Vector,
          Uₙₘ::Matrix, f1::Vector, f2::Vector) ->
          (::Vector, ::Vector, ::Vector)

Возращает три диагонали якоибана правой части сопряженной задачи.
Входные параметры полностью аналогичны [`adjointRP`](@ref).

# Return
Возвращает три вектора `dl`, `d`, `dl` элементов якоибана.
Нижняя и верхняя диагонали длины ``N-2``, главная — ``N-1``.
```
    [ d[1]  du[1]                         ]
    [ dl[1] d[2]  du[2]                   ]
    [ 0     dl[2] d[3] du[3]              ]
    [           ...  ...  ...             ]
    [                ...  ...     du[N-2] ]
    [                     dl[N-2] d[N-1]  ]
```

!!! warn
    Функция не входит в публичный API, поэтому размерность
    входных векторов отличается от аналогичных в [`solve_adjoint`](@ref).
     - `qₙ`, `y`    размера `N-1`,
     - `Xₙ`         размера `N+1`,
     - `Uₙₘ`        размера `N-1, M+1`
"""
function ARP_y(y::Vector, m::Int,
               Xₙ::Vector, N::Int,
               Tₘ::Vector, M::Int,
               ε::Real,
               ulₘ::Vector, urₘ::Vector,
               qₙ::Vector,
               Uₙₘ::Matrix, f1::Vector, f2::Vector)

    @assert length(Xₙ) == N+1
    @assert length(qₙ) == N-1
    @assert length(y) == N-1
    @assert m < length(ulₘ)

    d = zeros(N-1);
    dl = zeros(N-2);
    du = zeros(N-2);

    # Здесь нужно применить сдвиг индексов на +1 для Xₙ!
    # Xₙ[1] соответствует ПЕРВОМУ узлу нашей сетки — ``x_0``
    # Xₙ[2] — ``x_1``
    # Xₙ[3] — ``x_2``
    #
    # Для удобства, формулы записаны в следующем виде
    # d[1] = ... Xₙ[1+1+1] ...,
    # первая единица соответствует единице в d[1],
    # вторая — берем следующий элемент,
    # третья — совершаем сдвиг индексов для Xₙ
    d[1] = - 2ε / ( Xₙ[1+1+1] - Xₙ[1-1+1]) * ( -1 / (Xₙ[1+1+1] - Xₙ[1+1]) - 1 / (Xₙ[1+1] - Xₙ[1-1+1])) +
           qₙ[1]

    du[1] = - 2ε / ( Xₙ[1+1+1] - Xₙ[1-1+1]) * ( 1/ (Xₙ[1+1+1] - Xₙ[1+1]) ) +
            Uₙₘ[1, m] / (Xₙ[1+1+1] - Xₙ[1-1+1])

    # Xₙ[n  ] — соответствует ``x_{n-1}``
    # Xₙ[n+1] — соответствует ``x_{n  }``
    # Xₙ[n+2] — соответствует ``x_{n+1}``
    #
    # Для удобства, формулы записаны в следующем виде
    # d[n] = ... Xₙ[n-1+1],
    # первое слагаемое соответствует d[n],
    # второе — берем предыдущий элемент,
    # третье — совершаем сдвиг индексов для Xₙ
    for n in 2:N-2
        # Сдвиг индексов для поддиагональных элементов `dl`
        # см # Return документации функции
        dl[n - 1] = - 2ε / ( Xₙ[n+1+1] - Xₙ[n-1+1] ) * ( 1 / (Xₙ[n+1] - Xₙ[n-1+1])) -
                    Uₙₘ[n, m] / ( Xₙ[n+1+1] - Xₙ[n-1+1])

        d[n]      = - 2ε / ( Xₙ[n+1+1] - Xₙ[n-1+1]) * ( -1 / (Xₙ[n+1+1] - Xₙ[n+1]) - 1 / (Xₙ[n+1] - Xₙ[n-1+1])) +
                    qₙ[n]

        du[n]     = - 2ε / ( Xₙ[n+1+1] - Xₙ[n-1+1]) * ( 1/ (Xₙ[n+1+1] - Xₙ[n+1]) ) +
                    Uₙₘ[n, m] / (Xₙ[n+1+1] - Xₙ[n-1+1])
    end

    dl[N-2] = - 2ε / ( Xₙ[N+1] - Xₙ[N-2+1]) * ( 1 / (Xₙ[N-1+1] - Xₙ[N-2+1])) -
                Uₙₘ[N-1, m] / ( Xₙ[N+1] - Xₙ[N-2+1])
    d[N-1] =  - 2ε / ( Xₙ[N+1] - Xₙ[N-2+1]) * ( -1 / (Xₙ[N+1] - Xₙ[N-1+1]) - 1 / (Xₙ[N-1+1] - Xₙ[N-2+1])) +
                qₙ[N-1]

    return dl, d, du
end

@doc raw"""
    ∂ARP_∂y(y::Vector, m::Int,
          Xₙ::Vector, N::Int,
          Tₘ::Vector, M::Int,
          ε::Real,
          ulₘ::Vector, urₘ::Vector,
          qₙ::Vector,
          Uₙₘ::Matrix, f1::Vector, f2::Vector) -> Tridiagonal

Обертка фнукции [`ARP_y`](@ref), которая возвращает `Tridiagonal( ARP_y(...))`
трехдиагональную матрицы из векторов, которые возвращает `ARP_y`.
"""
function ∂ARP_∂y(y::Vector, m::Int,
               Xₙ::Vector, N::Int,
               Tₘ::Vector, M::Int,
               ε::Real,
               ulₘ::Vector, urₘ::Vector,
               qₙ::Vector,
               Uₙₘ::Matrix, f1::Vector, f2::Vector)
    return Tridiagonal( (ARP_y(y, m, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qₙ, Uₙₘ, f1, f2))... )
end

@doc raw"""
    ∂adjointRP_∂y(y::Vector, m::Int,
                   Xₙ::Vector, N::Int,
                   Tₘ::Vector, M::Int,
                   ε::Real,
                   ulₘ::Vector, urₘ::Vector,
                   qₙ::Vector,
                   Uₙₘ::Matrix, f1::Vector, f2::Vector,
                   w::Real)

Функция якобиана для `adjointRP`.
Все аргументы, размерности входных векторов такие же, как и у [`adjointRP`](@ref).

!!! warn
    Функция не входит в публичный API, поэтому размерность
    входных векторов отличается от аналогичных в [`solve_adjoint`](@ref).
     - `qₙ`, `y`    размера `N-1`,
     - `Xₙ`         размера `N+1`,
     - `Uₙₘ`        размера `N-1, M+1`
"""
function ∂adjointRP_∂y(y::Vector, m::Int,
                       Xₙ::Vector, N::Int,
                       Tₘ::Vector, M::Int,
                       ε::Real,
                       ulₘ::Vector, urₘ::Vector,
                       qₙ::Vector,
                       Uₙₘ::Matrix, f1::Vector, f2::Vector,
                       w::Real)
    return ForwardDiff.jacobian( z -> adjointRP(z, m, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qₙ, Uₙₘ, f1, f2, w), y)
end

@doc raw"""

!!! warning
    Массивы `Xₙ`, `Tₘ`, `u`, `f1`, `f2` Передаются **как есть**!
    Они переварачиваются внутри функции локально.
"""
function solve_adjoint(y₀::Vector, Xₙ::Array, N::Int,
                       Tₘ::Vector, M::Int,
                       ε::Real, ulₘ::Vector, urₘ::Vector,
                       qₙ::Vector,
                       Uₙₘ::Matrix, f1::Vector, f2::Vector,
                       RP::Function = adjointRP,
                       jac::Function = ∂ARP_∂y;
                       α::Complex = complex(0.5, 0.5),
                       w::Real = 0.00001)
    # Checking lengths
    # {{{
    if (typeof(Xₙ) <: Vector)
        length(Xₙ) == N+1 ||
        throw(ArgumentError("length(Xₙ) == $(length(Xₙ)), N == $(N) "*
                            "Массив Xₙ должен иметь размерность N+1."))
    elseif ( typeof(Xₙ) <: Matrix )
        size(Xₙ) == (N+1, M+1) ||
        throw(ArgumentError("size(Xₙ) == $(size(Xₙ)), N == $(N), M == $(M) " *
                            "Массив Xₙ должен иметь размерность (N+1, M+1)."))
    else
        throw(ArgumentError("Функция принимает `Xₙ` только в виде вектора или матрицы"))
    end

    if length(qₙ) != N+1
        throw(ArgumentError("""length(qₙ) == $(length(qₙ)), N == $(N)
        Массив qₙ должен иметь размерность N+1."""))
    end
    if length(y₀) != N+1
        throw(ArgumentError("""length(y) == $(length(y)), N == $(N)
        Массив `y` должен иметь размерность N+1."""))
    end

    if length(Tₘ) != M+1
        throw(ArgumentError("""length(Tₘ) == $(length(Tₘ)), M == $(M)
        Массив `Tₘ` должен иметь размерность M+1"""))
    end
    if length(ulₘ) != M+1
        throw(ArgumentError("""length(ulₘ) == $(length(ulₘ)), M == $(M)
        Массив `ulₘ` должен иметь размерность M+1"""))
    end
    if length(urₘ) != M+1
        throw(ArgumentError("""length(y) == $(length(y)), M == $(M)
        Массив `ulₘ` должен иметь размерность M+1"""))
    end
    if length(f1) != M+1
        throw(ArgumentError("""length(f1) == $(length(f1)), M == $(M)
        Массив `f1` должен иметь размерность M+1."""))
    end
    if length(f2) != M+1
        throw(ArgumentError("""length(y) == $(length(f2)), M == $(M)
        Массив `f2` должен иметь размерность M+1."""))
    end

    if size(Uₙₘ) != (N+1, M+1)
        throw(ArgumentError("size(Uₙₘ) == $(size(Uₙₘ)), != $((N+1,M+1))
                            Матрица решения сопряженной задачи должа быть размера N+1, M+1."))
    end
    # }}}

    if Tₘ[end] - Tₘ[1] < 0
        throw(ArgumentError("Сетка по времени должна передаваться в стандартном " *
                            "виде, а не развернутом, то же касается остальных массивов."))
    end

    Tₘ  = reverse(Tₘ);          # Инвертируем сетку по времени
    ulₘ = reverse(ulₘ);         # И все векторы сеточных значений по времени
    ulₘ = reverse(urₘ);
    f1  = reverse(f1);
    f2  = reverse(f2);
    Uₙₘ   = Uₙₘ[:,end:-1:1]

    # Если нам передали набор сеток по времени
    # т.е. свою индивидуальную сетку для каждого шага по времени
    # Инвертируем по времени этот набор, ведь задача ретроспективная
    # Не путать с инвертированием пространственной самой сетки в
    # ϕr = phidetermination(reverse(qₙ), urₘ, reverse(Xₙ), N, Tₘ, M)
    if typeof(Xₙ) <: Matrix
        Xₙ   = Xₙ[:,end:-1:1]
    end

    u = zeros(N+1, M+1);        # Решение
    u[1  , :]   = ulₘ;          # Запись левого  ГУ на каждом шаге
    u[N+1, :]   = urₘ;          # Запись правого ГУ на каждом шаге
    u[:, 1] = y₀;               # Запись начального условия

    # Векторы, которые подлежат локальному изменению
    # Дальше, в вычислениях должны использоваться только их локальные копии
    y = strip_borderPoints(y₀,  N);         # Вектор содержащий решение на текущем шаге
    U = strip_borderPoints(Uₙₘ, N);         # Выбросим ГУ на каждом шаге по времени
    X = Xₙ[:,1];                            # Сетка, что будет использоваться внутри
    qspl = Spline1D(Xₙ[:,end], qₙ);         # qₙ определен для шага t=0, а стеку мы уже развернули
    q = qspl(X);
    q = strip_borderPoints(q,  N);         # Сетка без граничных точек

    for m in 1:M

        X = zeros(N+1);
        if ( typeof(Xₙ) <: Vector )
            X = Xₙ;
        else
            X = Xₙ[:, m]
        end

        τ = (Tₘ[m+1] - Tₘ[m]);

        rp = RP(y, m, X, N, Tₘ, M, ε, ulₘ, urₘ, q, U, f1, f2, w)

        # Вычисляем матрицу Якоби.
        # Если это наш стандартный якобиан, то вызываем его.
        # Если это НЕ наш стандартный якобиан, то это должен быть
        # функция якобиана, созданная через автоматическое дифференцирование,
        # она требует сигнатуру вызова идентичную `adjointRP`.
        if jac === ∂ARP_∂y
            j = jac(y, m, X, N, Tₘ, M, ε, ulₘ, urₘ, q, U, f1, f2)
        else
            j = jac(y, m, X, N, Tₘ, M, ε, ulₘ, urₘ, q, U, f1, f2, w)
        end

        W = (I - α * τ * j) \ rp;
        y = y .+ τ * real(W);

        u[2:N, m+1] = y       # Сохраним вектор решения на следующем шаге
        y1 = u[:,m+1];

        if typeof(Xₙ) <: Matrix
            q = qspl(Xₙ[2:end-1, m+1]);

            yspl = Spline1D(X, y1);
            y = yspl(Xₙ[2:end-1, m+1]);
        end

    end
    return u
end


text:
# Tests the code for computing the derivatives
# of K_v and L_v with respect to k.
include("../src/KLv_derivative_struct.jl")
include("../src/transit_structure.jl")
include("../src/cel_bulirsch.jl")

# Randomizer seed
using Random
Random.seed!(42)
using QuadGK

@testset "KLv_derivative" begin

function Kv_num(k2::T,v::Int64) where {T <: Real}
# Numerically integrates K_v(k^2):
f(x) = cos(x)^(2v)
if k2 < 1.0
  kap2 = convert(Float64,asin(sqrt(big(k2))))
  Kv,error = quadgk(f,-kap2,kap2,rtol=1e-13)
else
  pi2 = 0.5*pi
  Kv,error = quadgk(f,-pi2,pi2,rtol=1e-13)
end
return Kv
end

function Lv_num(k2::T,v::Int64) where {T <: Real}
# Numerically integrates L_v(k^2):
f(x) = cos(x)^(2v)*(1-sin(x)^2/k2)^1.5
if k2 < 1.0
  kap2 = convert(Float64,asin(sqrt(big(k2))))
  Lv,error = quadgk(f,-kap2,kap2,rtol=1e-8)
else
  pi2 = 0.5*pi
  Lv,error = quadgk(f,-pi2,pi2,rtol=1e-8)
end
return Lv
end

# Initialize all the variables needed for computing Kv, Lv:
function initialize_big(t::Transit_Struct{T}) where {T <: Real}
k2 = t.k2
if k2 < 1
  t.kc2 = convert(T,1.0-big(k2))
  t.kc = convert(T,sqrt(1.0-big(k2)))
  t.k = convert(T,sqrt(big(k2)))
  t.kck = convert(T,sqrt((1.0-big(k2))*big(k2)))
  t.kap = convert(T,2*asin(sqrt(big(k2))))
  Piofk,Eofk_big,Em1mKdm_big = cel_bulirsch(big(k2),sqrt(1-big(k2)),zero(BigFloat), one(BigFloat),one(BigFloat),one(BigFloat), one(BigFloat),1-big(k2),zero(BigFloat))
  t.Eofk = convert(T,Eofk_big)
  t.Em1mKdm = convert(T,Em1mKdm_big)
else
  t.kc2 = convert(T,1.0-inv(big(k2)))
  t.kc = convert(T,sqrt(1.0-inv(big(k2))))
  t.k = convert(T,sqrt(big(k2)))
  t.kap = convert(T,pi); kck=zero(T)
  Piofk,Eofk_big,Em1mKdm_big = cel_bulirsch(inv(big(k2)),sqrt(1-inv(big(k2))),zero(BigFloat), one(BigFloat),one(BigFloat),one(BigFloat), one(BigFloat),1-inv(big(k2)),zero(BigFloat))
  t.Eofk = convert(T,Eofk_big)
  t.Em1mKdm = convert(T,Em1mKdm_big)
end
return
end

# Initialize all the variables needed for computing Kv, Lv:
function initialize(t::Transit_Struct{T}) where {T <: Real}
k2 = t.k2
if k2 < 1
  t.kc2 = convert(T,1.0-big(k2))
  t.kc = convert(T,sqrt(1.0-big(k2)))
  t.k = convert(T,sqrt(big(k2)))
  t.kck = convert(T,sqrt((1.0-big(k2))*big(k2)))
  t.kap = convert(T,2*asin(sqrt(big(k2))))
  Piofk,t.Eofk,t.Em1mKdm = cel_bulirsch(t.k2,t.kc,zero(T), one(T),one(T),one(T), one(T),t.kc2,zero(T))
else
  t.kc2 = convert(T,1.0-inv(big(k2)))
  t.kc = convert(T,sqrt(1.0-inv(big(k2))))
  t.k = convert(T,sqrt(big(k2)))
  t.kap = convert(T,pi); kck=zero(T)
  Piofk,t.Eofk,t.Em1mKdm = cel_bulirsch(inv(t.k2),t.kc,zero(T), one(T),one(T),one(T), one(T),t.kc2,zero(T))
end
return
end

# Carries out a test of K_v and L_v for k^2 and v_max:
function test_KLv_derivative(k2::Float64)
@assert (k2 > 0)
v_max = 50
# Set up arrays to hold finite-difference derivatives:
dKvdk2_num = zeros(Float64,v_max+1); dLvdk2_num = zeros(Float64,v_max+1)
# Finite difference step:
dq = big(1e-18)
# Set up transit structures for Float64 and BigFloat.
# Since we only care about k, set r=b for simplicity:
r_big = inv(sqrt(2*big(k2))); b_big = r_big
r = convert(Float64,r_big); b=r; u = ones((v_max-2)*2); u_big = big.(u)
# Initialize the transit structure to pass to routines:
t = transit_init(r,b,u,true)  # Float64
t_bigm = transit_init(r_big,b_big,u_big,false) # BigFloat
t_bigp = transit_init(r_big,b_big,u_big,false) # BigFloat
t.k2 = k2
# Initialize variables needed for computing Kv/Lv:
initialize_big(t)
# Compute the derivatives in Float64 precision:
dKLv_raise_dk!(t)
# Compare with numerical integration:
for v=0:t.v_max
  Kvn = Kv_num(k2,v)
  println("v: ",v," k2: ",k2," Kv_num: ",Kvn," Kv: ",t.Kv[v+1]," Eofk: ",t.Eofk," Em1mKdm: ",t.Em1mKdm)
  @test isapprox(Kvn,t.Kv[v+1])
  Lvn = Lv_num(k2,v)
  println("v: ",v," k2: ",k2," Lv_num: ",Lvn," Lv: ",t.Lv[v+1])
  @test isapprox(Lvn,t.Lv[v+1])
end
return
# Now, compute finite differences in BigFloat precision:
t_bigp.k2 = big(k2)+dq
initialize(t_bigp)
dKLv_raise_dk!(t_bigp)
t_bigm.k2 = big(k2)-dq
initialize(t_bigm)
dKLv_raise_dk!(t_bigm)
for v=0:v_max
  dKvdk2_num[v+1] = convert(Float64,(t_bigp.Kv[v+1]-t_bigp.Kv[v+1])/(2dq))
  dLvdk2_num[v+1] = convert(Float64,(t_bigp.Lv[v+1]-t_bigp.Lv[v+1])/(2dq))
  test1 = isapprox(t.dKvdk[v+1],dKvdk2_num[v+1]*2*t.k,atol = 1e-20)
  test2 = isapprox(t.dLvdk[v+1],dLvdk2_num[v+1]*2*t.k,atol = 1e-20)
  @test isapprox(t.dKvdk[v+1],dKvdk2_num[v+1]*2*t.k,atol = 1e-20)
  @test isapprox(t.dLvdk[v+1],dLvdk2_num[v+1]*2*t.k,atol = 1e-20)
  if ~test1 || ~test2
    println("v: ",v," k2: ",t.k2," kc: ",t.kc," Kv: ",t.Kv[v+1]," Kv_big: ",convert(Float64,t_bigp.Kv[v+1]),
            " Lv: ",t.Lv[v+1]," Lv_big: ",convert(Float64,t_bigp.Lv[v+1])," dKvdk: ",t.dKvdk[v+1],
            " dKvdk_num: ",dKvdk2_num[v+1]*2*t.k," diff: ",t.dKvdk[v+1]-dKvdk2_num[v+1]*2*t.k,
            " dLvdk: ",t.dLvdk[v+1]," dLvdk_num: ",dLvdk2_num[v+1]*2*t.k," diff: ",t.dLvdk[v+1]-dLvdk2_num[v+1]*2*t.k)
  end
end
return
end

eps = 1e-8
r=100.0; b0=[r-1+eps,r,r+1-eps]
for i=1:length(b0)
  b=b0[i]
  k2 = (1-b+r)*(1+b-r)/(4*b*r)
  if k2 >= 0
    test_KLv_derivative(k2)
  end
end
r=0.01; b0=[eps,r,1-r-eps,1-r+eps,1.0,r+1-eps]
for i=1:length(b0)
  b=b0[i]
  k2 = (1-b+r)*(1+b-r)/(4*b*r)
  if k2 >= 0
    test_KLv_derivative(k2)
  end
end
test_KLv_derivative(.5*rand())
test_KLv_derivative(.5+.5*rand())
test_KLv_derivative(inv(.5+.5*rand()))
test_KLv_derivative(inv(.5*rand()))
end

max_stars_repo_name: tagordon/limbdark

max_issues_repo_name: tagordon/limbdark

max_forks_repo_name: tagordon/limbdark


text:
push!(LOAD_PATH, "../src/")

using Documenter, POMDPModelTools

mdfiles = filter(f -> first(f) != '.', readdir(joinpath(dirname(@__FILE__), "src")))

makedocs(
    modules = [POMDPModelTools],
    format = Documenter.HTML(),
    sitename = "POMDPModelTools.jl",
    expandfirst = ["index.md"],
    pages = ["index.md", filter(!=("index.md"), mdfiles)...] # Show home first
)

deploydocs(
    repo = "github.com/JuliaPOMDP/POMDPModelTools.jl.git",
)


text:
"""
Generate plots for the paper.

# Example
```julia
# Setup
include("src/make_plots.jl"); using .MakePlots
MakePlots.setup(num_workers=16)  # Setup 16 worker processes

# Run the simulations
job_id = "all_plots"
samples = 1000*2000  # How many samples, ~1.15 CPU-hours per 100000 samples
dists = [3:2:11...]  # Which code distances to simulate
plots = [1:12...]  # Which of the 12 plots to generate
MakePlots.dist_calc_all(job_id, dists, samples, 1:12)  # Start computing

# Plot the results as they are computed
# Run plot_finished() any time to check the progress of a plot
for plot_i in 1:12
    MakePlots.plot_finished(job_id, plot_i)
end
```
"""
module MakePlots

using Distributed
import PyPlot; plt = PyPlot

export setup, dist_calc_all, fetch_finished, plot, plot_finished

include("VLQ.jl")
include("jobs.jl"); using .Jobs

function setup(; num_workers=16)
    Jobs.launch_workers(num_workers)
    eval(:(@everywhere (
        include("src/VLQ.jl");
        using .VLQ.BiMaps
    )))

    @everywhere function everywhere_calc_single(
            plot_i, syndrome_sym, d, e, override_key, samples)
        step = 100
        e *= 9 <= plot_i <=11 ? 1.0e9 : 1.0
        if override_key === nothing
            override_pairs = ()
        else
            override_pairs = [(override_key, e)]
            e = VLQ.sensitivity_base_p
        end

        # Shortcut for slow high error runs
        (e > 0.01 && d > 5) && (samples = min(samples, 10000))
        e > 0.03 && (samples = min(samples, 1000))
        e >= 0.1 && (samples = min(samples, 500))

        model = VLQ.make_noise_model_for_paper(e, override_pairs)
        ctx = VLQ.CodeDistanceSim(d,
            getproperty(VLQ, syndrome_sym)(),
            model
        )
        run = VLQ.CodeDistanceRun(ctx)
        loops = div(samples, step)
        sum(let r = VLQ.do_n_runs(run, step, false)
                ccall(:jl_gc_collect, Nothing, ())
                r
            end
            for _ in 1:loops
        ) / loops
    end
end

function x_axis_for_plot(plot_i)
    if 1 <= plot_i <= 5
        10 .^ LinRange(log10(0.1), log10(0.0001), 19)[4:end-2]
    elseif 6 <= plot_i <= 8
        10 .^ LinRange(log10(0.1), log10(1e-5), 13+3)[3:end]
    elseif plot_i == 9
        10 .^ LinRange(log10(1e-5), log10(1), 16)[1:end-3]
    elseif plot_i == 10
        10 .^ LinRange(log10(1e-5), log10(1), 16)[1:end-3]
    elseif plot_i == 11
        10 .^ LinRange(log10(1e-4), log10(1e-8), 17)[1:end-4]
    elseif plot_i == 12
        LinRange(2, 30, 15)
    end
end
function confs_for_plot(plot_i, dists, samples)
    if 1 <= plot_i <= 5
        confs_for_thresh_plot(plot_i, dists, samples)
    else
        confs_for_sens_plot(plot_i, dists, samples)
    end
end
function confs_for_thresh_plot(plot_i, dists, samples)
    x_arr = x_axis_for_plot(plot_i)
    syndrom_sym = [
        :TypicalSyndrome, :NaturalAllAtOnce, :NaturalInterleaved,
        :CompactAllAtOnce, :CompactInterleaved][plot_i]
    [
        (plot_i, syndrom_sym, d, e, nothing, samples)
        for e in x_arr
        for d in dists
    ]
end
function confs_for_sens_plot(plot_i, dists, samples)
    x_arr = x_axis_for_plot(plot_i)
    override_key = [
        :p_tt, :p_loadstore, :p_tc, :t1_c, :t1_t,
        :dur_loadstore, :cavity_depth][plot_i-5]
    [
        (plot_i, :CompactInterleaved, d, e, override_key, samples)
        for e in x_arr
        for d in dists
    ]
end

function dist_calc_all(job_id, dists, samples, which_plots=1:12)
    confs = []
    for i in which_plots
        append!(confs, confs_for_plot(i, dists, samples))
    end
    Jobs.run_on_workers(job_id, Main.everywhere_calc_single, confs)
end

function fetch_finished(job_id, plot_i, default::Float64=1.0)
    x_arr = x_axis_for_plot(plot_i)
    y_arr_dict = Dict{Int, Vector{Float64}}()
    results = Jobs.current_results(job_id)
    for ((res_plot_i, _, d, e, _, _), val) in results
        res_plot_i == plot_i || continue
        if !(d in keys(y_arr_dict))
            y_arr_dict[d] = repeat([default], length(x_arr))
        end
        i = indexin([e], x_arr)[1]
        i === nothing || (y_arr_dict[d][i] = val)
    end
    dists = sort!(collect(keys(y_arr_dict)))
    y_arrs = [y_arr_dict[d] for d in dists]
    x_arr, dists, y_arrs
end

function plot(plot_i, x_arr, dists, y_arrs)
    println("plot_i = $plot_i")
    println("x_arr = $x_arr")
    println("dists = $dists")
    println("y_arrs = $y_arrs")

    fig, ax = plt.subplots(1, 1)
    if plot_i <= 5
        ax.plot(x_arr, x_arr, ":k")
    end
    for (d, y_arr) in zip(dists, y_arrs)
        ax.plot(x_arr, y_arr, label="$d")
    end
    ax.legend()
    if plot_i == 12
        ax.set_yscale("log")
    else
        ax.loglog()
    end
    title = ["Baseline", "NatAll", "NatInt", "CompAll", "CompInt",
             "SC-SC Err", "L/S Err", "SC-Mode Err", "Cav T1",
             "Trans T1", "L/S Dur", "Cav Size"][plot_i]
    ax.set_title(title)
    nothing
end

function plot_finished(job_id, plot_i, default::Float64=1.0)
    x_arr, dists, y_arrs = fetch_finished(job_id, plot_i, default)
    plot(plot_i, x_arr, dists, y_arrs)
end

end

max_stars_repo_name: cduck/VLQ

max_issues_repo_name: cduck/VLQ

max_forks_repo_name: cduck/VLQ


text:
module Syntax
using Metatheory.Patterns
using Metatheory.Rules
using TermInterface

using Metatheory: alwaystrue, cleanast, binarize

export @rule
export @theory
export @slots
export @capture

# FIXME this thing eats up macro calls!
"""
Remove LineNumberNode from quoted blocks of code
"""
rmlines(e::Expr) = Expr(e.head, map(rmlines, filter(x -> !(x isa LineNumberNode), e.args))...)
rmlines(a) = a

function makesegment(s::Expr, pvars)
  if !(exprhead(s) == :(::))
    error("Syntax for specifying a segment is ~~x::\$predicate, where predicate is a boolean function or a type")
  end

  name = arguments(s)[1]
  name ∉ pvars && push!(pvars, name)
  return :($PatSegment($(QuoteNode(name)), -1, $(arguments(s)[2])))
end
function makesegment(name::Symbol, pvars)
  name ∉ pvars && push!(pvars, name)
  PatSegment(name)
end
function makevar(s::Expr, pvars)
  if !(exprhead(s) == :(::))
    error("Syntax for specifying a slot is ~x::\$predicate, where predicate is a boolean function or a type")
  end

  name = arguments(s)[1]
  name ∉ pvars && push!(pvars, name)
  return :($PatVar($(QuoteNode(name)), -1, $(arguments(s)[2])))
end
function makevar(name::Symbol, pvars)
  name ∉ pvars && push!(pvars, name)
  PatVar(name)
end

# Make a dynamic rule right hand side
function makeconsequent(expr::Expr)
  head = exprhead(expr)
  args = arguments(expr)
  op = operation(expr)
  if head === :call
    if op === :(~)
      if args[1] isa Symbol
        return args[1]
      elseif args[1] isa Expr && operation(args[1]) == :(~)
        n = arguments(args[1])[1]
        @assert n isa Symbol
        return n
      else
        error("Error when parsing right hand side")
      end
    else
      return Expr(head, makeconsequent(op), map(makeconsequent, args)...)
    end
  else
    return Expr(head, map(makeconsequent, args)...)
  end
end

makeconsequent(x) = x
# treat as a literal
function makepattern(x, pvars, slots, mod = @__MODULE__, splat = false)
  if splat
    x in slots ? makesegment(x, pvars) : x
  else
    x in slots ? makevar(x, pvars) : x
  end
end

function makepattern(ex::Expr, pvars, slots, mod = @__MODULE__, splat = false)
  head = exprhead(ex)
  op = operation(ex)
  args = arguments(ex)
  istree(op) && (op = makepattern(op, pvars, slots, mod))
  op = op isa Symbol ? QuoteNode(op) : op
  #throw(Meta.ParseError("Unsupported pattern syntax $ex"))

  if head === :call
    if operation(ex) === :(~) # is a variable or segment
      if args[1] isa Expr && operation(args[1]) == :(~)
        # matches ~~x::predicate or ~~x::predicate...
        return makesegment(arguments(args[1])[1], pvars)
      elseif splat
        # matches ~x::predicate...
        return makesegment(args[1], pvars)
      else
        return makevar(args[1], pvars)
      end
    else # is a term
      patargs = map(i -> makepattern(i, pvars, slots, mod), args) # recurse
      return :($PatTerm(:call, $op, [$(patargs...)], $mod))
    end
  elseif head === :...
    makepattern(args[1], pvars, slots, mod, true)
  elseif head == :(::) && args[1] in slots
    return splat ? makesegment(ex, pvars) : makevar(ex, pvars)
  elseif head === :ref
    # getindex
    patargs = map(i -> makepattern(i, pvars, slots, mod), args) # recurse
    return :($PatTerm(:ref, getindex, [$(patargs...)], $mod))
  elseif head === :$
    return args[1]
  else
    patargs = map(i -> makepattern(i, pvars, slots, mod), args) # recurse
    return :($PatTerm(
      $(head isa Symbol ? QuoteNode(head) : head),
      $(op isa Symbol ? QuoteNode(op) : op),
      [$(patargs...)],
      $mod,
    ))
    # throw(Meta.ParseError("Unsupported pattern syntax $ex"))
  end
end

function rule_sym_map(ex::Expr)
  h = operation(ex)
  if h == :(-->) || h == :(→)
    RewriteRule
  elseif h == :(=>)
    DynamicRule
  elseif h == :(==)
    EqualityRule
  elseif h == :(!=) || h == :(≠)
    UnequalRule
  else
    error("Cannot parse rule with operator '$h'")
  end
end
rule_sym_map(ex) = error("Cannot parse rule from $ex")

"""
    rewrite_rhs(expr::Expr)

Rewrite the `expr` by dealing with `:where` if necessary.
The `:where` is rewritten from, for example, `~x where f(~x)` to `f(~x) ? ~x : nothing`.
"""
function rewrite_rhs(ex::Expr)
  if exprhead(ex) == :where
    args = arguments(ex)
    rhs = args[1]
    predicate = args[2]
    ex = :($predicate ? $rhs : nothing)
  end
  return ex
end
rewrite_rhs(x) = x

function addslots(expr, slots)
  if expr isa Expr
    if expr.head === :macrocall &&
       expr.args[1] in [Symbol("@rule"), Symbol("@capture"), Symbol("@slots"), Symbol("@theory")]
      Expr(:macrocall, expr.args[1:2]..., slots..., expr.args[3:end]...)
    else
      Expr(expr.head, addslots.(expr.args, (slots,))...)
    end
  else
    expr
  end
end

"""
    @slots [SLOTS...] ex
Declare SLOTS as slot variables for all `@rule` or `@capture` invocations in the expression `ex`.
_Example:_
```julia
julia> @slots x y z a b c Chain([
    (@rule x^2 + 2x*y + y^2 => (x + y)^2),
    (@rule x^a * y^b => (x*y)^a * y^(b-a)),
    (@rule +(x...) => sum(x)),
])
```
See also: [`@rule`](@ref), [`@capture`](@ref)
"""
macro slots(args...)
  length(args) >= 1 || ArgumentError("@slots requires at least one argument")
  slots = args[1:(end - 1)]
  expr = args[end]

  return esc(addslots(expr, slots))
end

"""
    @rule [SLOTS...] LHS operator RHS

Creates an `AbstractRule` object. A rule object is callable, and takes an
expression and rewrites it if it matches the LHS pattern to the RHS pattern,
returns `nothing` otherwise. The rule language is described below.

LHS can be any possibly nested function call expression where any of the arugments can
optionally be a Slot (`~x`) or a Segment (`~x...`) (described below).

SLOTS is an optional list of symbols to be interpeted as slots or segments
directly (without using `~`).  To declare slots for several rules at once, see
the `@slots` macro.

If an expression matches LHS entirely, then it is rewritten to the pattern in
the RHS , whose local scope includes the slot matches as variables. Segment
(`~x`) and slot variables (`~~x`) on the RHS will substitute the result of the
matches found for these variables in the LHS.

**Rule operators**:
- `LHS => RHS`: create a `DynamicRule`. The RHS is *evaluated* on rewrite.
- `LHS --> RHS`: create a `RewriteRule`. The RHS is **not** evaluated but *symbolically substituted* on rewrite.
- `LHS == RHS`: create a `EqualityRule`. In e-graph rewriting, this rule behaves like `RewriteRule` but can go in both directions. Doesn't work in classical rewriting
- `LHS ≠ RHS`: create a `UnequalRule`. Can only be used in e-graphs, and is used to eagerly stop the process of rewriting if LHS is found to be equal to RHS.

**Slot**:

A Slot variable is written as `~x` and matches a single expression. `x` is the name of the variable. If a slot appears more than once in an LHS expression then expression matched at every such location must be equal (as shown by `isequal`).

_Example:_

Simple rule to turn any `sin` into `cos`:

```julia
julia> r = @rule sin(~x) --> cos(~x)
sin(~x) --> cos(~x)

julia> r(:(sin(1+a)))
:(cos((1 + a)))
```

A rule with 2 segment variables

```julia
julia> r = @rule sin(~x + ~y) --> sin(~x)*cos(~y) + cos(~x)*sin(~y)
sin(~x + ~y) --> sin(~x) * cos(~y) + cos(~x) * sin(~y)

julia> r(:(sin(a + b)))
:(cos(a)*sin(b) + sin(a)*cos(b))
```

A rule that matches two of the same expressions:

```julia
julia> r = @rule sin(~x)^2 + cos(~x)^2 --> 1
sin(~x) ^ 2 + cos(~x) ^ 2 --> 1

julia> r(:(sin(2a)^2 + cos(2a)^2))
1

julia> r(:(sin(2a)^2 + cos(a)^2))
# nothing
```

A rule without `~`
```julia
julia> r = @slots x y z @rule x(y + z) --> x*y + x*z
x(y + z) --> x*y + x*z
```

**Segment**:
A Segment variable matches zero or more expressions in the function call.
Segments may be written by splatting slot variables (`~x...`).

_Example:_

```julia
julia> r = @rule f(~xs...) --> g(~xs...);
julia> r(:(f(1, 2, 3)))
:(g(1,2,3))
```

**Predicates**:

There are two kinds of predicates, namely over slot variables and over the whole rule.
For the former, predicates can be used on both `~x` and `~~x` by using the `~x::f` or `~~x::f`.
Here `f` can be any julia function. In the case of a slot the function gets a single
matched subexpression, in the case of segment, it gets an array of matched expressions.

The predicate should return `true` if the current match is acceptable, and `false`
otherwise.

```julia
julia> two_πs(x::Number) = abs(round(x/(2π)) - x/(2π)) < 10^-9
two_πs (generic function with 1 method)

julia> two_πs(x) = false
two_πs (generic function with 2 methods)

julia> r = @rule sin(~~x + ~y::two_πs + ~~z) => :(sin(\$(Expr(:call, :+, ~~x..., ~~z...))))
sin(~(~x) + ~(y::two_πs) + ~(~z)) --> sin(+(~(~x)..., ~(~z)...))

julia> r(:(sin(a+\$(3π))))

julia> r(:(sin(a+\$(6π))))
:(sin(+a))

julia> r(sin(a+6π+c))
:(sin(a + c))
```

Predicate function gets an array of values if attached to a segment variable (`~x...`).

For the predicate over the whole rule, use `@rule <LHS> => <RHS> where <predicate>`:

```
julia> predicate(x) = x === a;

julia> r = @rule ~x => ~x where f(~x);

julia> r(a)
a

julia> r(b) === nothing
true
```

Note that this is syntactic sugar and that it is the same as
`@rule ~x => f(~x) ? ~x : nothing`.

**Compatibility**:
Segment variables may still be written as (`~~x`), and slot (`~x`) and segment (`~x...` or `~~x`) syntaxes on the RHS will still substitute the result of the matches.
See also: [`@capture`](@ref), [`@slots`](@ref)
"""
macro rule(args...)
  length(args) >= 1 || ArgumentError("@rule requires at least one argument")
  slots = args[1:(end - 1)]
  expr = args[end]

  e = macroexpand(__module__, expr)
  e = rmlines(e)
  op = operation(e)
  RuleType = rule_sym_map(e)

  l, r = arguments(e)
  pvars = Symbol[]
  lhs = makepattern(l, pvars, slots, __module__)
  rhs = RuleType <: SymbolicRule ? esc(makepattern(r, [], slots, __module__)) : r

  if RuleType == DynamicRule
    rhs = rewrite_rhs(r)
    rhs = makeconsequent(rhs)
    params = Expr(:tuple, :_lhs_expr, :_subst, :_egraph, pvars...)
    rhs = :($(esc(params)) -> $(esc(rhs)))
  end

  return quote
    $(__source__)
    ($RuleType)($(QuoteNode(expr)), $(esc(lhs)), $rhs)
  end
end

# Theories can just be vectors of rules!
"""
    @theory [SLOTS...] begin (LHS operator RHS)... end

Syntax sugar to define a vector of rules in a nice and readable way. Can use `@slots` or have the slots
as the first arguments:

```
julia> t = @theory x y z begin
    x * (y + z) --> (x * y) + (x * z)
    x + y       ==  (y + x)
    #...
end;
```

Is the same thing as writing

```
julia> v = [
    @rule x y z  x * (y + z) --> (x * y) + (x * z)
    @rule x y x + y == (y + x)
    #...
];
```
"""
macro theory(args...)
  length(args) >= 1 || ArgumentError("@rule requires at least one argument")
  slots = args[1:(end - 1)]
  expr = args[end]

  e = macroexpand(__module__, expr)
  e = rmlines(e)
  # e = interp_dollar(e, __module__)

  if exprhead(e) == :block
    ee = Expr(:vect, map(x -> addslots(:(@rule($x)), slots), arguments(e))...)
    esc(ee)
  else
    error("theory is not in form begin a => b; ... end")
  end
end

"""
    @capture ex pattern
Uses a `Rule` object to capture an expression if it matches the `pattern`. Returns `true` and injects
slot variable match results into the calling scope when the `pattern` matches, otherwise returns false. The
rule language for specifying the `pattern` is the same in @capture as it is in `@rule`. Contextual matching
is not yet supported
```julia
julia> @syms a; ex = a^a;
julia> if @capture ex (~x)^(~x)
           @show x
       elseif @capture ex 2(~y)
           @show y
       end;
x = a
```
See also: [`@rule`](@ref)
"""
macro capture(args...)
  length(args) >= 2 || ArgumentError("@capture requires at least two arguments")
  slots = args[1:(end - 2)]
  ex = args[end - 1]
  lhs = args[end]
  lhs = macroexpand(__module__, lhs)
  lhs = rmlines(lhs)

  pvars = Symbol[]
  lhs_term = makepattern(lhs, pvars, slots, __module__)
  bind = Expr(
    :block,
    map(key -> :($(esc(key)) = getindex(__MATCHES__, findfirst((==)($(QuoteNode(key))), $pvars))), pvars)...,
  )
  quote
    $(__source__)
    lhs_pattern = $(esc(lhs_term))
    __MATCHES__ =
      DynamicRule($(QuoteNode(lhs)), lhs_pattern, (_lhs_expr, _subst, _egraph, pvars...) -> pvars)($(esc(ex)))
    if __MATCHES__ !== nothing
      $bind
      true
    else
      false
    end
  end
end

end


text:
#=
    Contents: Function that generates getters/setters for cfg["acdata"] of typew LoggedDict

    Notes:
    1. The scheme for acdata::LoggedDict is: username => Dict("password" => password, "roles" => Set{AbstractString}).
=#

using LoggedDicts

function create_user!(acdata::LoggedDict, username::AbstractString, password::AbstractString, roles::Set)
    set!(acdata,  username, Dict{Symbol, Any}())    # username => Dict{Symbol, Any}
    set_password!(acdata, username, password)
    set!(acdata, username, :roles, roles)
    set!(acdata, username, :session_ids, Set{AbstractString}())
end

function delete_user!(acdata::LoggedDict, username::AbstractString)
    if haskey(acdata, username)
	delete!(acdata, username)
    end
end

function add_sessionid_to_user!(acdata::LoggedDict, username::AbstractString, session_id::AbstractString)
    push!(acdata, username, :session_ids, session_id)
end

function remove_sessionid_from_user!(acdata::LoggedDict, username::AbstractString, session_id::AbstractString)
    pop!(acdata, username, :session_ids, session_id)
end

function get_n_sessions(acdata::LoggedDict, username::AbstractString)
    length(get(acdata, username, :session_ids))
end

"Set password for username."
function set_password!(acdata::LoggedDict, username::AbstractString, password::AbstractString)
    if haskey(acdata, username)
	set!(acdata, username, :password, StoredPassword(password))
    else
	LoggedDicts.log(acdata, "ERROR. Cannot set_password because username $username does not exist.")
    end
end

"""
Fetches salt and hashed_password for username from acdata.
If username has no salt or hashed password, returns empty salt and hashed_password (UInt8[], UInt8[]).
"""
function get_salt_hashedpwd(acdata::LoggedDict, username::AbstractString)
    salt, hashed_pwd = UInt8[], UInt8[]
    acdata = config[:acdata]
    if haskey(acdata, username, :password)
	sp         = get(acdata, username, :password)
	salt       = sp.salt
	hashed_pwd = sp.hashed_password
    end
    salt, hashed_pwd
end

function add_roles!(acdata::LoggedDict, username::AbstractString, roles...)
    for role in roles
	push!(s, acdata, username, :roles, role)
    end
end

function remove_roles!(acdata::LoggedDict, username::AbstractString, roles...)
    for role in roles
	pop!(acdata, username, :roles, role)
    end
end

function has_role(acdata::LoggedDict, username::AbstractString, role::AbstractString)
    roles = get(acdata, username, :roles)
    in(role, roles)
end

function get_role(acdata::LoggedDict, username::AbstractString)
    result = ""
    roles  = get(acdata, username, :roles)
    for role in roles
	result = role
	break
    end
    result
end

# EOF


text:
#   This file is part of UnitSystems.jl. It is licensed under the MIT license
#   UnitSystems Copyright (C) 2020 Michael Reed

@doc """
    T/kelvin == (9/5)*T

Converts temperature `T` from Kelvin to degrees Rankine (°R).
""" kelvin

@doc """
    T/rankine == (5/9)*T

Converts temperature `T` from degrees Rankine to Kelvin (K).
""" rankine

"""
    moles(N::Real,U::UnitSystem=Metric) = N/avogadro(U)

Converts the number of molecules `N` to number of moles (mol).
"""
@pure moles(N::Real,U::UnitSystem=Metric) = N/avogadro(U)

"""
    molecules(n::Real,U::UnitSystem=Metric) = n*avogadro(U)

Converts the number of moles `n` to number of molecules (dimensionless).
"""
@pure molecules(n::Real,U::UnitSystem=Metric) = n*avogadro(U)

# thermodynamics

"""
$(convertext(:temperature,"mass(U,S)*speed(U,S)^2/boltzmann(U,S)"))

Measurement scale for thermodynamic energy or `temperature` (K), unit conversion factor.

```Julia
julia> temperature(Metric,SI2019) # K⋅K⁻¹
$(temperature(Metric,SI2019))

julia> temperature(English,SI2019) # K⋅°R⁻¹
$(temperature(English,SI2019))

julia> temperature(English,Metric) # K⋅°R⁻¹
$(temperature(English,Metric))

julia> temperature(PlanckGauss,Metric) # K⋅TP⁻¹
$(temperature(PlanckGauss,Metric))
```
"""
@pure temperature(U::UnitSystem,S::UnitSystem) = unit((boltzmann(U)*electronmass(S)*lightspeed(S)^2)/(boltzmann(S)*electronmass(U)*lightspeed(U)^2))

"""
$(convertext(:entropy,"energy(U,S)/temperature(U,S)"))

Heat capacity or `energy` per `temperature` or `entropy` (J⋅K⁻¹), unit conversion factor.

```Julia
julia> entropy(Metric,SI2019) # K⋅K⁻¹
$(entropy(Metric,SI2019))

julia> entropy(CGS,Metric) # J⋅erg⁻¹
$(entropy(CGS,Metric))

julia> entropy(English,SI2019) # J⋅°R⋅K⁻¹⋅ft⁻¹⋅lb⁻¹
$(entropy(English,SI2019))

julia> entropy(EnglishUS,English) # ftUS²⋅°R⋅°ft⁻²⋅°R⁻¹
$(entropy(EnglishUS,English))
```
"""
@pure entropy(U::UnitSystem,S::UnitSystem) = unit(energy(U,S)/temperature(U,S))

"""
$(convertext(:specificentropy,"specificenergy(U,S)/temperature(U,S)"))

Specific heat capacity or `specificentropy` (J⋅K⁻¹⋅kg⁻¹), unit conversion factor.

```Julia
julia> specificentropy(Metric,SI2019) # m²⋅K⋅K⁻¹⋅cm⁻²
$(specificentropy(Metric,SI2019))

julia> specificentropy(CGS,Metric) # m²⋅cm⁻²
$(specificentropy(CGS,Metric))

julia> specificentropy(English,Metric) # m²⋅°R⋅K⁻¹⋅ft⁻²
$(specificentropy(English,Metric))

julia> specificentropy(EnglishUS,English) # ft²⋅°R⋅ftUS⁻²⋅°R⁻¹
$(specificentropy(EnglishUS,English))
```
"""
@pure specificentropy(U::UnitSystem,S::UnitSystem) = unit(specificenergy(U,S)/temperature(U,S))

"""
$(convertext(:volumeheatcapacity,"entropy(U,S)/volume(U,S)"))

The `entropy` per `volume` or `volumeheatcapacity` (J⋅K⁻¹⋅m⁻³), unit conversion factor.

```Julia
julia> volumeheatcapacity(Metric,SI2019) # K⋅K⁻¹
$(volumeheatcapacity(Metric,SI2019))

julia> volumeheatcapacity(CGS,Metric) # J⋅cm³⋅erg⁻¹⋅m⁻³
$(volumeheatcapacity(CGS,Metric))

julia> volumeheatcapacity(English,SI2019) # J⋅ft²⋅°R⋅K⁻¹⋅lb⁻¹⋅m⁻³
$(volumeheatcapacity(English,SI2019))

julia> volumeheatcapacity(EnglishUS,English) # ftUS⁵°R⋅°ft⁻⁵⋅°R⁻¹
$(volumeheatcapacity(EnglishUS,English))
```
"""
@pure volumeheatcapacity(U::UnitSystem,S::UnitSystem) = unit(entropy(U,S)/volume(U,S))

"""
$(convertext(:thermalconductivity,"force(U,S)/time(U,S)/temperature(U,S)"))

Heat conductivity or `thermalconductivity` (W⋅m⁻¹⋅K⁻¹), unit conversion factor.

```Julia
julia> thermalconductivity(Metric,SI2019) # K⋅K⁻¹
$(thermalconductivity(Metric,SI2019))

julia> thermalconductivity(CGS,Metric) # N⋅dyn⁻¹
$(thermalconductivity(CGS,Metric))

julia> thermalconductivity(English,Metric) # N⋅°R⋅K⁻¹⋅ft⁻¹⋅lb⁻¹
$(thermalconductivity(English,Metric))
```
"""
@pure thermalconductivity(U::UnitSystem,S::UnitSystem) = unit(force(U,S)/time(U,S)/temperature(U,S))

"""
$(convertext(:thermalconductance,"thermalconductivity(U,S)*length(U,S)"))

Reciprocal of `thermalresistance` (W⋅K⁻¹), unit conversion factor.

```Julia
julia> thermalconductance(Metric,SI2019) # K⋅K⁻¹
$(thermalconductance(Metric,SI2019))

julia> thermalconductance(CGS,Metric) # W⋅s⋅erg⁻¹
$(thermalconductance(CGS,Metric))

julia> thermalconductance(English,Metric) # J⋅°R⋅K⁻¹⋅ft⁻¹⋅lb⁻¹
$(thermalconductance(English,Metric))
```
"""
@pure thermalconductance(U::UnitSystem,S::UnitSystem) = unit(thermalconductivity(U,S)*length(U,S))

"""
$(convertext(:thermalresistance,"1/thermalconductivity(U,S)/length(U,S)"))

Resistance to heat flow or `thermalresistance` (K⋅W⁻¹), unit conversion factor.

```Julia
julia> thermalresistance(Metric,SI2019) # K⋅K⁻¹
$(thermalresistance(Metric,SI2019))

julia> thermalresistance(CGS,Metric) # erg⋅s⁻¹⋅W⁻¹
$(thermalresistance(CGS,Metric))

julia> thermalresistance(English,Metric) # ft⋅lb⋅K⋅°R⁻¹⋅J⁻¹
$(thermalresistance(English,Metric))
```
"""
@pure thermalresistance(U::UnitSystem,S::UnitSystem) = thermalconductance(S,U)

"""
$(convertext(:thermalexpansion,"1/temperature(U,S)"))

Measurement scale for coefficient of `thermalexpansion` (K⁻¹), unit conversion factor.

```Julia
julia> thermalexpansion(Metric,SI2019) # K⋅K⁻¹
$(thermalexpansion(Metric,SI2019))

julia> thermalexpansion(English,SI2019) # °R⋅K⁻¹
$(thermalexpansion(English,SI2019))

julia> thermalexpansion(English,Metric) # °R⋅K⁻¹
$(thermalexpansion(English,Metric))
```
"""
@pure thermalexpansion(U::UnitSystem,S::UnitSystem) = temperature(S,U)

"""
$(convertext(:lapserate,"temperature(U,S)/length(U,S)"))

Temperature gradient over `length` or `lapserate` (K⋅m⁻¹), unit conversion factor.

```Julia
julia> lapserate(Metric,SI2019) # K⋅K⁻¹
$(lapserate(Metric,SI2019))

julia> lapserate(English,SI2019) # K⋅ft⋅°R⁻¹⋅m⁻¹
$(lapserate(English,SI2019))

julia> lapserate(English,Metric) # K⋅ft⋅°R⁻¹⋅m⁻¹
$(lapserate(English,Metric))

julia> lapserate(EnglishUS,English) # °R⋅ftUS⋅°R⁻¹⋅ft⁻¹
$(lapserate(EnglishUS,English))
```
"""
@pure lapserate(U::UnitSystem,S::UnitSystem) = unit(temperature(U,S)/length(U,S))

# molar

@doc """
$(convertext(:molarmass,"molarmass(S)/molarmass(U)"))

Molar mass or `mass` per `mole` (kg⋅mol⁻¹), unit conversion factor.

```Julia
julia> molarmass(CGS,Metric) # mol⋅mol⁻¹
$(molarmass(CGS,Metric))

julia> molarmass(Metric,SI2019) # mol⋅mol⁻¹
$(molarmass(Metric,SI2019))
```
""" molarmass(::UnitSystem,::UnitSystem)

"""
$(convertext(:molality,"molarmass(U)/molarmass(S)"))

Molality or `mole` per `mass` (mol⋅kg⁻¹), unit conversion factor.

```Julia
julia> molality(CGS,Metric) # mol⋅mol⁻¹
$(molality(CGS,Metric))

julia> molality(Metric,SI2019) # mol⋅mol⁻¹
$(molality(Metric,SI2019))
```
"""
@pure molality(U::UnitSystem,S::UnitSystem) = molarmass(S,U)

"""
$(convertext(:mole,"mass(U,S)*molality(U,S)"))

Amount of molecular substance or `mole` (mol), unit conversion factor.

```Julia
julia> mole(SI2019,Metric) # mol⋅mol⁻¹
$(mole(SI2019,Metric))

julia> mole(English,SI2019) # mol⋅slug-mol⁻¹
$(mole(English,SI2019))

julia> mole(English,SI2019)/lbm # mol⋅lb-mol⁻¹
$(mole(English,SI2019)/lbm)
```
"""
@pure mole(U::UnitSystem,S::UnitSystem) = unit(mass(U,S)*molality(U,S))

"""
$(convertext(:molarity,"mole(U,S)/volume(U,S)"))

Molar concentration or `mole` per `volume` or `molarity` (mol⋅m⁻³), unit conversion factor.

```Julia
julia> molarity(CGS,Metric) # cm³⋅m⁻³
$(molarity(CGS,Metric))

julia> molarity(English,SI2019) # ft³⋅m⁻³
$(molarity(English,SI2019))
```
"""
@pure molarity(U::UnitSystem,S::UnitSystem) = unit(mole(U,S)/volume(U,S))

"""
$(convertext(:molarvolume,"volume(U,S)/mole(U,S)"))

Occupied `volume` per `mole` or `molarvolume` (m³⋅mol⁻¹), unit conversion factor.

```Julia
julia> molarvolume(CGS,Metric) # m³⋅cm⁻³
$(molarvolume(CGS,Metric))

julia> molarvolume(English,SI2019) # m³⋅ft⁻³
$(molarvolume(English,SI2019))
```
"""
@pure molarvolume(U::UnitSystem,S::UnitSystem) = unit(volume(U,S)/mole(U,S))

"""
$(convertext(:molarentropy,"entropy(U,S)/mole(U,S)"))

Molar heat capacity or `entropy` per `mole` (J⋅K⁻¹⋅mol⁻¹), unit conversion factor.

```Julia
julia> molarentropy(CGS,Metric) # J⋅erg⁻¹
$(molarentropy(CGS,Metric))

julia> molarentropy(English,SI2019) # J⋅°R⋅slug-mol⋅ft⁻¹⋅lb⁻¹⋅K⁻¹⋅mol⁻¹
$(molarentropy(English,SI2019))
```
"""
@pure molarentropy(U::UnitSystem,S::UnitSystem) = unit(entropy(U,S)/mole(U,S))

"""
$(convertext(:molarenergy,"energy(U,S)/mole(U,S)"))

Gibbs free `energy` per `mole` or `molarenergy` (J⋅mol⁻¹), unit conversion factor.

```Julia
julia> molarenergy(CGS,Metric) # J⋅erg⁻¹
$(molarenergy(CGS,Metric))

julia> molarenergy(English,SI2019) # J⋅slug-mol⋅ft⁻¹⋅lb⁻¹⋅mol⁻¹
$(molarenergy(English,SI2019))
```
"""
@pure molarenergy(U::UnitSystem,S::UnitSystem) = unit(energy(U,S)/mole(U,S))

"""
$(convertext(:molarconductivity,"conductivity(U,S)*area(U,S)/mole(U,S)"))

Conductivity per `molarvolume` or `molarconductivity` (S⋅m²⋅mol⁻¹), unit conversion factor.

```Julia
julia> molarconductivity(EMU,Metric) # S⋅m²⋅abΩ⋅cm⁻²
$(molarconductivity(EMU,Metric))

julia> molarconductivity(ESU,Metric) # S⋅m²⋅statΩ⋅cm⁻²
$(molarconductivity(ESU,Metric))
```
"""
@pure molarconductivity(U::UnitSystem,S::UnitSystem) = unit(conductivity(U,S)*area(U,S)/mole(U,S))

"""
$(convertext(:molarsusceptibility,"specificsusceptibility(U,S)*molarmass(U,S)"))

Magnetic/electric molar mass `susceptibility` (m³⋅mol⁻¹), unit conversion factor.

```Julia
julia> molarsusceptibility(CGS,Metric) # m³⋅cm⁻³
$(molarsusceptibility(CGS,Metric))

julia> molarsusceptibility(Metric,SI2019) # m³⋅mol⋅mol⁻¹⋅cm⁻³
$(molarsusceptibility(Metric,SI2019))
```
"""
@pure molarsusceptibility(U::UnitSystem,S::UnitSystem) = unit(specificsusceptibility(U,S)*molarmass(U,S))

"""
$(convertext(:catalysis,"mole(U,S)/time(U,S)"))

Catalytic activity or `mole` per `time` or `catalysis` (kat, mol⋅s⁻¹), unit conversion factor.

```Julia
julia> catalysis(English,Metric) # kat⋅s⋅slug-mol⁻¹
$(catalysis(English,Metric))
```
"""
@pure catalysis(U::UnitSystem,S::UnitSystem) = unit(mole(U,S)/time(U,S))

"""
$(convertext(:specificity,"volume(U,S)/mole(U,S)/time(U,S)"))

Catalytic efficiency or `volume` per `mole` per `time` (m³⋅mol⁻¹⋅s⁻¹), unit conversion factor.

```Julia
julia> specificity(CGS,Metric) # m³⋅cm⁻³
$(specificity(CGS,Metric))

julia> specificity(English,Metric) # m³⋅slug-mol⋅mol⁻¹⋅ft⁻³
$(specificity(English,Metric))
```
"""
@pure specificity(U::UnitSystem,S::UnitSystem) = unit(volume(U,S)/mole(U,S)/time(U,S))

# photometrics

"""
$(convertext(:luminousflux,"luminousefficacy(U,S)*planck(U,S)/time(U,S)^2"))

Perceived power of light or `luminousflux` (lm, cd⋅sr), unit conversion factor.
"""
@pure luminousflux(U::UnitSystem,S::UnitSystem) = unit(frequency(U,S)^2*(luminousefficacy(S)*planckreduced(S))/(luminousefficacy(U)*planckreduced(U)))

"""
$(convertext(:luminance,"luminousflux(U,S)/area(U,S)"))

Luminous intensity per `area` or `luminance` (lx, lm⋅m⁻², cd⋅m⁻²⋅sr), unit conversion factor.

```Julia
julia> luminance(CGS,Metric) # lx⋅ph⁻¹
$(luminance(CGS,Metric))

julia> luminance(IAU,Metric) # lx⋅au²⋅lm⁻¹
$(luminance(IAU,Metric))

julia> luminance(English,Metric) # ft²⋅m⁻²
$(luminance(English,Metric))

julia> 1/10.76 # lx⋅fc⁻¹
$(1/10.76)
```
"""
@pure luminance(U::UnitSystem,S::UnitSystem) = unit(luminousflux(U,S)/area(U,S))

"""
$(convertext(:luminousenergy,"luminousflux(U,S)*time(U,S)"))

Perceived quantity of light or `luminousenergy` (lm⋅s, cd⋅s⋅sr), unit conversion factor.

```Julia
julia> luminousenergy(IAU,Metric) # s⋅day⁻¹
$(luminousenergy(IAU,Metric))
```
"""
@pure luminousenergy(U::UnitSystem,S::UnitSystem) = unit(frequency(U,S)*(luminousefficacy(S)*planckreduced(S))/(luminousefficacy(U)*planckreduced(U)))

"""
$(convertext(:luminousexposure,"luminance(U,S)*time(U,S)"))

Integrated `luminance` along `time` (lx⋅s, lm⋅s⋅m⁻², cd⋅s⋅m⁻²⋅sr), unit conversion factor.

```Julia
julia> luminousexposure(CGS,Metric) # lx⋅ph⁻¹
$(luminousexposure(CGS,Metric))

julia> luminousexposure(IAU,Metric) # s⋅au²⋅day⁻¹⋅m⁻²
$(luminousexposure(IAU,Metric))

julia> luminousexposure(English,Metric) # ft²⋅m⁻²
$(luminousexposure(English,Metric))
```
"""
@pure luminousexposure(U::UnitSystem,S::UnitSystem) = unit(luminance(U,S)*time(U,S))

@doc """
$(convertext(:luminousefficacy,"luminousefficacy(S)/luminousefficacy(U)"))

Ratio of `luminousflux` to `power` or `luminousefficacy` (lm⋅W⁻¹), unit conversion factor.

```Julia
julia> luminousefficacy(CGS,Metric) # erg⋅s⁻¹⋅W⁻¹
$(luminousefficacy(CGS,Metric))

julia> luminousefficacy(English,Metric) # ft⋅lb⋅s⁻¹⋅W⁻¹
$(luminousefficacy(English,Metric))
```
""" luminousefficacy(::UnitSystem,::UnitSystem)


text:
import .Zygote

struct ZygoteGradientLogDensity{L} <: ADGradientWrapper
    ℓ::L
end

"""
$(SIGNATURES)

Gradient using algorithmic/automatic differentiation via Zygote.

!!! NOTE
    Experimental, bug reports welcome.
"""
ADgradient(::Val{:Zygote}, ℓ) = ZygoteGradientLogDensity(ℓ)

Base.show(io::IO, ∇ℓ::ZygoteGradientLogDensity) = print(io, "Zygote AD wrapper for ", ∇ℓ.ℓ)

function logdensity_and_gradient(∇ℓ::ZygoteGradientLogDensity, x::AbstractVector)
    @unpack ℓ = ∇ℓ
    y, back = Zygote.forward(_logdensity_closure(ℓ), x)
    y, first(back(Zygote.sensitivity(y)))
end


text:
using MDPTools, NestedFixedPoint
#srand(122)

# 2 Model
bdisc = 0.95 # discount factor
truepar = [-1.9;1.;2.] # fixed costs, ..., entry costs
# State 1: Exogenous state
nX1 = 2
X1 = [0.;1.]
F1 = [sparse([1. 0.; 1. 0.]), sparse([0. 1.; 0. 1.])]

nX2 = 5
X2 = 1:nX2
M = [0.8 0.2 0.0 0.0 0.0;
	 0.2 0.6 0.2 0.0 0.0;
	 0.0 0.2 0.6 0.2 0.0;
	 0.0 0.0 0.2 0.6 0.2;
	 0.0 0.0 0.0 0.2 0.8;]

S = States(State(X1, F1),
           CommonState(X2, M))

Z = [zeros(nX2*2, 3), # don't enter
             [-ones(nX2) log(X2) -ones(nX2); # enter
			  -ones(nX2) log(X2)  zeros(nX2)]]

U = LinearUtility(Z, bdisc, truepar)

M, N, T = 2, 100, 1000
D = simulate(U, S, M, N, T)

n_transitions = zeros(nX2, nX2)
for n = 1:T*N-1
	if D.id[n] == D.id[n+1,1]
		n_transitions[D.xs[n,2], D.xs[n+1,2]] += 1
	end
end

estimated_F2 = n_transitions./sum(n_transitions,2)

estimated_S =  States(State(X1, F1), CommonState(X2, estimated_F2))

U.Θ*=0.0
am_nfxp = fit_nfxp(U, estimated_S, D)