text:
__precompile__()
module SpectralVis
    # Write your package code here.
    import Base.convert
    import Base: +, -, *, /, ∘

using Reexport
    @reexport using SSpline, Colors, Plots
    # environment types
    export  SPECENV, SpecEnvironment, set_specenv, set_extrap, set_limits, set_colormatch

    # spectral functions
    export  luminance_spec, reflectance_spec, transmittance_spec, cmfmat, normalize_spec, hardlimit_spec, reinterpolate, mbspec, mbcolor

    # cmf and cone fundamental functions
    export cmfunc, conefund, colorvisiondeficiency, vybrg

    # spectral types
    export  LSpec, RSpec, TSpec

    # multispectral types (CMFs and associated cone fundamentals)
    export  CIE31, LMS31, CIE31_J, CIE31_JV, CIE64, LMS64, CIE12_2, LMS06_2, CIE12_10, LMS06_10, VYBRG, CMatch, ConeFund, COpp, Spectrum, Multispectral

    # special functions
    export  normalize_spec, normalize_D_series, normalize_blackbody, adapt_spec, sconv, d_whitepoint, shift_spec, shift_cone, deactivate_cone, conefund, convmat

    # conversions
    export chromatic_adaptation, coneresponse, INGLING_TSOU, inglingtsou_opponent, OCS, mlms_ycc

    # spectral operators
    export ×, *, +

    # color matching function tables
    export  CMF1931, CMF1931_J, CMF1931_JV, CMF1964, CMF2012_2, CMF2012_10

    # cone fundamental tables
    export  LMS2006_2, LMS2006_10

    # spectra
    export  mars_sky_opportunity, mer_caltarget, pancam_L2, pancam_L3, pancam_Lpancam_L5, pancam_L6, pancam_L7, macbeth, munsell_specs, munsell_spec, munsell_keys,munsell_dict, wratten, wratten_filter, WrattenVis, WrattenIR, Wrattenfilters, mineral_specs, mineral, agfait_specs, agfait, vantablack

    # spectral generator tables
    export daylight_generator_table, b_c_series_table, f_series_table, gas_discharge_table, water_constants

    # spectral generators
    export  block_spec, led_spec, D_series_generator, D_series_illuminant, normalized_D_series, D_series_whitepoint, blackbody_illuminant, blackbody_whitepoint, D40_illuminant, D50_illuminant, D55_illuminant, D65_illuminant, D75_illuminant, D93_illuminant, illuminant_A, illuminant_B, illuminant_C, F_series_illuminant, gas_discharge_illuminant, sinusoidal_spd

    # matrix R, fundamental metamers, metameric blacks
    export  tristimulus, matrixA, matrixR, fundamental_metamer, metameric_black, cohen_pair

    # JND/color discrimination ellipse functions
    export ellipses, ellipse_parameters, ellipse, ellipse_a, ellipse_b, ellipse_θ, ellipse_center, ellipse_l0, ellipse_luminance, ellipse_g11, ellipse_g12, ellipse_g13, ellipse_g22, ellipse_g23, draw_ellipses

    # ellipse datasets
    export PGN_ellipses, WSS_ellipses, FMC_ellipses, WRB_ellipses, DLM_ellipses, B_ellipses, GF_ellipses, AR_ellipses, GW_ellipses

    # visualizations
    export  gamut_vis, gamutslice

    # doodles
    export  plottest, locustest, plotmunsell, spec_color, spec_idx, multiline, col, colv, sl, border, polar, getslice, sortslice

    include("cmf_tables.jl")
    include("spectral_types.jl")
    include("cmf_functions.jl")
    include("environment_types.jl")
    include("conversions.jl")
    include("spectral_functions.jl")
    include("special_functions.jl")
    include("spectral_generators.jl")
    include("spectral_operators.jl")
    include("matrix_R.jl")
    include("ellipse_functions.jl")

    #tables and other data
    include("spectral_generator_tables.jl")

    # Real world spectral data
    const datapath = "../data/"
    const specpath = datapath * "spectra/"
    const ellipsepath = datapath * "ellipses/"
    # MER Pancam data:
    # Mars Sky spectrum measured by Opportunity
    include(specpath * "spectrum_MarsSkyOpp.jl")
    # Spirit/Opportunity calibration target spectra
    include(specpath * "spectrum_MERCaltarget.jl")
    # Pancam color filter spectra
    include(specpath * "spectrum_Pancam.jl")
    include(specpath * "spectrum_MacBeth.jl")
    include(specpath * "spectrum_Wratten.jl")
    # Other spectra and reading functions
    include(specpath * "specread_munsell.jl")
    include(specpath * "specread_mineral.jl")
    include(specpath * "specread_agfait.jl")
    # MacAdam ellipses etc.
    include(ellipsepath * "ellipse_datasets.jl")
    # visualizations
    include("visualizations.jl")
    # doodles
    include("locustest.jl")
    include("lmsconversion.jl")
    include("col.jl")
end


text:
using Test
using Miletus

using Miletus.TermStructure
using Miletus.DayCounts

yc = ConstantYieldCurve(Actual365(), 0.01, :Continuous, :NoFrequency, Date("2016-11-01"))
core = CoreModel(40.00, yc, 0.0)
eucall = EuropeanCall(Date("2016-11-25"), SingleStock(), 41.00)
euput = EuropeanPut(Date("2016-11-25"), SingleStock(), 41.00)

σ = ivol(core, eucall, 1.5)
m = fit(GeomBMModel, core, eucall, 1.5)

@test m.volatility == σ
@test value(m, eucall) ≈ 1.5

coref = CoreForwardModel([Date("2016-11-25")], [42.0], yc)

σf = ivol(coref, eucall, 1.5)
mf = fit(GeomBMModel, coref, eucall, 1.5)

@test mf.volatility == σf
@test value(mf, eucall) ≈ 1.5

σ = ivol(core, euput, 1.5)
m = fit(GeomBMModel, core, euput, 1.5)

@test m.volatility == σ
@test value(m, euput) ≈ 1.5

let
    d1 = Dates.today()
    d2 = d1 + Day(365)
    c = EuropeanCall(d2, SingleStock(), 147.32)
    core = CoreModel(d1, 157.32, 0.01, 0.01)
    gbm = GeomBMModel(core, 0.1)
    v = ivol(core, c, 12, TianModel)
    gbm2 = GeomBMModel(core, v)
    @test isapprox(value(gbm2, c), 12, rtol = 1e-2)
    v = ivol(core, c, 12, CRRModel)
    gbm2 = GeomBMModel(core, v)
    @test isapprox(value(gbm2, c), 12, rtol = 1e-2)
    v = ivol(core, c, 12, JRModel)
    gbm2 = GeomBMModel(core, v)
    @test isapprox(value(gbm2, c), 12, rtol = 1e-2)
end

ext: jl


text:
module Parallel

using LightGraphs
using LightGraphs: sample, AbstractPathState, JohnsonState, BellmanFordState, FloydWarshallState
using Distributed: @distributed
using Base.Threads: @threads, nthreads, Atomic, atomic_add!, atomic_cas!
using SharedArrays: SharedMatrix, SharedVector, sdata
using ArnoldiMethod
using Random:shuffle
import SparseArrays: sparse
import Base: push!, popfirst!, isempty, getindex

include("shortestpaths/bellman-ford.jl")
include("shortestpaths/dijkstra.jl")
include("shortestpaths/floyd-warshall.jl")
include("shortestpaths/johnson.jl")
include("centrality/betweenness.jl")
include("centrality/closeness.jl")
include("centrality/pagerank.jl")
include("centrality/radiality.jl")
include("centrality/stress.jl")
include("distance.jl")
include("traversals/bfs.jl")
include("traversals/gdistances.jl")
include("traversals/greedy_color.jl")
include("utils.jl")
include("dominatingset/minimal_dom_set.jl")
include("independentset/maximal_ind_set.jl")
include("vertexcover/random_vertex_cover.jl")

end


text:
"""
costfunctions.jl
"""

import Base: copy, +

"""
Abstract type that represents a scalar-valued function that accepts a state and control
at a single knot point.
"""
abstract type CostFunction end

"""
Cost function of the form
```math
\\frac{1}{2} x^T Q x + \\frac{1}{2} u^T R u + u^T H x + q^T x + r^T u + c
```
"""
struct QuadraticCost{TQ,TR,TH,Tq,Tr,T} <: CostFunction
    Q::TQ
    R::TR
    H::TH
    q::Tq
    r::Tr
    c::T
    n::Int
    m::Int
    terminal::Bool
    Q_diag::Bool
    R_diag::Bool
    use_R::Bool
    use_H::Bool
    use_r::Bool
    const_grad::Bool
    const_hess::Bool
end

# constructors
function QuadraticCost(Q::TQ, R::TR, H::TH, q::Tq, r::Tr, c::T; use_R=true,
                       use_H=true, use_r=true, terminal=false, checks=true) where {
                           TQ,TR,TH,Tq,Tr,T}
    n = size(Q, 1)
    m = size(R, 1)
    if checks
        @assert length(q) == n
        if use_r
            @assert length(r) == m
        end
        if use_H
            @assert size(H) == (m, n)
        end
        if !ispossemidef(Array(Q))
            @warn "Q is not positive semidefinite:\n$(Q)"
        end
        if use_R && !isposdef(Array(R)) && !terminal
            @warn "R is not positive definite:\n$(R)"
        end
    end
    Q_diag = (TQ <: Diagonal)
    R_diag = (TR <: Diagonal)
    const_grad = false
    const_hess = true
    return QuadraticCost{TQ,TR,TH,Tq,Tr,T}(Q, R, H, q, r, c, n, m, terminal,
                                           Q_diag, R_diag, use_R, use_H, use_r, const_grad,
                                           const_hess)
end

# methods
Base.copy(cost::QuadraticCost{TQ,TR,TH,Tq,Tr,T}) where {TQ,TR,TH,Tq,Tr,T} = (
    QuadraticCost{TQ,TR,TH,Tq,Tr,T}(
        copy(cost.Q), copy(cost.R), copy(cost.H), copy(cost.q), copy(cost.r), cost.c,
        cost.terminal, cost.use_H, cost.Q_diag, cost.R_diag
))

@inline state_dim(cost::QuadraticCost) = cost.n
@inline control_dim(cost::QuadraticCost) = cost.m

# evaluation methods
function cost(cost_::QuadraticCost, X::AbstractVector, U::AbstractVector, k::Int; terminal=false)
    cost__ = 0.5 * X[k]' * cost_.Q * X[k] + dot(cost_.q, X[k]) + cost_.c
    if cost_.use_R
        cost__ += 0.5 * U[k]' * cost_.R * U[k]
    end
    if cost_.use_H
        cost__ += U[k]' * cost_.H * X[k]
    end
    if cost_.use_r
        cost__ += dot(cost_.r, U[k])
    end
    return cost__
end

function cost_derivatives!(E::QuadraticCost, cost::QuadraticCost, X::AbstractVector,
                           U::AbstractVector, k::Int; terminal=false)
    # E.Q
    E.Q .= cost.Q
    # E.q
    E.q .= cost.q
    mul!(E.q, cost.Q, X[k], 1., 1.)
    if cost.use_H
        mul!(E.q, Transpose(cost.H), U[k], 1., 1.)
    end
    # E.R
    if cost.use_R
        E.R .= cost.R

    else
        E.R .= 0
    end
    # E.r
    if cost.use_R
        mul!(E.r, cost.R, U[k])
    else
        E.r .= 0
    end
    if cost.use_r
        E.r .+= cost.r
    end
    if cost.use_H
        mul!(E.r, cost.H, X[k], 1., 1.)
    end
    # E.H
    if cost.use_H
        E.H .= cost.H
    else
        E.H .= 0
    end
    return nothing
end

# LQRCost
function LQRCost(Q::AbstractMatrix, xf::AbstractVector, R::AbstractMatrix, M, V;
                 use_R=true, terminal=false)
    n = size(Q, 1)
    m = size(R, 1)
    H = M(zeros(m, n))
    q = -Q * xf
    r = V(zeros(m))
    c = 0.5 * xf' * Q * xf
    return QuadraticCost(Q, R, H, q, r, c; use_R=use_R, use_H=false, use_r=false, terminal=terminal)
end

# """
# An abstract type that represents any [`CostFunction`](@ref) of the form
# ```math
# \\frac{1}{2} x^T Q x + \\frac{1}{2} u^T R u + u^T H x + q^T x + r^T u + c
# ```

# These types all support the following methods
# * [`is_diag`](@ref)
# * [`is_blockdiag`](@ref)
# * [`invert!`](@ref)

# As well as standard addition.
# """
# abstract type QuadraticCostFunction{n,m,T} <: CostFunction end

# """
#     is_diag(::QuadraticCostFunction)

# Determines if the hessian of a quadratic cost function is strictly diagonal.
# """
# is_diag(cost::QuadraticCostFunction) = is_blockdiag(cost) && cost.Q isa Diagonal && cost.R isa Diagonal

# """
#     is_diag(::QuadraticCostFunction)

# Determines if the hessian of a quadratic cost function is block diagonal (i.e. ``\\norm(H) = 0``).
# """
# is_blockdiag(::QuadraticCostFunction) = false
# state_dim(::QuadraticCostFunction{n}) where n = n
# control_dim(::QuadraticCostFunction{<:Any,m}) where m = m

# function (::Type{QC})(Q::AbstractArray, R::AbstractArray;
#         H::AbstractArray=SizedMatrix{size(R,1),size(Q,1)}(zeros(eltype(Q), size(R,1), size(Q,1))),
#         q::AbstractVector=(@SVector zeros(eltype(Q), size(Q,1))),
#         r::AbstractVector=(@SVector zeros(eltype(R), size(R,1))),
#         c::Real=zero(eltype(Q)), kwargs...) where QC <: QuadraticCostFunction
#     QC(Q, R, H, q, r, c; kwargs...)
# end

# """
#     stage_cost(costfun::CostFunction, x, u)
#     stage_cost(costfun::CostFunction, x)

# Calculate the scalar cost using `costfun` given state `x` and control `u`. If only the
# state is provided, it is assumed it is a terminal cost.
# """
# function stage_cost(cost::QuadraticCostFunction, x::AbstractVector, u::AbstractVector)
#     J = 0.5*u'cost.R*u + dot(cost.r,u) + stage_cost(cost, x)
#     if !is_blockdiag(cost)
#         J += u'cost.H*x
#     end
#     return J
# end

# function stage_cost(cost::QuadraticCostFunction, x::AbstractVector{T}) where T
#     0.5*x'cost.Q*x + dot(cost.q,x) + cost.c
# end

# """
#     gradient!(E::QuadraticCostFunction, costfun::CostFunction, x, u)
#     gradient!(E::QuadraticCostFunction, costfun::CostFunction, x)

# Evaluate the gradient of the cost function `costfun` at state `x` and control `u`, storing
#     the result in `E.q` and `E.r`. Return a `true` if the gradient is constant, and `false`
#     otherwise.
# """
# function gradient!(E::QuadraticCostFunction, cost::QuadraticCostFunction, x)
#     E.q .= cost.Q*x .+ cost.q
#     return false
# end

