text:
using DynamicPolynomials

using FileIO;

@with_kw mutable struct LagrangePoly
    p :: AbstractPolynomialLike
    grad_poly :: Union{Nothing, Vector{<:AbstractPolynomialLike}} = nothing
    res :: AbstractResult = init_res(Result)
    # TODO check if `res` is a reference (good) or copy (bad for memory)
    # if the latter, store only values at interpolation sites
end
Broadcast.broadcastable(lp :: LagrangePoly ) = Ref(lp);

@with_kw mutable struct LagrangeModel <: SurrogateModel
    n_vars :: Int;  # stored for convenience
    n_out :: Int;
    basis :: Vector{LagrangePoly} = LagrangePoly[];
    vals :: RVecArr = RVec[];
    fully_linear :: Bool = false;
    out_indices :: Vector{Int} = Int[];
end

@with_kw mutable struct LagrangeConfig <: SurrogateConfig
    degree :: Int = 2;
    θ_enlarge :: Real = 2;

    "Acceptance Parameter in first Poisedness-Algorithm."
    ε_accept :: Real = 1e-8;
    "Quality Parameter in Λ-Poisedness Algorithm."
    Λ :: Real = 1.5;

    allow_not_linear :: Bool = false;

    optimized_sampling :: Bool = true;

    # # if optimized_sampling = false, shall we try to use saved sites?
    save_path :: Union{AbstractString, Nothing} = nothing;
    io_lock :: Union{Nothing, ReentrantLock, Threads.SpinLock} = nothing;

    algo1_max_evals :: Union{Nothing,Int} = nothing;    # nothing = automatic
    algo2_max_evals :: Union{Nothing,Int} = nothing;

    algo1_solver :: Symbol = :LD_MMA
    algo2_solver :: Symbol = algo1_solver;

    max_evals :: Int64 = typemax(Int64);
end

# overwrite lock and unlock so we can use `nothing` as a lock
function Base.lock(::Nothing) end
function Base.unlock(::Nothing) end

max_evals( cfg::LagrangeConfig)::Int=cfg.max_evals;

fully_linear( lm::LagrangeModel ) :: Bool = lm.fully_linear;

combinable(::LagrangeConfig) ::Bool = true;

function combine(cfg1 :: LagrangeConfig, cfg2 :: LagrangeConfig) :: LagrangeConfig
    cfg1
end

@with_kw mutable struct LagrangeMeta <: SurrogateMeta
    interpolation_indices :: Vector{<:Union{Nothing,Int}} = Union{Nothing,Int}[];
    #lagrange_basis :: Vector{Any} = [];
end

function _init_model( cfg :: LagrangeConfig, objf :: AbstractObjective,
    mop :: AbstractMOP, id :: AbstractIterData, ac :: AbstractConfig ) :: Tuple{LagrangeModel, LagrangeMeta}
    lm = LagrangeModel(;
        n_vars = num_vars(objf),
        n_out = num_outputs(objf),
        out_indices = copy(output_indices( objf, mop)),
        basis = _canonical_basis( num_vars( mop ), cfg.degree ),
    );
    # prepare initial basis of polynomial space
    lmeta = LagrangeMeta();
    return update_model(lm, objf, lmeta, mop, id, ac;
        ensure_fully_linear = true)
end

function _eval_new_sites!( lm :: LagrangeModel, lmeta :: LagrangeMeta,
        mop :: AbstractMOP, id :: AbstractIterData )  :: Nothing

    basis_results = [lp.res for lp ∈ lm.basis];
    _eval_and_store_new_results!(id, basis_results, mop);
    #@show  [ get_id(res) for res ∈ basis_results ];
    lmeta.interpolation_indices = [ convert(Union{Int,Nothing},get_id(res)) for res ∈ basis_results ];
    nothing
end

function _set_gradient!(lp :: LagrangePoly) :: Nothing
    lp.grad_poly = differentiate.(lp.p, DynamicPolynomials.variables(lp.p));
    nothing
end

function _update_model_optimized( lm :: LagrangeModel, objf :: AbstractObjective, lmeta :: LagrangeMeta,
    mop :: AbstractMOP, id :: AbstractIterData, :: AbstractConfig;
    ensure_fully_linear :: Bool = false):: Tuple{LagrangeModel, LagrangeMeta}
    cfg = model_cfg(objf) :: LagrangeConfig;

    make_basis_poised!(
        lm.basis, id, mop ;
        Δ_factor = cfg.θ_enlarge, ε_accept = cfg.ε_accept,
        max_solver_evals = cfg.algo1_max_evals,
        solver = cfg.algo1_solver
    );

    if ensure_fully_linear || !cfg.allow_not_linear
        make_basis_lambda_poised!(
            lm.basis, id, mop;
            Δ_factor = cfg.θ_enlarge, Λ = cfg.Λ, max_solver_evals = cfg.algo2_max_evals,
            solver = cfg.algo2_solver
        );
        lm.fully_linear = true;
    end

    _eval_new_sites!( lm, lmeta, mop, id);
    _set_gradient!.(lm.basis);

    return lm, lmeta;
end

@memoize ThreadSafeDict function _get_unit_basis( n_vars :: Int, cfg :: LagrangeConfig ) :: Vector{LagrangePoly}
    return _get_unit_basis( Val(!isnothing( cfg.save_path ) ) , n_vars, cfg )
end

function _get_unit_basis( use_saved :: Val{true}, n_vars :: Int, cfg :: LagrangeConfig )
    load_successful = true;
    unit_sites = [];
    lock(cfg.io_lock)
    try
        file_dict = load(cfg.save_path);
        unit_sites = file_dict["sites"];
    catch
        stacktrace( catch_backtrace() )
        @warn "Could not load files from file. Calculating a poised set."
        load_successful = false;
    end

    if load_successful
        # check if the sites are any goodif length( unit_sites[1] ) != n_vars ||
        length( unit_sites ) != binomial( n_vars + cfg.degree, n_vars )
        if any( any( ( s .< 0 ) .| ( s .> 1 ) ) for s ∈ unit_sites )
            @warn "Loaded sites not suited for interpolation."
            load_successful = false;
        end
    end

    if !load_successful
        unit_basis = _get_unit_basis( Val(false), n_vars, cfg )
    else
        @logmsg loglevel4 "Loaded sites from $(cfg.save_path)"
        unit_basis = copy( _canonical_basis( n_vars, cfg.degree ) );
        for (i,x) ∈ enumerate(unit_sites)
            unit_basis[i].res = init_res( Result, x );
            _normalize!(unit_basis[i]);
            _orthogonalize!(unit_basis, i);
        end
    end

    unlock(cfg.io_lock)

    return unit_basis;
end

function _get_unit_basis( use_saved :: Val{false}, n_vars :: Int, cfg :: LagrangeConfig )
    unit_basis = _unit_basis( n_vars, cfg.degree;
        solver1 = cfg.algo1_solver, solver2 = cfg.algo2_solver,
        max_evals1 = cfg.algo1_max_evals, max_evals2 = cfg.algo2_max_evals,
        Λ = cfg.Λ
    );
    if !isnothing( cfg.save_path )
        #lock(cfg.io_lock)
        # better wrap the outer call in io_lock for parallelization
        try
            unit_sites = [ get_site(ub.res) for ub ∈ unit_basis ];
            save(cfg.save_path, Dict( "sites" => unit_sites ) );
        catch
            stacktrace( catch_backtrace() )
            @warn "Could not save unit sites."
        end
        #unlock(cfg.io_lock);
    end
    return unit_basis;
end

 function _update_model_unoptimized(
    lm :: LagrangeModel, objf :: AbstractObjective, lmeta :: LagrangeMeta,
    mop :: AbstractMOP, id :: AbstractIterData, :: AbstractConfig;
    ensure_fully_linear :: Bool = false ):: Tuple{LagrangeModel, LagrangeMeta}

    cfg = model_cfg(objf);
    unit_basis = _get_unit_basis( num_vars(mop), cfg ); # this is memoized
    lb_eff, ub_eff = local_bounds(mop, xᵗ(id), Δᵗ(id) * cfg.θ_enlarge);
    lm.basis = _scale_unit_basis( unit_basis, lb_eff, ub_eff );
    lm.fully_linear = true;

    _eval_new_sites!( lm, lmeta, mop, id);
    _set_gradient!.(lm.basis);
    return  lm, lmeta ;
end

function _scale_unit_basis( unit_basis :: Vector{LagrangePoly}, lb :: RVec, ub :: RVec )
    χ = variables( unit_basis[1].p );

    # combine the polyonmial from unit_basis (working on [0,1]^n)
    # with a poly that maps the current trust region to [0,1]^n
    scaling_poly = (χ .- lb) ./ (ub .- lb)

    scaled_basis = LagrangePoly[];
    for up ∈ unit_basis
        lp = LagrangePoly(;
            p = subs( up.p, χ => scaling_poly ),
            # we `unscale` the sites from [0,1]^n to correspond to sites in
            # current trust region
            res = init_res( Result, _unscale( get_site( up.res ), lb, ub ) )
        );
        push!(scaled_basis, lp);
    end
    return scaled_basis
end

function update_model( lm :: LagrangeModel, objf :: AbstractObjective, lmeta :: LagrangeMeta,
    mop :: AbstractMOP, id :: AbstractIterData, algo_config :: AbstractConfig;
    ensure_fully_linear :: Bool = false):: Tuple{LagrangeModel, LagrangeMeta}

    @logmsg loglevel3 "Building LagrangeModel with indices $(output_indices(objf, mop))."
    cfg = model_cfg(objf);
    if cfg.optimized_sampling
        @info "optimized"
        lm,lmeta = _update_model_optimized(lm,objf,lmeta,mop,id,algo_config; ensure_fully_linear);
    else
        lm,lmeta = _update_model_unoptimized(lm,objf,lmeta,mop,id,algo_config; ensure_fully_linear);
    end

    # TODO remove this
    #=
    oi = lm.out_indices #output_indices( objf, mop )
    for lp in lm.basis
        try
            @assert eval_poly( lp, get_site(lp.res) ) ≈ 1
            @assert all(get_value(lp.res)[oi] .≈ eval_models( lm, get_site(lp.res) ))
        catch e
            @warn "Imprecise Lagrange Models."
            @info get_value(lp.res)[oi] .- eval_models( lm, get_site(lp.res) )
        end
    end
    =#
    return lm, lmeta;
end

function improve_model( lm :: LagrangeModel, objf :: AbstractObjective, lmeta :: LagrangeMeta,
    mop :: AbstractMOP, id :: AbstractIterData, :: AbstractConfig; ensure_fully_linear :: Bool = false):: Tuple{LagrangeModel, LagrangeMeta}

    @logmsg loglevel3 "Performing an improvement step for LagrangeModel with indices $(output_indices(objf, mop))."

    cfg = model_cfg( objf );
    make_basis_lambda_poised!( lm.basis, id, mop;
        Δ_factor = cfg.θ_enlarge, Λ = cfg.Λ,
        max_solver_evals = cfg.algo2_max_evals, solver = cfg.algo2_solver
    );
    lm.fully_linear = true;
    _eval_new_sites!( lm, lmeta, mop, id);
    _set_gradient!.(lm.basis)
    return lm, lmeta
end

function eval_models( lm :: LagrangeModel, x̂ :: RVec, ℓ :: Int ) :: Real
    sum( get_value(lp.res)[lm.out_indices[ ℓ ]] * eval_poly(lp.p, x̂) for lp ∈ lm.basis)
end

function eval_models( lm :: LagrangeModel, x̂ :: RVec ) :: RVec
    sum( get_value(lp.res)[ lm.out_indices ] * eval_poly(lp.p, x̂) for lp ∈ lm.basis)
end

function get_gradient( lm :: LagrangeModel, x̂ :: RVec, ℓ :: Int ) :: RVec
    sum( get_value(lp.res)[lm.out_indices[ ℓ ]] * eval_poly(lp.grad_poly, x̂) for lp ∈ lm.basis )
end

function get_jacobian( lm :: LagrangeModel, x̂ :: RVec ) :: RMat
    grad_evals = [ eval_poly(lp.grad_poly, x̂) for lp ∈ lm.basis ];
    return Matrix(transpose( hcat(
        [ sum( get_value(lm.basis[i].res)[ lm.out_indices[ ℓ ] ] * grad_evals[i]
            for i = eachindex(grad_evals) )
                for ℓ = 1 : lm.n_out ]...
    )));
end

#@memoize
function _unit_basis( n_vars :: Int, degree :: Int ;
        solver1 :: Symbol, solver2 :: Symbol, max_evals1 :: Union{Nothing,Int},
        max_evals2 :: Union{Nothing,Int}, Λ :: Real ) :: Vector{LagrangePoly}
    lb = zeros(n_vars);
    ub = ones(n_vars);

    basis = _canonical_basis( n_vars, degree );

    candidates = [init_res(Result, lb .+ .5 .* (ub .- lb) ),];
    _make_poised!( basis, lb, ub, candidates; ε_accept = 0, max_solver_evals = max_evals1, solver = solver1 );
    _make_lambda_poised!( basis, lb, ub; Λ = Λ, max_solver_evals = max_evals2, solver = solver2);
    return basis
end

function _make_poised!( basis :: Vector{LagrangePoly}, lb :: RVec, ub :: RVec, candidates :: Vector{<:AbstractResult};
    ε_accept ::Real, solver :: Symbol, max_solver_evals :: Union{Nothing,Int}, loop_start :: Int = 1 ) :: Nothing
    n_vars = length(lb);
    if isnothing( max_solver_evals)
        max_solver_evals = 300*(n_vars+1);   # TODO is sensible?
    end
    p = length( basis )
    for i = loop_start : p
        X = get_site.( candidates );
        l_max, j = isempty( X ) ? (0, 0) : findmax( abs.( eval_poly.( basis[i], X) ) );
        if l_max > ε_accept
            @logmsg loglevel4 " 1.$(i)) Recycling a point from the database."
            res = candidates[j];
            deleteat!( candidates, j);
        else
            @logmsg loglevel4 " 1.$(i)) Computing a poised point by Optimization."
            opt = NLopt.Opt( solver , n_vars )
            opt.lower_bounds = lb;
            opt.upper_bounds = ub;
            opt.maxeval = max_solver_evals;
            opt.xtol_rel = 1e-3;
            opt.max_objective = (x,g) -> abs( eval_poly( basis[i], x) )
            x₀ = _rand_box_point(lb, ub);
            (_, x, ret) = NLopt.optimize(opt, x₀)

            res = init_res(Result, x);
        end
        basis[i].res = res;
        _normalize!(basis[i]);
        _orthogonalize!(basis, i);
    end
    nothing;