# function gradient!(E::QuadraticCostFunction, cost::QuadraticCostFunction, x, u)
#     gradient!(E, cost, x)
#     E.r .= cost.R*u .+ cost.r
#     if !is_blockdiag(cost)
#         E.q .+= cost.H'u
#         E.r .+= cost.H*x
#     end
#     return false
# end

# """
#     hessian!(E::QuadraticCostFunction, costfun::CostFunction, x, u)
#     hessian!(E::QuadraticCostFunction, costfun::CostFunction, x)

# Evaluate the hessian of the cost function `costfun` at state `x` and control `u`, storing
#     the result in `E.Q`, `E.R`, and `E.H`. Return a `true` if the hessian is constant, and `false`
#     otherwise.
# """
# function hessian!(E::QuadraticCostFunction, cost::QuadraticCostFunction, x)
#     if is_diag(cost)
#         for i = 1:length(x); E.Q[i,i] = cost.Q[i,i] end
#     else
#         E.Q .= cost.Q
#     end
#     return true
# end

# function hessian!(E::QuadraticCostFunction, cost::QuadraticCostFunction, x, u)
#     hessian!(E, cost, x)
#     if is_diag(cost)
#         for i = 1:length(u); E.R[i,i] = cost.R[i,i]; end
#     else
#         E.R .= cost.R
#     end
#     if !is_blockdiag(cost)
#         E.H .= cost.H
#     end
#     return true
# end

# function Base.copy(c::QC) where QC<:QuadraticCostFunction
#     QC(copy(c.Q), copy(c.R), H=copy(c.H), q=copy(c.q), r=copy(c.r), c=c.c,
#         terminal=c.terminal, checks=false)
# end

# # Other methods
# """
#     set_LQR_goal!(cost::QuadraticCostFunction, xf)
#     set_LQR_goal!(cost::QuadraticCostFunction, xf, uf)

# Change the reference state and control for an LQR tracking cost.
# Only changes `q` and `r`, and not the constant term `c`. If `uf` is
# not passed in, it isn't changed.
# """
# function set_LQR_goal!(cost::QuadraticCostFunction, xf)
#     cost.q .= -cost.Q * xf
#     return nothing
# end

# function set_LQR_goal!(cost::QuadraticCostFunction, xf, uf)
#     set_LQR_goal!(cost, xf)
#     cost.r .= -cost.R * uf
#     return nothing
# end

# function +(c1::QuadraticCostFunction, c2::QuadraticCostFunction)
#     @assert state_dim(c1) == state_dim(c2)
#     @assert control_dim(c1) == control_dim(c2)
#     n,m = state_dim(c1), control_dim(c1)
#     H1 = c1 isa DiagonalCost ? zeros(m,n) : c1.H
#     H2 = c2 isa DiagonalCost ? zeros(m,n) : c2.H
#     QC = promote_type(typeof(c1), typeof(c2))
#     QC(c1.Q + c2.Q, c1.R + c2.R, H1 + H2,
#        c1.q + c2.q, c1.r + c2.r, c1.c + c2.c,
#        checks=false, terminal=c1.terminal && c2.terminal)
# end

# """
#     invert!(Ginv, cost::QuadraticCostFunction)

# Invert the hessian of the cost function, storing the result in `Ginv`. Performs the inversion
#     efficiently, depending on the structure of the Hessian (diagonal or block diagonal).
# """
# function invert!(Ginv, cost::QuadraticCostFunction{n,m}) where {n,m}
#     ix = 1:n
#     iu = n .+ (1:m)
#     if is_diag(cost)
#         for i = 1:n; Ginv[i,i] = inv(cost.Q[i,i]); end
#         if !cost.terminal
#             for i = 1:m; Ginv[i+n,i+n] = inv(cost.R[i,i]); end
#         end
#     elseif is_blockdiag(cost)
#         Ginv[ix,ix] .= inv(SizedMatrix{n,n}(cost.Q))
#         if !cost.terminal
#             Ginv[iu,iu] .= inv(SizedMatrix{m,m}(cost.R))
#         end
#     else
#         G1 = [cost.Q cost.H']
#         G2 = [cost.H cost.R ]
#         G = [G1; G2]
#         Ginv .= inv(G)
#     end
# end

# #######################################################
# #              COST FUNCTION INTERFACE                #
# #######################################################

# """
#     DiagonalCost{n,m,T}

# Cost function of the form
# ```math
# \\frac{1}{2} x^T Q x + \\frac{1}{2} u^T R u + q^T x + r^T u + c
# ```
# where ``Q`` and ``R`` are positive semi-definite and positive definite diagonal matrices,
# respectively, and ``x`` is `n`-dimensional and ``u`` is `m`-dimensional.

# # Constructors
#     DiagonalCost(Qd, Rd, q, r, c; kwargs...)
#     DiagonalCost(Q, R, q, r, c; kwargs...)
#     DiagonalCost(Qd, Rd; [q, r, c, kwargs...])
#     DiagonalCost(Q, R; [q, r, c, kwargs...])

# where `Qd` and `Rd` are the diagonal vectors, and `Q` and `R` are matrices.

# Any optional or omitted values will be set to zero(s). The keyword arguments are
# * `terminal` - A `Bool` specifying if the cost function is terminal cost or not.
# * `checks` - A `Bool` specifying if `Q` and `R` will be checked for the required definiteness.
# """
# struct DiagonalCost{n,m,T} <: QuadraticCostFunction{n,m,T}
#     Q::Diagonal{T,SVector{n,T}}
#     R::Diagonal{T,SVector{m,T}}
#     q::MVector{n,T}
#     r::MVector{m,T}
#     c::T
#     terminal::Bool
#     function DiagonalCost(Qd::StaticVector{n}, Rd::StaticVector{m},
#                           q::StaticVector{n},  r::StaticVector{m},
#                           c::Real; terminal::Bool=false, checks::Bool=true) where {n,m}
#         T = promote_type(typeof(c), eltype(Qd), eltype(Rd), eltype(q), eltype(r))
#         if checks
#             if any(x->x<0, Qd)
#                 @warn "Q needs to be positive semi-definite."
#             elseif any(x->x<=0, Rd) && !terminal
#                 @warn "R needs to be positive definite."
#             end
#         end
#         new{n,m,T}(Diagonal(SVector(Qd)), Diagonal(SVector(Rd)), q, r, T(c), terminal)
#     end
# end

# Base.copy(c::DiagonalCost) =
#     DiagonalCost(c.Q.diag, c.R.diag, copy(c.q), copy(c.r), c.c,
#         checks=false, terminal=c.terminal)

# # strip away the H
# @inline function (::Type{<:DiagonalCost})(Q::AbstractArray, R::AbstractArray, H::AbstractArray,
#         q::AbstractVector, r::AbstractVector, c; kwargs...)
#     DiagonalCost(Q, R, q, r, c; kwargs...)
# end

# # convert to form for inner constructor
# function DiagonalCost(Q::AbstractArray, R::AbstractArray,
#         q::AbstractVector, r::AbstractVector, c; kwargs...)
#     n,m = length(q), length(r)
#     Qd = SVector{n}(diag(Q))
#     Rd = SVector{m}(diag(R))
#     DiagonalCost(Qd, Rd, convert(MVector{n},q), convert(MVector{m},r), c; kwargs...)
# end

# # Pass in vectors for Q and R
# function DiagonalCost(Q::AbstractVector, R::AbstractVector,
#         q::AbstractVector, r::AbstractVector, c::Real; kwargs...)
#     DiagonalCost(Diagonal(Q), Diagonal(R), q, r, c; kwargs...)
# end

# function DiagonalCost(Q::AbstractVector, R::AbstractVector; kwargs...)
#     DiagonalCost(Diagonal(Q), Diagonal(R); kwargs...)
# end

# is_blockdiag(::DiagonalCost) = true
# is_diag(::DiagonalCost) = true

# function LinearAlgebra.inv(cost::DiagonalCost)
#     return DiagonalCost(inv(cost.Q), inv(cost.R), cost.q, cost.r, cost.c, terminal=cost.terminal)
# end

# function Base.:\(cost::DiagonalCost, z::AbstractKnotPoint)
#     x = cost.Q\state(z)
#     u = cost.R\control(z)
#     return StaticKnotPoint([x;u], z._x, z._u, z.dt, z.t)
# end

# function change_dimension(cost::DiagonalCost, n::Int, m::Int, ix, iu)
#     Qd = zeros(n)
#     Rd = zeros(m)
#     q = zeros(n)
#     r = zeros(m)
#     Qd[ix] = diag(cost.Q)
#     Rd[iu] = diag(cost.R)
#     q[ix] = cost.q
#     r[iu] = cost.r
#     DiagonalCost(Qd, Rd, q, r, cost.c, terminal=cost.terminal, checks=false)
# end

# """
#     QuadraticCost{n,m,T,TQ,TR}

# Cost function of the form
# ```math
# \\frac{1}{2} x^T Q x + \\frac{1}{2} u^T R u + u^T H x + q^T x + r^T u + c
# ```
# where ``R`` must be positive definite, ``Q`` and ``Q_f`` must be positive semidefinite.

# The type parameters `TQ` and `TR` specify the type of ``Q`` and ``R``.

# # Constructor
#     QuadraticCost(Q, R, H, q, r, c; kwargs...)
#     QuadraticCost(Q, R; H, q, r, c, kwargs...)

# Any optional or omitted values will be set to zero(s). The keyword arguments are
# * `terminal` - A `Bool` specifying if the cost function is terminal cost or not.
# * `checks` - A `Bool` specifying if `Q` and `R` will be checked for the required definiteness.
# """
# mutable struct QuadraticCost{n,m,T,TQ,TR} <: QuadraticCostFunction{n,m,T}
#     Q::TQ                     # Quadratic stage cost for states (n,n)
#     R::TR                     # Quadratic stage cost for controls (m,m)
#     H::Matrix{T}   # Quadratic Cross-coupling for state and controls (m,n)
#     q::Vector{T}           # Linear term on states (n,)
#     r::Vector{T}           # Linear term on controls (m,)
#     c::T                      # constant term
#     terminal::Bool
#     zeroH::Bool
#     function (::Type{QC})(Q::TQ, R::TR, H::TH,
#             q::Tq, r::Tr, c::Real; checks=true, terminal=false) where {TQ,TR,TH,Tq,Tr,QC<:QuadraticCost}
#         @assert size(Q,1) == length(q)
#         @assert size(R,1) == length(r)
#         @assert size(H) == (length(r), length(q))
#         if checks
#             if !isposdef(Array(R)) && !terminal
#                 @warn "R is not positive definite"
#             end
#             if !ispossemidef(Array(Q))
#                 @warn "Q is not positive semidefinite"
#             end
#         end
#         zeroH = norm(H,Inf) ≈ 0
#         m,n = size(H)
#         T = promote_type(eltype(Q), eltype(R), eltype(H), eltype(q), eltype(r), typeof(c))
#         new{n,m,T,TQ,TR}(Q,R,H,q,r,c,terminal,zeroH)
#     end
#     function QuadraticCost{n,m,T,TQ,TR}(qcost::QuadraticCost) where {n,m,T,TQ,TR}
#         new{n,m,T,TQ,TR}(qcost.Q, qcost.R, qcost.H, qcost.q, qcost.r, qcost.c, qcost.terminal, qcost.zeroH)
#     end
# end

# state_dim(cost::QuadraticCost) = length(cost.q)
# control_dim(cost::QuadraticCost) = length(cost.r)
# is_blockdiag(cost::QuadraticCost) = cost.zeroH

# function QuadraticCost{T}(n::Int,m::Int; terminal=false) where T
#     Q = Matrix(one(T)*I,n,n)
#     R = Matrix(one(T)*I,m,m)
#     H = zeros(T,m,n)
#     q = zeros(T,n)
#     r = zeros(T,m)
#     c = zero(T)
#     QuadraticCost(Q,R,H,q,r,c, checks=false, terminal=terminal)
# end

# QuadraticCost(cost::QuadraticCost) = cost

# function LinearAlgebra.inv(cost::QuadraticCost)
#     if norm(cost.H,Inf) ≈ 0
#         QuadraticCost(inv(cost.Q), inv(cost.R), cost.H, cost.q, cost.r, cost.c,
#             checks=false, terminal=cost.terminal)
#     else
#         m,n = size(cost.H)
#         H1 = [cost.Q cost.H']
#         H2 = [cost.H cost.R ]
#         H = [H1; H2]
#         Hinv = inv(H)
#         ix = 1:n
#         iu = n .+ (1:m)
#         Q = SizedMatrix{n,n}(Hinv[ix, ix])
#         R = SizedMatrix{m,m}(Hinv[iu, iu])
#         H = SizedMatrix{m,n}(Hinv[iu, ix])
#         QuadraticCost(Q,R,H, cost.q, cost.r, cost.c,
#             checks=false, terminal=cost.terminal)
#     end
# end

# @inline function (::Type{<:QuadraticCost})(dcost::DiagonalCost)
#     QuadraticCost(dcost,Q, dcost.R, q=dcost.q, r=dcost.r, c=dcost.c, terminal=dcost.terminal)
# end

# function Base.convert(::Type{<:QuadraticCost}, dcost::DiagonalCost)
#     QuadraticCost(dcost.Q, dcost.R, q=dcost.q, r=dcost.r, c=dcost.c, terminal=dcost.terminal)
# end

# Base.promote_rule(::Type{<:QuadraticCostFunction}, ::Type{<:QuadraticCostFunction}) = QuadraticCost

# function Base.promote_rule(::Type{<:QuadraticCost{n,m,T1,Q1,R1}},
#                            ::Type{<:QuadraticCost{n,m,T2,Q2,R2}}) where {n,m,T1,T2,Q1,Q2,R1,R2}
#     T = promote_type(T1,T2)
#     function diag_type(T1,T2,n)
#         elT = promote_type(eltype(T1), eltype(T2))
#         if T1 == T2
#             return T1
#         elseif (T1 <: Diagonal) && (T2 <: Diagonal)
#             return Diagonal{elT,MVector{n,elT}}
#         else
#             return SizedMatrix{n,n,elT,2}
#         end
#     end
#     Q = diag_type(Q1,Q2,n)
#     R = diag_type(R1,R2,m)
#     QuadraticCost{n,m,T, Q, R}
# end

# @inline Base.convert(::Type{QC}, cost::QuadraticCost) where QC <: QuadraticCost = QC(cost)

# """
#     LQRCost(Q, R, xf, [uf; kwargs...])

# Convenience constructor for a `QuadraticCostFunction` of the form:
# ```math
# \\frac{1}{2} (x-x_f)^T Q (x-xf) + \\frac{1}{2} (u-u_f)^T R (u-u_f)
# ```

# If ``Q`` and ``R`` are diagonal, the output will be a `DiagonalCost`, otherwise it will
# be a `QuadraticCost`.
# """
# function LQRCost(Q::AbstractArray, R::AbstractArray,
#         xf::AbstractVector, uf=(@SVector zeros(size(R,1))); kwargs...)
#     H = @SMatrix zeros(size(R,1),size(Q,1))
#     q = -Q*xf
#     r = -R*uf
#     c = 0.5*xf'*Q*xf + 0.5*uf'R*uf
#     return DiagonalCost(Q, R, H, q, r, c, kwargs...)
# end