end

function make_basis_poised!( basis :: Vector{LagrangePoly}, id :: AbstractIterData, mop :: AbstractMOP;
    Δ_factor :: Real = 1, ε_accept :: Real = 1e-8, solver :: Symbol = :GN_AGS, max_solver_evals :: Union{Int,Nothing} = nothing )
    @logmsg loglevel4 "Finding a poised set..."

    x = xᵗ( id );
    N = length(x);

    # always include current iterate
    @assert xᵗ_index( id ) isa XInt
    basis[1].res = get_result( id, xᵗ_index(id));
    _normalize!(basis[1]);
    _orthogonalize!(basis, 1);

    # find all points in database in current trust region
    lb_eff, ub_eff = local_bounds(mop, x, Δᵗ(id) * Δ_factor);
    box_indices = find_points_in_box(
        id, lb_eff, ub_eff;
        exclude_indices = [xᵗ_index(id)]
    );

    _make_poised!( basis, lb_eff, ub_eff, get_result.(id, box_indices); ε_accept, solver, max_solver_evals, loop_start = 2 );
    nothing
end

function _make_lambda_poised!( basis :: Vector{LagrangePoly}, lb :: RVec, ub :: RVec;
    Λ :: Real, max_solver_evals :: Union{Int,Nothing}, solver :: Symbol ) :: Nothing

    N = length(lb);
    if isnothing( max_solver_evals)
        max_solver_evals = 300*(N+1);   # TODO is sensible?
    end

    iter_counter = 1;
    while !isinf(Λ)
        # determine a swap index iₖ, 1 ≤ iₖ ≤ length(basis):
        # maximize absolute value of each basis polynomial within trust region
        # if we find a maximizer xᵢ with |lᵢ| > Λ then we can replace the
        # site with index i with xᵢ
        # iₖ is set so as to favor the argmax of all |lᵢ|
        # or to correspond to unevaluated sites and not the current iterate if possible
        iₖ = -1;
        xₖ = similar(lb);
        Λ_max = -Inf;
        swap_id = nothing;
        for (i,lp) = enumerate(basis)
            opt = NLopt.Opt( solver, N);
            opt.lower_bounds = lb;
            opt.upper_bounds = ub;
            opt.maxeval = max_solver_evals;
            opt.xtol_rel = 1e-3;
            opt.max_objective = (x,g) -> abs( eval_poly( lp, x) )
            x₀ = _rand_box_point(lb, ub);
            (abs_lᵢ, xᵢ, _) = NLopt.optimize(opt, x₀);

            if abs_lᵢ > Λ # found a potential swap candidate
                res_id = get_id(lp.res)
                if isa(swap_id, XInt) || abs_lᵢ > Λ_max || isnothing( res_id )
                    Λ_max = abs_lᵢ
                    iₖ = i;
                    xₖ[:] = xᵢ;
                    swap_id = res_id
                    if isnothing(res_id)
                        # break if we have a suitable swapping point that is not yet
                        # in the database (i.e. not evaluated yet)
                        break;
                    end
                end
            end
        end#for
        if iₖ > 0
            # perform a point swap and normalize
            @logmsg loglevel4 " 2.$(iter_counter)) Replacing point at index $(iₖ)."
            basis[iₖ].res = init_res(Result, xₖ);
            _normalize!(basis[iₖ])
            _orthogonalize!(basis, iₖ);
       else
            # finished, there is no poly with abs value > Λ
            return nothing
        end#if
        iter_counter += 1;
    end#while
    nothing
end

function make_basis_lambda_poised!( basis :: Vector{LagrangePoly}, id :: AbstractIterData, mop :: AbstractMOP;
        Δ_factor :: Real = 1, Λ :: Real = 1.5, max_solver_evals :: Union{Int,Nothing} = nothing,
        solver :: Symbol = :GN_AGS ) :: Nothing
    @logmsg loglevel4 "Making the set $(Λ)-poised..."

    lb_eff, ub_eff = local_bounds(mop, xᵗ(id), Δᵗ(id) * Δ_factor);
    _make_lambda_poised!( basis, lb_eff, ub_eff; Λ, max_solver_evals, solver );
end

# helper function to easily evaluate polynomial
function eval_poly(p :: AbstractPolynomialLike, x :: RVec)
    return p( variables(p) => x )
end

function eval_poly( lp :: LagrangePoly, x :: RVec )
    eval_poly( lp.p, x )
end

function _orthogonalize!( basis :: Vector{LagrangePoly}, i :: Int ) :: Nothing
    p = length(basis)
    y = get_site( basis[i].res );
    for j=1:p
        if j≠i
            basis[j].p -= (eval_poly( basis[j], y )* basis[i].p )
        end
    end
    nothing
end

function _normalize!( lp :: LagrangePoly ) :: Nothing
    x = get_site(lp.res);
    @assert !isempty(x);
    lp.p /= eval_poly( lp, x);
    nothing
end

# helper function to easily evaluate polynomial array for gradient
function eval_poly(p::Vector{<:AbstractPolynomialLike}, x :: RVec)
    [g( variables(p) => x ) for g ∈ p]
end

# helper function for initial monomial basis
@doc """
Return array of solution vectors [x_1, …, x_len] to the equation
``x_1 + … + x_len = rhs``
where the variables must be non-negative integers.
"""
function non_negative_solutions( rhs :: Int, len :: Int )
    if len == 1
        return rhs
    else
        solutions = [];
        for i = 0 : rhs
            for shorter_solution ∈ non_negative_solutions( i, len - 1)
                push!( solutions, [ rhs-i; shorter_solution ] )
            end
        end
        return solutions
    end
end

@doc "Factorial of a multinomial."
_multifactorial( arr :: Vector{Int} ) =  prod( factorial(α) for α in arr )

# @memoize
# TODO memoization caused a lot of trouble here in parallel usage.
# if important: use IdDicts and Memoization.jl
function _canonical_basis( n_vars :: Int, degree :: Int ) :: Vector{LagrangePoly}
    basis = LagrangePoly[];
    @polyvar χ[1:n_vars]
    for d = 0 : degree
        for multi_exponent ∈ non_negative_solutions( d, n_vars )
            poly = 1 / _multifactorial( multi_exponent ) * prod( χ.^multi_exponent );
            push!(basis, LagrangePoly(; p = poly ));
        end
    end
    return basis
end

max_stars_repo_path: src/archive/LagrangeModel copy.jl

max_issues_repo_path: src/archive/LagrangeModel copy.jl

max_forks_repo_path: src/archive/LagrangeModel copy.jl


text:
module VectorEngine

using GPUArrays
using Adapt

using GPUCompiler
using LLVM
using LLVM.Interop

import Core: LLVMPtr

include("types.jl")
include("pointer.jl")

include("veda/VEDA.jl")
using .VEDA

# Device and Memory management
include("stream.jl")
include("memory.jl")

# Device sources must load _before_ the compiler infrastructure
# because of generated functions.
include(joinpath("device", "tools.jl"))
include(joinpath("device", "pointer.jl"))
include(joinpath("device", "globals.jl"))
include(joinpath("device", "output.jl"))
include(joinpath("device", "memory.jl"))
include(joinpath("device", "array.jl"))
include(joinpath("device", "runtime.jl"))
include(joinpath("device", "llvm.jl"))
include(joinpath("device", "strings.jl"))

# Compiler infrastructure
include(joinpath("compiler", "compiler.jl"))
include(joinpath("compiler", "execution.jl"))
include(joinpath("compiler", "reflection.jl"))

# array essentials
include("array.jl")

# High-level functionality
include("vectorize.jl")
include("random.jl")
include("utils.jl")
end # module

ext: jl


text:
function force(A::AbstractSystem, p::EmpiricalPotential)
    f = fill(zeros(3), length(A))
    nnlist = neighborlist(A, p.rcutoff)

    for ii in 1:length(A)
        j = nnlist.j[ii]
        r = nnlist.r[ii]
        R = nnlist.R[ii]
        for (jj, rj, Rj) in zip(j, r, R)
            fo = force(Rj, rj, p)
            f[ii] += fo
            f[jj] -= fo
        end
    end
    [SVector{3}(fi) for fi in f]
end


text:
using BDF, Base.Test

#test split BDF at two time points
origFilePath = joinpath(dirname(@__FILE__), "Newtest17-256.bdf")
tSplitFilePath1 = joinpath(dirname(@__FILE__), "Newtest17-256_1.bdf")
tSplitFilePath2 = joinpath(dirname(@__FILE__), "Newtest17-256_2.bdf")
tSplitFilePath3 = joinpath(dirname(@__FILE__), "Newtest17-256_3.bdf")
tSplitFilePathComb = joinpath(dirname(@__FILE__), "Newtest17-256_comb.bdf")

bdfHeader = readBDFHeader(origFilePath)
dats, evtTab, trigs, statusChan = readBDF(origFilePath)

splitBDFAtTime(origFilePath, [20; 40])
dats1, evtTab1, trigs1, statusChan1 = readBDF(tSplitFilePath1)
dats2, evtTab2, trigs2, statusChan2 = readBDF(tSplitFilePath2)
dats3, evtTab3, trigs3, statusChan3 = readBDF(tSplitFilePath3)

writeBDF(tSplitFilePathComb, [dats1 dats2 dats3], [trigs1; trigs2; trigs3], [statusChan1; statusChan2; statusChan3], bdfHeader["sampRate"][1])
datsComb, evtTabComb, trigsComb, statusChanComb = readBDF(tSplitFilePathComb)

rm(tSplitFilePath1)
rm(tSplitFilePath2)
rm(tSplitFilePath3)
rm(tSplitFilePathComb)

@test isequal(dats, datsComb)
@test isequal(evtTab, evtTabComb)
@test isequal(trigs, trigsComb)
@test isequal(statusChan, statusChanComb)


text:
using .NetworkResilience
using Statistics

function main(t)

	ind = parse(Int64, ARGS[1])

	configs = []
	n = 67
	for j=1:n
		for i=1:n-j
			ne = n - i - j
			push!(configs, [i j ne])
		end
	end
	# break1 = length(configs)

	# #Make the file name
	# directory = "swing_coupling_q0/"
	# indexdict = Dict("50" => ind, "100" => ind-break1)
	# subdirectorydict = Dict("50" => "50tri/", "100" => "100tri/")
	# ntot = sum(configs[ind])
	# thisid = indexdict[string(ntot)]
	# thissubdir = subdirectorydict[string(ntot)]
	# fname = "data/" * directory * thissubdir * string(thisid) * ".txt"

	#Compute the cascade
	ns = configs[ind][1]
	nd = configs[ind][2]
	ne = configs[ind][3]
	ensemblesize = 3
	tol = 5e-4
	n = ns + nd + ne
	q = 0.0
	I = 1.0
	D = 1.0
	k = 4
	Ptot = 1.0
	κ = 5.5
	criticalcouplings = criticalcouplingbisection(ensemblesize, tol, n, ns, nd, q, k, κ, I, D)
	println(criticalcouplings)
	# open(fname, "w") do f
	# 	for z=1:ensemblesize
	# 		x = criticalcouplings[z]
	# 		write(f, "$x \n")
	# 	end
	# end

end

main(ARGS)

max_stars_repo_name: o-smith/GridResilience

max_issues_repo_name: o-smith/GridResilience

max_forks_repo_name: o-smith/GridResilience


text:
#hellwarththeory.jl

"""
Implementation of Hellwarth's extended treatment of Osaka's finite temperature version of Feynman's variational technique applied to the Polaron model; generalising it from the case of a single phonon mode to the case of multiple phonon modes.

And Hellwarth 1999:
https://doi.org/10.1103/PhysRevB.60.299

"""

using Optim
using QuadGK

function singlemode_variation(α, β; v = 7.2, w = 6.5, verbose::Bool = false)

    initial = [v, w]
    lower = [0.0, 0.0]
    upper = [Inf, v]

    f(x) = F(x[1], x[2], β, α)

    solution = Optim.optimize(
        Optim.OnceDifferentiable(f, initial; autodiff = :forward),
        lower,
        upper,
        initial,
        Fminbox(BFGS()),
    )

    if Optim.converged(solution) == false
        print("\tWARNING: Failed to converge to v,w soln? : ", Optim.converged(solution))
    end

    if verbose
        println()
        show(solution)
    end

    v, w = Optim.minimizer(solution)
    return v, w
end


text:
## PCA vs LDA
## Height Weight data
# Grab data from https://www.kaggle.com/mustafaali96/weight-height
# https://opendata.stackexchange.com/questions/7793/age-weight-and-height-dataset
# https://www.reddit.com/r/datasets/comments/1dii3t/data_sets_for_heights_and_weights/

# https://vincentarelbundock.github.io/Rdatasets/datasets.html

# a related tutorial
# http://goelhardik.github.io/2016/10/04/fishers-lda/

##
using Plots
##
# https://vincentarelbundock.github.io/Rdatasets/csv/carData/Davis.csv
using CSV
df = CSV.read(download("https://vincentarelbundock.github.io/Rdatasets/csv/carData/Davis.csv"))
X = Float64.([df[:weight] df[:height]])

##
g = df[:sex]

##
scatter(X[:,1],X[:,2])
xlabel!("Weight")
ylabel!("Height")

## Filter the data
goodpts = X[:,1] .<= 115
Xf = X[goodpts,:]
gf = g[goodpts]

##
theme(:dark)
scatter(Xf[:,1],Xf[:,2],color=2)
xlabel!("Weight")
ylabel!("Height")

## Save a picture
pyplot()
scatter(Xf[:,1],Xf[:,2],color=2,size=(300,300), legend=false)
xlabel!("Weight")
ylabel!("Height")
savefig("weight-height-1.pdf")