# function LQRCost(Q::Diagonal{<:Any,<:SVector{n}},R::Diagonal{<:Any,<:SVector{m}},
#         xf::AbstractVector, uf=(@SVector zeros(m)); kwargs...) where {n,m}
#     q = -Q*xf
#     r = -R*uf
#     c = 0.5*xf'*Q*xf + 0.5*uf'R*uf
#     return DiagonalCost(Q, R, q, r, c; kwargs...)
# end

# ############################################################################################
# #                        QUADRATIC QUATERNION COST FUNCTION
# ############################################################################################

# struct QuadraticQuatCost{T,N,M,N4} <: CostFunction
#     Q::Diagonal{T,SVector{N,T}}
#     R::Diagonal{T,SVector{M,T}}
#     q::SVector{N,T}
#     r::SVector{M,T}
#     c::T
#     w::T
#     q_ref::SVector{4,T}
#     q_ind::SVector{4,Int}
#     Iq::SMatrix{N,4,T,N4}
#     function QuadraticQuatCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}},
#             q::SVector{N,T}, r::SVector{M,T}, c::T, w::T,
#             q_ref::SVector{4,T}, q_ind::SVector{4,Int}) where {T,N,M}
#         Iq = @MMatrix zeros(N,4)
#         for i = 1:4
#             Iq[q_ind[i],i] = 1
#         end
#         Iq = SMatrix{N,4}(Iq)
#         return new{T,N,M,N*4}(Q, R, q, r, c, w, q_ref, q_ind, Iq)
#     end
# end

# state_dim(::QuadraticQuatCost{T,N,M}) where {T,N,M} = N
# control_dim(::QuadraticQuatCost{T,N,M}) where {T,N,M} = M

# function QuadraticQuatCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}};
#         q=(@SVector zeros(N)), r=(@SVector zeros(M)), c=zero(T), w=one(T),
#         q_ref=(@SVector [1.0,0,0,0]), q_ind=(@SVector [4,5,6,7])) where {T,N,M}
#     QuadraticQuatCost(Q, R, q, r, c, q_ref, q_ind)
# end

# function stage_cost(cost::QuadraticQuatCost, x::SVector, u::SVector)
#     stage_cost(cost, x) + 0.5*u'cost.R*u + cost.r'u
# end

# function stage_cost(cost::QuadraticQuatCost, x::SVector)
#     J = 0.5*x'cost.Q*x + cost.q'x + cost.c
#     q = x[cost.q_ind]
#     dq = cost.q_ref'q
#     J += cost.w*min(1+dq, 1-dq)
# end

# function gradient(cost::QuadraticQuatCost{T,N,M}, x::SVector, u::SVector) where {T,N,M}
#     Qx = cost.Q*x + cost.q
#     q = x[cost.q_ind]
#     dq = cost.q_ref'q
#     if dq < 0
#         Qx += cost.w*cost.Iq*cost.q_ref
#     else
#         Qx -= cost.w*cost.Iq*cost.q_ref
#     end
#     Qu = cost.R*u + cost.r
#     return Qx, Qu
# end

# function hessian(cost::QuadraticQuatCost, x::SVector{N}, u::SVector{M}) where {N,M}
#     Qxx = cost.Q
#     Quu = cost.R
#     Qux = @SMatrix zeros(M,N)
#     return Qxx, Quu, Qux
# end

# function QuatLQRCost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}}, xf,
#         uf=(@SVector zeros(M)); w=one(T), quat_ind=(@SVector [4,5,6,7])) where {T,N,M}
#     r = -R*uf
#     q = -Q*xf
#     c = 0.5*xf'Q*xf + 0.5*uf'R*uf
#     q_ref = xf[quat_ind]
#     return QuadraticQuatCost(Q, R, q, r, c, w, q_ref, quat_ind)
# end

# function change_dimension(cost::QuadraticQuatCost, n, m)
#     n0,m0 = state_dim(cost), control_dim(cost)
#     Q_diag = diag(cost.Q)
#     R_diag = diag(cost.R)
#     q = cost.q
#     r = cost.r
#     if n0 != n
#         dn = n - n0  # assumes n > n0
#         pad = @SVector zeros(dn)
#         Q_diag = [Q_diag; pad]
#         q = [q; pad]
#     end
#     if m0 != m
#         dm = m - m0  # assumes m > m0
#         pad = @SVector zeros(dm)
#         R_diag = [R_diag; pad]
#         r = [r; pad]
#     end
#     QuadraticQuatCost(Diagonal(Q_diag), Diagonal(R_diag), q, r, cost.c, cost.w,
#         cost.q_ref, cost.q_ind)
# end

# function (+)(cost1::QuadraticQuatCost, cost2::QuadraticCost)
#     @assert state_dim(cost1) == state_dim(cost2)
#     @assert control_dim(cost1) == control_dim(cost2)
#     @assert norm(cost2.H) ≈ 0
#     QuadraticQuatCost(cost1.Q + cost2.Q, cost1.R + cost2.R,
#         cost1.q + cost2.q, cost1.r + cost2.r, cost1.c + cost2.c,
#         cost1.w, cost1.q_ref, cost1.q_ind)
# end

# (+)(cost1::QuadraticCost, cost2::QuadraticQuatCost) = cost2 + cost1

#
#
# struct ErrorQuadratic{Rot,N,M} <: CostFunction
#     model::RigidBody{Rot}
#     Q::Diagonal{Float64,SVector{12,Float64}}
#     R::Diagonal{Float64,SVector{M,Float64}}
#     r::SVector{M,Float64}
#     c::Float64
#     x_ref::SVector{N,Float64}
#     q_ind::SVector{4,Int}
# end
#
#
# state_dim(::ErrorQuadratic{Rot,N,M}) where {Rot,N,M} = N
# control_dim(::ErrorQuadratic{Rot,N,M}) where {Rot,N,M} = M
#
# function ErrorQuadratic(model::RigidBody{Rot}, Q::Diagonal{T,<:SVector{12}},
#         R::Diagonal{T,<:SVector{M}},
#         x_ref::SVector{N}, u_ref=(@SVector zeros(T,M)); r=(@SVector zeros(T,M)), c=zero(T),
#         q_ind=(@SVector [4,5,6,7])) where {T,N,M,Rot}
#     r += -R*u_ref
#     c += 0.5*u_ref'R*u_ref
#     return ErrorQuadratic{Rot,N,M}(model, Q, R, r, c, x_ref, q_ind)
# end
#
# function stage_cost(cost::ErrorQuadratic, x::SVector)
#     dx = state_diff(cost.model, x, cost.x_ref)
#     return 0.5*dx'cost.Q*dx + cost.c
# end
#
# function stage_cost(cost::ErrorQuadratic, x::SVector, u::SVector)
#     stage_cost(cost, x) + 0.5*u'cost.R*u + cost.r'u
# end
#
# function cost_expansion(cost::ErrorQuadratic{Rot}, model::AbstractModel,
#         z::KnotPoint{T,N,M}, G) where {T,N,M,Rot<:UnitQuaternion}
#     x,u = state(z), control(z)
#     model = cost.model
#     Q = cost.Q
#     q = orientation(model, x)
#     q_ref = orientation(model, cost.x_ref)
#     dq = SVector(q_ref\q)
#     err = state_diff(model, x, cost.x_ref)
#     dx = @SVector [err[1],  err[2],  err[3],
#                     dq[1],   dq[2],   dq[3],   dq[4],
#                    err[7],  err[8],  err[9],
#                    err[10], err[11], err[12]]
#     G = state_diff_jacobian(model, dx) # n × dn
#
#     # Gradient
#     dmap = inverse_map_jacobian(model, dx) # dn × n
#     Qx = G'dmap'Q*err
#     Qu = cost.R*u
#
#     # Hessian
#     ∇jac = inverse_map_∇jacobian(model, dx, Q*err)
#     Qxx = G'dmap'Q*dmap*G + G'∇jac*G + ∇²differential(model, x, dmap'Q*err)
#     Quu = cost.R
#     Qux = @SMatrix zeros(M,N-1)
#     return Qxx, Quu, Qux, Qx, Qu
# end
#
# function cost_expansion(cost::ErrorQuadratic, model::AbstractModel,
#         z::KnotPoint{T,N,M}, G) where {T,N,M}
#     x,u = state(z), control(z)
#     model = cost.model
#     q = orientation(model, x)
#     q_ref = orientation(model, cost.x_ref)
#     err = state_diff(model, x, cost.x_ref)
#     dx = err
#     G = state_diff_jacobian(model, dx) # n × n
#
#     # Gradient
#     dmap = inverse_map_jacobian(model, dx) # n × n
#     Qx = G'dmap'cost.Q*err
#     Qu = cost.R*u + cost.r
#
#     # Hessian
#     Qxx = G'dmap'cost.Q*dmap*G
#     Quu = cost.R
#     Qux = @SMatrix zeros(M,N)
#     return Qxx, Quu, Qux, Qx, Qu
# end
#
# function change_dimension(cost::ErrorQuadratic, n, m)
#     n0,m0 = state_dim(cost), control_dim(cost)
#     Q_diag = diag(cost.Q)
#     R_diag = diag(cost.R)
#     r = cost.r
#     if n0 != n
#         dn = n - n0  # assumes n > n0
#         pad = @SVector zeros(dn) # assume the new states don't have quaternions
#         Q_diag = [Q_diag; pad]
#     end
#     if m0 != m
#         dm = m - m0  # assumes m > m0
#         pad = @SVector zeros(dm)
#         R_diag = [R_diag; pad]
#         r = [r; pad]
#     end
#     ErrorQuadratic(cost.model, Diagonal(Q_diag), Diagonal(R_diag), r, cost.c,
#         cost.x_ref, cost.q_ind)
# end
#
# function (+)(cost1::ErrorQuadratic, cost2::QuadraticCost)
#     @assert control_dim(cost1) == control_dim(cost2)
#     @assert norm(cost2.H) ≈ 0
#     @assert norm(cost2.q) ≈ 0
#     if state_dim(cost2) == 13
#         rm_quat = @SVector [1,2,3,4,5,6,8,9,10,11,12,13]
#         Q2 = Diagonal(diag(cost2.Q)[rm_quat])
#     else
#         Q2 = cost2.Q
#     end
#     ErrorQuadratic(cost1.model, cost1.Q + Q2, cost1.R + cost2.R,
#         cost1.r + cost2.r, cost1.c + cost2.c,
#         cost1.x_ref, cost1.q_ind)
# end
#
# (+)(cost1::QuadraticCost, cost2::ErrorQuadratic) = cost2 + cost1
#
#
#
#


text:
# Obtained from https://en.wikipedia.org/wiki/Telegrapher%27s_equations#Lossy_transmission_line
"
    telegrapher(L,C,R=0,G=0,n=100)

returns an array of differential equations that can be
used to solve for the travelling voltages and currents as a function of time.
All quantities are expected to be per-length quantities (in S.I. units).

The problem is solved by discretizing the transmission line into n sections. If
the solution is too coarse then try changing n.
"
function telegrapher(L, C, R=0, G=0 ;n=100, vs=t->1, is=t->0)
    if L == 0 || C == 0
        error("Both the inductance and capacitance need to be non-zero.")
    end
    return function(t,u,du)
        B2(t,view(u,n+1:2n),view(du,1:n))
        scale!(view(du,1:n), -1/L)

        B1(t,view(u,1:n),view(du,n+1:2n))
        scale!(view(du,n+1:2n), -1/C)
    end
end
"
    tline_solve(f,L,C,R=0,G=0,t=(0,1),n=100)

Solves for the time-domain response of a transmission line to an incident
voltage waveform f(t).

### Arguments
u0: initial voltage applied to the line\n
t: solution interval\n
n: number of tline discretizations\n
"
function tline_solve(u0,vs,is,L,C,R=0,G=0;t=(0.,1.),n=100,dt=1e-3)
    prob = DifferentialEquations.ODEProblem(telegrapher(L,C,R,G,n=n,vs=vs),u0,t)
    #solver = DifferentialEquations.Euler()
    solver = DifferentialEquations.CVODE_BDF()
    sol = DifferentialEquations.solve(
                                      prob,
                                      solver,
                                      reltol=1e-8,
                                      abstol=1e-8,
                                      dt=dt
                                     )
    return sol
end


text:
"""
    boys_quadgk(x, n::Int; rtol=1e-12)

Calculates Boys function with Gaussian Quadrature using QuadGK.jl.
Accuracy can be tuned with `rtol` argument.
This is accurate, but slow.
"""
function boys_quadgk(x, n::Int; rtol=1e-12)
    tmp = quadgk( 0, 1; rtol=rtol) do t
        t^(2n) * exp( -x*t^2 )
    end
    return tmp[1]
end

"""
    boys_gamma(x, n::Int)

Evaluate Boys function with gamma function.
This is accurate.
"""
function boys_gamma(x, n::Int)
    upper = gamma( n + 1//2 ) - gamma( n + 1//2, x )
    lower = 2x^( n + 1//2 )
    return upper / lower
end

"""
    boys_rec_largex(x, n::Int)

For large `x` > 30. Calculate using recursion.
Uses approximations.
This is fast, but loses accuracy on high `n` values.

Based on:
A. K. H. Weiss, C. Ochsenfeld J. Comput.
Chem. 2015, 36, 1390–1398. DOI: 10.1002/jcc.23935
"""
function boys_rec_largex(x, n::Int)
    if n == 0
        return √π / (2√x)
    else
        return ( n - 1//2 ) * boys_rec_largex(x, n-1) / x
    end
end


text:
### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ 006a6698-f9dd-11ea-3f64-152c31c749ef
using Pkg, DrWatson

# ╔═╡ 0090af24-f9dd-11ea-0eb4-eb1758ba5442
begin
	using Distributions
	using StatsPlots, Plots
	using Turing
	using Logging
	using LaTeXStrings
nd

# ╔═╡ 9453eb8e-f9dc-11ea-1865-f168816cdb46
md"## Clip-02-06-07.jl"

# ╔═╡ 009dc1c8-f9dd-11ea-193f-3f1baf487786
md"## snippet 2.6"

# ╔═╡ 009f1ad2-f9dd-11ea-2abc-0579fc5057d2
@model globethrowing(w, l) = begin
    p ~ Uniform(0, 1)
    w ~ Binomial(w + l, p)
end

# ╔═╡ 00a7e61c-f9dd-11ea-028c-ddfdbdd620cb
m = globethrowing(6, 3);

# ╔═╡ 00a87104-f9dd-11ea-2ef2-bf5a7c31c607
r = quap(m)

# ╔═╡ ea69e428-f9e0-11ea-0c97-e3b70306a220
begin
	p_grid = range(0, 1, length = 20)
	prior = ones(20)
	likelihood = pdf.(Binomial.(9, p_grid), 6)
	posterior = likelihood .* prior
	posterior ./= sum(posterior)

	x = 0:0.01:1
	w = 6
	l = 3

	f = plot( x, pdf.(Normal(mean(r.coef.p), √r.vcov[1]), x ), lab="quap")
	plot!( x, pdf.(Beta( w+1 , l+1 ) , x ), lab="exact", leg=:topleft, title="n = $(w+l)")
end

# ╔═╡ 00dd7c78-f9dd-11ea-29ee-b9fb0b25bd56
md"## End clip-02-06-07t.jl"

# ╔═╡ Cell order:
# ╟─9453eb8e-f9dc-11ea-1865-f168816cdb46
# ╠═006a6698-f9dd-11ea-3f64-152c31c749ef
# ╠═0090af24-f9dd-11ea-0eb4-eb1758ba5442
# ╟─009dc1c8-f9dd-11ea-193f-3f1baf487786
# ╠═009f1ad2-f9dd-11ea-2abc-0579fc5057d2
# ╠═00a7e61c-f9dd-11ea-028c-ddfdbdd620cb
# ╠═00a87104-f9dd-11ea-2ef2-bf5a7c31c607
# ╠═ea69e428-f9e0-11ea-0c97-e3b70306a220
# ╟─00dd7c78-f9dd-11ea-29ee-b9fb0b25bd56


text:
function J2Gravity(Body::T, PosRsph::Vector{Float64}) where {T <: abstractCelestialBody}
  #Convert δ (latitude) into co-latitude α
  α = pi/2 - PosRsph[3];
  Re_r = mean_radius(Body)/PosRsph[1]
  J2planet = j2(Body)
  muplanet = mu(Body)

  #Gravity Field Terms
  g = zeros(3)
  g[3] = -muplanet/PosRsph[1]^2*(1- 1.5*J2planet*Re_r^2*(3*cos(α)^2 - 1))
  g[2] = 0.0
  g[1] = -(3*muplanet/PosRsph[1]^2*Re_r^2*sin(α)*cos(α)*(g[3] + J2planet))
  return g
end


text:
"""
	H_matrix(U, input_state, partition_occupancy_vector)
	H_matrix(interf::Interferometer, i::Input, o::FockDetection)
	H_matrix(interf::Interferometer, i::Input, subset_modes::ModeOccupation)

H matrix for a partition defined by `partition_occupancy_vector`, see definition
in the article below.

!!! note
	Conventions follow the author's [Boson bunching is not
	maximized by indistinguishable particles](https://arxiv.org/abs/2203.01306)
	which are the ones compatible with Tichy's conventions (Shshnovitch has a
	different one for the evolution of the creation operators).
"""
function H_matrix(U, input_state::Vector, partition_occupancy_vector::Vector)

	part = occupancy_vector_to_partition(partition_occupancy_vector)
	input_modes = occupancy_vector_to_mode_occupancy(input_state)

	number_photons = sum(input_state)
	if number_photons == 0
		error("no input photons")
	end

	if size(U,1) < number_photons
		@warn "more photons than modes (trivially gives the identity as photons are conserved)"
	end

	H = Matrix{ComplexF64}(undef,number_photons,number_photons)

	for i in 1: number_photons
		for j in 1:number_photons
			H[i,j] = sum([conj(U[l, input_modes[i]]) * U[l,input_modes[j]] for l in part])
		end
	end

	H

end

H_matrix(interf::Interferometer, i::Input, o::FockDetection) = H_matrix(interf.U, i.r.state, o.s)

H_matrix(interf::Interferometer, i::Input, subset_modes::ModeOccupation) = isa_subset(subset_modes) ? H_matrix(interf.U, i.r.state, subset_modes.state) : error("invalid subset")

"""

	full_bunching_probability(interf::Interferometer, i::Input, subset_modes::Subset)
	full_bunching_probability(interf::Interferometer, i::Input, mo::ModeOccupation)

Computes the probability that all n photons end up in the subset of chosen
output modes following.

!!! note "Reference"
	[Universality of Generalized Bunching and Efficient Assessment of Boson Sampling](https://arxiv.org/abs/1509.01561)
"""
function full_bunching_probability(interf::Interferometer, i::Input, subset_modes::Subset)

	return clean_proba(permanent(H_matrix(interf,i,subset_modes) .* transpose(i.G.S)))

end

function full_bunching_probability(interf::Interferometer, i::Input, mo::ModeOccupation)

	return clean_proba(permanent(H_matrix(interf,i,mo) .* transpose(i.G.S)))

end

#
# """
#
# 	bunching_events(input_state::Input, sub::Subset)
#
# generates the output configurations corresponding to a full
# bunching in the subset_modes
# """
# function bunching_events(input_state::Input, sub::Subset)
#
# 	#photon_distribution_in_subset_modes =
# 	all_mode_configurations(input_state, sub, only_photon_number_conserving = false)
#
# 	######### convert to output ModeOccupations
#
# end

"""
	bunching_probability_brute_force_bosonic(U, input_state, output_state; print_output = false)
	bunching_probability_brute_force_bosonic(interf::Interferometer, i::Input, subset_modes::ModeOccupation)

Bosonic bunching probability by direct summation of all possible cases

`bunching_event_proba` gives the probability to get the event of ``[1^n 0^(m-n)]``.
"""
function bunching_probability_brute_force_bosonic(U, input_state, output_state; print_output = false)

    n = sum(input_state)
    m = size(U,1)

    bunching_proba = 0
    bunching_proba_array = []
    bunching_event_proba = nothing

    print_output ? println("bunching probabilities : ") : nothing

    for t in reverse.(Iterators.product(fill(0:n,n)...))[:]
        if sum(t) == n # cases where t is physical
            output_state = zeros(Int,m)
            output_state[1:n] .= t[1:n]
            this_proba = process_probability(U, input_state, output_state)
            bunching_proba += this_proba
            push!(bunching_proba_array,[t, this_proba])
            print_output ? println("output = ", t, " p = ", this_proba) : nothing

            if output_state[1:n] == ones(Int, n)
                bunching_event_proba = this_proba
            end
        end
    end

    H = H_matrix(U,input_state,partition)
    @test bunching_proba ≈ real(permanent(H))

    return bunching_proba, bunching_proba_array, bunching_event_proba

end

bunching_probability_brute_force_bosonic(interf::Interferometer, i::Input, subset_modes::ModeOccupation) = bunching_probability_brute_force_bosonic(interf.U, i.r.state, subset_modes.r.state; print_output = false)


text:
Pkg.add("Images")
Pkg.add("Colors")
Pkg.add("ImageMagick")

if VERSION > v"0.4.5"
 Pkg.add("Combinatorics")
end


text:
struct SequenceGraphLink
    source::NodeID
    destination::NodeID
    dist::Int64
end

source(l::SequenceGraphLink) = l.source
destination(l::SequenceGraphLink) = l.destination
distance(l::SequenceGraphLink) = l.dist

is_forward_link(l::SequenceGraphLink, n::NodeID) = source(l) == -n

@inline function Base.isequal(a::SequenceGraphLink, b::SequenceGraphLink)
    return source(a) == source(b) && destination(a) == destination(b)
end


text:
# A stochastic extension to the SIR and SEIR model

"""
We consider the 2 dimensional system of SDEs

dX_t = F(X_t) dt + G(X_t) dW_t

where X_t = (S_t; I_t) and F:R^2 -> R^2 and G(⋅) ∈ R^4.

The component R for the standard SIR model need not be considered.
"""

using DifferentialEquations
###########################################
#First some testing
f(du,u,p,t) = du .= 1.01u
function g(du,u,p,t)
  du[1,1] = 0.3u[1]
  du[1,2] = 0.6u[1]
  du[2,1] = 1.2u[2]
  du[2,2] = 0.2u[2]
end
prob = SDEProblem(f,g,ones(2),(0.0,1.0),noise_rate_prototype=zeros(2,2))
sol = solve(prob,LambaEM())
plot(sol)

############################################################
beta = 0.25; gamma = 0.07; N = 7000;
function f(du,u,p,t)
    du[1] = -beta * u[2]*u[1]/N;
    du[2] = beta*u[1]*u[2]/N - gamma*u[2]
end

function g(du,u,p,t)
    du[1,1] = -sqrt(beta*u[1]*u[2]/N)
    du[1,2] = 0
    du[2,1] = sqrt(beta*u[1]*u[2]/N)
    du[2,2] = -sqrt(gamma*u[2])
  end

prob = SDEProblem(f,g,[N-15;15],(0.0,200.0),noise_rate_prototype=zeros(2,2))
sol = solve(prob,LambaEM())

plot(sol)


text:
using AWS
using AWS: AWSCredentials
using AWS: AWSServices
using AWS.AWSExceptions: AWSException, InvalidFileName, NoCredentials, ProtocolNotDefined
using AWS.AWSMetadata:
    ServiceFile,
    _clean_documentation,
    _filter_latest_service_version,
    _generate_low_level_definition,
    _generate_high_level_definition,
    _generate_high_level_definitions,
    _get_service_files,
    _get_service_and_version,
    _get_function_parameters,
    _clean_uri,
    _format_name,
    _splitline,
    _wraplines,
    _validindex
using Base64
using Compat: mergewith
using Dates
using GitHub
using HTTP
using IniFile: Inifile
using JSON
using OrderedCollections: LittleDict, OrderedDict
using MbedTLS: digest, MD_SHA256, MD_MD5
using Mocking
using Pkg
using Random
using Retry
using Suppressor
using Test
using UUIDs
using XMLDict

Mocking.activate()

include("patch.jl")

const TEST_MINIO = begin
    all(k -> haskey(ENV, k), ("MINIO_ACCESS_KEY", "MINIO_SECRET_KEY", "MINIO_REGION_NAME"))
end

aws = AWSConfig()

function _now_formatted()
    return lowercase(Dates.format(now(Dates.UTC), dateformat"yyyymmdd\THHMMSSsss\Z"))
end

@testset "AWS.jl" begin
    include("AWSExceptions.jl")
    include("AWSMetadataUtilities.jl")
    include("test_pkg.jl")
    include("utilities.jl")
    include("AWSConfig.jl")

    backends = [AWS.HTTPBackend, AWS.DownloadsBackend]
    @testset "Backend: $(nameof(backend))" for backend in backends
        AWS.DEFAULT_BACKEND[] = backend()
        include("AWS.jl")
        include("AWSCredentials.jl")
        include("issues.jl")

        if TEST_MINIO
            include("minio.jl")
        end
    end
end


text:
# See https://github.com/TuringLang/Turing.jl/issues/1199
ChainRulesCore.@non_differentiable push!!(
    vi::VarInfo, vn::VarName, r, dist::Distribution, gidset::Set{Selector}
)

ChainRulesCore.@non_differentiable updategid!(
    vi::AbstractVarInfo, vn::VarName, spl::Sampler
)

# https://github.com/TuringLang/Turing.jl/issues/1595
ZygoteRules.@adjoint function dot_observe(
    spl::Union{SampleFromPrior,SampleFromUniform},
    dists::AbstractArray{<:Distribution},
    value::AbstractArray,
    vi,
)
    function dot_observe_fallback(spl, dists, value, vi)
        increment_num_produce!(vi)
        return sum(map(Distributions.loglikelihood, dists, value)), vi
    end
    return ZygoteRules.pullback(__context__, dot_observe_fallback, spl, dists, value, vi)
end

max_stars_repo_name: TuringLang/MicroPPL

max_issues_repo_name: TuringLang/MicroPPL

max_forks_repo_name: TuringLang/MicroPPL


text:
import Distances.evaluate

"The top level distance type"
abstract type AbstractDistance <: Distances.Metric end
"All subtypes of this type operates on rational transfer functions"
abstract type AbstractRationalDistance <: AbstractDistance end
"All subtypes of this type operates on signals"
abstract type AbstractSignalDistance <: AbstractDistance end
"All subtypes of this type operates on the roots of rational transfer functions"
abstract type AbstractRootDistance <: AbstractRationalDistance end
"All subtypes of this type operates on the coefficients of rational transfer functions"
abstract type AbstractCoefficientDistance <: AbstractRationalDistance end
abstract type AbstractWelchDistance <: AbstractSignalDistance end

Base.Broadcast.broadcastable(p::AbstractDistance) = Ref(p)

"A Union that represents a collection of distances"
const DistanceCollection = Union{Tuple, Vector{<:AbstractDistance}}

struct Identity end
struct Log end
Base.Broadcast.broadcastable(p::Identity) = Ref(p)
Base.Broadcast.broadcastable(p::Log) = Ref(p)
magnitude(d) = Identity()

evaluate(d::DistanceCollection,x,y;kwargs...) = sum(evaluate(d,x,y;kwargs...) for d in d)
Base.:(+)(d::AbstractDistance...) = d

(d::AbstractDistance)(x,y,args...;kwargs...) = evaluate(d, x, y,args...;kwargs...)
(d::DistanceCollection)(x,y,args...;kwargs...) = evaluate(d, x, y,args...;kwargs...)

"""
    SymmetricDistance{D} <: AbstractDistance

Evaluates `d(x,y) - 0.5(d(x,x) + d(y,y))`
"""
struct SymmetricDistance{D} <: AbstractDistance
    d::D
end

function evaluate(d::SymmetricDistance, x, y, args...; kwargs...)
    evaluate(d.d, x, y, args...; kwargs...) -
    0.5 * (
        evaluate(d.d, x, x, args...; kwargs...) +
        evaluate(d.d, y, y, args...; kwargs...)
    )
end

"""
    CoefficientDistance{D, ID} <: AbstractCoefficientDistance

Distance metric based on model coefficients

# Arguments:
- `domain::D`: [`Discrete`](@ref) or [`Continuous`](@ref)
- `distance::ID = SqEuclidean()`: Inner distance between coeffs
"""
CoefficientDistance
@kwdef struct CoefficientDistance{D,ID} <: AbstractCoefficientDistance
    domain::D
    distance::ID = SqEuclidean()
end

"""
    ModelDistance{D <: AbstractDistance} <: AbstractSignalDistance

A model distance operates on signals and works by fitting an LTI model to the signals before calculating the distance. The distance between the LTI models is defined by the field `distance`. This is essentially a wrapper around the inner distance that handles the fitting of a model to the signals. How the model is fit is determined by `fitmethod`.

# Arguments:
- `fitmethod::`[`FitMethod`](@ref): [`LS`](@ref), [`TLS`](@ref) or [`PLR`](@ref)
- `distance::D`: The inner distance between the models

# Example:
```julia
using SpectralDistances
innerdistance = OptimalTransportRootDistance(domain=Continuous(), β=0.005, p=2)
dist = ModelDistance(TLS(na=30), innerdistance)
```
"""
struct ModelDistance{D <: AbstractDistance, FM <: FitMethod} <: AbstractSignalDistance
    fitmethod::FM
    distance::D
end

"""
    EuclideanRootDistance{D, A, F1, F2} <: AbstractRootDistance

Simple euclidean distance between roots of transfer functions

# Arguments:
- `domain::D`: [`Discrete`](@ref) or [`Continuous`](@ref)
- `assignment::A =` [`SortAssignement`](@ref)`(imag)`: Determines how roots are assigned. An alternative is `HungarianAssignement`
- `transform::F1 = identity`: DESCRIPTION
- `weight` : A function used to calculate weights for the induvidual root distances. A good option is [`residueweight`](@ref)
- `p::Int = 2` : Order of the distance
"""
EuclideanRootDistance
@kwdef struct EuclideanRootDistance{D,A,F1,F2} <: AbstractRootDistance
    domain::D = Continuous()
    assignment::A = SortAssignement(imag)
    transform::F1 = identity
    weight::F2 = unitweight
    p::Int = 2
end

"""
    OptimalTransportRootDistance{D, F1, F2, S} <: AbstractRootDistance

The Sinkhorn distance between roots. The weights are provided by `weight`, which defaults to [`residueweight`](@ref).

# Arguments:
- `domain::D = Continuous()`: [`Discrete`](@ref) or [`Continuous`](@ref)
- `transform::F1 = identity`: Probably not needed.
- `weight::F2 = `[`simplex_residueweight`](@ref): A function used to calculate weights for the induvidual root distances.
- `β::Float64 = 0.01`: Amount of entropy regularization
- `p::Int = 2` : Order of the distance
- `divergence::S = nothing`: A divergence that penalizes creation and destruction of mass.
"""
OptimalTransportRootDistance
@kwdef struct OptimalTransportRootDistance{
    D,
    F1,
    F2,
    S<:Union{UnbalancedOptimalTransport.AbstractDivergence,Nothing},
} <: AbstractRootDistance
    domain::D = Continuous()
    transform::F1 = identity
    weight::F2 = simplex_residueweight
    β::Float64 = 0.01
    p::Int = 2
    divergence::S = nothing
end

@deprecate SinkhornRootDistance(args...;kwargs...) OptimalTransportRootDistance(args...;kwargs...)

"""
    HungarianRootDistance{D, ID <: Distances.PreMetric, F} <: AbstractRootDistance

Similar to [`EuclideanRootDistance`](@ref) but does the pole assignment using the Hungarian method.

# Arguments:
- `domain::D = Continuous()`: [`Discrete`](@ref) or [`Continuous`](@ref)
- `distance::ID = SqEuclidean()`: Inner distance
- `transform::F = identity`: If provided, this Function transforms all roots before the distance is calculated
"""
HungarianRootDistance
@kwdef struct HungarianRootDistance{D,ID <: Distances.PreMetric,F} <: AbstractRootDistance
    domain::D = Continuous()
    distance::ID = SqEuclidean()
    transform::F = identity
end

"""
    KernelWassersteinRootDistance{D, F, DI} <: AbstractRootDistance

A kernel version of the root distance

# Arguments:
- `domain::D = Continuous()`: [`Discrete`](@ref) or [`Continuous`](@ref)
- `λ::Float64 = 1.0`: Kernel precision, lower value means wider kernel.
- `transform::F = identity`: If provided, this Function transforms all roots before the distance is calculated
"""
KernelWassersteinRootDistance
@kwdef struct KernelWassersteinRootDistance{D,F,DI} <: AbstractRootDistance
    domain::D = Continuous()
    λ::Float64   = 1.
    transform::F = identity
    distance::DI = SqEuclidean()
end

"""
    DiscretizedRationalDistance{WT, DT} <: AbstractRationalDistance

This distance discretizes the spectrum before performing the calculations.

# Arguments:
- `w::WT = LinRange(0.01, 0.5, 300)`: Frequency set
- `distmat::DT = distmat_euclidean(w, w)`: DESCRIPTION
"""
DiscretizedRationalDistance
@kwdef struct DiscretizedRationalDistance{WT,DT} <: AbstractRationalDistance
    w::WT = LinRange(0.01, 0.5, 300)
    distmat::DT = distmat_euclidean(w,w)
end

"""
    WelchOptimalTransportDistance{DT, AT <: Tuple, KWT <: NamedTuple} <: AbstractWelchDistance

Calculates the Wasserstein distance between two signals by estimating a Welch periodogram of each.

# Arguments:
- `distmat::DT`: you may provide a matrix array for this
- `args::AT = ()`: Options to the Welch function
- `kwargs::KWT = NamedTuple()`: Options to the Welch function
- `p::Int = 2` : Order of the distance
"""
WelchOptimalTransportDistance
@kwdef struct WelchOptimalTransportDistance{DT,AT <: Tuple, KWT <: NamedTuple} <: AbstractWelchDistance
    distmat::DT = nothing
    args::AT = ()
    kwargs::KWT = NamedTuple()
    p::Int = 2
end

"""
    WelchLPDistance{AT <: Tuple, KWT <: NamedTuple, F} <: AbstractWelchDistance

Lᵖ distance between welch spectra, `mean(abs(x1-x2)^p)`.

#Arguments:
- `args::AT = ()`: These are sent to `welch_pgram`
- `kwargs::KWT = NamedTuple()`: These are sent to `welch_pgram`
- `p::Int = 2`: Order of the distance
- `normalized::Bool = true`: Normlize spectrum to sum to 1 (recommended)
- `transform::F = identity`: Optional function to apply to the spectrum, example `log1p` or `sqrt`. Must not produce negative values, so `log` is not a good idea. The function is applied like this: `transform.(x1)`.
"""
WelchLPDistance
@kwdef struct WelchLPDistance{AT <: Tuple, KWT <: NamedTuple, F} <: AbstractWelchDistance
    args::AT = ()
    kwargs::KWT = NamedTuple()
    p::Int = 2
    normalized::Bool = true
    transform::F = identity
end

"""
    OptimalTransportHistogramDistance{DT} <: AbstractDistance

Optimal transport between two histograms. If you pass two vectors to this distance two histograms will be computed automatically (`fit(Histogram,x)`). Pass two histograms to get better manual control.

# Arguments:
- `p::Int = 1`: order
"""
OptimalTransportHistogramDistance
@kwdef struct OptimalTransportHistogramDistance <: AbstractDistance
    p::Int = 1
end

"""
    DiscreteGridTransportDistance{DT} <: AbstractDistance

Optimal transport between two measures on a common discrete grid.
"""
DiscreteGridTransportDistance
struct DiscreteGridTransportDistance{D,TD,T,V} <: AbstractDistance
    inner::D
    distmat::Matrix{TD}
    C::Matrix{T}
    c1::V
    c2::V
end

"""
    DiscreteGridTransportDistance(Cityblock(), Float32, n, n)

The constructor takes an inner distance, the numeric type and the length of the two measures.
"""
function DiscreteGridTransportDistance(inner::Distances.PreMetric, T, n, m)
    distmat = inner.(one(T):n, (one(T):m)')
    C = zeros(T, n,m)
    c1 = zeros(T, n)
    c2 = zeros(T, m)
    DiscreteGridTransportDistance(inner,distmat,C,c1,c2)
end

"""
    CationalOptimalTransportDistance{DT, MT} <: AbstractRationalDistanCe

calculates the Wasserstein distance using the closed-form sold*c for (d,c) in zip(D,C)s and inverse cumulative functions.

C ArgumentC:
- `domain::DT = Continuous()`: [`Discrete`](@ref) or [`Continuous`](@red*c for (d,c) in zip(D,C)r
- `magnitude::MT = Identity()`:
- `interval = (-(float(π)), float(π))`: Integration interval
"""
RationalOptimalTransportDistance
@kwdef struct RationalOptimalTransportDistance{DT,MT} <: AbstractRationalDistance
    domain::DT = Continuous()
    p::Int = 1
    magnitude::MT = Identity()
    interval = (-float(π,),float(π))
end
magnitude(d::RationalOptimalTransportDistance) = d.magnitude

"""
    RationalCramerDistance{DT} <: AbstractRationalDistance

Similar to `RationalOptimalTransportDistance` but does not use inverse functions.

# Arguments:
- `domain::DT = Continuous()`: [`Discrete`](@ref) or [`Continuous`](@ref)
- `p::Int = 2`: order
- `interval = (-(float(π)), float(π))`: Integration interval
"""
RationalCramerDistance
@kwdef struct RationalCramerDistance{DT} <: AbstractRationalDistance
    domain::DT = Continuous()
    p::Int = 2
    interval = (-float(π,),float(π))
end

"""
    BuresDistance <: AbstractDistance

Distance between pos.def. matrices
"""
struct BuresDistance <: AbstractDistance
end

"""
    EnergyDistance <: AbstractSignalDistance

`std(x1) - std(x2)`
This distance can be added to a loss function to ensure that the energy in the two signals is the same. Some of the optimal transport-based distances are invariant to the energy in the signal, requiring this extra cost if that invariance is not desired. Combining distances is done by putting two or more in a tuple.
    Usage: `combined_loss = (primary_distance, EnergyDistance())`
"""
struct EnergyDistance <: AbstractSignalDistance
end

"""
    domain(d::AbstractDistance)
    $(TYPEDSIGNATURES)

Return the domain of the distance
"""
domain(d) = d.domain
domain(d::ModelDistance) = domain(d.distance)

"""
    domain_transform(d::AbstractDistance, e)

Change domain of roots
"""
domain_transform(d::AbstractDistance, e) = domain_transform(domain(d), e)

transform(d::AbstractRootDistance) = d.transform
transform(d::AbstractRootDistance, x) = d.transform(x)

"""
    distmat_euclidean(e1::AbstractVector, e2::AbstractVector, p = 2)

The euclidean distance matrix between two vectors of complex numbers
"""
distmat_euclidean(e1::AbstractVector,e2::AbstractVector,p=2) = @fastmath abs.(e1 .- transpose(e2)).^p # Do not replace with call to mutating version, Zygote limitation

"""
    distmat_euclidean!(D, e1::AbstractVector, e2::AbstractVector, p = 2) = begin

In-place version
"""
distmat_euclidean!(D, e1::AbstractVector,e2::AbstractVector,p=2) = @fastmath D .= abs.(e1 .- transpose(e2)).^p

# distmat(dist,e1,e2) = (n = length(e1[1]); [dist(e1[i],e2[j]) for i = 1:n, j=1:n])

"""
    $(TYPEDSIGNATURES)

Compute the symmetric, pairwise distance matrix using the specified distance.

- `normalize`: set to true to normalize distances such that the diagonal is zero. This is useful for distances that are not true distances due to `d(x,y) ≠ 0` such as the [`OptimalTransportRootDistance`](@ref)

This function uses multiple threads if available and copies the distance object to each thread in case it stores an internal cache.
"""
function distmat(dist,e::AbstractVector; normalize=false, kwargs...)
    n = length(e)
    T = typeof(evaluate(dist,e[1],e[1];kwargs...))
    D = zeros(T,n,n)
    dists = [deepcopy(dist) for _ in 1:nthreads()]
    @showprogress 1 "distmat" for i = 1:n
        D[i,i] = evaluate(dist,e[i],e[i];kwargs...) # Note we do calc this distance since it's nonzero for regularized distances.
        @threads for j=i+1:n
            D[i,j] = evaluate(dists[threadid()], e[i], e[j]; kwargs...)
            (D[j,i] = D[i,j])
        end
    end
    normalize && symmetrize!(D)
    Symmetric(D)
end

function symmetrize!(D)
    d = diag(D)
    for i = 1:size(D,1)
        x = 0.5d[i]
        D[:,i] .-= x
        D[i,:] .-= x
    end
    D
end

function symmetrize!(D::SparseMatrixCSC)
    d = deepcopy(diag(D))
    N = length(d)
    for k = 1:N
        x = 0.5d[k]
        for i = 1:N
            D[i,k] == 0 && continue
            D[i,k] -= x
            D[k,i] -= x
        end
    end
    D
end

# distmat_logmag(e1::AbstractArray,e2::AbstractArray) = distmat_euclidean(logmag.(e1), logmag.(e2))

function preprocess_roots(d::AbstractRootDistance, e::AbstractRoots)
    e2 = domain_transform(d,e)
    T = similar_type(e2)
    T(d.transform.(e2))
end

function preprocess_roots(d::AbstractRootDistance, m::AbstractModel)
    e = roots(domain(d), m)
    preprocess_roots(d, e)
end

function evaluate(d::AbstractRootDistance,w1::AbstractModel,w2::AbstractModel; kwargs...)
    # if isderiving()
    #     # w1,w2 = change_precision(Float64, w1), change_precision(Float64, w2)
    #     return evaluate(d, preprocess_roots(d,w1), preprocess_roots(d,w2); solver=sinkhorn_log, tol=0.01) # workaround for https://github.com/FluxML/Zygote.jl/issues/584, when that has been closed, this branch can be removed
    # end
    r1 = preprocess_roots(d,w1)
    r2 = preprocess_roots(d,w2)
    evaluate(d, r1, r2; kwargs...)
end
function evaluate(d::AbstractRootDistance,w1::ARMA,w2::ARMA; kwargs...)
    d1 = evaluate(d, preprocess_roots(d,pole(domain(d),w1)), preprocess_roots(d,pole(domain(d),w2)))
    d2 = evaluate(d, preprocess_roots(d,tzero(domain(d),w1)), preprocess_roots(d,tzero(domain(d),w2)))
    d1 + d2
end

function evaluate(d::HungarianRootDistance, e1::AbstractRoots{CT}, e2::AbstractRoots{CT}; kwargs...)::real(CT) where CT
    e1,e2 = toreim(e1), toreim(e2)
    n     = length(e1[1])
    dist  = d.distance
    dm    = [dist(e1[1][i],e2[1][j]) + dist(e1[2][i],e2[2][j]) for i = 1:n, j=1:n]
    _,c   = hungarian(dm)
    c
end

# function kernelsum(e1,e2,λ)
#     s = 0.
#     λ = -λ
#     @inbounds for i in eachindex(e1)
#         s += exp(λ*abs(e1[i]-e2[i])^2)
#         for j in i+1:length(e2)
#             s += 2exp(λ*abs(e1[i]-e2[j])^2)
#         end
#     end
#     s / length(e1)^2
# end
#
# function logkernelsum(e1,e2,λ)
#     s = 0.
#     λ = -λ
#     le2 = logreim.(e2)
#     @inbounds for i in eachindex(e1)
#         le1 = logreim(e1[i])
#         s += exp(λ*abs2(le1-le2[i]))
#         for j in i+1:length(e2)
#             s += 2exp(λ*abs2(le1-le2[j]))
#         end
#     end
#     s / length(e1)^2
# end

function evaluate(d::EuclideanRootDistance, e1::AbstractRoots,e2::AbstractRoots; kwargs...)
    length(e1) == 0 && return zero(real(eltype(e1)))
    w1,w2 = d.weight(e1), d.weight(e2)
    # sum(eachindex(e1)) do i
    #     i1 = I1[i]
    #     i2 = I2[i]
    #     (w1[i1]*w2[i2])^β*abs(w1[i1]*e1[i1]-w2[i2]*e2[i2])^p
    # end # below is workaround for zygote #314
    I1,I2 = d.assignment(e1, e2)
    evaluate(d,e1[I1],e2[I2],w1[I1],w2[I2])
end

function evaluate(d::EuclideanRootDistance, e1,e2, w1, w2; kwargs...)
    p = d.p
    β = (1-p)/2
    l = zero(real(eltype(e1)))
    @inbounds @fastmath for i in 1:length(e1)
        l += (w1[i]*w2[i])^β*abs(w1[i]*e1[i]-w2[i]*e2[i])^p
    end
    real(l) # Workaround for expanding to complex for Zygote support
end

function evaluate(
    d::OptimalTransportRootDistance,
    e1::AbstractRoots,
    e2::AbstractRoots;
    solver = sinkhorn_log!,
    kwargs...,
)
    D = distmat_euclidean(e1, e2, d.p)
    w1 = d.weight(e1)
    w2 = d.weight(e2)
    # C     = solver(D,SVector{length(w1)}(w1),SVector{length(w2)}(w2); β=d.β, kwargs...)[1]

    if d.divergence !== nothing
        return OT!(
            d.divergence,
            D,
            DiscreteMeasure(w1),
            DiscreteMeasure(w2),
            d.β;
            max_iters = get(kwargs, :iters, 100_000),
            tol = get(kwargs, :tol, 1e-8),
            warn = true
        )
    end

    C = solver(D, w1, w2; β = d.β, kwargs...)[1]
    if any(isnan, C) && !isderiving()
        println("Nan in OptimalTransportRootDistance, increasing precision")
        C = solver(big.(D), big.(w1), big.(w2); β = d.β, kwargs...)[1]
        any(isnan, C) &&
        error("No solution found by solver $(solver), check your input and consider increasing β ($(d.β)).")
        C = eltype(D).(C)
    end
    dot(C,D)
end

function evaluate(d::AbstractRootDistance, a1::AbstractVector{<: Real},a2::AbstractVector{<: Real}; kwargs...)
    evaluate(d, AutoRoots(domain(d), hproots(rev(a1))), AutoRoots(domain(d), hproots(rev(a2))); kwargs...)
end

# function eigval_dist_wass_logmag_defective(d::KernelWassersteinRootDistance, e1::AbstractRoots,e2::AbstractRoots)
#     λ     = d.λ
#     error("this does not yield symmetric results")
#     # e1 = [complex((logmag(magangle(e1)))...) for e1 in e1]
#     # e2 = [complex((logmag(magangle(e2)))...) for e2 in e2]
#     # e1    = logreim.(e1)
#     # e2    = logreim.(e2)
#     # e1    = sqrtreim.(e1)
#     # e2    = sqrtreim.(e2)
#     dm1   = logkernelsum(e1,e1,λ)
#     dm2   = logkernelsum(e2,e2,λ)
#     dm12  = logkernelsum(e1,e2,λ)
#     dm1 - 2dm12 + dm2
# end
function evaluate(d::KernelWassersteinRootDistance, e1::AbstractRoots,e2::AbstractRoots; kwargs...)
    λ     = d.λ
    # dm1   = exp.(.- λ .* distmat(d.distance, e1,e1))
    # dm2   = exp.(.- λ .* distmat(d.distance, e2,e2))
    # dm12  = exp.(.- λ .* distmat(d.distance, e1,e2))

    D = distmat_euclidean(e1,e1)
    @avx D .= exp.(.- λ .* D)
    c = mean(D)
    distmat_euclidean!(D,e2,e2)
    @avx D .= exp.(.- λ .* D)
    c += mean(D)
    distmat_euclidean!(D,e1,e2)
    @avx D .= exp.(.- λ .* D)
    c -= 2mean(D)
    c
end

evaluate(d::AbstractCoefficientDistance,w1::AbstractModel,w2::AbstractModel; kwargs...) = evaluate(d, coefficients(domain(d),w1), coefficients(domain(d),w2); kwargs...)

function evaluate(d::CoefficientDistance,w1::AbstractArray,w2::AbstractArray; kwargs...)
    evaluate(d.distance,w1,w2; kwargs...)
end

function evaluate(d::ModelDistance,X,Xh; kwargs...)
    w = fitmodel(d.fitmethod, X)
    wh = fitmodel(d.fitmethod, Xh)
    evaluate(d.distance, w, wh; kwargs...)
end

# function batch_loss(bs::Int, loss, X, Xh; kwargs...)
#     l = zero(eltype(Xh))
#     lx = length(X)
#     n_batches = length(X)÷bs
#     inds = 1:bs
#     # TODO: introduce overlap for smoother transitions  #src
#     for i = 1:n_batches
#         l += loss(X[inds],Xh[inds]; kwargs...)
#         inds = inds .+ bs
#     end
#     l *= bs
#     residual_inds = inds[1]:lx
#     lr = length(residual_inds)
#     lr > 0 && (l += loss(X[residual_inds],Xh[residual_inds]; kwargs...)*lr)
#     l /= length(X)
#     l / n_batches
# end

evaluate(d::EnergyDistance,X::AbstractArray,Xh::AbstractArray; kwargs...) = (std(X)-std(Xh))^2 + (mean(X)-mean(Xh))^2

"""
    precompute(d::AbstractDistance, As, threads=true)

Perform computations that only need to be donce once when several pairwise distances are to be computed

# Arguments:
- `As`: A vector of models
- `threads`: Us multithreading? (true)
"""
function precompute(d::DiscretizedRationalDistance, As::AbstractArray{<:AbstractModel}, threads=true)
    mapfun = threads ? tmap : map
    mapfun(As) do A1
        w = d.w
        tm = tf(m1)
        b1 = bode(tm, w.*2π)[1] |> vec
        b1 .= abs2.(b1)
        b1 ./= sum(b1)
        b1
    end
end

function evaluate(d::DiscretizedRationalDistance, m1::AbstractModel, m2::AbstractModel; kwargs...)
    w = d.w
    m1 = tf(m1)
    b1,_,_ = bode(m1, w.*2π) .|> vec
    b1 .= abs2.(b1)
    # b1 .-= (minimum(b1) - 1e-9) # reg to ensure no value is exactly 0
    b1 ./= sum(b1)
    m2 = tf(m2)
    b2,_,_ = bode(m2, w.*2π) .|> vec
    b2 .= abs2.(b2)
    # b2 .-= (minimum(b2) - 1e-9)
    b2 ./= sum(b2)
    evaluate(d, b1, b2; kwargs...)
end

function evaluate(d::DiscretizedRationalDistance, b1, b2; kwargs...)
    plan = discrete_grid_transportplan(b1, b2)
    cost = dot(plan, d.distmat)
end

function evaluate(
    d::WelchOptimalTransportDistance,
    w1::DSP.Periodograms.TFR,
    w2::DSP.Periodograms.TFR;
    kwargs...,
)
    if d.distmat === nothing
        return discrete_grid_transportcost(s1(w1.power), s1(w2.power), d.p, 1e-3; inplace=true)
    else
        D = d.distmat
        C = discrete_grid_transportplan(s1(w1.power), s1(w2.power), 1e-3; inplace=true)
        return dot(C, D)
    end
end

function evaluate(d::WelchLPDistance, w1::DSP.Periodograms.TFR, w2::DSP.Periodograms.TFR;  kwargs...)
    x1,x2 = d.normalized ? (s1(w1.power),s1(w2.power)) : (w1.power, w2.power)
    mean(abs.(d.transform.(x1)-d.transform.(x2)).^d.p)
end

function evaluate(d::AbstractWelchDistance, x1, x2; kwargs...)
    evaluate(d, welch_pgram(x1, d.args...; d.kwargs...), welch_pgram(x2, d.args...; d.kwargs...))
end

centers(x) = 0.5*(x[1:end-1] + x[2:end])
function evaluate(d::OptimalTransportHistogramDistance, x1, x2; solver=IPOT, kwargs...)
    p  = d.p
    h1 = fit(Histogram,x1)
    h2 = fit(Histogram,x2)
    distmat = [abs(e1-e2)^p for e1 in centers(h1.edges[1]), e2 in centers(h2.edges[1])]
    plan = solver(distmat, s1(h1.weights), s1(h2.weights); kwargs...)[1]
    # plan = sinkhorn_plan_log(distmat, b1, b2; ϵ=1/10, rounds=300)
    cost = dot(plan, distmat)
end

function evaluate(d::DiscreteGridTransportDistance, x1, x2; kwargs...)
    inner = d.inner
    D = d.distmat
    (length(x1),length(x2)) == size(D) || throw(ArgumentError("Sizes do not match, should be (length(x1),length(x2)) == size(D), but got $((length(x1),length(x2))), $(size(D))"))
    c1,c2 = d.c1, d.c2
    d.C .= 0
    @avx c1 .= x1 ./ sum(x1)
    @avx c2 .= x2 ./ sum(x2)
    if typeof(inner) ∈ (Cityblock, Euclidean)
        return discrete_grid_transportcost(c1, c2, 1, inplace=true)
    elseif inner isa SqEuclidean
        return discrete_grid_transportcost(c1, c2, 2, inplace=true)
    end
    C = discrete_grid_transportplan(c1, c2, g=d.C, inplace=true)
    dot(d.C, D)
end

"""
    discrete_grid_transportplan(x::AbstractVector{T}, y::AbstractVector{T}, tol=sqrt(eps(T)); inplace=false) where T

Calculate the optimal-transport plan between two vectors that are assumed to have the same support, with sorted support points.
- ´inplace` if true, `x` is overwritten.
- It's possible to supply keyword argument `g` to provide preallocated memory for the transport plan. Make sure it's all zeros if you do.
"""
function discrete_grid_transportplan(x::AbstractVector{T},y::AbstractVector{T},tol=sqrt(eps(T)); g = zeros(T,length(x), length(x)), inplace=false) where T
    n  = length(x)
    length(y) == n || throw(ArgumentError("Inputs must have the same length"))
    if !inplace
        x  = copy(x)
    end
    yf = zero(T)
    i = j = 1
    @inbounds while j <= n && i <= n
        needed = y[j] - yf
        available = x[i]
        if available >= needed
            g[i,j] = needed
            x[i] -= needed
            yf = zero(T)
            j += 1
        else
            g[i,j] = available
            yf += available
            i += 1
        end
    end
    if j < n || i < n
        takenfromy = yf#min(needed,available)
        @views sumleft = (sum(x[i:end]),sum(y[j:end])-takenfromy)
        if sumleft[1] > tol || sumleft[2] > tol
            error("Not all indices were covered (n,i,j) = $((n,i,j)), sum left (x,y) = $(sumleft)")
        end
    end
    g
end

"""
    discrete_grid_transportcost(x::AbstractVector{T}, y::AbstractVector{T}, p = 2, tol=sqrt(eps(T)); inplace=false) where T

Calculate the optimal-transport cost between two vectors that are assumed to have the same support, with sorted and equdistant support points.

The calculated cost corresponds to the following
```julia
D = [abs2((i-j)/(n-1)) for i in 1:n, j in 1:n] # note the 1/(n-1)
Γ = discrete_grid_transportplan(x, y)
dot(Γ, D)
```
- ´p´ is the power of the distance
- ´inplace` if true, `x` is overwritten.
"""
function discrete_grid_transportcost(x::AbstractVector{T},y::AbstractVector{T}, p = 2, tol=sqrt(eps(T)); inplace=false) where T
    n  = length(x)
    length(y) == n || throw(ArgumentError("Inputs must have the same length"))
    if !inplace
        x  = copy(x)
    end
    yf = zero(T)
    i = j = 1
    cost = zero(T)
    @inbounds while j <= n && i <= n
        needed = y[j] - yf
        available = x[i]
        if available >= needed
            cost += (abs(i-j)/(n-1))^p*needed # the -1 is because the largest difference is n-1
            x[i] -= needed
            yf = zero(T)
            j += 1
        else
            cost += (abs(i-j)/(n-1))^p*available # the -1 is because the largest difference is n-1
            yf += available
            i += 1
        end
    end
    if j < n || i < n
        takenfromy = yf#min(needed,available)
        @views sumleft = (sum(x[i:end]),sum(y[j:end])-takenfromy)
        if sumleft[1] > tol || sumleft[2] > tol
            error("Not all indices were covered (n,i,j) = $((n,i,j)), sum left (x,y) = $(sumleft)")
        end
    end
    cost
end

# function Base.inv(f::AbstractVector)
#     n = length(f)
#     r = LinRange(0, f[end], n)
#     map(r) do r
#         findfirst(>=(r), f)
#     end
# end

Base.inv(f::Function, interval) = x->finv(f, x, interval)
# Base.inv(f::Function,fp) = x->finv(f, fp, x)

function finv(f, z, interval)
    f(interval[1])-z > 0 && return interval[1]
    f(interval[2])-z < 0 && return interval[2]
    w = Roots.fzero(w->f(w)-z, interval...)
end

# function finv(f, fp, z)
#     w = Roots.fzero(w->f(w)-z, fp, π)
# end

"""
    ∫(f, a, b) = begin

Integrate `f` between `a` and `b`
"""
∫(f,a,b)::Float64 = quadgk(f,a,b; atol=1e-12, rtol=1e-7)[1]

"""
    c∫(f, a, b; kwargs...)

Cumulative integration of `f` between `a` and `b`

# Arguments:
- `kwargs`: are sent to the DiffEq solver
"""
function c∫(f,a,b;kwargs...)
    fi    = (u,p,t) -> f(t)
    tspan = (a,b)
    prob  = ODEProblem(fi,0.,tspan)
    retry = false
    local sol
    try
        sol   = solve(prob,AutoTsit5(Rosenbrock23());reltol=1e-12,abstol=1e-45,kwargs...)
    catch
        retry = true
    end
    retry = retry || length(sol) <= 1 || abs(sol[end]-1) > 0.1
    if retry # In this case the solver failed due to numerical issues
        @warn "Spectral integration failed, increasing precision"
        prob  = ODEProblem(fi,big(0.),big.(tspan))
        sol   = solve(prob,AutoVern9(Rodas5());reltol=1e-16,abstol=1e-45,kwargs...)
    end
    sol
end

@inline ControlSystems.evalfr(::Discrete, m::Identity, w, a::AbstractArray, b::AbstractVector) =
        (n=length(a);m=length(b);abs2(sum(j->b[j]*cis(w*(m-j)), 1:m)/sum(j->a[j]*cis(w*(n-j)), 1:n)))
@inline ControlSystems.evalfr(::Continuous, m::Identity, w, a::AbstractArray, b::AbstractVector) =
        (n=length(a);m=length(b);abs2(sum(j->b[j]*(im*w)^(m-j), 1:m)/sum(j->a[j]*(im*w)^(n-j), 1:n)))

@inline ControlSystems.evalfr(::Discrete, m::Identity, w, a::AbstractArray, scale::Number=1) =
        (n=length(a);abs2(scale/sum(j->a[j]*cis(w*(n-j)), 1:n)))
@inline ControlSystems.evalfr(::Continuous, m::Identity, w, a::AbstractArray, scale::Number=1) =
        (n=length(a);abs2(scale/sum(j->a[j]*(im*w)^(n-j), 1:n)))
@inline ControlSystems.evalfr(::Discrete, m::Log, w, a::AbstractArray, scale::Number=1) =
        (n=length(a);-log(abs2(sum(j->a[j]*cis(w*(n-j)), 1:n))) + scale/(2π))
@inline ControlSystems.evalfr(::Continuous, m::Log, w, a::AbstractArray, scale::Number=1) =
        (n=length(a);-log(abs2(sum(j->a[j]*(im*w)^(n-j), 1:n))) + scale/(2π))

@inline ControlSystems.evalfr(d, m, w, a::ARMA, scale=1) = evalfr(d,m,w,denvec(d, a), scale*denvec(d,a))
@inline ControlSystems.evalfr(d, m, w, a::AR, scale=1) = evalfr(d,m,w,denvec(d, a), scale)
@inline ControlSystems.evalfr(r::AbstractRoots, m, w, scale=1) = evalfr(domain(r),m,w,roots2poly(r), scale)

function invfunctionbarrier(sol1::T1,sol2::T2,p,interval) where {T1,T2}
    σ1    = sol1(interval[2]) # The total energy in the spectrum
    σ2    = sol2(interval[2]) # The total energy in the spectrum
    abs(σ1-1) > 0.05 && @warn "Cumulative spectral energy not equal to 1, σ1 = $(σ1)"
    abs(σ2-1) > 0.05 && @warn "Cumulative spectral energy not equal to 1, σ2 = $(σ2)"
    F1(w) = sol1(w)/σ1
    F2(w) = sol2(w)/σ2
    ∫(z->abs(inv(F1, interval)(z) - inv(F2, interval)(z))^p, 0, 1)
end

function closed_form_wass(d::RationalOptimalTransportDistance,sol1,sol2)
    p,interval = d.p, d.interval
    invfunctionbarrier(sol1,sol2,p,interval)
end

function functionbarrier(sol1::T1,sol2::T2,p,interval) where {T1,T2}
    σ1    = sol1(interval[2]) # The total energy in the spectrum
    σ2    = sol2(interval[2]) # The total energy in the spectrum
    abs(σ1-1) > 0.05 && @warn "Cumulative spectral energy not equal to 1, σ1 = $(σ1)"
    abs(σ2-1) > 0.05 && @warn "Cumulative spectral energy not equal to 1, σ2 = $(σ2)"
    F1(w) = sol1(w)/σ1
    F2(w) = sol2(w)/σ2
    ∫(interval...) do w
        abs(F1(w)-F2(w))^p
    end
end

function evaluate(d::Union{RationalOptimalTransportDistance, RationalCramerDistance}, A1::AbstractModel, A2::AbstractModel)
    sc   = d.interval[1] == 0 ? sqrt(2) : 1
    e1   = sc*sqrt(A1.b/spectralenergy(domain(d), A1))
    f1   = w -> evalfr(domain(d), magnitude(d), w, A1, e1)
    e2   = sc*sqrt(A2.b/spectralenergy(domain(d), A2))
    f2   = w -> evalfr(domain(d), magnitude(d), w, A2, e2)
    sol1 = c∫(f1,d.interval...)
    sol2 = c∫(f2,d.interval...)
    d isa RationalOptimalTransportDistance && d.p > 1 && (return closed_form_wass(d,sol1,sol2))
    evaluate(d, sol1, sol2)
end

function evaluate(d::Union{RationalOptimalTransportDistance, RationalCramerDistance}, sol1, sol2)
    d isa RationalOptimalTransportDistance && d.p > 1 && (return closed_form_wass(d,sol1,sol2))
    functionbarrier(sol1,sol2,d.p,d.interval)
end

function evaluate(d::Union{RationalOptimalTransportDistance, RationalCramerDistance}, A1::AbstractModel, P::DSP.Periodograms.TFR)
    @assert d.interval[1] == 0 "Integration interval must start at 0 to compare against a periodogram"
    p    = d.p
    e    = sqrt(2)sqrt(A1.b/spectralenergy(domain(d), A1))
    f    = w -> evalfr(domain(d), magnitude(d), w, A1, e)
    w    = P.freq .* (2π)
    plan = discrete_grid_transportplan(s1(f.(w)), s1(sqrt.(P.power)))
    c    = 0.
    for i in axes(plan,1), j in axes(plan,2)
        c += abs(w[i]-w[j])^p * plan[i,j]
    end
    @assert c >= 0 "Distance turned out to be negative :/ c=$c"
    c
end

function precompute(d::Union{RationalOptimalTransportDistance, RationalCramerDistance}, As::AbstractArray{<:AbstractModel}, threads=true)
    mapfun = threads ? tmap : map
    sc   = d.interval[1] == 0 ? sqrt(2) : 1
    mapfun(As) do A1
        e1 = sc/sqrt(spectralenergy(domain(d), A1))
        f1   = w -> evalfr(domain(d), magnitude(d), w, A1, e1)
        sol1 = c∫(f1,d.interval...)
    end
end

# function closed_form_log_wass(a1,a2,p=1)
#     n     = length(a1)
    # f1    = w -> -log(abs(evalfr(Discrete(), w, a1))) + 1/(2π)
#     f2    = w -> -log(abs(evalfr(Discrete(), w, a2))) + 1/(2π)
#     endpoint = min(∫(f1,0,2π), ∫(f2,0,2π))
#     F1(w) = ∫(f1,0,w)
#     F2(w) = ∫(f2,0,w)
#     ∫(z->abs(inv(F1)(z) - inv(F2)(z))^p, 1e-9, endpoint-1e-9)
# end

"""
    stabilize(sys)

takes a statespace system and reflects all eigenvalues to be in the left half plane
"""
function stabilize(sys)
    A = sys.A
    e = eigen(A)
    e.values .= reflect(ContinuousRoots(e.values))
    sys.A .= real((e.vectors) * Diagonal(e.values) * inv(e.vectors))
    sys
end

function evaluate(d::BuresDistance, A1::AbstractModel, A2::AbstractModel)
    evaluate(d,tf(A1),tf(A2))
end

function evaluate(d::BuresDistance, A1::LTISystem, A2::LTISystem)
    sys1 = (ss((A1))) # stabilize
    sys2 = (ss((A2)))
    X1   = gram(sys1,:c)
    X2   = gram(sys2,:c)
    evaluate(d,(sys1.C*X1*sys1.C'),(sys2.C*X2*sys2.C'))
    # evaluate(d,inv(X1),inv(X2))
    # 2π*tr(sys.C*X*sys.C')
end

function evaluate(d::BuresDistance, A::AbstractMatrix, B::AbstractMatrix)
    sA = sqrt(A)
    tr(A+B - 2sqrt(sA*B*sA))
end


text:
using Documenter, TravelingSalesmanExact

makedocs(;
    modules=[TravelingSalesmanExact],
    authors="Eric P. Hanson",
    repo="https://github.com/ericphanson/TravelingSalesmanExact.jl/blob/{commit}{path}#L{line}",
    sitename="TravelingSalesmanExact.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://ericphanson.github.io/TravelingSalesmanExact.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/ericphanson/TravelingSalesmanExact.jl",
)


text:
# AUTO GENERATED FILE - DO NOT EDIT

export ''_fefferycontextmenu

"""
    ''_fefferycontextmenu(;kwargs...)

A FefferyContextMenu component.

Keyword arguments:
- `id` (optional)
- `className` (optional)
- `loading_state` (optional)
- `setProps` (optional): Dash-assigned callback that should be called to report property changes
to Dash, to make them available for callbacks.
- `style` (optional)
"""
function ''_fefferycontextmenu(; kwargs...)
        available_props = Symbol[:id, :className, :loading_state, :style]
        wild_props = Symbol[]
        return Component("''_fefferycontextmenu", "FefferyContextMenu", "feffery_utils_components", available_props, wild_props; kwargs...)
end

max_stars_repo_name: CNFeffery/feffery-utils-components

max_issues_repo_name: CNFeffery/feffery-utils-components

max_forks_repo_name: CNFeffery/feffery-utils-components


text:
using DataFrames
using SQLite
using Queryverse
using Gadfly

## Exemplos de statements comumente utilizados, para os quais desejamos as funções pré-compiladas,
## de forma a evitar

## lendo arquivo de entrada
db = SQLite.DB("bd_dados_qsa_cnpj.db") # carrega o arquivo SQLite com as tabelas

## cria índices com colunas úteis para consultas
SQLite.createindex!(db, "cnpj_dados_cadastrais_pj", "idx_situacao_cadastral", "situacao_cadastral",
    unique=false, ifnotexists=true)

SQLite.createindex!(db, "cnpj_dados_cadastrais_pj", "idx_identificador_matriz_filial", "identificador_matriz_filial",
    unique=false, ifnotexists=true)

## executa a consulta base e retorna um result set, que será depois usado
# para trazer as empresas de interesse
dados_pj_rs =
    DBInterface.execute(db,
"
    SELECT
        pj.cnpj,
        pj.nome_fantasia,
        pj.situacao_cadastral as situacao,
        case pj.porte_empresa
            when '00' then '?'
            when '01' then 'MICRO'
            when '03' then 'PEQUENA'
            when '05' then 'DEMAIS'
        end as porte,
        case pj.opcao_pelo_simples
            when '' then 'N'
            when '0' then 'N'
            when '5' then 'S'
            when '7' then 'S'
            when '6' then 'S - Excluido'
            when '8' then 'S - Excluido'
        end as simples,
        case pj.opcao_pelo_mei
            when 'S' then 'S'
            else 'N'
        end as MEI,
        nat_jur.nm_natureza_juridica as nat_juridica,
        pj.data_inicio_atividade as dat_ini_atividade,
        pj.cnae_fiscal as cod_cnae_principal,
        cnae.nm_cnae as cnae_principal,
        cnae.nm_divisao as divisao_cnae_principal,
        cnae.nm_classe as classe_cnae_principal,
        cnae.nm_grupo as grupo_cnae_principal,
        pj.bairro, pj.codigo_municipio as cod_municipio, pj.municipio, pj.uf
    FROM `cnpj_dados_cadastrais_pj` pj
        INNER JOIN `tab_cnae` cnae on cnae.cod_cnae = pj.cnae_fiscal
        INNER JOIN `tab_natureza_juridica` nat_jur on nat_jur.cod_subclass_natureza_juridica = pj.codigo_natureza_juridica
    WHERE
        situacao_cadastral='02' -- somente ativas
        and identificador_matriz_filial='1' -- somente matrizes
    ORDER BY data_inicio_atividade, cnpj
    ;
") # consulta base

# OBS: o result set somente suporta ser consumido uma única vez (single pass, forware only).
# Para usá-lo de novo, é preciso resetá-lo (rodar a query de novo) com: SQLite.reset!(dados_pj_rs)

## tendo o result set inicial, podemos executar as análises de interesse

## Análise 1: quantos cartórios temos no Brasil?
SQLite.reset!(dados_pj_rs) # para recarregar o result set, se necessário
dados_cartorios = dados_pj_rs |>
    @filter(_.cnae_principal=="Cartórios") |>
    DataFrame

print("O Brasil possui $(nrow(dados_cartorios)) cartórios!")

# mas onde eles estão localizados? são muitos porque temos muitas cidades?
dados_cartorios_uf = dados_cartorios |>
    @groupby(_.uf) |> @map({uf=key(_), n_cartorios=length(_)}) |> @orderby_descending(_.n_cartorios)

dados_cartorios_uf_municipios = dados_cartorios |>
    @select(:uf, :municipio) |> @unique() |>
    @groupby(_.uf) |> @map({uf=key(_), n_municipios=length(_)}) |> @orderby_descending(_.n_municipios)

comparacao_cartorios_uf = dados_cartorios_uf |>
    @join(dados_cartorios_uf_municipios, _.uf, _.uf, {_.uf, _.n_cartorios, __.n_municipios}) |>
    @orderby_descending(_.n_cartorios) |> DataFrame

Statistics::cor(
    convert(Array{Float64,1}, comparacao_cartorios_uf.n_cartorios),
    convert(Array{Float64,1}, comparacao_cartorios_uf.n_municipios)
)

Gadfly.push_theme(plot_style())
set_default_plot_size(21cm, 10cm)

plot(dados_cartorios_uf, y=:uf, x=:n_cartorios, label=:n_cartorios,
    Geom.bar(orientation=:horizontal), Geom.label(position=:right),
    Guide.title("Top UFs com cartório"))

max_stars_repo_name: llrt/eda-qsa-rfb-jl

max_issues_repo_name: llrt/eda-qsa-rfb-jl

max_forks_repo_name: llrt/eda-qsa-rfb-jl


text:
function size(cur_matrix::FourBlockMatrix)
  cur_dim = cur_matrix.n + cur_matrix.m
  (cur_dim, cur_dim)
end


text:
# Auxiliary functions for TimeModeling module
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: September 2016
#

export source, add_pml, remove_pml, get_computational_nt, smooth10, damp_boundary, calculate_dt, setup_grid, setup_3D_grid
export convertToCell, reshape_for_devito, reshape_from_devito, limit_model_to_receiver_area, extend_gradient, plot_geometry_map

function limit_model_to_receiver_area(srcGeometry::Geometry,recGeometry::Geometry,model::Model,buffer;pert=[])
	# Restrict full velocity model to area that contains either sources and receivers

	# scan for minimum and maximum x and y source/receiver coordinates
	min_x = minimum([vec(recGeometry.xloc[1]); vec(srcGeometry.xloc[1])])
	max_x = maximum([vec(recGeometry.xloc[1]); vec(srcGeometry.xloc[1])])
	min_y = minimum([vec(recGeometry.yloc[1]); vec(srcGeometry.yloc[1])])
	max_y = maximum([vec(recGeometry.yloc[1]); vec(srcGeometry.yloc[1])])

	# add buffer zone if possible
	min_x = max(model.o[1],min_x-buffer)
	max_x = min(model.o[1] + model.d[1]*(model.n[1]-1),max_x+buffer)
	min_y = max(model.o[2],min_y-buffer)
	max_y = min(model.o[2] + model.d[2]*(model.n[2]-1),max_y+buffer)

	# extract part of the model that contains sources/receivers
	nx_min = Int(round(min_x/model.d[1])) + 1
	nx_max = Int(round(max_x/model.d[1])) + 1
	ny_min = Int(round(min_y/model.d[2])) + 1
	ny_max = Int(round(max_y/model.d[2])) + 1
	ox = (nx_min - 1)*model.d[1]
	oy = (ny_min - 1)*model.d[2]
	oz = model.o[3]

	# Extract relevant model part from full domain
	n_orig = model.n
	model.m = model.m[nx_min:nx_max,ny_min:ny_max,:]
	model.o = (ox,oy,oz)
	model.n = size(model.m)
	if isempty(pert)
		return model
	else
		pert = reshape(pert,n_orig)[nx_min:nx_max,ny_min:ny_max,:]
		return model,vec(pert)
	end
end

function extend_gradient(model_full::Model,model::Model,gradient::Array)
	# Extend gradient back to full model size
	full_gradient = zeros(Float32,model_full.n)
	nx_start = Int((model.o[1] - model_full.o[1])/model.d[1] + 1)
	nx_end = nx_start + model.n[1] - 1
	ny_start = Int((model.o[2] - model_full.o[2])/model.d[2] + 1)
	ny_end = ny_start + model.n[2] - 1
	full_gradient[nx_start:nx_end,ny_start:ny_end,:] = gradient
	return full_gradient
end

function ccallsyms(ccallargs, n, argsyms)
	# create argument list for ccall
    if n > 0
        if length(argsyms) == n
            ccallargs = Any[ccallargs..., argsyms...]
        else
            for i = 1:length(argsyms)-1
                push!(ccallargs, argsyms[i])
            end
            for i = 1:n-length(argsyms)+1
                push!(ccallargs, Expr(:getindex, argsyms[end], i))
            end
        end
    end
    ccallargs
end

function ccallexpr(lib::String, ccallsym::Symbol, argtypes::Array{DataType,1}, argsyms::Array{Any,1})
    ccallargs = Any[Libdl.dlsym(Libdl.dlopen(lib),ccallsym), Cint, Expr(:tuple, argtypes...)]
    ccallargs = ccallsyms(ccallargs, length(argtypes), argsyms)
    eval(:(ccall($(ccallargs...))))
end

"""
Processing of input varibales for the ccall from the symbolic values obtained from devito
Some convention is applied here for the keywords
forward wavefield : u
adjoint wavefield : v
square slowness : m
boundary damping : damp
density : rho
source : src
source position : src_coords
data : rec
data locations : rec_coords
model has to be an input t infer the number of dimensions and size of the domain
use as :

op = devito[:operator]()
args = process_args(op, model; u=u, rec=...)
"""
function process_args(op, kwargs_dict)
	# kwargs_dict = Dict(kwargs)
	args = op[:arguments]()
	nargs = length(args)
	TypesIn = []
	ArgsIn = Array{Any}(length(args))
	for i=1:length(args)
		if args[i][2] == nothing
			loc = string(args[i][1])
			TypesIn = [TypesIn..., Ref{Cfloat}]
			if haskey(kwargs_dict, Symbol(args[i][1]))
				ArgsIn[i] = kwargs_dict[Symbol(args[i][1])]
			else
				@printf("Unrecognized key %s", Symbol(args[i][1]))
				error("Unknown key for inputs, check the inputs")
			end
		else
			TypesIn = [TypesIn..., Cint]
			ArgsIn[i] = args[i][2]
		end
	end
	return TypesIn, ArgsIn
end

function add_pml(array,nb)
	# extent model domain for absorbing boundaries
	if length(array)>1
		pad_list =[(nb, nb),(nb, nb)]	# x,z
		if ndims(array) == 3
			pad_list =[(nb, nb), (nb, nb), (nb, nb)]	# x,y,z
		end
		return np.pad(array, pad_list, "edge")
	else
		return array
	end
end

function remove_pml(array,nb)
# Cut off PML from model with Python dimensions
    if ndims(array) == 3
       return array[nb+1:end-nb, nb+1:end-nb, nb+1:end-nb]	# x,y,z
    else
       return array[nb+1:end-nb, nb+1:end-nb]	# x,z
    end
end

function damp_boundary(damp, model)
	# add absorbing boundaries to model

	# Number of dimensions
    num_dim = ndims(damp)

	# Damping coefficients
	if	num_dim == 2
		dampcoeffX = 1.5 * log(1.0 / 0.001) / (model.nb * model.d[1])
    	dampcoeffZ = 1.5 * log(1.0 / 0.001) / (model.nb * model.d[2])
	else
		dampcoeffX = 1.5 * log(1.0 / 0.001) / (model.nb * model.d[1])
		dampcoeffY = 1.5 * log(1.0 / 0.001) / (model.nb * model.d[2])
    	dampcoeffZ = 1.5 * log(1.0 / 0.001) / (model.nb * model.d[3])
	end

	if num_dim == 2
		for i=1:model.nb	# add x pml
        	pos = abs((model.nb-i)/float(model.nb))
        	val = dampcoeffX * (pos - sin(2*pi*pos)/(2*pi))
        	damp[i, :] += val
        	damp[end-i+1, :] += val
		end
    	for i=1:model.nb	# add z pml
       		pos = abs((model.nb-i)/float(model.nb))
       		val = dampcoeffZ * (pos - sin(2*pi*pos)/(2*pi))
        	damp[:, i] += val
        	damp[:, end-i+1] += val
		end
	else
		for i=1:model.nb	# add x pml
        	pos = abs((model.nb-i)/float(model.nb))
        	val = dampcoeffX * (pos - sin(2*pi*pos)/(2*pi))
            damp[i, :, :] += val
            damp[end-i+1, :, :] += val
		end
		for i=1:model.nb	# add y pml
        	pos = abs((model.nb-i)/float(model.nb))
        	val = dampcoeffY * (pos - sin(2*pi*pos)/(2*pi))
            damp[:, i, :] += val
            damp[:, end-i+1, :] += val
		end
	   	for i=1:model.nb	# add z pml
       		pos = abs((model.nb-i)/float(model.nb))
       		val = dampcoeffZ * (pos - sin(2*pi*pos)/(2*pi))
            damp[:, :, i] += val
            damp[:, :, end-i+1] += val
		end
	end
	return damp
end

"""
    convertToCell(x)

Convert an array `x` to a cell array (`Array{Any,1}`) with `length(x)` entries,\\
where the i-th cell contains the i-th entry of `x`.

"""
function convertToCell(x)
	n = length(x)
	y = Array{Any}(n)
	for j=1:n
		y[j] = x[j]
	end
	return y
end

# 1D source time function
"""
    source(tmax, dt, f0)

Create seismic Ricker wavelet of length `tmax` (in milliseconds) with sampling interval `dt` (in milliseonds)\\
and central frequency `f0` (in kHz).

"""
function source(tmax, dt, f0)
	t0 = 0.
	nt = Int(trunc((tmax - t0)/dt + 1))
	t = linspace(t0,tmax,nt)
	r = (pi * f0 * (t - 1 / f0))
	q = zeros(Float32,nt,1)
	q[:,1] = (1. - 2.*r.^2.).*exp.(-r.^2.)
	return vec(q)
end

function calculate_dt(n,d,o,v; epsilon=0)
	model = ct.IGrid(o, d, n, v, epsilon=epsilon)
	return dt = model[:critical_dt]
end

"""
    get_computational_nt(srcGeometry, recGeoemtry, model)

Estimate the number of computational time steps. Required for calculating the dimensions\\
of the matrix-free linear modeling operators. `srcGeometry` and `recGeometry` are source\\
and receiver geometries of type `Geometry` and `model` is the model structure of type \\
`Model`.

"""
function get_computational_nt(srcGeometry, recGeometry, model::ModelTTI)
	# Determine number of computational time steps
	if typeof(srcGeometry) == GeometryOOC
		nsrc = length(srcGeometry.container)
	else
		nsrc = length(srcGeometry.xloc)
	end
	nt = Array{Any}(nsrc)
	dtComp = calculate_dt(model.n,model.d,model.o,sqrt.(1./model.m); epsilon=model.epsilon)
	for j=1:nsrc
		ntRec = Int(trunc(recGeometry.dt[j]*(recGeometry.nt[j]-1))) / dtComp
		ntSrc = Int(trunc(srcGeometry.dt[j]*(srcGeometry.nt[j]-1))) / dtComp
		nt[j] = max(Int(trunc(ntRec)), Int(trunc(ntSrc)))
	end
	return nt
end

function get_computational_nt(srcGeometry, recGeometry, model::Union{Model,ModelRho})
	# Determine number of computational time steps
	if typeof(srcGeometry) == GeometryOOC
		nsrc = length(srcGeometry.container)
	else
		nsrc = length(srcGeometry.xloc)
	end
	nt = Array{Any}(nsrc)
	dtComp = calculate_dt(model.n,model.d,model.o,sqrt.(1./model.m))
	for j=1:nsrc
		ntRec = Int(trunc(recGeometry.dt[j]*(recGeometry.nt[j]-1))) / dtComp
		ntSrc = Int(trunc(srcGeometry.dt[j]*(srcGeometry.nt[j]-1))) / dtComp
		nt[j] = max(Int(trunc(ntRec)), Int(trunc(ntSrc)))
	end
	return nt
end

function setup_grid(geometry,n, origin)
	# 3D grid
	if length(n)==3
		if length(geometry.xloc[1]) > 1
			source_coords = Array{Float32,2}([vec(geometry.xloc[1]) vec(geometry.yloc[1]) vec(geometry.zloc[1])])
		else
			source_coords = Array{Float32,2}([geometry.xloc[1] geometry.yloc[1] geometry.zloc[1]])
		end
		orig = Array{Float32}([origin[1] origin[2] origin[3]])
	else
	# 2D grid
		if length(geometry.xloc[1]) > 1
			source_coords = Array{Float32,2}([vec(geometry.xloc[1]) vec(geometry.zloc[1])])
		else
			source_coords = Array{Float32,2}([geometry.xloc[1] geometry.zloc[1]])
		end
		orig = Array{Float32}([origin[1] origin[2]])
	end
	return source_coords .- orig
end

function setup_3D_grid(xrec::Array{Any,1},yrec::Array{Any,1},zrec::Array{Any,1})
	# Take input 1d x and y coordinate vectors and generate 3d grid. Input are cell arrays
	nsrc = length(xrec)
	xloc = Array{Any}(nsrc)
	yloc = Array{Any}(nsrc)
	zloc = Array{Any}(nsrc)
	for i=1:nsrc
		nxrec = length(xrec[i])
		nyrec = length(yrec[i])

		xloc[i] = zeros(nxrec*nyrec)
		yloc[i] = zeros(nxrec*nyrec)
		zloc[i] = zeros(nxrec*nyrec)

		idx = 1

		for k=1:nyrec
			for j=1:nxrec
				xloc[i][idx] = xrec[i][j]
				yloc[i][idx] = yrec[i][k]
				zloc[i][idx] = zrec[i]
				idx += 1
			end
		end
	end
	return xloc, yloc, zloc
end

function setup_3D_grid(xrec,yrec,zrec)
# Take input 1d x and y coordinate vectors and generate 3d grid. Input are arrays/linspace
	nxrec = length(xrec)
	nyrec = length(yrec)

	xloc = zeros(nxrec*nyrec)
	yloc = zeros(nxrec*nyrec)
	zloc = zeros(nxrec*nyrec)
	idx = 1
	for k=1:nyrec
		for j=1:nxrec
			xloc[idx] = xrec[j]
			yloc[idx] = yrec[k]
			zloc[idx] = zrec
			idx += 1
		end
	end
	return xloc, yloc, zloc
end

function smooth10(velocity,shape)
	# 10 point smoothing function
	out = ones(Float32,shape)
	nz = shape[end]
	if length(shape)==3
		out[:,:,:] = velocity[:,:,:]
		for a=5:nz-6
			out[:,:,a] = sum(velocity[:,:,a-4:a+5],3) / 10
		end
	else
		out[:,:] = velocity[:,:]
		for a=5:nz-6
			out[:,a] = sum(velocity[:,a-4:a+5],2) / 10
		end
	end
	return out
end

# Vectorization of single variable (not defined in Julia)
vec(x::Float64) = x;
vec(x::Float32) = x;
vec(x::Int64) = x;
vec(x::Int32) = x;

function plot_geometry_map(model,srcGeometry,recGeometry,shot_no;colormap="viridis")
	close("all")
	map = zeros(Float32,model.n[1],model.n[2])

	xmin = minimum(vec(recGeometry.xloc[shot_no]))
	xmax = maximum(vec(recGeometry.xloc[shot_no]))
	ymin = minimum(vec(recGeometry.yloc[shot_no]))
	ymax = maximum(vec(recGeometry.yloc[shot_no]))

	xmin_n = Int(round((xmin-model.o[1])/model.d[1] + 1))
	xmax_n = Int(round((xmax-model.o[1])/model.d[1] + 1))
	ymin_n = Int(round((ymin-model.o[2])/model.d[2] + 1))
	ymax_n = Int(round((ymax-model.o[2])/model.d[2] + 1))

	src_x = Int(round((srcGeometry.xloc[shot_no][1]-model.o[1])/model.d[1] + 1))
	src_y = Int(round((srcGeometry.yloc[shot_no][1]-model.o[2])/model.d[2] + 1))

	map[xmin_n:xmax_n,ymin_n:ymax_n] = 1.
	#map[src_x,src_y] = -1.
	imshow(map; cmap=colormap); colorbar()
	xlabel("y-coordinate")
	ylabel("x-coordinate")
	plot(src_y,src_x,"x")
end

wrap_retry(f,n) = retry(f;n=n)

max_stars_repo_name: SINBADconsortium/opesciSLIM

max_issues_repo_name: SINBADconsortium/opesciSLIM

max_forks_repo_name: SINBADconsortium/opesciSLIM


text:
mutable struct SDEOptions{tTypeNoUnits,tType,F2,F3,F4,F5,F6,tstopsType,discType,ECType,SType,MI,A,R,D,
           tcache,savecache,disccache}
  maxiters::MI
  save_everystep::Bool
  adaptive::Bool
  abstol::A
  reltol::R
  gamma::tTypeNoUnits
  qmax::tTypeNoUnits
  qmin::tTypeNoUnits
  failfactor::tTypeNoUnits
  dtmax::tType
  dtmin::tType
  internalnorm::F2
  save_idxs::SType
  tstops::tstopsType
  saveat::tstopsType
  d_discontinuities::discType
  tstops_cache::tcache
  saveat_cache::savecache
  d_discontinuities_cache::disccache
  userdata::ECType
  progress::Bool
  progress_steps::Int
  progress_name::String
  progress_message::F5
  timeseries_errors::Bool
  dense_errors::Bool
  beta1::tTypeNoUnits
  beta2::tTypeNoUnits
  delta::D
  qoldinit::tTypeNoUnits
  dense::Bool
  save_on::Bool
  save_start::Bool
  save_end::Bool
  save_noise::Bool
  callback::F3
  isoutofdomain::F4
  unstable_check::F6
  verbose::Bool
  calck::Bool
  force_dtmin::Bool
  advance_to_tstop::Bool
  stop_at_next_tstop::Bool
end


text:
# Activation funcitons `y = σ(x)` and the corresponding gradient functions `y' = ∇σ(x, y)`.

export
    ∇tanh,
    logistic, ∇logistic,
    relu, ∇relu,
    softplus, ∇softplus,
    ∇cos,
    log_logistic, ∇log_logistic

# tanh: built-in
@inline function ∇tanh(x::Float64, y::Float64)::Float64
    1.0 - y ^ 2
end

@inline function logistic(x::Float64)::Float64
    1.0 / (1.0 + exp(-x))
end

@inline function ∇logistic(x::Float64, y::Float64)::Float64
    y * (1.0 - y)
end

@inline function relu(x::Float64)::Float64
    max(x, 0.0)
end

@inline function ∇relu(x::Float64, y::Float64)::Float64
    x >= 0.0 ? 1.0 : 0.0
end

@inline function softplus(x::Float64)::Float64
    log(1.0 + exp(x))
end

@inline function ∇softplus(x::Float64, y::Float64)::Float64
    logistic(x)
end

# cos: built-in
@inline function ∇cos(x::Float64, y::Float64)::Float64
    -sin(x)
end

@inline function log_logistic(x::Float64)::Float64
    x - softplus(x)
end

@inline function ∇log_logistic(x::Float64, y::Float64)::Float64
    1.0 / (1.0 + exp(x))
end

max_stars_repo_name: gyk/TrivialSolutions

max_issues_repo_name: gyk/TrivialSolutions

max_forks_repo_name: gyk/TrivialSolutions