## Partition by classes
s1 = gf .== "F"
s2 = gf .== "M"
X1 = Xf[s1,:]
X2 = Xf[s2,:]
gr()
scatter(X1[:,1],X1[:,2])
scatter!(X2[:,1],X2[:,2])

## Save a picture
pyplot()
scatter(X1[:,1],X1[:,2],size=(300,300), legend=false)
scatter!(X2[:,1],X2[:,2])
savefig("weight-height-2.pdf")
##
using Statistics
m = mean(Xf,dims=1)
m1 = mean(X1,dims=1)
m2 = mean(X2,dims=1)
n1 = size(X1,1)
n2 = size(X2,1)

B = (m1-m)'*n1*(m1-m) + (m2-m)'*n2*(m2-m)
W = (X1.-m1)'*(X1.-m1) + (X2.-m2)'*(X2.-m2)

## The matrices B and W add up to the total covariance
using LinearAlgebra
C = (Xf.-m)'*(Xf.-m)
norm(B+W - C)

##
lams,V = eigen(B,W)

##
v = V[:,2]
p1 = histogram((X1)*v,nbins=10)
histogram!((X2)*v,nbins=10)
## Save
plot!(size=(300,300),legend=false)
savefig("weight-height-lda.pdf")
##

## Compare to PCA
lams,V = eigen(C)
v2 = V[:,2]
p2 = histogram(X1*v2,nbins=10)
histogram!(X2*v2,nbins=10)
plot(p1,p2)
##
## Save
plot(p2)
plot!(size=(300,300),legend=false)
savefig("weight-height-pca.pdf")
##
X1c = X1.-m
X2c = X2.-m
scatter(X1c[:,1],X1c[:,2], label="")
scatter!(X2c[:,1],X2c[:,2], label="")
Plots.abline!(v[2]/v[1],0,label="LDA")
Plots.abline!(v2[2]/v2[1],0,label="PCA")
##
plot!(size=(300,300),legend=false)
savefig("weight-height-lda-pca.pdf")


text:
# This file is a part of ROOTFramework.jl, licensed under the MIT License (MIT).

using Cxx

cxxinclude("TDirectory.h")
cxxinclude("TFile.h")

import Base: cd, show, write, open, close

export current_tdirectory
export path

export TDirectoryRef, TDirectoryPtr, TDirectoryInst
const TDirectoryRef = rcpp"TDirectory"
const TDirectoryPtr = pcpp"TDirectory"
const TDirectoryInst = Union{TDirectoryRef, TDirectoryPtr}

export TFile, TFilePtr, TFileInst
const TFile = cxxt"TFile"
const TFilePtr = pcpp"TFile"
const TFileInst = Union{TFile, TFilePtr}

export ATDirectoryInst
const ATDirectoryInst = Union{TDirectoryInst, TFileInst}

current_tdirectory() = begin
    lock(gROOTMutex()) do
        icxx""" (TDirectory*)TDirectory::CurrentDirectory(); """
    end
end

show(io::IO, x::TDirectoryRef) =
    print(io, "TDirectoryRef(\"$(escape_string(path(x)))\")")

show(io::IO, x::TDirectoryPtr) =
    print(io, "TDirectoryPtr(\"$(escape_string(path(x)))\")")

function Base.keys(tdir::ATDirectoryInst)
    tdir_ptr = as_pointer(tdir)
    names = icxx"""
        std::vector<std::string> names;

        TIter nextobj($tdir_ptr->GetList());
        TObject *obj = nullptr;
        while ((obj = (TObject *) nextobj())) {
            // std::cout << obj->GetName() << std::endl;
            names.push_back(obj->GetName());
        }

        TIter nextkey($tdir_ptr->GetListOfKeys());
        TKey *key = nullptr;
        while ((key = (TKey *) nextkey())) {
            // std::cout << key->GetName() << std::endl;
            names.push_back(key->GetName());
        }
        names;
    """

    unique(map(String, names))
end

Base.haskey(tdir::ATDirectoryInst, key::String) = !isnullptr(@cxx tdir->FindKey(pointer(key)))

function Base.getindex(tdir::ATDirectoryInst, key::String)
    ptr = @cxx tdir->Get(pointer(key))
    isnullptr(ptr) ? throw(KeyError(key)) : dyncast(ptr)
end

open(::Type{TFile}, fname::AbstractString, mode::AbstractString = "", ftitle::AbstractString = "", compress::Integer = 1) = begin
    lock(gROOTMutex()) do
        prev_dir = current_tdirectory()
        try
            @cxx TFile(pointer(fname), pointer(mode), pointer(ftitle), Int32(compress))
        finally
            cd(prev_dir)
        end
    end
end

show(io::IO, x::TFile) =
    print(io, "TFile(\"$(escape_string(path(x)))\")")

show(io::IO, x::TFilePtr) =
    print(io, "TFilePtr(\"$(escape_string(path(x)))\")")

write(tdir::ATDirectoryInst, obj::CppValue, name::AbstractString = "", option::Integer = 0, bufsize::Integer = 0) =
    write(tdir, pointer_to(obj), name, option, bufsize)

write(tdir::ATDirectoryInst, obj::CppPtr, name::AbstractString = "", option::Integer = 0, bufsize::Integer = 0) = begin
    cd(tdir) do
       @cxx obj->Write(isempty(name) ? Ptr{UInt8}(0) : pointer(name), Int32(option), Int32(bufsize))
   end
end

close(tfile::TFileInst) = begin
    tfile_ptr = as_pointer(tfile)
    icxx"""
        if ($tfile_ptr->IsWritable()) {
            $tfile_ptr->Write();
            $tfile_ptr->Close();
        }
    """
    nothing
end

cd(dir::ATDirectoryInst) = begin
    lock(gROOTMutex()) do
        @cxx dir->cd()
    end
end

cd(f::Function, dir::ATDirectoryInst) = begin
    prev_dir = current_tdirectory()
    try
        cd(dir)
        f()
    finally
        cd(prev_dir)
    end
end

path(dir::ATDirectoryInst) = unsafe_string(@cxx dir->GetPath())


text:
#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*-----------------------------------------------------------------------
#* Author: Tiago Os򱨯rio
#* $Id$
#*---------------------------------------------------------------------
type Lead_lag
	Lead_lag()=begin
		[
			:(diff(aux)= 0/"s"),
		],
		[
			"",
		],
		new(
			positive(Dict{Symbol,Any}(
				:Brief=>"model gain"
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"lead time constant",
				:Unit=>"s"
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"lag time constant",
				:Unit=>"s"
			)),
			DanaReal(Dict{Symbol,Any}(
				:Brief=>"input signal"
			)),
			DanaReal(Dict{Symbol,Any}(
				:Brief=>"output signal"
			)),
			DanaReal(Dict{Symbol,Any}(
				:Brief=>"internal variable"
			)),
			[
				:(alpha*diff(aux)=gain*input-aux),
				:(output=beta*diff(aux) + aux),
			],
			[
				"Calculate variable aux","Calculate output",
			],
			[:gain,:beta,:alpha,],
			[:input,:output,:aux,]
		)
	end
	gain::positive
	beta::positive
	alpha::positive
	input::DanaReal
	output::DanaReal
	aux::DanaReal
	equations::Array{Expr,1}
	equationNames::Array{String,1}
	initials::Array{Expr,1}
	initialNames::Array{String,1}
	parameters::Array{Symbol,1}
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export Lead_lag
function setEquationFlow(in::Lead_lag)
	addEquation(1)
	addEquation(2)
end
function initial(in::Lead_lag)
	addEquation(1)
end
function atributes(in::Lead_lag,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=false
	fields[:Icon]="icon/Lead_Lag"
	fields[:Brief]="Model Lead lag."
	fields[:Info]="== Inputs ==
* One input signal.

== Outputs ==
* One output signal.
"
	drive!(fields,_)
	return fields
end
Lead_lag(_::Dict{Symbol,Any})=begin
	newModel=Lead_lag()
	newModel.attributes=atributes(newModel,_)
	newModel
end


text:
mutable struct HMM
    dict::AhoCorasickAutomaton{Ti}
    words::Vector{String}
    usr_words::Int
    tags::Vector{String}
    pmp::Dict{String,Int}
    hpr::Vector{Tv}
    h2h::Matrix{Tv}
    h2v::Vector{Dict{Int, Tv}}
    INF::Vector{Tv}
    usrh2v::Vector{Dict{Int, Tv}}
    usrINF::Vector{Tv}
    usrw::Tv
end

function HMM(corpus, dictpath, usrw)
        ispath(dictpath) || @warn "dictpath is not a valid path"
        @assert usrw < 1.

        tags, words = KongYiji.postags(corpus), KongYiji.words(corpus)
        np, pmp = length(tags), KongYiji.dict(tags)

        hpr, h2h, h2v, INF = fill(Tv(0), np), fill(Tv(0), (np, np)), [Dict{Int, Tv}() for _ in 1:np], fill(Tv(0), np)
        usrh2v, usrINF = [Dict{Int, Tv}() for _ in 1:np], fill(Tv(0), np)
        for sent in sents(corpus)
                pp = 0
                for (ip, iw) in sent
                        if pp == 0 hpr[ip] += 1 else h2h[pp,ip] += 1 end
                        pp = ip
                        h2v[ip][iw] = get(h2v[ip], iw, 0) + 1
                end
        end

        wmp = KongYiji.dict(words)

        if ispath(dictpath)
                cd(() -> begin
                                for tag in tags
                                        if !ispath(tag); continue; end
                                        ip = get(pmp, tag, 0)
                                        for file in readdir(tag)
                                                for word in eachline(joinpath(tag, file))
                                                        iw = getid!(wmp, word)
                                                        if ip > 0; usrh2v[ip][iw] = get(usrh2v[ip], iw, 0) + 1; end
                                                end
                                        end
                                end
                        end, dictpath)
        end

        dict = AhoCorasickAutomaton{Ti}(wmp)
        usr_words = length(wmp) - length(words)
        words = KongYiji.vec(wmp)
        normalize!(hpr)
        normalize!(h2h)
        normalize!(h2v, INF)
        log!(hpr)
        log!(h2h)
        normalize!(usrh2v, usrINF)

        HMM(dict, words, usr_words, tags, pmp, hpr, h2h, h2v, INF, usrh2v, usrINF, usrw)
end

function (hmm::HMM)(xs::Vector{String}; recover=true, withpos=false)
        nc_max = mapreduce(ncodeunits, max, xs)
        np = length(hmm.hpr)
        @assert np > 0
        dp = fill(Tv(-Inf), (nc_max + 1, np))
        pr = fill((1, 0), (nc_max + 1, np))
        return [hmm(x, dp, pr, recover, withpos) for x in xs]
end

function (hmm::HMM)(x::String; recover=true, withpos=false)
        return hmm([x]; recover=recover, withpos=withpos)[1]
end

function (hmm::HMM)(x::String, dp::Matrix{Tv}, pr::Matrix{Tuple{Int, Int}}, recover, withpos)::Vector{String}
        origin = x
        x = mask(x)
        chrs = codeunits(x) #todo slow?
        vtxs = collect(eachmatch(hmm.dict, chrs))
        sort!(vtxs)
        nv, nc, np = length.((vtxs, chrs, hmm.hpr))
        for i = 1:nc + 1, j in 1:np dp[i,j] = -Inf end
        pv = 1
        dp[1, :] = hmm.hpr
        pre_i = 1
        for i in 1:nc + 1
                if dp[i,1] != -Inf pre_i = i end
                while pv <= nv && vtxs[pv].s < i pv = pv + 1 end
                if !(pv <= nv && vtxs[pv].s == i || i == nc + 1) continue end
                if dp[i, 1] == -Inf
                        #@show i, pre_i
                        for pi = 1:np, pj = 1:np
                                maybe = dp[pre_i, pj] + hmm.h2h[pj, pi] + logph2v(hmm, pj, 0)
                                if maybe > dp[i, pi] dp[i, pi] = maybe; pr[i, pi] = (pre_i, pj) end
                        end
                        #pre_i = i
                end
                while pv <= nv && vtxs[pv].s == i
                        vtx = vtxs[pv]
                        j = i + length(vtx)
                        for pi = 1:np
                                for pj = 1:np
                                        maybe = dp[i, pi] + hmm.h2h[pi, pj] + logph2v(hmm, pi, vtx.i)
                                        if maybe > dp[j,pj] dp[j,pj] = maybe; pr[j,pj] = (i, pi) end
                                end
                        end
                        pv = pv + 1
                end
        end
        #@show hmm
        #@show maximum(dp[nc + 1, :])
        #@show dp'
        #@show pr'
        #@show vtxs
        v = (nc + 1, argmax(dp[nc + 1,:]))
        words = fill("", 0)
        while v[1] != 1
                pv = pr[v[1],v[2]]
                word = x[pv[1]:prevind(x, v[1])]
                push!(words, word)
                v = pv
        end
        reverse!(words)
        if recover; words = demask(words, origin); end
        ret = words

        if withpos
                v = (nc + 1, argmax(dp[nc + 1,:]))
                postags = fill("", 0)
                while v[1] != 1
                        pv = pr[v[1],v[2]]
                        postag = hmm.tags[pv[2]]
                        push!(postags, postag)
                        v = pv
                end
                reverse!(postags)
                ret = [join(p, "/") for p in zip(words, postags)]
        end

        return ret
end

logph2v(hmm::HMM, p, v) = ((1. - hmm.usrw) * get(hmm.h2v[p], v, hmm.INF[p]) + hmm.usrw * get(hmm.usrh2v[p], v, hmm.usrINF[p])) |> log

function ==(a::HMM, b::HMM)
        return all(fname -> getfield(a, fname) == getfield(b, fname), fieldnames(HMM))
end


text:
include("sparsestcut.jl")

##
mat = matread("Netscience_xy.mat")
A = mat["A"]

n = size(A, 1)
T = Int(ceil(10 * log(n)^2))

min_S, min_expansion = cut_matching(A, T)
@show min_S
@show min_expansion
@show compute_expansion(A, min_S)
##
using MatrixMarket

A = MatrixMarket.mmread("ca-AstroPh.mtx")
n = size(A, 1)
T = Int(ceil(10 * log(n)^2))

min_S, min_expansion = cut_matching(A, T)
@show min_S
@show min_expansion
@show compute_expansion(A, min_S)
##

max_stars_repo_name: luotuoqingshan/Sparsest-Cut

max_issues_repo_name: luotuoqingshan/Sparsest-Cut

max_forks_repo_name: luotuoqingshan/Sparsest-Cut


text:
module SchedWCTGraph

mutable struct SchedWCTinstance
    m::Int64 # number of machines
    weight::Vector{Int64}
    procTime::Vector{Int64}
    releaseTime::Vector{Int64}
    deadline::Vector{Int64}
end

SchedWCTinstance() = SchedWCTinstance(0,[],[],[],[])

function readSchedWCTInstance(fileName::String, m::Int64)
    instance = SchedWCTinstance()
    open(fileName) do f
        allFile = read(f, String)
        splitFile = split(allFile)
        #println(splitFile)
        n = parse(Int64,splitFile[2])
        weight = zeros(Int64, n)
        procTime = zeros(Int64, n)
        releaseTime = zeros(Int64, n)
        deadline = typemax(Int64)*ones(Int64,n)
        for i=1:n
            #println("i=$i")
            weight[i] = parse(Int64,splitFile[1+3*i])
            procTime[i] = parse(Int64,splitFile[2+3*i])
        end
        instance = SchedWCTinstance(m,weight, procTime, releaseTime, deadline)
    end
    return instance
end

function constructSchedWCTGraph(instance::SchedWCTinstance)
    n = length(instance.weight)
    m = instance.m
    itemOrder = sort(collect(1:n), by = i -> (instance.weight[i] / instance.procTime[i]), rev=true)
    #itemOrder = sort(collect(1:n))
    println("itemOrder: ", itemOrder)
    for i in itemOrder
        println("$i: weight: $(instance.weight[i]), processing time: $(instance.procTime[i]), ratio: $(instance.weight[i] / instance.procTime[i])")
    end
    # compute latest completion time based on property #3 from A Branch-and-Price Algorithm for Parallel Machine
    # Scheduling Using ZDDs and Generic Branching (Daniel Kowalczyk, Roel Leus)
    HMax = sum(instance.procTime) / m + (m-1)/m*maximum(instance.procTime)
    println("HMax = $HMax")
end

end

#SchedWCTGraph.readSchedWCTInstance("Data/scheduling_WCT/instances_wct/20_3/instance1_20_3_0.txt")
#SchedWCTGraph.constructSchedWCTGraph(x)


text:
using Plots, Statistics, Distributions
initial=[]
NList=[1000,10000,100000,1000000]
final_res=[]
for N in NList
    res=[]
    for i in 1:N
        num=rand(0:9)
        push!(initial, num)
    end
    for i in 2:N
        if initial[i-1]==4
            push!(res, initial[i])
        end
    end
    push!(final_res, res)
end
p1=histogram(final_res[1], bins=15, alpha=0.5, xlabel="The Random Number", ylabel="Number",c = :red, label="N=1000")
p2=histogram(final_res[2], bins=15, alpha=0.5, xlabel="The Random Number", ylabel="Number",c = :green,label="N=10000")
p3=histogram(final_res[3], bins=15, alpha=0.5, xlabel="The Random Number", ylabel="Number",c = :yellow,label="N=100000")
p4=histogram(final_res[4], bins=15, alpha=0.5, xlabel="The Random Number", ylabel="Number",c = :blue,label="N=1000000")
plot(p1,p2,p3,p4, layout=4,dpi=600,guidefontsize=9, legendfont=6, legend = :outertop)
savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter6\\6.2\\Figs\\TheFrequencyCruve.png")


text:
# this script is used to generate the ascii plots used in the window docstrings

using UnicodePlots
using DSP

open(joinpath(@__DIR__, "src", "winplots.jl"), "w") do io

    println(io, """
    #####################################################
    # THIS FILE IS AUTO-GENERATED
    # to re-generate use the `gen_winplots.jl` script
    #####################################################
    """)

    for (winfunc, args) in [
        (rect, ()),
        (hanning, ()),
        (hamming, ()),
        (tukey, (0.4)),
        (cosine, ()),
        (lanczos, ()),
        (triang, ()),
        (bartlett, ()),
        (gaussian, (0.2)),
        (bartlett_hann, ()),
        (blackman, ()),
        (kaiser, (3)),
        (dpss, (2, 1)),
        ]

        ymax = winfunc == dpss ? 0.2 : 1.0
        n = winfunc == triang ? 7 : 69

        fname = split(string(winfunc), ".")[end]
        println(io, "const $(fname)_winplot = padplot(\"\"\"")
        # convert Nx1 matrices to vectors with [:] - necessary for dpss
        print(io, lineplot(winfunc(n, args...)[:], ylim=[0,ymax], xlim=[1,n], width=70, canvas=BlockCanvas))
        println(io, "\"\"\")\n")
    end
end


text:
using SparseArrays: rowvals, nzrange, spzeros, spdiagm, findnz, rowvals, nonzeros
using FunctionWrappers: FunctionWrapper
using PolynomialRoots: roots
using QuantumOptics: SparseOperator, embed

export hamiltonian

"""
    hamiltonian(
            T::Trap; timescale::Real=1e-6, lamb_dicke_order::Union{Vector{Int},Int}=1,
            rwa_cutoff::Real=Inf, displacement="truncated", time_dependent_eta=false
        )
Constructs the Hamiltonian for `T` as a function of time. Return type is a function
`h(t::Real, ψ)` that, itself, returns a `QuantumOptics.SparseOperator`.

**args**
* `timescale`: e.g. a value of 1e-6 will take time to be in ``\\mu s``
* `lamb_dicke_order`: Only consider terms that change the phonon number by up to this value.
    If this is an `Int`, then the cutoff is applied to all modes. If this is a `Vector{Int}`,
    then `lamb_dicke_order[i]` is applied to the iᵗʰ mode, according to the order in
    `T.basis`.
    Note: this isn't quite the same thing as the Lamb-Dicke approximation since setting
    `lamb_dicke_order=1` will retain, for example, terms proportional to ``a^\\dagger a ``.
* `rwa_cutoff`: drop terms in the Hamiltonian that oscillate faster than this cutoff.
* `displacement`: This can be either `"truncated"`(default) or `"analytic"`.

   When an atom is irradiated, both the atom's energy and its momentum will generally be
   affected. For an atom in a harmonic potential, the exchange of momentum can be modeled as
   a displacement operation ``D(α=iηe^{-iνt}) = exp[αa^† - α^*a]``, where ``η`` is the
   Lamb-Dicke parameter, which can be described equivalently as either being proportional to
   the square root of the ratio of the recoil frequency with the ground state energy of the
   atom's motion or as the ratio of the spread of the ground state wavefunction to the
   wavelength of the laser.

   When `"truncated"` is selected, the matrix elements of ``D(α)`` are computed by
   constructing ``α^* a, αa^†`` in a truncated basis (according to the dimension specified in
   your model) and then exponentiating their difference. This has the advantage, amongst
   other things, of guaranting unitarity.

   If `"analytic"` is selected, then the matrix elements are computed assuming an infinite-
   dimensional Hilbert space.

   For small displacements (``η ≪ N``, where ``N`` is the dimension of the motion's Hilbert
   space), both of these methods will be good approximations.
* `time_dependent_eta::Bool`: In addition to impacting the vibrational subspace directly, a
   change in the trap frequency, ``δν``, will also change the Lamb-Dicke parameter. Since
   typically ``δν≪ν``, this effect will be small ``η ≈ η₀(1 + δν/2ν)`` and doesn't warrant
   the additional computational resources needed to calculate and update it in time. In this
   case, we can set `time_dependent_eta=false` (default), which will set ``η(t) = η₀``.

"""
function hamiltonian(
        T::Trap; timescale::Real=1e-6, lamb_dicke_order::Union{Vector{Int},Int}=1,
        rwa_cutoff::Real=Inf, displacement::String="truncated", time_dependent_eta::Bool=false
    )
    hamiltonian(
            T, T.configuration, timescale, lamb_dicke_order, rwa_cutoff, displacement,
            time_dependent_eta
        )
end

#############################################################################################
# Hamiltonian for a linear configuration of ions
#############################################################################################

# At each time step, this function updates in-place the 2D array describing the full system
# Hamiltonian.
function hamiltonian(
        T::Trap, configuration::LinearChain, timescale::Real,
        lamb_dicke_order::Union{Vector{Int},Int}, rwa_cutoff::Real, displacement::String,
        time_dependent_eta::Bool
    )
    b, indxs, cindxs = _setup_base_hamiltonian(
            T, timescale, lamb_dicke_order, rwa_cutoff, displacement, time_dependent_eta
        )
    aui, gbi, gbs, bfunc, δνi, δνfuncs = _setup_fluctuation_hamiltonian(T, timescale)
    S = SparseOperator(get_basis(T))
    function f(t, ψ)  # a two argument function is required in the QuantumOptics solvers
        @inbounds begin
            @simd for i in 1:length(indxs)
                bt_i, conj_bt_i = b[i](t)::Tuple{ComplexF64,ComplexF64}
                @simd for j in 1:length(indxs[i])
                    i1, i2 = indxs[i][j]
                    S.data[i1, i2] = bt_i
                    if i1 != i2
                        S.data[i2, i1] = conj(bt_i)
                        if length(cindxs[i]) != 0
                            flag = cindxs[i][1][1]
                            i3, i4 = cindxs[i][j+1]
                            if flag == -1
                                S.data[i3, i4] = -conj_bt_i
                                S.data[i4, i3] = -conj(conj_bt_i)
                            else
                                S.data[i3, i4] = conj_bt_i
                                S.data[i4, i3] = conj(conj_bt_i)
                            end
                        end
                    end
                end
            end
            if length(gbi) == 0 && length(δνi) == 0
                return S
            else
                @simd for indx in aui
                    S.data[indx, indx] = complex(0.)
                end
                @simd for i in 1:length(gbi)
                    zeeman_t = bfunc(t)::Float64
                    @simd for j in 1:length(gbi[i])
                        indx = gbi[i][j]
                        S.data[indx, indx] += zeeman_t * gbs[i][j]
                    end
                end
                @simd for i in 1:length(δνi)
                    δν_t = δνfuncs[i](t)::Float64
                    @simd for n in 1:length(δνi[i])
                        @simd for indx in δνi[i][n]
                            S.data[indx, indx] += n * δν_t
                        end
                    end
                end
            end
        end
        return S
    end
    return f
end

#=
The purpose of the hamiltonian function is to evaluate a vector of time-dependent functions
and use the returned values to update, in-place, a pre-allocated array.

The pre-allocated array holds the full Hamiltonian -- a tensor product defined over all of
the individual ion and vibrational mode subspaces -- at a particular point in time.

However, we don't know a priori the exact form of the Hilbert space or the details of the
Hamiltonian's time dependence, since it will be defined by the user.

The _setup_hamiltonian function extracts this user-defined information from a <:Trap struct
and converts it into two vector of vectors of indices, corresponding to redundant (see note
below) matrix elements of the Hamiltonian, and a matched vector of time-dependent functions
for updating these elements.

------- Note -------
Since the terms of the Hamiltonian will always be of the form of a single ion operator
tensored with a single vibrational mode operator, there will be a lot of redundancy in the
Hamiltonian's matrix elements. E.g.

                                     [ σ₊ ⊗ D(α(t))      0         ]
             H = 𝐼 ⊗ σ₊ ⊗ D(α(t)) =  [       0        σ₊ ⊗ D(α(t)) ]

So to avoid unnecessarily evaluating functions more than once, _setup_hamiltonian also
returns a vector of vectors of indices that keep track of this redundancy.

Also, we have: <m|D(α)|n> = (-1)^(n-m) × conjugate(<n|D(α)|m>). We keep track of this in an
additional vector of vectors of indices.

Finally, since we require the Hamiltonian to be Hermitian, h[i, j] = conj(h[j, i]), this
function does not keeps track of only one of these pairs.
=#
function _setup_base_hamiltonian(
        T, timescale, lamb_dicke_order, rwa_cutoff, displacement, time_dependent_eta
    )
    rwa_cutoff *= timescale
    modes = reverse(get_vibrational_modes(T.configuration))
    L = length(modes)
    νlist = Tuple([mode.ν for mode in modes])
    mode_dims = [mode.N+1 for mode in modes]

    ions = reverse(T.configuration.ions)
    Q = prod([ion.shape[1] for ion in ions])
    ion_arrays = [spdiagm(0 => [true for _ in 1:ion.shape[1]]) for ion in ions]

    ηm, Δm, Ωm = _ηmatrix(T), _Δmatrix(T, timescale), _Ωmatrix(T, timescale)
    lamb_dicke_order = _check_lamb_dicke_order(lamb_dicke_order, L)
    ld_array, rows, vals = _ld_array(mode_dims, lamb_dicke_order, νlist, timescale)
    if displacement == "truncated" && time_dependent_eta
        rootlist = map(x->real.(roots(_He(x))), mode_dims)
    end

    indxs_dict = Dict()
    repeated_indices = Vector{Vector{Tuple{Int64,Int64}}}(undef, 0)
    conj_repeated_indices = Vector{Vector{Tuple{Int64,Int64}}}(undef, 0)
    functions = FunctionWrapper[]
    work_eta = zeros(Float64, L)
    local ts, ion_rows, ion_cols, ion_idxs, ion_reps, rn

    # iterate over ions and lasers
    for n in eachindex(ions), m in eachindex(T.lasers)
        if m ≡ 1
            rn = length(ions) - n + 1
            ts = _transitions(ions[n])
            C = sum([i*real.(sigma(ions[n], reverse(ts[i])...).data) for i in 1:length(ts)])
            if length(ions) == 1
                K = C
            else
                K = kron(ion_arrays[1:n-1]..., C, ion_arrays[n+1:length(ions)]...)
            end
            ion_rows, ion_cols, ion_vals = findnz(K)
            ion_idxs = sortperm(real.(ion_vals))
            ion_reps = Int(Q / size(C, 1))
        end
        ηnm = view(ηm, rn, m, :)
        function ηlist(t)
            for i in 1:L
                η = ηnm[i]
                typeof(η) <: Number ? work_eta[i] = ηnm[i] : work_eta[i] = ηnm[i](t)
            end
            work_eta
        end
        if displacement == "truncated" && !time_dependent_eta
            D_arrays = []
            for (i, mode) in enumerate(modes)
                push!(D_arrays, real.(displace(mode, ηlist(0)[i]).data))
            end
        end

        # iterate over ion-laser transitions
        for (ti, tr) in enumerate(ts)
            Δ, Ω = Δm[rn, m][ti], Ωm[rn, m][ti]
            Δ_2π = Δ/2π
            typeof(Ω) <: Number && continue  # e.g. the laser doesn't shine on this ion
            locs = view(ion_idxs, (ti - 1) * ion_reps + 1 : ti * ion_reps)
            for j in 1:prod(mode_dims); for i in nzrange(ld_array, j)
                ri = rows[i]
                ri<j && continue
                cf = vals[i]
                pflag = abs(Δ_2π + cf) > rwa_cutoff
                nflag = abs(Δ_2π - cf) > rwa_cutoff
                (pflag && nflag) && continue
                rev_indxs = false
                idxs = _inv_get_kron_indxs((rows[i], j), mode_dims)
                for l in 1:L
                    (idxs[1][l] ≠ idxs[2][l] && typeof(ηnm[l]) <: Number) && @goto cl
                end
                s_ri = []
                s_cri = []
                for loc in locs
                    if !pflag
                        push!(s_ri, (Q*(ri - 1) + ion_rows[loc], Q*(j - 1) + ion_cols[loc]))
                        if !nflag && ri≠j
                            push!(s_cri,
                                    (Q*(j - 1) + ion_rows[loc], Q*(ri - 1) + ion_cols[loc]))
                        end
                    elseif !nflag
                        push!(s_ri, (Q*(j - 1) + ion_rows[loc], Q*(ri - 1) + ion_cols[loc]))
                        rev_indxs = true
                    end
                end
                if rev_indxs
                    idxs = reverse(idxs)
                end
                if length(s_cri) > 0
                    parity = sum(map(x -> isodd(abs(x[1] - x[2])), zip(idxs...)))
                    pushfirst!(s_cri, (-1 * isodd(parity), 0))
                end

                # push information to top-level lists/ construct time-dep function
                if displacement == "truncated" && !time_dependent_eta
                    D = Tuple([D_arrays[i][idxs[1][i], idxs[2][i]] for i in 1:L])
                elseif displacement == "analytic" && !time_dependent_eta
                    D = Tuple([_Dnm_cnst_eta(ηlist(0)[i], idxs[1][i], idxs[2][i]) for i in 1:L])
                elseif displacement == "truncated"
                    pflist = [_pf(mode_dims[i], idxs[1][i], idxs[2][i]) for i in 1:L]
                end
                row, col = s_ri[1]
                if haskey(indxs_dict, s_ri[1])
                    # this will happen when multiple lasers address the same transition
                    functions[indxs_dict[row, col]] =
                        let
                            a = functions[indxs_dict[row, col]]
                            if !time_dependent_eta
                                FunctionWrapper{Tuple{ComplexF64,ComplexF64},Tuple{Float64}}(
                                    t -> a(t) .+ _D_cnst_eta(
                                        Ω(t), Δ, νlist, timescale, idxs, D, t, L))
                            elseif displacement == "analytic"
                                FunctionWrapper{Tuple{ComplexF64,ComplexF64},Tuple{Float64}}(
                                    t -> a(t) .+ _D(
                                        Ω(t), Δ, ηlist(t), νlist, timescale, idxs, t, L))
                            elseif displacement == "truncated"
                                FunctionWrapper{Tuple{ComplexF64,ComplexF64},Tuple{Float64}}(
                                    t -> a(t) .+ _Dtrunc(
                                        Ω(t), Δ, ηlist(t), νlist, rootlist, mode_dims, idxs,
                                        pflist, timescale, L, t))
                            end
                        end
                else
                    if !time_dependent_eta
                        f = FunctionWrapper{Tuple{ComplexF64,ComplexF64},Tuple{Float64}}(
                                t -> _D_cnst_eta(Ω(t), Δ, νlist, timescale, idxs, D, t, L))
                    elseif displacement == "analytic"
                        f = FunctionWrapper{Tuple{ComplexF64,ComplexF64},Tuple{Float64}}(
                                t -> _D(Ω(t), Δ, ηlist(t), νlist, timescale, idxs, t, L))
                    elseif displacement == "truncated"
                        f = FunctionWrapper{Tuple{ComplexF64,ComplexF64},Tuple{Float64}}(
                                    t -> _Dtrunc(
                                        Ω(t), Δ, ηlist(t), νlist, rootlist, mode_dims, idxs,
                                        pflist, timescale, L, t))
                    end
                    push!(functions, f)
                    push!(repeated_indices, s_ri)
                    push!(conj_repeated_indices, s_cri)
                    indxs_dict[row, col] = length(repeated_indices)
                end
                @label cl
            end; end
        end
    end
    functions, repeated_indices, conj_repeated_indices
end

# δν(t) × aᵀa terms for Hamiltonian. This function returns an array of functions
# δν_functions = [2π×ν.δν(t)×timescale for ν in modes]. It also returns an array of arrays
# of arrays of indices, δν_indices, such that δν_indices[i][j] lists all diagonal elements
# of the full 2D system matrix upon which have been mapped the jth diagonal element of the
# ith mode.
function _setup_δν_hamiltonian(T, timescale)
    N = length(T.configuration.ions)
    modes = get_vibrational_modes(T.configuration)
    δν_indices = Vector{Vector{Vector{Int64}}}(undef, 0)
    δν_functions = FunctionWrapper[]
    τ = timescale
    for l in eachindex(modes)
        δν = modes[l].δν
        mode = modes[l]
        (mode._cnst_δν && δν(0) == 0) && continue
        push!(δν_functions, FunctionWrapper{Float64,Tuple{Float64}}(
                    t -> @fastmath 2π * δν(t) * τ
                ))
        δν_indices_l = Vector{Vector{Int64}}(undef, 0)
        mode_op = number(mode)
        A = embed(get_basis(T), [N+l], [mode_op]).data
        mode_dim = mode.shape[1]
        for i in 1:mode_dim-1
            indices = [x[1] for x in getfield.(findall(x->x.==complex(i, 0), A), :I)]
            push!(δν_indices_l, indices)
        end
        push!(δν_indices, δν_indices_l)
    end
    δν_indices, δν_functions
end

# Hamiltonian terms for global Bfield fluctuations encoded in T.δB. If T.δB=0, this function
# returns a collection of empty arrays. Otherwise it iterates over the selected levels for
# each ion and creates an array (global_B_scales), which encodes the magnetic field
# susceptibility of each level. It also returns an array of indices (global_B_indices), that
# keeps track of the indices of the full 2D array that represents the tensored system,
# corresponding to the energy of each level. Finally it returns a function (bfunc) encoding
# the time-dependence of δB. When the system is integrated, the Hamiltonian terms will be
# updated at each time step by by bfunc(t) times the individual susceptibilities.
function _setup_global_B_hamiltonian(T, timescale)
    ions = T.configuration.ions
    global_B_indices = Vector{Vector{Int64}}(undef, 0)
    global_B_scales = Vector{Float64}(undef, 0)
    δB = T.δB
    τ = timescale
    bfunc = FunctionWrapper{Float64,Tuple{Float64}}(t -> 2π * δB(t * τ))
    if T._cnst_δB && δB(0) == 0
        return global_B_indices, global_B_scales, bfunc
    end
    for n in eachindex(ions)
        transitions_list = []
        for (ti, tr) in enumerate(_transitions(ions[n]))
            for level in tr
                (level in transitions_list) && continue
                push!(transitions_list, level)
                ion_op = sigma(ions[n], level)
                A = embed(get_basis(T), [n], [ion_op]).data
                indices = [x[1] for x in getfield.(findall(x->x.==complex(1, 0), A), :I)]
                push!(global_B_indices, indices)
                # zeeman_shift(1, ions[n].level_structure[level]) is the Zeeman shift of
                # ions[n].level_structure[level] in units of δB.
                push!(global_B_scales, τ * zeeman_shift(1, ions[n].level_structure[level]))
            end
        end
    end
    global_B_indices, global_B_scales, bfunc
end

# This mostly just strings together the results from _setup_global_B_hamiltonian and
# _setup_δν_hamiltonian for use in the hamiltonian function. The one additional task
# performed is the creation of an array of indices (all_unique_indices), which keeps track
# of all the diagonal indices affected by δν and/or δB, which is useful in hamiltonian().
function _setup_fluctuation_hamiltonian(T, timescale)
    gbi, gbs, bfunc = _setup_global_B_hamiltonian(T, timescale)
    δνi, δνfuncs = _setup_δν_hamiltonian(T, timescale)
    all_unique_indices = convert(Vector{Int64}, _flattenall(unique([gbi; δνi])))
    all_unique_indices, gbi, gbs, bfunc, δνi, δνfuncs
end

#############################################################################################
# internal functions
#############################################################################################

# https://gist.github.com/ivirshup/e9148f01663278ca4972d8a2d9715f72
function _flattenall(a::AbstractArray)
    while any(x -> typeof(x)<:AbstractArray, a)
        a = collect(Iterators.flatten(a))
    end
    a
end

# A 3D array of Lamb-Dicke parameters for each combination of ion, laser and mode. Modes are
# populated in reverse order.
function _ηmatrix(T)
    ions = T.configuration.ions
    vms = get_vibrational_modes(T.configuration)
    lasers = T.lasers
    (N, M, L) = map(x -> length(x), [ions, lasers, vms])
    ηnml = Array{Any}(undef, N, M, L)
    for n in 1:N, m in 1:M, l in 1:L
        δν = vms[l].δν
        ν = vms[l].ν
        eta = get_η(vms[l], lasers[m], ions[n], scaled=true)
        if eta == 0
            ηnml[n,m,L-l+1] = 0
        else
            ηnml[n,m,L-l+1] = FunctionWrapper{Float64,Tuple{Float64}}(t->eta / √(ν + δν(t)))
        end
    end
    ηnml
end

# Returns an array of vectors. The rows and columns of the array refer to ions and lasers,
# respectively. For each row/column we have a vector of detunings from the laser frequency
# for each ion transition. We need to separate this calculation from _Ωmatrix to implement
# RWA easily.
function _Δmatrix(T, timescale)
    ions = T.configuration.ions
    lasers = T.lasers
    (N, M) = length(ions), length(lasers)
    B = T.B
    ∇B = T.∇B
    Δnmkj = Array{Vector}(undef, N, M)
    for n in 1:N, m in 1:M
        transitions = _transitions(ions[n])
        Btot = B + ∇B * ions[n].position
        v = Vector{Float64}(undef, 0)
        for t in transitions
            L1 = ions[n].selected_level_structure[t[1]]
            L2 = ions[n].selected_level_structure[t[2]]
            stark_shift = ions[n].stark_shift[t[1]] - ions[n].stark_shift[t[2]]
            ωa = (abs(L1.E  + zeeman_shift(Btot, L1) - (L2.E + zeeman_shift(Btot, L2)))
                  + stark_shift)
            push!(v, 2π * timescale * ((c / lasers[m].λ) + lasers[m].Δ - ωa))
        end
        Δnmkj[n, m] = v
    end
    Δnmkj
end

# Returns an array of vectors. the rows and columns of the array refer to ions and lasers,
# respectively. For each row/column we have a vector of coupling strengths between the laser
# and all allowed electronic ion transitions.
function _Ωmatrix(T, timescale)
    ions = T.configuration.ions
    lasers = T.lasers
    (N, M) = length(ions), length(lasers)
    Ωnmkj = Array{Vector}(undef, N, M)
    for n in 1:N, m in 1:M
        E = lasers[m].E
        phase = lasers[m].ϕ
        (γ, ϕ) = map(x -> rad2deg(acos(ndot(T.Bhat, x))), [lasers[m].ϵ, lasers[m].k])
        transitions = _transitions(ions[n])
        s_indx = findall(x -> x[1] == n, lasers[m].pointing)
        if length(s_indx) == 0
            Ωnmkj[n, m] = [0 for _ in 1:length(transitions)]
            continue
        else
            s = lasers[m].pointing[s_indx[1]][2]
        end
        v = []
        for t in transitions
            Ω0 = 2π * timescale * s * ions[n].selected_matrix_elements[t](1.0, γ, ϕ) / 2.0
            if Ω0 == 0
                push!(v, 0)
            else
                push!(v, FunctionWrapper{ComplexF64,Tuple{Float64}}(
                        t -> Ω0 * E(t) * exp(-im * phase(t)))
                    )
            end
        end
        Ωnmkj[n, m] = v
    end
    Ωnmkj
end

_transitions(ion) = collect(keys(ion.selected_matrix_elements))

# Returns a tuple correpsonding to: [σ₊(t)]_ij ⋅ [D(ξ(t))]_ij, [σ₊(t)]_ji ⋅ [D(ξ(t))]_ji.
# [D(ξ(t))]_ij is calculated assuming an infinite dimensional Hilbert space for the HO.
function _D(Ω, Δ, η, ν, timescale, n, t, L)
    d = complex(1, 0)
    for i in 1:L
        d *= _Dnm(1im * η[i] * exp(im * 2π * ν[i] * timescale * t), n[1][i], n[2][i])
    end
    g = Ω * exp(-1im * t * Δ)
    g * d, g * conj(d)
end

# Returns a tuple correpsonding to: [σ₊(t)]_ij ⋅ [D(ξ(t))]_ij, [σ₊(t)]_ji ⋅ [D(ξ(t))]_ji.
# [D(ξ(t))]_ij is calculated assuming an infinite dimensional Hilbert space for the HO.
# As opposed to _D, in this case, we assume η(t) = η₀, which allows us to precompute _Dnm.
# This precomputation is performed externally to the function and fed in as the argument `D`.
function _D_cnst_eta(Ω, Δ, ν, timescale, n, D, t, L)
    d = complex(1, 0)
    for i in 1:L
        d *= D[i] * exp(1im * (n[1][i] - n[2][i]) * (2π * ν[i] * timescale * t + π/2))
    end
    g = Ω * exp(-1im * t * Δ)
    g * d, g * conj(d)
end

# Consider: T = X₁ ⊗ X₂ ⊗ ... ⊗ X_n (Xᵢ ∈ ℝ{dims[i]×dims[i]}), and indices:
# indxs[1], indxs[2], ..., indsx[N] = (i1, j1), (i2, j2), ..., (iN, jN).
# This function returns (k, l) such that: T[k, l] = X₁[i1, j1] * X₂[i2, j2] *...* X_N[iN, jN]
function _get_kron_indxs(indxs::Vector{Tuple{Int64,Int64}}, dims::Vector{Int64})
    L = length(indxs)
    rowcol = Int64[0, 0]
    @simd for i in 0:L-1
        if i == 0
            @inbounds rowcol .+= indxs[L-i]
        else
            @inbounds rowcol .+= (indxs[L-i] .- 1) .* prod(view(dims, 1:i))
        end
    end
    rowcol
end

# The inverse of _get_kron_indxs. If T = X₁ ⊗ X₂ ⊗ X₃ and X₁, X₂, X₃ are M×M, N×N and L×L
# dimension matrices, then we should input dims=(M, N, L).
function _inv_get_kron_indxs(indxs, dims)
    row, col = indxs
    N = length(dims)
    ret_rows = Array{Int64}(undef, N)
    ret_cols = Array{Int64}(undef, N)
    for i in 1:N
        tensor_N = prod(dims[i:N])
        M = tensor_N ÷ dims[i]
        rowflag = false
        colflag = false
        for j in 1:dims[i]
            jM = j * M
            if !rowflag && row <= jM
                @inbounds ret_rows[i] = j
                row -= jM - M
                rowflag = true
            end
            if !colflag && col <= jM
                @inbounds ret_cols[i] = j
                col -= jM - M
                colflag = true
            end
            rowflag && colflag && break
        end
    end
    Tuple(ret_rows), Tuple(ret_cols)
end

# similar to _Dnm, but meant to be used when η is assumed constant in ξ=iηe^(i2πνt)
function _Dnm_cnst_eta(ξ::Number, n::Int, m::Int)
    n < m && return _Dnm_cnst_eta(ξ, m, n) * (-1)^isodd(abs(n-m))
    n -= 1; m -= 1
    s = 1.0
    for i in m+1:n
        s *= i
    end
    ret = sqrt(1 / s) * ξ^(n-m) * exp(-abs2(ξ) / 2.0) * _alaguerre(abs2(ξ), m, n-m)
    isnan(ret) && return 1.0 * (n==m)
    ret
end

# If lamb_dicke_order is <: Int, this constructs a constant vector with this value of length
# L (i.e. same lamb_dicke_order for all modes). Otherwise lamb_dicke_order is reversed and
# returned.
function _check_lamb_dicke_order(lamb_dicke_order, L)
    if typeof(lamb_dicke_order) <: Int
        return [lamb_dicke_order for _ in 1:L]
    else
        @assert(length(lamb_dicke_order) == L,
            "if typeof(lamb_dicke_order)<:Vector, then length of lamb_dicke_order must ",
            "equal number of modes"
        )
        reverse(lamb_dicke_order)
    end
end

function _ld_array(mode_dims, lamb_dicke_order, νlist, timescale)
    a = [spzeros(Float16, d, d) for d in mode_dims]
    @inbounds for (i, d) in enumerate(mode_dims)
        for k in 1:d, l in 1:k
            if k-l <= lamb_dicke_order[i]
                val = (l - k) * νlist[i] * timescale
                a[i][k, l] = exp(val)
                l ≠ k && @inbounds a[i][l, k] = exp(-val)
            end
        end
    end
    length(a)==1 ? ld_array = a[1] : ld_array = kron(a...)
    ld_array, rowvals(ld_array), log.(nonzeros(ld_array))
end


text:
# Can't use replace(Array, Pair) because of #5
function replaceshape(A::Node, pair::Pair)
    dualspaces = [i == first(pair)' ? last(pair)' : i for i in A.shape]
    return tuple([i == first(pair) ? last(pair) : i for i in dualspaces]...)
end

Assignment = Union{Pair{VariableTensor{rank}, T},
    Pair{T1,T2},
    Pair{FreeIndex{T1}, FixedIndex{T1}}} where {rank, T<:AbstractTensor{rank}, T1 <: AbstractVectorSpace, T2 <: AbstractVectorSpace}

function _assign(visited, A::Tensor{T, rank}, pair::Pair{T1, T2}) where {T, rank, T1<:AbstractVectorSpace, T2<:AbstractVectorSpace}
    return Tensor(A.value, replaceshape(A, pair)...)
end

function _assign(visited, A::ConstantTensor{rank}, pair::Pair{T1, T2}) where {rank, T1<:AbstractVectorSpace, T2<:AbstractVectorSpace}
    return ConstantTensor(A.value, replaceshape(A, pair)...)
end

function _assign(visited, A::ZeroTensor{rank}, pair::Pair{T1, T2}) where {rank, T1<:AbstractVectorSpace, T2<:AbstractVectorSpace}
    return ZeroTensor(replaceshape(A, pair)...)
end

function _assign(visited, A::DeltaTensor{rank}, pair::Pair{T1, T2}) where {rank, T1<:AbstractVectorSpace, T2<:AbstractVectorSpace}
    return DeltaTensor(replaceshape(A, pair)...)
end

function _assign(visited, A::VariableTensor{rank}, pair::Pair{T1, T2}) where {rank, T1<:AbstractVectorSpace, T2<:AbstractVectorSpace}
    return VariableTensor(A.name, replaceshape(A, pair)...)
end

function _assign(visited, indices::Indices, pair::Pair{T1, T2}) where {T1<:AbstractVectorSpace, T2<:AbstractVectorSpace}
    return Indices([i.V == first(pair) ? updatevectorspace(i, last(pair)) : i for i in indices.indices]...)
end

function _assign(visited, indices::Indices, pair::Pair{FreeIndex{T}, FixedIndex{T}}) where {T<:AbstractVectorSpace}
    return Indices([i == first(pair) ? last(pair) : i for i in indices.indices]...)
end

function _assign(visited, A::VariableTensor, pair::Pair{VariableTensor{rank},
        T}) where {rank, T<:AbstractTensor{rank}}
    return A == first(pair) ? last(pair) : A
end

function _assign(visited, A::ConstantTensor, pair::Pair{VariableTensor{0}, S}) where {S<: Scalar}
    return A.value == first(pair) ? ConstantTensor(last(pair), A.shape...) : A
end

function _assign(visited, s::VariableTensor{0}, pair::Pair{VariableTensor{0}, S}) where {S<: Scalar}
    return s == first(pair) ? last(pair) : s
end

function _assign(visited, A::Tensor, pair::Pair{VariableTensor{0}, S}) where {S<: Scalar}
    function assigninarray(x, pair)
        if last(pair) isa ConstantTensor && x == first(pair)
            return last(pair).value
        elseif last(pair) isa ZeroTensor && x == first(pair)
            return 0
        else
            return assign(x, pair)
        end
    end
    return Tensor(assigninarray.(A.value, pair), A.shape...)
end

function _assign(visited, x::Number, pair::Assignment, children::Vararg{Node})
    return x
end

function _assign(visited, x::Node, pair::Assignment, children::Vararg{Node})
    return updatechildren(x, children...)
end

"""
    assign(node::Node, pair::Pair{VariableTensor{rank}, T}) where {rank, T<:AbstractTensor{rank}}
    assign(node::Node, pair::Pair{T, N}) where {T<:AbstractTensor{0}, N<:Number}
    assign(node::Node, pair::Pair{T1, T2}) where {T1<:AbstractVectorSpace, T2<:AbstractVectorSpace}
    assign(node::Node, pair::Pair{FreeIndex{T}, Union{Int, FixedIndex{T}}}) where {T<:AbstractVectorSpace}

Create a new node from node, where every value first(pair) is replaced with
last(pair). Can be used on tensors, indices and vector spaces.
"""
function assign(node::Node, pair::Pair{VariableTensor{rank}, T}) where {rank, T<:AbstractTensor{rank}}
    # Check that vector spaces match
    A, B = pair
    for i in 1:length(A.shape)
        if A.shape[i] != B.shape[i]
            throw(DomainError((A.shape[i], B.shape[i]), string("Invalid reshaping of ", A.shape, " to ", B.shape)))
        end
    end
    return _callassign(node, pair)
end

function assign(node::Node, pair::Pair{T1, T2}, visited) where {T1<:AbstractVectorSpace, T2<:AbstractVectorSpace}
    # Check that vector spaces match
    if dim(first(pair)) != dim(last(pair)) && dim(first(pair)) !== nothing
        throw(DomainError((T1,T2), "Cannot assign different dimension vector spaces"))
    end
    return _callassign(node, pair, visited)
end

function assign(node::Node, pair::Pair{T1, T2}) where {T1<:AbstractVectorSpace, T2<:AbstractVectorSpace}
    return assign(node, pair, nothing)
end

function assign(node::Node, pair::Pair{T, N}, visited) where {T<:AbstractTensor{0}, N<:Number}
    return _callassign(node, first(pair)=>Tensor(last(pair)), visited)
end

function assign(node::Node, pair::Pair{T, N}) where {T<:AbstractTensor{0}, N<:Number}
    return assign(node, pair, nothing)
end

function assign(node::Node, pair::Pair{FreeIndex{T}, FixedIndex{T}}) where {T<:AbstractVectorSpace}
    if first(pair).V != last(pair).V
        throw(Dimensionmismatch(string(first(pair).V, "!=", last(pair).V)))
    end
    return _callassign(node, pair)
end

function assign(node::Node, pair::Pair{FreeIndex{T}, Int}) where {T<:AbstractVectorSpace}
    return _callassign(node, first(pair)=>FixedIndex(first(pair).V, last(pair)))
end

function assign(node::Node, pair::Assignment)
    return _callassign(node, pair)
end

function _callassign(node::Node, pair::Assignment, visited)
    return traversal(node, x-> x, _assign, nothing, pair, visited)
end

function _callassign(node::Node, pair::Pair)
    return _callassign(node, pair, nothing)
end

function assign(i::Number, pair::Assignment)
    return i
end

max_stars_repo_name: Sohva/oolong

max_issues_repo_name: Sohva/oolong

max_forks_repo_name: oolongtensor/oolong


text:
using Documenter, Query

makedocs(
	modules = [Query],
	sitename = "Query.jl",
	pages = [
		"Introduction" => "index.md",
		"Getting Started" => "gettingstarted.md",
		"Standalone Query Commands" => "standalonequerycommands.md",
		"LINQ Style Query Commands" => "linqquerycommands.md",
		"Data Sources" => "sources.md",
		"Data Sinks" => "sinks.md",
		"Experimental Features" => "experimental.md",
		"Internals" => "internals.md"]
)

deploydocs(
    repo = "github.com/queryverse/Query.jl.git"
)


text:
@with_kw mutable struct TTCIntersectionDriver <: DriverModel{LonAccelDirection}
    a::LonAccelDirection = LonAccelDirection(0., 1)
    navigator::RouteFollowingIDM = RouteFollowingIDM()
    intersection::Vector{Lane{Float64}} = Lane{Float64}[]
    intersection_pos::VecSE2 = VecE2(0., 0., 0.)
    ttc_threshold::Float64 = 4.5
    horizon::Float64 = 20.0 # distance from the intersection to start the logic
    stop_delta::Float64 = 0. # precision for stopping slightly before the line
    accel_tol::Float64 = 1e-2 # if |a| < accel_tol then a = 0.
    priorities::Dict{Tuple{LaneTag, LaneTag}, Bool} = Dict{Tuple{LaneTag, LaneTag}, Bool}()
    # states
    priority::Bool = false
    stop::Bool = false # switch to true when stopped at the intersection, stays true until the end
end

get_name(::TTCIntersectionDriver) = "TTCIntersectionDriver"
Base.rand(model::TTCIntersectionDriver) = model.a

function AutomotiveDrivingModels.set_desired_speed!(model::TTCIntersectionDriver, v_des::Float64)
    model.navigator.v_des = v_des
    model
end

function AutomotiveDrivingModels.reset_hidden_state!(model::TTCIntersectionDriver)
    model.priority = false
    model.stop = false
    model
end

function AutomotiveDrivingModels.observe!(model::TTCIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int)
    ego = scene[findfirst(egoid, scene)]
    AutomotiveDrivingModels.observe!(model.navigator, scene, roadway, egoid) # set the direction
    dir = model.navigator.dir
    a_lon =0.
    a_lon_idm = model.navigator.a
    passed = has_passed(model, scene, roadway, egoid)
    is_engaged = engaged(model, scene, roadway, egoid)
    right_of_way = model.priorities[(model.navigator.route[1].tag,model.navigator.route[end].tag)]
    is_clogged = is_intersection_clogged(model, scene, roadway, egoid)
    ttc = ttc_check(model, scene, roadway, egoid)
    # if !model.stop
    #     model.stop = isapprox(ego.state.v, 0.)
    # end

    if isempty(model.intersection) || passed
        a_lon = a_lon_idm
    elseif !passed
        if right_of_way
            if is_clogged && !passed && is_engaged && !isapprox(ego.state.v, 0.)
                # println("Vehicle $egoid : emergency break")
                a_lon = -model.navigator.d_max
            else
                a_lon = a_lon_idm
            end
        else # left turn
            if !ttc && !is_engaged  # before left turn
                a_lon = min(a_lon_idm, stop_at_end(model, ego, roadway))
            elseif is_clogged && !passed && is_engaged && !isapprox(ego.state.v, 0.) #!ttc && !passed && is_engaged || (is_clogged && is_engaged)
                # println("Vehicle $egoid : emergency break")
                a_lon = -model.navigator.d_max
            elseif ttc
                a_lon = a_lon_idm
            end
        end
    end
    # println("veh $egoid | clogged $is_clogged | passed $passed | engaged $is_engaged | stop $(isapprox(ego.state.v, 0.))| right of way $right_of_way")
    # println("veh $egoid: ttc $ttc")
    # println(" ID ", egoid, " stop ", model.stop, " priority ", model.priority)
    model.a = LonAccelDirection(a_lon, dir) # add noise to break ties #XXX remove with priorities
    model
end

function engaged(model::TTCIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int)
    ego = scene[findfirst(egoid, scene)]
    lane = get_lane(roadway, ego)
    inter_width = 7.5 #todo parameterized
    if normsquared(VecE2(model.intersection_pos - ego.state.posG)) < inter_width^2
        return true
    end
    return false
    # if isempty(lane.entrances)
    #     return false
    # end
    # return true
end

function has_passed(model::TTCIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int)
    ego = scene[findfirst(egoid, scene)]
    lane = get_lane(roadway, ego)
    inter_to_car = ego.state.posG - model.intersection_pos
    car_vec = get_front(ego) - ego.state.posG
    has_passed = dot(inter_to_car, car_vec) > 0. || lane ∈ get_exit_lanes(roadway)
    return has_passed
end

"""
Check if the time to collision is above some threshold
"""
function ttc_check(model::TTCIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int)
    min_ttc = Inf
    inter_width = 6.0 #todo parameterized
    otherid = -1
    ego = scene[findfirst(egoid, scene)]
    for veh in scene
        if veh.id != egoid && veh.def.class != AgentClass.PEDESTRIAN && !is_behind(ego, veh, roadway)
            posF = veh.state.posF
            int_x, int_y, int_θ = model.intersection_pos
            lane = get_lane(roadway, veh)
            int_proj = Frenet(proj(model.intersection_pos, lane, roadway, move_along_curves=false), roadway)
            if normsquared(VecE2(model.intersection_pos - veh.state.posG)) < inter_width^2 # vehicle is in the middle
                ttc = 0.
            else
                ttc = (int_proj.s - posF.s)/veh.state.v
            end
            if 0 <= ttc < min_ttc
                otherid = veh.id
                min_ttc = ttc
            end
        end
    end
    # println("veh id ", egoid, "min_ttc ", min_ttc, " threshold ", model.ttc_threshold, " other id ", otherid)
    if 0 <= min_ttc < model.ttc_threshold
        return false
    else
        return true
    end
end

function is_intersection_clogged(model::TTCIntersectionDriver, scene::Scene, roadway::Roadway, egoid::Int64)
    inter_width = 5.0 #todo parameterized
    ego = scene[findfirst(egoid, scene)]
    for veh in scene
        if veh.id == egoid
            continue
        end
        # println("veh id $(veh.id) : ", sqrt(normsquared(VecE2(model.intersection_pos - veh.state.posG))))
        if normsquared(VecE2(model.intersection_pos - veh.state.posG)) < inter_width^2 && !is_behind(ego, veh, roadway)
            return true
        end
    end
    return false
end

function is_behind(veh1::Vehicle, veh2::Vehicle, roadway::Roadway)
    lane1 = get_lane(roadway, veh1)
    # project veh2 on lane2
    veh2_l1 = Frenet(proj(veh2.state.posG, lane1, roadway, move_along_curves=false), roadway)
    return veh2_l1.s < veh1.state.posF.s
end

max_stars_repo_name: MaximeBouton/AutomotivePOMDPs

max_issues_repo_name: MaximeBouton/AutomotivePOMDPs

max_forks_repo_name: MaximeBouton/AutomotivePOMDPs


text:
module StringDistances

##############################################################################
##
## Export
##
##############################################################################
using DataStructures
import Base: eltype, length, iterate, ==, hash, isless, convert, show, @deprecate
import Distances: evaluate, Hamming, hamming, PreMetric, SemiMetric
export
evaluate,
compare,
Hamming,
Levenshtein,
DamerauLevenshtein,
Jaro,
RatcliffObershelp,
QGram,
Cosine,
Jaccard,
SorensenDice,
Overlap,
Winkler,
Partial,
TokenSort,
TokenSet,
TokenMax,
qgram,
find_best,
find_all

##############################################################################
##
## include
##
##############################################################################
include("utils.jl")
include("edit.jl")
include("qgram.jl")
include("compare.jl")
include("find.jl")

end

##############################################################################
##
## Some memo about Strings

# length: number of characters
# ncodeunits: Return the number of code units in a string (aking to index of vector). Not all such indices  are valid – they may not be the start of a character,.
# sizeof:  Size, in bytes, of the string str. Equal to the number of code units in str  multiplied by the size, in bytes, of one code unit in str.

# lastindex: Return the last index of a collection
# nextinds(s, i):  return the index of the start of the character whose encoding starts after index i
# nextind(s, 0, N): return the index of the Nth character of s (or, if there are less than N characters, return ncodeunits(str) + (N - length(s))

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


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder

name = "Pango"
version = v"1.47.0"

# Collection of sources required to build Pango
sources = [
    ArchiveSource("http://ftp.gnome.org/pub/GNOME/sources/pango/$(version.major).$(version.minor)/pango-$(version).tar.xz",
                  "730db8652fc43188e03218c3374db9d152351f51fc7011b9acae6d0a6c92c367"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/pango-*/
mkdir build && cd build
meson --cross-file="${MESON_TARGET_TOOLCHAIN}" \
    -Dintrospection=false \
    ..
ninja -j${nproc}
ninja install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = filter!(p -> arch(p) != "armv6l", supported_platforms(; experimental=true))

# The products that we will ensure are always built
products = [
    LibraryProduct(["libpango", "libpango-1", "libpango-1.0"], :libpango),
    LibraryProduct(["libpangocairo", "libpangocairo-1", "libpangocairo-1.0"], :libpangocairo),
    LibraryProduct(["libpangoft2", "libpangoft2-1", "libpangoft2-1.0"], :libpangoft),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency("FriBidi_jll"),
    Dependency("FreeType2_jll"),
    Dependency("Glib_jll", v"2.68.1"; compat="2.68.1"),
    Dependency("Fontconfig_jll"),
    Dependency("HarfBuzz_jll", v"2.8.1"; compat="2.8.1"),
    Dependency("Cairo_jll", v"1.16.1"; compat="1.16.1"),
    BuildDependency("Xorg_xorgproto_jll"; platforms=filter(p->Sys.islinux(p)||Sys.isfreebsd(p), platforms)),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat="1.6")

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
module TrainingSetGenerator

#Declare linear function, which separates observations in 2D space
f(x) = 8*x+3.5

#Generating training data with labels for given N lenght
function generate_data(n)
  training_set = []
  for i in 1:n
    x = rand(0:200)
    y = rand(0:200)
    bias = rand(0:200)
    lineY = f(x)
    if y>lineY
      label = -1
    else
      label = 1
    end
    input_values = [x,y, bias]
    push!(training_set, [input_values, label])
  end
  return training_set
end

end

max_stars_repo_name: AlexAdamenko/Perceptron

max_issues_repo_name: AlexAdamenko/Perceptron

max_forks_repo_name: AlexAdamenko/Perceptron


text:
#! /usr/bin/julia

# Rosetta Code, Strip a set of characters from a string

function fullstrip(s::String, r::String)
    replace(s, Set(r), "")
end

tests = "She was a soul stripper. She took my heart!"
testr = "aei"

println(tests, " => ", fullstrip(tests, testr))

max_stars_repo_name: MichaeLeroy/rosetta-code

max_issues_repo_name: MichaeLeroy/rosetta-code

max_forks_repo_name: MichaeLeroy/rosetta-code


text:
"""
    BlueStyle()

Formatting style based on [BlueStyle](https://github.com/invenia/BlueStyle)
and [JuliaFormatter#283](https://github.com/domluna/JuliaFormatter.jl/issues/283).

!!! note
    This style is still work-in-progress, and does not yet implement all of the
    BlueStyle guide.

Configurable options with different defaults to [`DefaultStyle`](@ref) are:
- `always_use_return` = true
- `short_to_long_function_def` = true
- `whitespace_ops_in_indices` = true
- `remove_extra_newlines` = true
- `always_for_in` = true
- `import_to_using` = true
- `pipe_to_function_call` = true
- `whitespace_in_kwargs` = false
- `annotate_untyped_fields_with_any` = false
"""
struct BlueStyle <: AbstractStyle
    innerstyle::Union{Nothing,AbstractStyle}
end
BlueStyle() = BlueStyle(nothing)

@inline getstyle(s::BlueStyle) = s.innerstyle === nothing ? s : s.innerstyle

function options(style::BlueStyle)
    return (;
        always_use_return = true,
        short_to_long_function_def = true,
        whitespace_ops_in_indices = true,
        remove_extra_newlines = true,
        always_for_in = true,
        import_to_using = true,
        pipe_to_function_call = true,
        whitespace_in_kwargs = false,
        annotate_untyped_fields_with_any = false,
        conditional_to_if = true,
    )
end

function nestable(::BlueStyle, cst::CSTParser.EXPR)
    is_assignment(cst) && is_iterable(cst[end]) && return false
    return nestable(DefaultStyle(), cst)
end

function p_do(bs::BlueStyle, cst::CSTParser.EXPR, s::State)
    style = getstyle(bs)
    t = FST(cst, nspaces(s))
    add_node!(t, pretty(style, cst[1], s), s)
    add_node!(t, Whitespace(1), s)
    add_node!(t, pretty(style, cst[2], s), s, join_lines = true)
    if cst[3].fullspan != 0
        add_node!(t, Whitespace(1), s)
        add_node!(t, pretty(style, cst[3], s), s, join_lines = true)
    end
    if cst[4].typ === CSTParser.Block
        s.indent += s.opts.indent
        n = pretty(style, cst[4], s, ignore_single_line = true)
        add_node!(t, n, s, max_padding = s.opts.indent)
        s.indent -= s.opts.indent
    end
    add_node!(t, pretty(style, cst.args[end], s), s)
    t
end

function p_call(bs::BlueStyle, cst::CSTParser.EXPR, s::State)
    style = getstyle(bs)
    t = p_call(DefaultStyle(style), cst, s)
    if !parent_is(cst, is_function_or_macro_def)
        separate_kwargs_with_semicolon!(t)
    end
    t
end

function p_binaryopcall(
    bs::BlueStyle,
    cst::CSTParser.EXPR,
    s::State;
    nonest = false,
    nospace = false,
)
    style = getstyle(bs)
    t = FST(cst, nspaces(s))
    op = cst[2]
    nonest = nonest || op.kind === Tokens.COLON
    if cst.parent.typ === CSTParser.Curly &&
       op.kind in (Tokens.ISSUBTYPE, Tokens.ISSUPERTYPE) &&
       !s.opts.whitespace_typedefs
        nospace = true
    elseif op.kind === Tokens.COLON
        nospace = true
    end
    nospace_args = s.opts.whitespace_ops_in_indices ? false : nospace

    if is_opcall(cst[1])
        n = pretty(style, cst[1], s, nonest = nonest, nospace = nospace_args)
    else
        n = pretty(style, cst[1], s)
    end

    if op.kind === Tokens.COLON &&
       s.opts.whitespace_ops_in_indices &&
       !is_leaf(cst[1]) &&
       !is_iterable(cst[1])
        paren = FST(CSTParser.PUNCTUATION, -1, n.startline, n.startline, "(")
        add_node!(t, paren, s)
        add_node!(t, n, s, join_lines = true)
        paren = FST(CSTParser.PUNCTUATION, -1, n.startline, n.startline, ")")
        add_node!(t, paren, s, join_lines = true)
    else
        add_node!(t, n, s)
    end

    nrhs = nest_rhs(cst)
    nrhs && (t.nest_behavior = AlwaysNest)
    nest = (nestable(style, cst) && !nonest) || nrhs

    if op.fullspan == 0
        # Do nothing - represents a binary op with no textual representation.
        # For example: `2a`, which is equivalent to `2 * a`.
    elseif op.kind === Tokens.EX_OR
        add_node!(t, pretty(style, op, s), s, join_lines = true)
    elseif (
        is_number(cst[1]) ||
        op.kind === Tokens.FWDFWD_SLASH ||
        op.kind === Tokens.CIRCUMFLEX_ACCENT
    ) && op.dot
        add_node!(t, Whitespace(1), s)
        add_node!(t, pretty(style, op, s), s, join_lines = true)
        nest ? add_node!(t, Placeholder(1), s) : add_node!(t, Whitespace(1), s)
    elseif op.kind !== Tokens.IN && (
        nospace || (
            op.kind !== Tokens.ANON_FUNC && CSTParser.precedence(op) in (
                CSTParser.ColonOp,
                CSTParser.RationalOp,
                CSTParser.PowerOp,
                CSTParser.DeclarationOp,
                CSTParser.DotOp,
            )
        )
    )
        add_node!(t, pretty(style, op, s), s, join_lines = true)
    else
        add_node!(t, Whitespace(1), s)
        add_node!(t, pretty(style, op, s), s, join_lines = true)
        nest ? add_node!(t, Placeholder(1), s) : add_node!(t, Whitespace(1), s)
    end

    if is_opcall(cst[3])
        n = pretty(style, cst[3], s, nonest = nonest, nospace = nospace_args)
    else
        n = pretty(style, cst[3], s)
    end

    if op.kind === Tokens.COLON &&
       s.opts.whitespace_ops_in_indices &&
       !is_leaf(cst[3]) &&
       !is_iterable(cst[3])
        paren = FST(CSTParser.PUNCTUATION, -1, n.startline, n.startline, "(")
        add_node!(t, paren, s, join_lines = true)
        add_node!(t, n, s, join_lines = true)
        paren = FST(CSTParser.PUNCTUATION, -1, n.startline, n.startline, ")")
        add_node!(t, paren, s, join_lines = true)
    else
        add_node!(t, n, s, join_lines = true)
    end

    if nest
        # for indent, will be converted to `indent` if needed
        insert!(t.nodes, length(t.nodes), Placeholder(0))
    end

    t
end

function p_return(bs::BlueStyle, cst::CSTParser.EXPR, s::State)
    style = getstyle(bs)
    t = FST(cst, nspaces(s))
    add_node!(t, pretty(style, cst[1], s), s)
    if cst[2].fullspan != 0
        for a in cst.args[2:end]
            add_node!(t, Whitespace(1), s)
            add_node!(t, pretty(style, a, s), s, join_lines = true)
        end
    elseif cst[2].kind === Tokens.NOTHING
        add_node!(t, Whitespace(1), s)
        no = FST(CSTParser.IDENTIFIER, -1, t.endline, t.endline, "nothing")
        add_node!(t, no, s, join_lines = true)
    end
    t
end


text:
int_rules_4_3_4_1 = @theory begin

    #= ::Subsection::Closed:: =#
    #= 4.3.4.1*(a+b*tan)^m*(A+B*tan+C*tan^2) =#
    @apply_utils Antiderivative((~(a') + ~(b') * tan(~(e') + ~(f') * ~x)) ^ ~(m') * (~A + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2), ~x) => (~A / (~b * ~f)) * Subst(Antiderivative((~a + ~x) ^ ~m, ~x), ~x, ~b * tan(~e + ~f * ~x)) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~C, ~m], ~x) && EqQ(~A, ~C)
    @apply_utils Antiderivative((~(a') + ~(b') * cot(~(e') + ~(f') * ~x)) ^ ~(m') * (~A + ~(C') * cot(~(e') + ~(f') * ~x) ^ 2), ~x) => (-(~A) / (~b * ~f)) * Subst(Antiderivative((~a + ~x) ^ ~m, ~x), ~x, ~b * cot(~e + ~f * ~x)) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~C, ~m], ~x) && EqQ(~A, ~C)
    @apply_utils Antiderivative((~(a') + ~(b') * tan(~(e') + ~(f') * ~x)) ^ ~(m') * (~(A') + ~(B') * tan(~(e') + ~(f') * ~x) + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2), ~x) => (1 / (~b) ^ 2) * Antiderivative((~a + ~b * tan(~e + ~f * ~x)) ^ (~m + 1) * Simp((~b * ~B - ~a * ~C) + ~b * ~C * tan(~e + ~f * ~x), ~x), ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~B, ~C, ~m], ~x) && EqQ((~A * (~b) ^ 2 - ~a * ~b * ~B) + (~a) ^ 2 * ~C, 0)
    @apply_utils Antiderivative((~(a') + ~(b') * tan(~(e') + ~(f') * ~x)) ^ ~(m') * (~(A') + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2), ~x) => (-(~C) / (~b) ^ 2) * Antiderivative((~a + ~b * tan(~e + ~f * ~x)) ^ (~m + 1) * (~a - ~b * tan(~e + ~f * ~x)), ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~C, ~m], ~x) && EqQ(~A * (~b) ^ 2 + (~a) ^ 2 * ~C, 0)
    @apply_utils Antiderivative((~(a') + ~(b') * tan(~(e') + ~(f') * ~x)) ^ ~(m') * (~(A') + ~(B') * tan(~(e') + ~(f') * ~x) + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2), ~x) => (-(((~a * ~A + ~b * ~B) - ~a * ~C)) * tan(~e + ~f * ~x) * (~a + ~b * tan(~e + ~f * ~x)) ^ ~m) / (2 * ~a * ~f * ~m) + (1 / (2 * (~a) ^ 2 * ~m)) * Antiderivative((~a + ~b * tan(~e + ~f * ~x)) ^ (~m + 1) * Simp(((~b * ~B - ~a * ~C) + ~a * ~A * (2 * ~m + 1)) - (~b * ~C * (~m - 1) + (~A * ~b - ~a * ~B) * (~m + 1)) * tan(~e + ~f * ~x), ~x), ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~B, ~C], ~x) && (NeQ((~A * (~b) ^ 2 - ~a * ~b * ~B) + (~a) ^ 2 * ~C, 0) && (LeQ(~m, -1) && EqQ((~a) ^ 2 + (~b) ^ 2, 0)))
    @apply_utils Antiderivative((~(a') + ~(b') * tan(~(e') + ~(f') * ~x)) ^ ~(m') * (~(A') + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2), ~x) => (-((~a * ~A - ~a * ~C)) * tan(~e + ~f * ~x) * (~a + ~b * tan(~e + ~f * ~x)) ^ ~m) / (2 * ~a * ~f * ~m) + (1 / (2 * (~a) ^ 2 * ~m)) * Antiderivative((~a + ~b * tan(~e + ~f * ~x)) ^ (~m + 1) * Simp((-(~a) * ~C + ~a * ~A * (2 * ~m + 1)) - (~b * ~C * (~m - 1) + ~A * ~b * (~m + 1)) * tan(~e + ~f * ~x), ~x), ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~C], ~x) && (NeQ(~A * (~b) ^ 2 + (~a) ^ 2 * ~C, 0) && (LeQ(~m, -1) && EqQ((~a) ^ 2 + (~b) ^ 2, 0)))
    @apply_utils Antiderivative((~A + ~(B') * tan(~(e') + ~(f') * ~x) + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2) / (~(a') + ~(b') * tan(~(e') + ~(f') * ~x)), ~x) => (((~a * ~A + ~b * ~B) - ~a * ~C) * ~x) / ((~a) ^ 2 + (~b) ^ 2) + (((~A * (~b) ^ 2 - ~a * ~b * ~B) + (~a) ^ 2 * ~C) / ((~a) ^ 2 + (~b) ^ 2)) * Antiderivative((1 + tan(~e + ~f * ~x) ^ 2) / (~a + ~b * tan(~e + ~f * ~x)), ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~B, ~C], ~x) && (NeQ((~a) ^ 2 + (~b) ^ 2, 0) && EqQ((~A * ~b - ~a * ~B) - ~b * ~C, 0))
    @apply_utils Antiderivative((~A + ~(B') * tan(~(e') + ~(f') * ~x) + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2) / tan(~(e') + ~(f') * ~x), ~x) => ~B * ~x + ~A * Antiderivative(1 / tan(~e + ~f * ~x), ~x) + ~C * Antiderivative(tan(~e + ~f * ~x), ~x) <-- FreeQ([~e, ~f, ~A, ~B, ~C], ~x) && NeQ(~A, ~C)
    @apply_utils Antiderivative((~A + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2) / tan(~(e') + ~(f') * ~x), ~x) => ~A * Antiderivative(1 / tan(~e + ~f * ~x), ~x) + ~C * Antiderivative(tan(~e + ~f * ~x), ~x) <-- FreeQ([~e, ~f, ~A, ~C], ~x) && NeQ(~A, ~C)
    @apply_utils Antiderivative((~A + ~(B') * tan(~(e') + ~(f') * ~x) + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2) / (~(a') + ~(b') * tan(~(e') + ~(f') * ~x)), ~x) => ((((~a * ~A + ~b * ~B) - ~a * ~C) * ~x) / ((~a) ^ 2 + (~b) ^ 2) - (((~A * ~b - ~a * ~B) - ~b * ~C) / ((~a) ^ 2 + (~b) ^ 2)) * Antiderivative(tan(~e + ~f * ~x), ~x)) + (((~A * (~b) ^ 2 - ~a * ~b * ~B) + (~a) ^ 2 * ~C) / ((~a) ^ 2 + (~b) ^ 2)) * Antiderivative((1 + tan(~e + ~f * ~x) ^ 2) / (~a + ~b * tan(~e + ~f * ~x)), ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~B, ~C], ~x) && (NeQ((~A * (~b) ^ 2 - ~a * ~b * ~B) + (~a) ^ 2 * ~C, 0) && (NeQ((~a) ^ 2 + (~b) ^ 2, 0) && NeQ((~A * ~b - ~a * ~B) - ~b * ~C, 0)))
    @apply_utils Antiderivative((~A + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2) / (~a + ~(b') * tan(~(e') + ~(f') * ~x)), ~x) => ((~a * (~A - ~C) * ~x) / ((~a) ^ 2 + (~b) ^ 2) - ((~b * (~A - ~C)) / ((~a) ^ 2 + (~b) ^ 2)) * Antiderivative(tan(~e + ~f * ~x), ~x)) + (((~a) ^ 2 * ~C + ~A * (~b) ^ 2) / ((~a) ^ 2 + (~b) ^ 2)) * Antiderivative((1 + tan(~e + ~f * ~x) ^ 2) / (~a + ~b * tan(~e + ~f * ~x)), ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~C], ~x) && (NeQ((~a) ^ 2 * ~C + ~A * (~b) ^ 2, 0) && (NeQ((~a) ^ 2 + (~b) ^ 2, 0) && NeQ(~A, ~C)))
    @apply_utils Antiderivative((~(a') + ~(b') * tan(~(e') + ~(f') * ~x)) ^ ~m * (~(A') + ~(B') * tan(~(e') + ~(f') * ~x) + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2), ~x) => (((~A * (~b) ^ 2 - ~a * ~b * ~B) + (~a) ^ 2 * ~C) * (~a + ~b * tan(~e + ~f * ~x)) ^ (~m + 1)) / (~b * ~f * (~m + 1) * ((~a) ^ 2 + (~b) ^ 2)) + (1 / ((~a) ^ 2 + (~b) ^ 2)) * Antiderivative((~a + ~b * tan(~e + ~f * ~x)) ^ (~m + 1) * Simp((~b * ~B + ~a * (~A - ~C)) - ((~A * ~b - ~a * ~B) - ~b * ~C) * tan(~e + ~f * ~x), ~x), ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~B, ~C], ~x) && (NeQ((~A * (~b) ^ 2 - ~a * ~b * ~B) + (~a) ^ 2 * ~C, 0) && (LtQ(~m, -1) && NeQ((~a) ^ 2 + (~b) ^ 2, 0)))
    @apply_utils Antiderivative((~(a') + ~(b') * tan(~(e') + ~(f') * ~x)) ^ ~m * (~(A') + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2), ~x) => ((~A * (~b) ^ 2 + (~a) ^ 2 * ~C) * (~a + ~b * tan(~e + ~f * ~x)) ^ (~m + 1)) / (~b * ~f * (~m + 1) * ((~a) ^ 2 + (~b) ^ 2)) + (1 / ((~a) ^ 2 + (~b) ^ 2)) * Antiderivative((~a + ~b * tan(~e + ~f * ~x)) ^ (~m + 1) * Simp(~a * (~A - ~C) - (~A * ~b - ~b * ~C) * tan(~e + ~f * ~x), ~x), ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~C], ~x) && (NeQ(~A * (~b) ^ 2 + (~a) ^ 2 * ~C, 0) && (LtQ(~m, -1) && NeQ((~a) ^ 2 + (~b) ^ 2, 0)))
    @apply_utils Antiderivative((~(a') + ~(b') * tan(~(e') + ~(f') * ~x)) ^ ~(m') * (~(A') + ~(B') * tan(~(e') + ~(f') * ~x) + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2), ~x) => (~C * (~a + ~b * tan(~e + ~f * ~x)) ^ (~m + 1)) / (~b * ~f * (~m + 1)) + Antiderivative((~a + ~b * tan(~e + ~f * ~x)) ^ ~m * Simp((~A - ~C) + ~B * tan(~e + ~f * ~x), ~x), ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~B, ~C, ~m], ~x) && (NeQ((~A * (~b) ^ 2 - ~a * ~b * ~B) + (~a) ^ 2 * ~C, 0) && Not(LeQ(~m, -1)))
    @apply_utils Antiderivative((~(a') + ~(b') * tan(~(e') + ~(f') * ~x)) ^ ~(m') * (~(A') + ~(C') * tan(~(e') + ~(f') * ~x) ^ 2), ~x) => (~C * (~a + ~b * tan(~e + ~f * ~x)) ^ (~m + 1)) / (~b * ~f * (~m + 1)) + (~A - ~C) * Antiderivative((~a + ~b * tan(~e + ~f * ~x)) ^ ~m, ~x) <-- FreeQ([~a, ~b, ~e, ~f, ~A, ~C, ~m], ~x) && (NeQ(~A * (~b) ^ 2 + (~a) ^ 2 * ~C, 0) && Not(LeQ(~m, -1)))
end

max_stars_repo_path: src/rules/4 Trig functions/4.3 Tangent/4.3.4.1 (a+b tan)^m (A+B tan+C tan^2).jl


text:
module ConditionalComponent

println("\nRectifier: Demonstrating conditional components")

using Modia
using Modia.Electric

# Desired:
#   using ModiaMath: plot
#
# In order that these packages need not to be defined in the user environment, they are included via Modia:
using Modia.ModiaMath: plot

@model ConditionalLoad begin
    R=Resistor(R=1)
    r=Resistor(R=0.01)
    C=Capacitor(C=1,start=0)
    D=IdealDiode()
    V=SineVoltage(V=5,freqHz=1.5, offset=0, startTime=0)
    extraLoad = true
    R2=if extraLoad; Resistor(R=1) end
@equations begin
    connect(V.p, D.p)
    connect(D.n, R.p)
    connect(R.n, r.p)
    connect(r.n, V.n)
    connect(C.n, R.n)
    connect(C.p, R.p)
    connect(C.n, R2.n)
    connect(C.p, R2.p)
    end
end

result = simulate(ConditionalLoad, 2, logTranslation=true)
plot(result, ("C.v", "V.v"), heading="Conditional", figure=12)

@model NoExtraLoad begin
    @extends ConditionalLoad(extraLoad = false)
end

result = simulate(NoExtraLoad, 2) # , logTranslation=true) # Problem with logging nothing in Julia 1.0
plot(result, ("C.v", "V.v"), heading="Conditional", figure=12)

end


text:
using SymbolicUtils, Test
using TermInterface

TermInterface.issym(s::Symbol) = true
Base.nameof(s::Symbol) = s

for f ∈ [:+, :-, :*, :/, :^]
    @eval begin
        Base.$f(x::Union{Expr, Symbol}, y::Number) = Expr(:call, $f, x, y)
        Base.$f(x::Number, y::Union{Expr, Symbol}) = Expr(:call, $f, x, y)
        Base.$f(x::Union{Expr, Symbol}, y::Union{Expr, Symbol}) = (Expr(:call, $f, x, y))
    end
end

Base.zero(t::Expr) = 0
TermInterface.symtype(::Expr) = Real
TermInterface.symtype(::Symbol) = Real

ex = 1 + (:x - 2)

@test simplify(ex) == -1 + :x
@test simplify(:a * (:b + -1 * :c) + -1 * (:b * :a + -1 * :c * :a), expand=true) == 0