text:
__precompile__()
"""
    src.jl

1. julia <= 1.2
2. StatsBase <= 0.32.0
"""
module src
    using StatsBase, Random
    export ModelBasic, EmptyHouseholdPolicy, HouseholdPolicy

# ======================= TYPES
    abstract type AbstractDGEModel <: Any end
    abstract type AbstractPolicyFunctions <: Any end

    # ---------- MODEL
    mutable struct ModelBasic <: AbstractDGEModel
        N::Float64  # total population size
        BOTTOM_VAL::Float64 # bottom value of val func
        AGG_SEED::Int # seed used in aggregation
        AGG_SIMU_SIZE::Int # size of simulation in aggregation
        AGG_DROP_FIRST_SIZE::Int # num of obs to drop in aggregation simulation (to drop non-statinary obs)
        # ------ ECON PARS
        θ::Float64 # cpaital income share
        δ::Float64 # depreciation rate
        β::Float64 # util disc factor
        μ::Float64 # CRRA risk aversion
        # ------ HH VAL FUNC ITER PARS
        Na::Int # num of asset grid units
        Nl::Int # num of labor endow grid units
        NaNl::Int # Na * Nl
        alim::Tuple{Float64,Float64} # asset state space
        agrid::Vector{Float64} # asset grid space
        lgrid::Vector{Float64} # labor endow grid space
        # ------ IDIO LAB ENDOW PROC - MARKOV CHAIN
        P::Matrix{Float64} # one-step trnasition matrix
        π::Vector{Float64} # stationary distribution of labor endow
        # ------ HOUSEHOLD POLICY FUNCTIONS
        HHPolicy::AbstractPolicyFunctions  # household policy functions
        # ------ OTHER ECON DATA
        r::Float64 # net interest rate
        w::Float64 # wage rate
        K::Float64 # agg capital
        L::Float64 # labor factior
        C::Float64 # agg consumption
        Y::Float64 # agg output
        # ------------- CONSTRUCTOR
        """
            ModelBasic( ; N::Float64 = 1.0,
                BOTTOM_VAL::Float64 = -6.66E36, AGG_SEED::Float64 = time(),
                AGG_SIMU_SIZE::Int = 10000, AGG_DROP_FIRST_SIZE::Int = 1000,
                Na::Int = 200, alim::Tuple{Float64, Float64} = (0.0,4.5),
                lgrid::Vector{Float64} = [0.6703,1.4918], P::Matrix{Float64} = [0.8 0.2;0.2 0.8], π::Vector{Float64} = [0.5,0.5],
                HHPolicy::AbstractPolicyFunctions = EmptyHouseholdPolicy(),
                θ::Float64 = 0.36, δ::Float64 = 0.08, β::Float64 = 0.96, μ::Float64 = 2.0 )
        """
        function ModelBasic( ; N::Float64 = 1.0,
            BOTTOM_VAL::Float64 = -6.66E36, AGG_SEED::Int = floor(Int,1E6 * rand()),
            AGG_SIMU_SIZE::Int = 10000, AGG_DROP_FIRST_SIZE::Int = 1000,
            Na::Int = 200, alim::Tuple{Float64, Float64} = (0.0,4.5),
            lgrid::Vector{Float64} = [0.6703,1.4918], P::Matrix{Float64} = [0.8 0.2;0.2 0.8], π::Vector{Float64} = [0.5,0.5],
            HHPolicy::AbstractPolicyFunctions = EmptyHouseholdPolicy(),
            θ::Float64 = 0.36, δ::Float64 = 0.08, β::Float64 = 0.96, μ::Float64 = 2.0 )
            # ----------
            local agrid = Vector{Float64}(LinRange(alim[1],alim[2], Na))

            new(N,BOTTOM_VAL,AGG_SEED,AGG_SIMU_SIZE,AGG_DROP_FIRST_SIZE,
                θ,δ,β,μ,
                Na, length(lgrid), Na * length(lgrid), alim, agrid, lgrid,
                P,π,
                HHPolicy,
                0.1, 1.0, 1.0, 1.0, 1.0, 1.0 )
        end # ModelBasic()
    end

    # ---------- HOUSEHOLD POLICY (CLOSURE)
    """
        EmptyHouseholdPolicy <: AbstractPolicyFunctions

    an empty type of household policies;
    used to initialize models.
    """
    struct EmptyHouseholdPolicy <: AbstractPolicyFunctions
    end # EmptyHouseholdPolicy
    # --------------
    """
        HouseholdPolicy <: AbstractPolicyFunctions

    household policy functions type
    """
    struct HouseholdPolicy <: AbstractPolicyFunctions
        Na::Int  # num of asset grid units
        Nl::Int  # num of labor endow grid units
        NaNl::Int  # Na * Nl
        agrid::Vector{Float64} # asset grid
        lgrid::Vector{Float64} # labor endow grid
        aOpt::Matrix{Float64} # decision matrix of asset, Na * Nl
        cOpt::Matrix{Float64} # decision matrix of consumption, Na * Nl
        vMat1::Matrix{Float64} # val func iter result - old guess
        vMat2::Matrix{Float64} # val func iter result - updated
        # --------- CONSTRUCTOR
        """
            HouseholdPolicy(agrid::Vector{Float64}, lgrid::Vector{Float64}, aOpt::Matrix{Float64}, cOpt::Matrix{Float64}, vMat1::Matrix{Float64}, vMat2::Matrix{Float64})
        """
        function HouseholdPolicy(agrid::Vector{Float64}, lgrid::Vector{Float64}, aOpt::Matrix{Float64}, cOpt::Matrix{Float64}, vMat1::Matrix{Float64}, vMat2::Matrix{Float64})
            local Na = length(agrid)
            local Nl = length(lgrid)
            new(Na,Nl,Na*Nl, agrid,lgrid, aOpt,cOpt, vMat1,vMat2)
        end # HouseholdPolicy()
    end # HouseholdPolicy
    # ------ APPENDED METHOD
    """
        eval_hhpolicy_interp(HHP::HouseholdPolicy, val_athis::Float64, idx_lthis::Int)

    evaluates `a(t+1)` at given `val_athis` (asset value of `a(t)`) and `idx_lthis` (index of labor endow state);
    returns a `Tuple{Float64,Float64}` consisting of
    `a(t+1)` and `c(t)` (in order);
    using linear interpolation.
    """
    function eval_hhpolicy_interp(HHP::HouseholdPolicy, val_athis::Float64, idx_lthis::Int)
        local anext = interp1(HHP.agrid, HHP.aOpt[:,idx_lthis], val_athis)
        local cthis = interp1(HHP.agrid, HHP.cOpt[:,idx_lthis], val_athis)
        return (anext,cthis)::Tuple{Float64, Float64}
    end # eval_hhpolicy_interp()

# ======================= NUMERICAL METHODS

    # ------- LINEAR INTERPOLATION (REFACTOR MATLAB interp1)
    function interp1(X::Vector{Float64}, Y::Vector{Float64}, xeval::Float64)
        # require X ascd sorted, at least has 2 ele
        local idx = -1::Int
        for j in 1:(length(X)-1)
            if X[j] <= xeval <= X[j+1]
                idx = j
                break
            end # if
        end # for j
        local res = Y[idx] + (xeval - X[idx])/(X[idx+1] - X[idx] + eps()) * (Y[idx+1] - Y[idx])
        return res::Float64
    end # interp1

    # ------- THREE-POINT GOLDEN-SECTION (HEER & MAUSSNER, 2008)
    function golden_max_3points(f::Function, LLB::Float64, LB::Float64, RB::Float64 ; TOL::Float64 = 1E-6)
        local r1 = 0.61803399
        local r2 = 1 - r1
        local x0 = LLB
        local x3 = RB

        if abs(RB-LB) <= abs(LB-LLB)
            x1 = LB; x2 = LB + r2 * (RB-LB)
        else
            x2 = LB; x1 = LB - r2 * (LB-LLB)
        end # if

        f1 = - f(x1)
        f2 = - f(x2)

        # Searching
        local counter = 1;
        while (  abs(x3-x0) > (TOL * (abs(x1)+abs(x2)))  ) | (counter < 200)
            if f2 < f1
                x0=x1; # Update the very lower bound
                x1=x2;
                x2=r1*x1+r2*x3; # new left golden position
                f1=f2;
                f2= - f(x2);
            else
                x3=x2;
                x2=x1;
                x1=r1*x2+r2*x0;
                f2=f1;
                f1= - f(x1);
            end
            counter = counter +1;
        end

        # Post-convergence
        if f1<=f2
            return x1::Float64
        else
            return x2::Float64
        end

        return NaN::Float64
    end # golden_max_3points

# ======================= MODEL-SPECIFIC METHODS

    # ------- UTIL
    ufunc(c::Float64, μ::Float64) = begin
        return ( (c + 1E-8) ^ (1 - μ) / (1 - μ) )::Float64
    end # ufunc
    # ------- INTER-PERIOD BUDGET
    get_c(lthis::Float64, athis::Float64, anext::Float64, w::Float64) = begin
        return (w * lthis + athis - anext)::Float64
    end # get_c
    # ------- DISCOUNTED VAL FUNC (t+1) EVAL
    function vfunc(idx_lthis::Int, val_anext::Float64, mat_vnext::Matrix{Float64}, M::ModelBasic)
        res = 0.0;
        for j in 1:M.Nl
            res += M.P[idx_lthis, j] * interp1(M.agrid, mat_vnext[:,j], val_anext)
        end # for j
        return (M.β * res)::Float64
    end # vfunc
    # ------- BELLMAN BEFORE MAXIMIZATION
    function bellman(val_athis::Float64, idx_lthis::Int, val_anext::Float64, mat_vnext::Matrix{Float64}, M::ModelBasic)
        local cthis = get_c(M.lgrid[idx_lthis], val_athis, val_anext, M.w)
        if cthis < 0
            # (vthis, cthis)
            return ( M.BOTTOM_VAL, 0.0 )::Tuple{Float64, Float64}
        else
            local vthis = ufunc(cthis, M.μ) + vfunc(idx_lthis, val_anext, mat_vnext, M)
            return ( vthis, cthis )::Tuple{Float64, Float64}
        end # if
    end # bellman

# ======================= SIMULATION AGGREGATION METHODS

    # ------------- AGG: VIA SIMU HH PATH
    """
        agg_hh_simu(M::ModelBasic, T::Int, seed::Int ; a0::Float64 = 1.0, drop_first::Int = 100)

    aggregation to get agg variables via simulating a household path (long enough);
    returns a tuple `Tuple{Float64, Float64}` consisting of aggregate capital and aggregate consumption.
    """
    function agg_hh_simu(M::ModelBasic, T::Int, seed::Int ; a0::Float64 = 1.0, drop_first::Int = 100)
        # init
        local avgK = a0
        local avgC = 0.0
        local nowK = a0
        local nowC = 0.0
        # sampling from stationary distribution of labor endowment
        # NOTE: we expect there will not be more than 2^8 states, so use `Int8` to
        #       improve gc performance when T is very large
        #       (e.g. >= 100,000,000, or equivalent run-time >= 1s)
        vidx_L = sample( MersenneTwister(seed), Vector{Int8}(1:M.Nl), Weights(M.π), T )

        # simulation
        for t in 1:T
            nowK, nowC = eval_hhpolicy_interp(M.HHPolicy, nowK, Int(vidx_L[t]) )
            if t > drop_first
                avgK = avgK + nowK
                avgC = avgC + nowC
            end # if t
        end # t

        # transformation & return
        avgK = avgK / (T - drop_first) * M.N
        avgC = avgC / (T - drop_first) * M.N

        return (avgK, avgC)::Tuple{Float64, Float64}
    end # agg_hh_simu()

# ======================= SOLVERS: HOUSEHOLD POLICY FUNCTIONS

    # ----------- HOUSEHOLD POLICY FUNCTIONS SEARCHING
    """
        solve_hh_decision_rule!(M::ModelBasic ; max_iter::Int = 700, rtol::Float64 = 1E-2, print_err::Bool = true, step_size::Float64 = 0.7, if_golden::Bool = true)

    solves household policy functions, using val func iter;
    returns a `HouseholdPolicy` instance.

    where `if_golden` controls if to use golden-section to precise-ize val func (if not, just use grid points of asset)
    """
    function solve_hh_decision_rule!(M::ModelBasic ;
        max_iter::Int = 700, rtol::Float64 = 1E-2, print_err::Bool = true, step_size::Float64 = 0.7, if_golden::Bool = true)

        # init
        local vMat1 = zeros(M.Na, M.Nl) # old guess of val func
        local vMat2 = zeros(M.Na, M.Nl) # updated guess of val func
        local aOpt = zeros(M.Na, M.Nl) # opt asset, saving levels
        local cOpt = zeros(M.Na, M.Nl) # opt consumption, saving levels
        # init guess of val func
        for jl in 1:M.Nl, ja in 1:M.Na
            vMat1[ja,jl] = ufunc( get_c(M.lgrid[jl], M.agrid[ja], M.agrid[1], M.w), M.μ )
        end # for jl,ja

        # val func iter
        for iter in 1:max_iter
            # computes V(a,l)
            for jl in 1:M.Nl, ja in 1:M.Na
                # lambda
                local afunc_bellman(val_anext::Float64) = bellman( M.agrid[ja], jl, val_anext, vMat1, M )[1]
                # eval temp val func space for current (a,l), then find max
                # then, find grid max
                local tmploc = findmax([ afunc_bellman(M.agrid[x]) for x in 1:M.Na ])[2]

                # dispatching
                if if_golden  # if choose to precise-ize result
                    if (tmploc == 1) | (tmploc == M.Na)
                        aOpt[ja,jl] = M.agrid[tmploc]
                    else
                        # other wise, use golden section
                        local LLB = M.agrid[tmploc - 1]
                        local LB = M.agrid[tmploc]
                        local RB = M.agrid[tmploc + 1]
                        aOpt[ja,jl] = golden_max_3points(afunc_bellman, LLB, LB, RB, TOL = 1E-6)
                    end # if tmplocs
                else
                    aOpt[ja,jl] = M.agrid[tmploc]  # save opt a grid
                end # if if_golden

                # eval at opt asset, save opt val func & consumption
                # NOTE: here we do not use afunc_bellman() because it only returns the 1st return
                vMat2[ja,jl], cOpt[ja,jl] = bellman( M.agrid[ja], jl, aOpt[ja,jl], vMat1, M )
            end # for jl, ja

            # check convergence
            local vFuncErr = findmax(abs.( vMat2 ./ ( vMat1 .+ eps() ) .- 1.0 ) )[1]
            print_err && println("- Round ",iter,", Rel Err: ",vFuncErr)

            if vFuncErr < rtol
                println("\t- HH Converged in Rel Err Meaning in round ",iter)
                break
            elseif iter == max_iter
                println("\t- HH Max Iter Reached")
            else
                vMat1 = (1.0 .- step_size) .* vMat1 .+ step_size .* vMat2  # update
            end # if isapprox()

        end # for iter

        # Modify model instance, add in HouseholdPolicy
        M.HHPolicy = HouseholdPolicy( M.agrid, M.lgrid, aOpt, cOpt, vMat1, vMat2 )

        return nothing
    end # solve_hh_decision_rule

# ======================= SOLVERS: STEADY STATE

    # ------------ SS SEARCH
    """
        solve_ss_gs!(M::ModelBasic, K_guess::Float64 ;
            max_iter_ss::Int = 500, max_iter_hh::Int = 500,
            rtol_ss::Float64 = 1E-3, rtol_hh::Float64 = 1E-3,
            step_size_ss::Float64 = 0.7, step_size_hh::Float64 = 0.7,
            print_err_ss::Bool = true, print_err_hh::Bool = false,
            if_golden::Bool = true,
            seed_agg::Int = M.AGG_SEED, simu_size_agg::Int = M.AGG_SIMU_SIZE )

    solves steady state for given `ModelBasic`.s
    """
    function solve_ss_gs!(M::ModelBasic, K_guess::Float64 ;
        max_iter_ss::Int = 500, max_iter_hh::Int = 500,
        rtol_ss::Float64 = 1E-3, rtol_hh::Float64 = 1E-3,
        step_size_ss::Float64 = 0.7, step_size_hh::Float64 = 0.7,
        print_err_ss::Bool = true, print_err_hh::Bool = false,
        if_golden::Bool = true,
        seed_agg::Int = M.AGG_SEED, simu_size_agg::Int = M.AGG_SIMU_SIZE )
        # ------------

        # init guess
        M.K = K_guess
        # init labor factor
        M.L = M.N * sum(M.π .* M.lgrid)

        # fixed-point iter
        for iter in 1:max_iter_ss
            # FIRM DEPT
            M.Y = (M.K)^(M.θ) * (M.L)^(1-M.θ);  # output
            M.r = (M.θ) * (M.K / M.L)^(1-M.θ) - M.δ;  # net interest rate
            M.w = (1-M.θ) * (M.K / M.L)^(M.θ);  # wage rate

            # SOCIAL SECURITY (PASSED)

            # HH DECISION MAKING
            print_err_ss && println("\t- Solving HH ...")
            solve_hh_decision_rule!(M,
                max_iter = max_iter_hh, rtol = rtol_hh,
                print_err = print_err_hh, step_size = step_size_hh, if_golden = if_golden )

            # AGGREGATION
            Knew, M.C = agg_hh_simu(M, simu_size_agg, seed_agg ; a0 = M.K, drop_first = M.AGG_DROP_FIRST_SIZE)
            # ROUND SUMMARY
            if print_err_ss
                println("\t- agg K: ",Knew)
                println("\t- agg L: ",M.L)
                println("\t- agg Y: ",M.Y)
                println("\t- agg C: ",M.C)
                println("\t- agg S: ",M.Y - M.C)
                println("\t- r    : ",M.r)
                println("\t- w    : ",M.w)
            end # if print_err_ss

            # CHECK CONVERGENCE
            local IterErr = abs( Knew / ( M.K + eps() ) - 1.0 )
            print_err_ss && println("- Round ",iter," Rel Err: ",IterErr)

            if IterErr < rtol_ss
                println("- SS Converged in Rel Err Meaning")
                break
            elseif iter == max_iter_ss
                println("- SS Max Iter Reached")
            else
                (M.r < 0.001) && (Knew = 0.05)  # binding error fluctuation
                # update
                M.K = (1.0 - step_size_ss) * M.K + step_size_ss * Knew
            end # if

        end # for iter

        return nothing
    end # solve_ss_gs()

end # src


text:
using BenchmarkTools
const BT=BenchmarkTools.DEFAULT_PARAMETERS;
BT.seconds = 10; BT.samples = 10_000; BT.time_tolerance = 1.0e-10;

using LinearAlgebra, Statistics
using Quadmath, RandomNumbers
using PairArithmetic

ext: jl


text:
if true

	p = Pattern()
	@test get_size(p) == 0
	set_size(p, 3)
	@test get_size(p) == 3

	for i = 0 : 2, j = i+1 : 2
		@test get_edge(p, i, j) == DSL_EDGETYPE_NONE
		@test get_edge(p, j, i) == DSL_EDGETYPE_NONE
		@test has_directed_path(p, i, j) == false
		@test has_directed_path(p, j, i) == false
		@test has_incoming_edge(p, i) == false
		@test has_outgoing_edge(p, i) == false
	end

	@test has_cycle(p) == false
	@test is_DAG(p) == true

	set_edge(p, 0, 1, DSL_EDGETYPE_DIRECTED)
	set_edge(p, 1, 2, DSL_EDGETYPE_DIRECTED)
	@test is_DAG(p) == true
	@test has_directed_path(p, 0, 2) == true
	@test has_incoming_edge(p, 1) == true
	@test has_outgoing_edge(p, 1) == true

	set_edge(p, 2, 0, DSL_EDGETYPE_DIRECTED)
	@test is_DAG(p) == false

	@test get_parents(p, 0) == [2]
	@test get_parents(p, 1) == [0]
	@test get_parents(p, 2) == [1]
	@test get_children(p, 0) == [1]
	@test get_children(p, 1) == [2]
	@test get_children(p, 2) == [0]
	@test get_adjacent_nodes(p, 0) == [1,2]
	@test get_adjacent_nodes(p, 1) == [0,2]
	@test get_adjacent_nodes(p, 2) == [0,1]

	set_edge(p, 0, 2, DSL_EDGETYPE_DIRECTED)
	@test get_children(p, 0) == [1,2]

	p = Pattern()
	set_size(p, 3)
	set_edge(p, 0, 1, DSL_EDGETYPE_DIRECTED)
	set_edge(p, 0, 2, DSL_EDGETYPE_DIRECTED)
	set_edge(p, 1, 2, DSL_EDGETYPE_DIRECTED)
	net = Network()

	mat = round(Int, 3*rand(1000,3))
	dset = Dataset()
	add_int_var(dset, "A")
	add_int_var(dset, "B")
	add_int_var(dset, "C")
	set_number_of_records(dset, 1000)
	dset[0,0:999] = mat[:,1]
	dset[1,0:999] = mat[:,2]
	dset[2,0:999] = mat[:,3]

	to_network(p, dset, net)

	@test get_number_of_nodes(net) == 3
	@test get_parents(p, 0) == []
	@test get_parents(p, 1) == [0]
	@test get_parents(p, 2) == [0,1]
	@test get_children(p, 0) == [1,2]
	@test get_children(p, 1) == [2]
	@test get_children(p, 2) == []

	# test learning pattern
	# pat = Pattern()
	# n = 1000
	# mat = Array(Int, n, 2)
	# for i = 1 : n
	# 	a = rand(0:3)
	# 	b = clamp(a+rand(-1:1), 0, 3)
	# 	mat[i,:] = [a,b]
	# end

	# learn!(pat, Smile.matrix_to_dataset(mat), DSL_PC )

	# println("size: ", get_size(pat))
	# println("isdag: ", is_DAG(pat))
	# println("children 0: ", get_children(pat, 0))
	# println("children 1: ", get_children(pat, 1))

end


text:
function handle_error(r::Response)
    r.status == 403 && error("403 Forbidden request to MapLight API - Probably you need an authentication key")
    r.status != 200 && error("Error querying MapLight API: $(r.status)")

    r.data == "\"Missing required arguments.\"" && error("Missing required arguments in MapLight API request")
end


text:
# ============================================================
# Structs of Initial & Boundary Conditions
# ============================================================

"""
$(TYPEDEF)

Initial & boundary conditions with no distribution function

# Fields

$(FIELDS)
"""
mutable struct IB{T} <: AbstractCondition
    fw::Function
    bc::T
    p::NamedTuple
end

IB(fw, bc = nothing) = IB{typeof(bc)}(fw, bc, NamedTuple())

(m::IB)(args...) = begin
    m.fw(args..., m.p), m.bc(args..., m.p)
end

"""
$(TYPEDEF)

Initial & boundary conditions with 1 distribution function

# Fields

$(FIELDS)
"""
mutable struct IB1F{T} <: AbstractCondition
    fw::Function
    ff::Function
    bc::T
    p::NamedTuple
end

(m::IB1F)(args...) = begin
    m.fw(args..., m.p), m.ff(args..., m.p), m.bc(args..., m.p)
end

#=
function IB1F(fw, vs::AbstractVelocitySpace, gas::AbstractProperty)
    bc = function (args...)
        w = fw(args...)
        prim = begin
            if ndims(w) == 1
                conserve_prim(w, gas.γ)
            else
                mixture_conserve_prim(w, gas.γ)
            end
        end

        return prim
    end

    ff = function (args...)
        prim = bc(args...)

        M = begin
            if !isdefined(vs, :v)
                if ndims(prim) == 1
                    maxwellian(vs.u, prim)
                else
                    mixture_maxwellian(vs.u, prim)
                end
            elseif !isdefined(vs, :w)
                if ndims(prim) == 1
                    maxwellian(vs.u, vs.v, prim)
                else
                    mixture_maxwellian(vs.u, vs.v, prim)
                end
            else
                if ndims(prim) == 1
                    maxwellian(vs.u, vs.v, vs.w, prim)
                else
                    mixture_maxwellian(vs.u, vs.v, vs.w, prim)
                end
            end
        end

        return M
    end

    return IB1F{typeof(bc)}(fw, ff, bc)
end
=#

"""
$(TYPEDEF)

Initial & boundary conditions with 2 distribution functions

# Fields

$(FIELDS)
"""
mutable struct IB2F{T} <: AbstractCondition
    fw::Function
    ff::Function
    bc::T
    p::NamedTuple
end

(m::IB2F)(args...) = begin
    m.fw(args..., m.p), m.ff(args..., m.p)[1], m.ff(args..., m.p)[2], m.bc(args..., m.p)
end

#=
function IB2F(fw, vs::AbstractVelocitySpace, gas::AbstractProperty)
    bc = function (args...)
        w = fw(args...)
        prim = begin
            if ndims(w) == 1
                conserve_prim(w, gas.γ)
            else
                mixture_conserve_prim(w, gas.γ)
            end
        end

        return prim
    end

    ff = function (args...)
        prim = bc(args...)

        if !isdefined(vs, :v)
            if ndims(prim) == 1
                H = maxwellian(vs.u, prim)
                B = H .* gas.K / 2 / prim[end]
            else
                H = mixture_maxwellian(vs.u, prim)
                B = zero(H)
                for j in axes(B, 2)
                    B[:, j] = H[:, j] * gas.K / (2.0 * prim[end, j])
                end
            end
        elseif !isdefined(vs, :w)
            if ndims(prim) == 1
                H = maxwellian(vs.u, vs.v, prim)
                B = H .* gas.K / 2 / prim[end]
            else
                H = mixture_maxwellian(vs.u, vs.v, prim)
                B = zero(H)
                for j in axes(B, 3)
                    B[:, :, j] = H[:, :, j] * gas.K / (2.0 * prim[end, j])
                end
            end
        else
            if ndims(prim) == 1
                H = maxwellian(vs.u, vs.v, vs.w, prim)
                B = H .* gas.K / 2 / prim[end]
            else
                H = mixture_maxwellian(vs.u, vs.v, vs.w, prim)
                B = zero(H)
                for j in axes(B, 4)
                    B[:, :, :, j] = H[:, :, :, j] * gas.K / (2.0 * prim[end, j])
                end
            end
        end

        return H, B
    end

    return IB2F{typeof(bc)}(fw, ff, bc)
end
=#

"""
$(TYPEDEF)

Initial & boundary conditions with 3 distribution functions

# Fields

$(FIELDS)
"""
mutable struct IB3F{T} <: AbstractCondition
    fw::Function
    ff::Function
    fE::Function
    fB::Function
    fL::Function
    bc::T
    p::NamedTuple
end

"""
$(TYPEDEF)

Initial & boundary conditions with 4 distribution functions

# Fields

$(FIELDS)
"""
mutable struct IB4F{T} <: AbstractCondition
    fw::Function
    ff::Function
    fE::Function
    fB::Function
    fL::Function
    bc::T
    p::NamedTuple
end


text:
function test_circuit_codes( p::Parameters )
  funcs = default_funcs( p.numinputs )
  c = random_chromosome( p, funcs )
  cc = circuit_code( c )
  # See https://en.wikibooks.org/wiki/Introducing_Julia/Strings_and_characters#Streams
  iobuffer = IOBuffer()
  print_build_chromosome(iobuffer,c)
  c_str = String(take!(iobuffer))
  ctc = code_to_circuit(cc,p)
  iobuffer = IOBuffer()
  print_build_chromosome(iobuffer,c)
  ctc_str = String(take!(iobuffer))
  if c_str != ctc_str
    println("error c != ctc ")
    println("c_str:   ",c_str)
    println("ctc_str: ",ctc_str)
  end
end

function ttt(j)
  for i = 1:4
   println(j)
   r = j % 2
   j ÷= 5
   println("i: ",i,"  r: ",r,"  j: ",j)
   r = j % 2
   j ÷= 2
   println("i: ",i,"  r: ",r,"  j: ",j)
   r = j % 5
   j ÷= 2
   println("i: ",i,"  r: ",r,"  j: ",j)
  end
end

function intc( j::Integer, p::Parameters )
  result = Int64[]
  for i = 1:p.numinteriors
    jmod = j % 2
    push!( result, jmod )
    println("i: ",i,"  j % 2: ",jmod,"  j: ",j)
    j ÷= 2
    jmod = j % 2
    push!( result, jmod )
    println("i: ",i,"  j % 2: ",jmod,"  j: ",j)
    j ÷= 2
    jmod = j % 5
    push!( result, jmod )
    println("i: ",i,"  j % 5: ",jmod,"  j: ",j)
    j ÷= 5
    println("res: ",result)
  end
  transpose(result)
end
function int_to_circuit_code( c_int::Integer, p::Parameters )
  c_int = Int128(c_int)
  c_code = zeros(Int64,3*p.numinteriors)
  k = 3*p.numinteriors
  for i = p.numinteriors:-1:1
    multiplier = min(p.numlevelsback,i-1+p.numinputs)
    println("i: ",i,"  multiplier: ",multiplier)
    for j = p.nodearity:-1:1
      c_int_mod= c_int % multiplier
      c_int ÷= multiplier
      println("i: ",i,"  j: ",j,"  c_int: ",c_int,"  c_int_mod: ",c_int_mod)
      c_code[k] = c_int_mod
      k -= 1
    end
    multiplier = length(funcs)
    println("i: ",i,"  multiplier: ",multiplier)
    c_int_mod= c_int % multiplier
    c_int ÷= multiplier
    println("i: ",i,"  c_int: ",c_int,"  c_int_mod: ",c_int_mod)
    c_code[k] = c_int_mod
    k -= 1
  end
  c_code
end


text:
using Test
using SymArrays
using SymArrays: symarrlength, _sub2grp, which_symgrp
using TensorOperations
using Random
using CUDA
CUDA.allowscalar(false)

@testset "SymArrays.jl" begin
    @testset "SymArray" begin
        @test symarrlength((3,6,4,3),(3,2,1,3)) == 8400

        S = SymArray{(2,),Float64}(5,5)
        @test nsymgrps(S) == 1
        @test symgrps(S) == (2,)
        @test _sub2grp(S,2,5) == _sub2grp(S,5,2)
        @test _sub2grp(S,3,4) != _sub2grp(S,5,3)
        # _sub2grp has to have same number of arguments as size of N
        @test_throws MethodError _sub2grp(S,3,5,1)
        # indexing the array allows having additional 1s at the end
        @test S[3,5,1] == S[3,5]
        @test_throws BoundsError S[3,5,3]
        # this calculation gives an index for data that is within bounds, but should not be legal
        # make sure this is caught
        @test_throws BoundsError S[0,6]
        @test (S[1,5] = 2.; S[1,5] == 2.)
        @test_throws BoundsError S[0,6] = 2.

        S = SymArray{(3,1,2,2),Float64}(3,3,3,2,4,4,4,4)
        @test nsymgrps(S) == 4
        @test symgrps(S) == (3,1,2,2)
        @test size(S) == (3,3,3,2,4,4,4,4)
        @test length(S) == 2000
        # iterating over all indices should give only the distinct indices,
        # i.e., give the same number of terms as the array length
        @test sum(1 for s in S) == length(S)
        @test sum(1 for I in eachindex(S)) == length(S)

        # calculating the linear index when iterating over Cartesian indices should give sequential access to the array
        @test 1:length(S) == [(LinearIndices(S.data)[_sub2grp(S,Tuple(I)...)...] for I in eachindex(S))...]

        @testset "_sub2grp" begin
            # test that permuting exchangeable indices accesses the same array element
            i1  = _sub2grp(S,1,2,3,1,4,3,2,1)
            @test _sub2grp(S,2,3,1,1,4,3,2,1) == i1
            @test _sub2grp(S,2,3,1,1,3,4,2,1) == i1
            @test _sub2grp(S,2,3,1,1,3,4,1,2) == i1
            @test _sub2grp(S,2,1,3,1,4,3,1,2) == i1
            # make sure that swapping independent indices gives a different array element
            @test _sub2grp(S,2,1,3,1,4,1,3,2) != i1
            @test _sub2grp(S,1,2,3,2,4,3,2,1) != i1
            @test _sub2grp(S,1,2,3,1,1,3,2,4) != i1

            SI = eachindex(S)
            @test first(SI) == CartesianIndex(1,1,1,1,1,1,1,1)

            NN = 4
            maxNdim = 40
            for Ndim = 1:maxNdim
                S = SymArray{(Ndim,),Float64}(ntuple(i->NN,Ndim)...)
                Is = ntuple(i->NN,Ndim)
                @test _sub2grp(S,Is...) == size(S.data)
                Is = ntuple(i->1,Ndim)
                @test _sub2grp(S,Is...) == (1,)
            end
        end

        A = rand(5,5)
        @test A' != A
        @test SymArray{(1,1)}(A) == A
        @test SymArray{(2,)}(A) != A

        A = A + A'
        # the standard equality test uses iteration over the arrays, (a,b) in zip(A,B),
        # which only accesses the actually different indices in SymArrays
        @test SymArray{(2,)}(A) != A
        # broadcasting goes over all indices
        @test all(SymArray{(2,)}(A) .== A)

        # views on existing vector-like types
        x = rand(3)
        S = SymArray{(2,)}(x,2,2)
        @test S.data === x
        fill!(S,8.)
        @test all(S .== 8.)

        S2 = SymArray{(2,)}(0*x,2,2)
        copyto!(S2,S)
        @test S == S2

        x = 5:10
        S = SymArray{(2,)}(x,3,3)
        @test S.data === x

        A = rand(10)
        S = SymArray{(1,)}(A,10)
        @test S.data === A

        # but if called without sizes, the copy constructor should be used
        S = SymArray{(1,)}(A)
        @test S.data !== A
    end

    @testset "which_symgrp" begin
        S = SymArray{(3,2),Float64}(5,5,5,3,3)
        @test which_symgrp(S,1) == (1,3)
        @test which_symgrp(S,3) == (1,3)
        @test which_symgrp(S,4) == (2,2)
        @test which_symgrp(S,5) == (2,2)
    end

    @testset "Contractions" begin
        @testset "Manual" begin
            N,M,O = 3, 5, 8
            for T in (Float64,ComplexF64)
                A = rand(T,N)

                for U in (Float64,ComplexF64)
                    for (n,dims) in enumerate([(N,M,O),(O,N,M),(M,O,N)])
                        B = rand(U,dims...)
                        n==1 && (@tensor D1[j,k] := A[i]*B[i,j,k])
                        n==2 && (@tensor D1[j,k] := A[i]*B[j,i,k])
                        n==3 && (@tensor D1[j,k] := A[i]*B[j,k,i])
                        D2 = contract(A,B,Val(n))
                        @test D1 ≈ D2
                        if T==U
                            contract!(D2,A,B,Val(n))
                            @test D1 ≈ D2
                        else
                            @test_throws MethodError contract!(D2,A,B,Val(n))
                        end
                    end

                    S = SymArray{(3,),U}(N,N,N)
                    rand!(S.data)
                    B = collect(S)
                    @test contract(A,S,Val(1)) == contract(A,S,Val(2))
                    @test contract(A,S,Val(1)) == contract(A,S,Val(3))
                    for n = 1:3
                        @test collect(contract(A,S,Val(n))) ≈ contract(A,B,Val(n))
                    end

                    S = SymArray{(2,1),U}(N,N,N)
                    rand!(S.data)
                    B = collect(S)
                    @test contract(A,S,Val(1)) == contract(A,S,Val(2))
                    for n = 1:3
                        @test collect(contract(A,S,Val(n))) ≈ contract(A,B,Val(n))
                    end
                    @test !(collect(contract(A,S,Val(1))) ≈ collect(contract(A,S,Val(3))))

                    S = SymArray{(2,),U}(N,N)
                    rand!(S.data)
                    B = collect(S)
                    @test contract(A,S,Val(1)) ≈ contract(A,S,Val(2))
                    for n = 1:2
                        @test collect(contract(A,S,Val(n))) ≈ contract(A,B,Val(n))
                    end
                end
            end
        end
        @testset "Generated" begin
            # use small dimension sizes here so the tests do not take too long
            N,M,O = 4, 5, 6
            arrTypes = has_cuda_gpu() ? (Array,CuArray) : (Array,)

            for arrType in arrTypes
                for T in (Float64,ComplexF64)
                    A = rand(T,N,M,O) |> arrType
                    for U in (Float64,ComplexF64)
                        S = SymArray{(2,3,1),U}(arrType,M,M,N,N,N,O)
                        @test storage_type(S) <: arrType

                        rand!(S.data)
                        # first collect GPU->CPU, then SymArray -> Array
                        B = collect(collect(S)) |> arrType
                        @test storage_type(B) <: arrType

                        TU = promote_type(U,T)
                        res21 = SymArray{(1,1,1,3,1),TU}(arrType,N,O,M,N,N,N,O)
                        contract!(res21,A,S,Val(2),Val(1))
                        @tensor res21_tst[i,k,l,m,n,o,p] := A[i,j,k] * B[j,l,m,n,o,p]
                        @test collect(collect(res21)) ≈ collect(res21_tst)
                        if T==U
                            contract!(res21_tst,A,B,Val(2),Val(1))
                            @test collect(collect(res21)) ≈ collect(res21_tst)
                        end

                        res13 = SymArray{(1,1,2,2,1),TU}(arrType,M,O,M,M,N,N,O)
                        contract!(res13,A,S,Val(1),Val(3))
                        @tensor res13_tst[j,k,l,m,n,o,p] := A[i,j,k] * B[l,m,i,n,o,p]
                        @test collect(collect(res13)) ≈ collect(res13_tst)
                        if T==U
                            contract!(res13_tst,A,B,Val(1),Val(3))
                            @test collect(collect(res13)) ≈ collect(res13_tst)
                        end

                        # dimension 3, 4, and 5 should be equivalent
                        contract!(res13,A,S,Val(1),Val(4))
                        @test collect(collect(res13)) ≈ collect(res13_tst)
                        contract!(res13,A,S,Val(1),Val(5))
                        @test collect(collect(res13)) ≈ collect(res13_tst)
                    end

                    A = rand(T,N) |> arrType
                    S = SymArray{(2,3,1),T}(arrType,N,N,N,N,N,N)
                    rand!(S.data)
                    # first collect GPU->CPU, then SymArray -> Array
                    B = collect(collect(S)) |> arrType

                    res11 = SymArray{(1,3,1),T}(arrType,N,N,N,N,N)
                    contract!(res11,A,S,Val(1),Val(1))
                    @tensor res11_tst[j,k,l,m,n] := A[i] * B[i,j,k,l,m,n]
                    @test collect(collect(res11)) ≈ collect(res11_tst)
                    contract!(res11_tst,A,B,Val(1),Val(1))
                    @test collect(collect(res11)) ≈ collect(res11_tst)

                    res13 = SymArray{(2,2,1),T}(arrType,N,N,N,N,N)
                    contract!(res13,A,S,Val(1),Val(3))
                    @tensor res13_tst[j,k,l,m,n] := A[i] * B[j,k,i,l,m,n]
                    @test collect(collect(res13)) ≈ collect(res13_tst)
                    contract!(res13_tst,A,B,Val(1),Val(3))
                    @test collect(collect(res13)) ≈ collect(res13_tst)

                    # dimension 3, 4, and 5 should be equivalent
                    contract!(res13_tst,A,B,Val(1),Val(4))
                    @test collect(collect(res13)) ≈ collect(res13_tst)
                    contract!(res13_tst,A,B,Val(1),Val(5))
                    @test collect(collect(res13)) ≈ collect(res13_tst)

                    res16 = SymArray{(2,3),T}(arrType,N,N,N,N,N)
                    contract!(res16,A,S,Val(1),Val(6))
                    @tensor res16_tst[j,k,l,m,n] := A[i] * B[j,k,l,m,n,i]
                    @test collect(collect(res16)) ≈ collect(res16_tst)
                    contract!(res16_tst,A,B,Val(1),Val(6))
                    @test collect(collect(res16)) ≈ collect(res16_tst)

                    # check that contraction from SymArray to Array works
                    A = rand(T,M,O) |> arrType
                    S = SymArray{(1,2,1),T}(arrType,N,M,M,O)
                    rand!(S.data)
                    B = collect(collect(S)) |> arrType

                    res12 = zeros(T,O,N,M,O) |> arrType
                    res12_tst = zeros(T,O,N,M,O) |> arrType
                    contract!(res12,A,S,Val(1),Val(2))
                    contract!(res12_tst,A,B,Val(1),Val(2))
                    @test collect(res12) ≈ collect(res12_tst)
                end
            end
        end
    end
end


text:
using HDF5: attributes, h5open, ishdf5, File
using TimeZones: TimeZone, ZonedDateTime, localzone

"""
    readhspy(filename::String, name=nothing)

Read a HyperSpectrum from a HyperSpy-style HDF5 file.  The `name` argument allows the user to
optionally specify an experiment to load.
"""
function readhspy(filename::String, name=nothing)
    function read_attributes(fid)
        cs = attributes(fid)
        Dict(map(k->k=>read(cs,k),keys(cs))...)
    end
    function assign(props, prop, val)
        if !ismissing(val) props[prop]=val end
    end
    function read_stage(stage_attr)
        Dict(
            :X => get(stage_attr, "x", 0.0)*0.1,
            :Y => get(stage_attr, "y", 0.0)*0.1,
            :Z => get(stage_attr, "z", 0.0)*0.1,
            :R => deg2rad(get(stage_attr, "rotation", 0.0)),
            :T => deg2rad(get(stage_attr, "tilt_alpha", 0.0)),
            :B => deg2rad(get(stage_attr, "tilt_beta", 0.0)),
        )
    end
    function read_eds(props, det_attr)
        assign(props, :Azimuthal, get(det_attr, "azimuth_angle", missing)*deg2rad(1.0))
        assign(props, :Elevation, get(det_attr, "elevation_angle", missing)*deg2rad(1.0))
        assign(props, :HSResolution, get(det_attr, "energy_resolution_MnKa", missing))
        assign(props, :LiveTime, get(det_attr, "live_time", missing))
        assign(props, :RealTime, get(det_attr, "real_time", missing))
    end
    function read_datetime(general)
        tz = get(general, "time_zone", missing)
        res = nothing
        try
            dt = get(general, "date", "1.1.1900")
            res = DateTime(dt, dateformat"dd.mm.yyyy")
        catch
            @warn "Error parsing date: $dt not of format dd.mm.yyyy"
            res = Date("1.1.1970", dateformat("dd.mm.yy"))
        end
        try
            tm = get(general, "time", missing)
            if !ismissing(tm)
                res += Time(tm, dateformat"HH:MM:SS")
            end
        catch
            @warn "Error parsing time: $tm not of format HH:MM:SS"
        end
        try
            if !ismissing(tz)
                res = ZonedDateTime(res, TimeZone(tz))
            end
        catch
            @warn "Error parsing timezone: $tz is not a standard timezone."
        end
        res
    end
    #=
        ├── Acquisition_instrument
        │   ├── SEM
        │   │   ├── Detector
        │   │   │   ├── detector_type
        │   │   │   └── EDS
        │   │   │       ├── azimuth_angle (º)
        │   │   │       ├── elevation_angle (º)
        │   │   │       ├── energy_resolution_MnKa (eV)
        │   │   │       ├── live_time (s)
        │   │   │       └── real_time (s)
        │   │   ├── beam_current (nA)
        │   │   ├── beam_energy (keV)
        │   │   ├── probe_area (nm²)
        │   │   ├── convergence_angle (mrad)
        │   │   ├── magnification
        │   │   ├── microscope
        │   │   ├── Stage
        │   │   │   ├── rotation (º)
        │   │   │   ├── tilt_alpha (º)
        │   │   │   ├── tilt_beta (º)
        │   │   │   ├── x (mm)
        │   │   │   ├── y (mm)
        │   │   │   └── z (mm)
        │   │   └── working_distance (mm)
        │   └── TEM
        │       ├── Detector
        │       │   ├── EDS
        │       │   │   ├── azimuth_angle (º)
        │       │   │   ├── elevation_angle (º)
        │       │   │   ├── energy_resolution_MnKa (eV)
        │       │   │   ├── live_time (s)
        │       │   │   └── real_time (s)
        │       │   └── EELS
        │       │       ├── aperture (mm)
        │       │       ├── collection_angle (mrad)
        │       │       ├── dwell_time (s)
        │       │       ├── exposure (s)
        │       │       ├── frame_number
        │       │       └── spectrometer
        │       ├── Biprism
        │       │   ├── azimuth_angle (º)
        │       │   ├── position
        │       │   └── voltage (V)
        │       ├── acquisition_mode
        │       ├── beam_current (nA)
        │       ├── beam_energy (keV)
        │       ├── probe_area (nm²)
        │       ├── camera_length (mm)
        │       ├── convergence_angle (mrad)
        │       ├── magnification
        │       ├── microscope
        │       └── Stage
        │           ├── rotation (º)
        │           ├── tilt_alpha (º)
        │           ├── tilt_beta (º)
        │           ├── x (mm)
        │           ├── y (mm)
        │           └── z (mm)
        ├── General
        │   ├── authors
        │   ├── date
        │   ├── doi
        │   ├── original_filename
        │   ├── notes
        │   ├── time
        │   ├── time_zone
        │   └── title
        ├── Sample
        │   ├── credits
        │   ├── description
        │   ├── elements
        │   ├── thickness
        │   └── xray_lines
        └── Signal
            ├── FFT
            │   └── shifted
            ├── Noise_properties
            │   ├── Variance_linear_model
            │   │   ├── correlation_factor
            │   │   ├── gain_factor
            │   │   ├── gain_offset
            │   │   └── parameters_estimation_method
            │   └── variance
            ├── binned
            ├── quantity
            ├── signal_type
            └── signal_origin
    =#
    h5open(filename,"r") do hdf5
        ff = read_attributes(hdf5)
        @assert get(ff, "file_format", "") == "HyperSpy" "This file does not appear to be a HyperSpy file."
        ver = parse(Float64, get(ff, "file_format_version", 0.0))
        (ver < 1.2) && @warn "This reader may not be compatible with version $ver of HyperSpy data files."
        props = Dict{Symbol,Any}()
        props[:Filename] = filename
        if haskey(hdf5, "Experiments")
            exps = keys(hdf5["Experiments"])
            exp = something(name, first(exps))
            exp_key = "Experiments/$exp"
            sig_key = "$exp_key/metadata/Signal"
            if haskey(hdf5, sig_key)
            #=  └── Signal
                    ├── FFT
                    │   └── shifted # N/A
                    ├── Noise_properties # N/A
                    │   ├── Variance_linear_model
                    │   │   ├── correlation_factor
                    │   │   ├── gain_factor
                    │   │   ├── gain_offset
                    │   │   └── parameters_estimation_method
                    │   └── variance
                    ├── binned
                    ├── quantity
                    ├── signal_type
                    └── signal_origin =#
                sig = read_attributes(hdf5[sig_key])
                st = get(sig, "signal_type", "Unknown")
                if st == "EDS_TEM"
                    props[:HSSignalType] = "EDS_TEM"
                elseif st == "EDS_SEM"
                    props[:HSSignalType] = "EDS_SEM"
                elseif st == "EELS"
                    props[:HSSignalType] = "EELS"
                else
                    @warn "Unsupported signal type $st which loading a HyperSpy file."
                    assign(props, :HSSignalType, st)
                end
                @assert get(sig, "recorded_by", "spectrum") == "spectrum" "Orderings other than `spectrum` are not supported"
                assign(props, :HSSignalOrigin, get(sig, "signal_origin", missing))
                assign(props, :HSQuantity, get(sig, "quantity",missing))
                assign(props, :HSBinned, get(sig, "binned",missing))
            end
            sem_key = "$exp_key/metadata/Acquisition_instrument/SEM"
            if haskey(hdf5, sem_key)
                @assert get(props, :HSSignalType, "EDS_SEM")=="EDS_SEM"
            #=  ├── Acquisition_instrument
                │   ├── SEM
                │   │   ├── Detector
                │   │   │   ├── detector_type
                │   │   │   └── EDS
                │   │   │       ├── azimuth_angle (º)
                │   │   │       ├── elevation_angle (º)
                │   │   │       ├── energy_resolution_MnKa (eV)
                │   │   │       ├── live_time (s)
                │   │   │       └── real_time (s)
                │   │   ├── beam_current (nA)
                │   │   ├── beam_energy (keV)
                │   │   ├── probe_area (nm²)
                │   │   ├── convergence_angle (mrad)
                │   │   ├── magnification
                │   │   ├── microscope
                │   │   ├── Stage
                │   │   │   ├── rotation (º)
                │   │   │   ├── tilt_alpha (º)
                │   │   │   ├── tilt_beta (º)
                │   │   │   ├── x (mm)
                │   │   │   ├── y (mm)
                │   │   │   └── z (mm)
                │   │   └── working_distance (mm) =#
                assign(props, :InstrumentType, "SEM")
                sem_attr=read_attributes(hdf5[sem_key])
                assign(props, :ProbeCurrent, get(sem_attr, "beam_current", missing)) # nA
                assign(props, :BeamEnergy, get(sem_attr, "beam_energy", missing)*1000.0) # keV
                assign(props, :HSProbeArea, get(sem_attr, "probe_area", missing))
                assign(props, :HSConvergenceAngle, get(sem_attr, "convergence_angle", missing) * 0.001) # mrad
                assign(props, :HSMagnification, get(sem_attr, "magnification", missing))
                assign(props, :Instrument,  get(sem_attr, "microscope", missing))
                assign(props, :WorkingDistance, get(sem_attr,"working_distance", missing)*0.1) # mm
                if haskey(hdf5[sem_key],"Stage")
                    props[:Stage] = read_stage(read_attributes(hdf5["$sem_key/Stage"]))
                end
                det_key = "$sem_key/Detector"
                if haskey(hdf5, det_key)
                    det_attr = read_attributes(hdf5[det_key])
                    assign(props, :HSDetectorType, get(det_attr,"detector_type",missing))
                    eds_key = "$det_key/EDS"
                    if haskey(hdf5, eds_key)
                        read_eds(props, read_attributes(hdf5[eds_key]))
                    end
                end
            end
            tem_key = "$exp_key/metadata/Acquisition_instrument/TEM"
            if haskey(hdf5, tem_key)
            #=  ├── Acquisition_instrument
                │   └── TEM
                │       ├── Detector
                │       │   ├── EDS
                │       │   │   ├── azimuth_angle (º)
                │       │   │   ├── elevation_angle (º)
                │       │   │   ├── energy_resolution_MnKa (eV)
                │       │   │   ├── live_time (s)
                │       │   │   └── real_time (s)
                │       │   └── EELS
                │       │       ├── aperture (mm)
                │       │       ├── collection_angle (mrad)
                │       │       ├── dwell_time (s)
                │       │       ├── exposure (s)
                │       │       ├── frame_number
                │       │       └── spectrometer
                │       ├── Biprism
                │       │   ├── azimuth_angle (º)
                │       │   ├── position
                │       │   └── voltage (V)
                │       ├── acquisition_mode
                │       ├── beam_current (nA)
                │       ├── beam_energy (keV)
                │       ├── probe_area (nm²)
                │       ├── camera_length (mm)
                │       ├── convergence_angle (mrad)
                │       ├── magnification
                │       ├── microscope
                │       └── Stage
                │           ├── rotation (º)
                │           ├── tilt_alpha (º)
                │           ├── tilt_beta (º)
                │           ├── x (mm)
                │           ├── y (mm)
                │           └── z (mm) =#
                assign(props, :InstrumentType, "TEM")
                tem_attr=read_attributes(hdf5[tem_key])
                assign(props, :ProbeCurrent, get(tem_attr, "beam_current", missing)) # nA
                assign(props, :BeamEnergy, get(tem_attr, "beam_energy", missing)*1000.0) # keV
                assign(props, :HSProbeArea, get(tem_attr, "probe_area", missing))
                assign(props, :HSConvergenceAngle, get(tem_attr, "convergence_angle", missing) * 0.001 )# mrad
                assign(props, :HSMagnification, get(tem_attr, "magnification", missing))
                assign(props, :Instrument,  get(tem_attr, "microscope", missing))
                assign(props, :HSCameraLength, get(tem_attr, "camera_length", missing)*0.1) # mm
                assign(props, :HSTEMTiltStage, get(tem_attr, "tilt_stage", missing) * deg2rad(1.0))
                assign(props, :HSAcquisitionMode, get(tem_attr, "acquisition_mode", missing))
                if haskey(hdf5[tem_key],"Stage")
                    props[:Stage] = read_stage(read_attributes(hdf5["$tem_key/Stage"]))
                end
                det_key = "$tem_key/Detector"
                if haskey(hdf5, det_key)
                    det_attr = read_attributes(hdf5[det_key])
                    assign(props, :HSDetectorType, get(det_attr,"detector_type",missing))
                    eds_key = "$det_key/EDS"
                    @assert get(props, :HSSignalType, "EDS_TEM")=="EDS_TEM"
                    if haskey(hdf5, eds_key)
                        read_eds(props, read_attributes(hdf5[eds_key]))
                    end
                    eels_key = "$det_key/EELS"
                    if haskey(hdf5, eels_key)
                        @assert get(props, :HSSignalType, "EELS")=="EELS"
                        eels_attr = read_attributes(hdf5[eels_attr])
                        assign(props, :HSAperture, get(eels_attr,"aperture",missing))
                        assign(props, :HSCollectionAngle, get(eels_attr,"collection_angle",missing) * 0.001)
                        assign(props, :HSDwellTime, get(eels_attr,"dwell_time",missing))
                        assign(props, :HSExposure, get(eels_attr,"exposure",missing))
                        assign(props, :HSFrameNumber, get(eels_attr,"frame_number",missing))
                        assign(props, :HSSpectrometer, get(eels_attr,"spectrometer",missing))
                    end
                end
            end
            #= │       ├── Biprism
               │       │   ├── azimuth_angle (º)
               │       │   ├── position
               │       │   └── voltage (V) =#
            bi_key = "$tem_key/Biprism"
            if haskey(hdf5, bi_key)
                bi_attr = read_attributes(hdf5[bi_key])
                assign(props, :HSBiprismAzAngle, get(bi_key, "azimuth_angle", missing) * deg2rad(1.0))
                assign(props, :HSBiprismPosition, get(bi_key, "position", missing))
                assign(props, :HSBiprismVoltage, get(bi_key, "voltage", missing))
            end
            #=  ├── General
                    │   ├── authors
                    │   ├── date
                    │   ├── doi
                    │   ├── original_filename
                    │   ├── notes
                    │   ├── time
                    │   ├── time_zone
                    │   └── title =#
            gen_key = "$exp_key/metadata/General"
            if haskey(hdf5, gen_key)
                general = read_attributes(hdf5[gen_key])
                assign(props, :Owner, get(general, "authors", missing))
                assign(props, :AcquisitionTime, read_datetime(general))
                assign(props, :Description, get(general, "notes", missing))
                assign(props, :Name, get(general, "title", missing))
                assign(props, :DOI, get(general, "doi", missing))
                assign(props, :Filename, get(general, "original_filename", missing))
                assign(props, :HSNotes, get(general, "notes", missing))
            end
            samp_key = "$exp_key/metadata/Sample"
            if haskey(hdf5, samp_key)
                #=  ├── Sample
                    │   ├── credits
                    │   ├── description
                    │   ├── elements
                    │   ├── thickness
                    │   └── xray_lines =#
                samp_attr = read_attributes(hdf5[samp_key])
                assign(props, :Sample, get(samp_attr, "description", missing))
                assign(props, :HSCredits, get(samp_attr, "credits", missing))
                elms = get(samp_attr, "elements", missing)
                if !ismissing(elms)
                    props[:Elements] = map(e->parse(Element, e), elms)
                end
                assign(props, :HSThickness, get(samp_attr, "thickness", missing))
                assign(props, :HSXRayLines, get(samp_attr, "xray_lines", missing))
            end
            scale, offset = 10.0, 0.0
            axisnames, fov, offsets = String[], Float64[], Float64[]
            i=0
            while haskey(hdf5["$exp_key"],"axis-$i")
                axis = read_attributes(hdf5["$exp_key/axis-$i"])
                unitname = uppercase(get(axis, "units", "eV"))
                if unitname == "KEV"
                    unit = 1000.0
                elseif unitname == "EV"
                    unit = 1.0
                elseif unitname == "MM"
                    unit = 0.1
                elseif unitname == "NM"
                    unit = 1.0e-7
                else
                    unit = 1.0
                end
                push!(axisnames, uppercase(get(axis, "name", "axis-$i")))
                scale = get(axis, "scale", 10.0/unit) * unit
                offset = get(axis, "offset", 0.0) * unit
                push!(fov, scale * get(axis, "size", 1))
                push!(offsets, offset)
                i+=1
            end
            data = read(hdf5, "$exp_key/data")
            axisnames = reverse(axisnames[1:end-1])
            fov = reverse(fov[1:end-1])
            offsets = reverse(offsets[1:end-1])
            # Now swap the first two axes because column major vs row major...
            if length(axisnames)>1
                newdims = collect(1:length(axisnames)+1)
                newdims[2], newdims[3] = newdims[3], newdims[2]
                data = permutedims(data, newdims)
                axisnames[1], axisnames[2] = axisnames[2], axisnames[1]
                fov[1], fov[2] = fov[2], fov[1]
                offsets[1], offsets[2] = offsets[2], offsets[1]
            end
            if haskey(props, :HSResolution)
                props[:Detector] = simpleEDS(size(data,1), scale, offset, props[:HSResolution])
            end
            return HyperSpectrum(LinearEnergyScale(offset, scale), props, data, #
                axisnames=axisnames, fov=fov, offset=offsets)
        end
    end
end

function write_hspy(filename::String, hspec::HyperSpectrum, mode=:SEM)
    h5open(filename) do hdf5
        write_hspy(hdf5, hspec, mode)
    end
end

function write_hspy(hdf5::File, hspec::HyperSpectrum, mode=:SEM)
    function assign_attr(gr, att, prop, tran=identity)
        if haskey(hspec,prop)
            attributes(gr)[att] = tran(hspec[prop])
        end
    end
    groot = create_group(hdf5, "Experiments/$(hspec[:Name])")
    for (i, aa) in enumerate(AxisArrays.axes(h.counts))
        idx = ndims(h.counts)-i
        gax = create_group(groot, "axis-$idx")
        attributes(gax)["name"] = String(axisnames(h.counts)[idx])
        if i==1
            attributes(gax)["offset"] = first(aa.val)
            attributes(gax)["scale"] = step(aa.val)
            attributes(gax)["size"] = size(h.counts,i)
            attributes(gax)["units"] = "eV"
        else
            attributes(gax)["offset"] = first(aa.val)*1.0e7
            attributes(gax)["scale"] = step(aa.val)*1.0e7
            attributes(gax)["size"] = size(h.counts,i)
            attributes(gax)["units"] = "nm"
        end
    end
    geds = create_group(ginst, "Acquisition_instrument/"*String(mode)*"/Detector/EDS")
    assign_attr(geds, "azimuth_angle", :Azimuthal, rad2deg)
    assign_attr(geds, "elevation_angle", :TakeOffAngle, rad2deg)
    assign_attr(geds, "energy_resolution_MnKa", :Resolution)
    ggen = create_group(groot, "General")
    if haskey(attributes(gdet), :AcquisitionTime)
        attributes(ggen)["date"] = "dd.mm.yyyy"
        attributes(ggen)["time"] = "hh.mm.ss"
        attributes(ggen)["time_zone"] = "+5.0"
    end
    gsamp = create_group(groot, "Sample")
    assign_attr(gsamp, "description", :Description)
    els = elms(hspec)
    if length(els)>0
        attributes(gsamp)["elements"] = join(symbol.(els),",")
    end
    gsamp = create_group(groot, "Signal")
    attributes(gsamp)["binned"] = true
    attributes(gsamp)["recorded_by"] = "spectrum"
    attributes(gsamp)["signal_type"] = "EDS_"*String(mode)
    groot["data"] = hspec.counts
end

"""
    ishspy(filename::String)::Bool

Is the file a HyperSpy-style HDF5 file?
"""
function ishspy(filename::String)::Bool
    return ishdf5(filename) && h5open(filename,"r") do h
        a = attributes(h)
        haskey(a,"file_format") && read(a["file_format"])=="HyperSpy"
    end
end

max_stars_repo_name: NicholasWMRitchie/NeXLSpectrum

max_issues_repo_name: NicholasWMRitchie/NeXLSpectrum

max_forks_repo_name: NicholasWMRitchie/NeXLSpectrum


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

struct NetworkInterfacesApi <: SwaggerApi
    client::Swagger.Client
end

"""

Creates or updates a network interface.
Param: resourceGroupName::String (required)
Param: networkInterfaceName::String (required)
Param: parameters::NetworkInterface (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: NetworkInterface
"""
function networkInterfacesCreateOrUpdate(_api::NetworkInterfacesApi, resourceGroupName::String, networkInterfaceName::String, parameters, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "PUT", NetworkInterface, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/networkInterfaces/{networkInterfaceName}", ["azure_auth"], parameters)
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "networkInterfaceName", networkInterfaceName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Deletes the specified network interface.
Param: resourceGroupName::String (required)
Param: networkInterfaceName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: Nothing
"""
function networkInterfacesDelete(_api::NetworkInterfacesApi, resourceGroupName::String, networkInterfaceName::String, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "DELETE", Nothing, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/networkInterfaces/{networkInterfaceName}", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "networkInterfaceName", networkInterfaceName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Gets information about the specified network interface.
Param: resourceGroupName::String (required)
Param: networkInterfaceName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Param: expand::String
Return: NetworkInterface
"""
function networkInterfacesGet(_api::NetworkInterfacesApi, resourceGroupName::String, networkInterfaceName::String, api_version::String, subscriptionId::String; expand=nothing, _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", NetworkInterface, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/networkInterfaces/{networkInterfaceName}", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "networkInterfaceName", networkInterfaceName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_param(_ctx.query, "\$expand", expand)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Gets all route tables applied to a network interface.
Param: resourceGroupName::String (required)
Param: networkInterfaceName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: EffectiveRouteListResult
"""
function networkInterfacesGetEffectiveRouteTable(_api::NetworkInterfacesApi, resourceGroupName::String, networkInterfaceName::String, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "POST", EffectiveRouteListResult, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/networkInterfaces/{networkInterfaceName}/effectiveRouteTable", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "networkInterfaceName", networkInterfaceName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Gets all network interfaces in a resource group.
Param: resourceGroupName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: NetworkInterfaceListResult
"""
function networkInterfacesList(_api::NetworkInterfacesApi, resourceGroupName::String, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", NetworkInterfaceListResult, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/networkInterfaces", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Gets all network interfaces in a subscription.
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: NetworkInterfaceListResult
"""
function networkInterfacesListAll(_api::NetworkInterfacesApi, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", NetworkInterfaceListResult, "/subscriptions/{subscriptionId}/providers/Microsoft.Network/networkInterfaces", ["azure_auth"])
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Gets all network security groups applied to a network interface.
Param: resourceGroupName::String (required)
Param: networkInterfaceName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: EffectiveNetworkSecurityGroupListResult
"""
function networkInterfacesListEffectiveNetworkSecurityGroups(_api::NetworkInterfacesApi, resourceGroupName::String, networkInterfaceName::String, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "POST", EffectiveNetworkSecurityGroupListResult, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/networkInterfaces/{networkInterfaceName}/effectiveNetworkSecurityGroups", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "networkInterfaceName", networkInterfaceName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

export networkInterfacesCreateOrUpdate, networkInterfacesDelete, networkInterfacesGet, networkInterfacesGetEffectiveRouteTable, networkInterfacesList, networkInterfacesListAll, networkInterfacesListEffectiveNetworkSecurityGroups


text:
# Visualize results on a particular problem that are saved in the results folder (names should match)

using Statistics
using PyPlot
using HDF5
using DataFrames
using CSV

include("Benchmarks.jl")

# OPTIONAL: write results in a csv file
updateCSV = "on"

# NAME OF THE INSTANCE (should match the stored data in the results/ folder)
PB = "PB1"
tweak = 1
T = 20000
N = 100

# INCLUDE ADSWITCH?
withAdSwitch = "off"
NAd = 50 # number of repetitions used for AdSwitch (can be different from N)

fname="results/$(PB)_tweak_$(tweak)_"

# POLICIES TO INCLUDE (check that you have data for them)
policies=["Oracle","klUCB","EXP3SOpt","SWKLUCB","DKLUCB","MKLUCB","MKLUCB2","CUSUMKLUCB","CUSUMKLUCB2","GLRKLUCB","GLRKLUCBConstant","GLRKLUCBLoc","GLRKLUCBLocConstant"]
names=["Oracle","klUCB","EXP3SOpt","SWKLUCB","DKLUCB","MKLUCB","MKLUCB2","CUSUMKLUCB","CUSUMKLUCB2","GLRKLUCB","GLRKLUCBConstant","GLRKLUCBLoc","GLRKLUCBLocConstant"]

# sub-sample of policies for the regret curves
#policies=["Oracle","klUCB","EXP3SOpt","SWKLUCB","DKLUCB","MKLUCB2","CUSUMKLUCB2","GLRKLUCB","GLRKLUCBLoc"]
#names =["Oracle","klUCB","EXP3.S","SW-klUCB","D-klUCB","M-klUCB","CUSUM-klUCB","GLR-klUCB Global","GLR-klUCB Local"]

lP = length(policies)

df=DataFrame(Algo= String[],Regret=Float32[],Stdev=Float32[],NbRestarts=Float32[])
clf()

# plot the problem
name = "$(fname)_$(policies[1])_T_$(T)_N_$(N)"
Problem = h5read(name,"Problem")
clf()
PlotProblem(Problem)
figure()

# display results and plot the regret
for imeth in 1:lP
    policy = policies[imeth]
    policyname = names[imeth]
    name = "$(fname)_$(policy)_T_$(T)_N_$(N)"
    Regret=h5read(name,"Regret")
    Restarts=h5read(name,"Restarts")
    tsave=floor.(Int,collect(range(1,T,length=500)))
    # print stuff
    regfinal = mean(Regret[:,end])
    stdevfinal = std(Regret[:,end])
    restarts = mean(Restarts)
    push!(df,("$(policy)",regfinal,stdevfinal,restarts))
    print("Mean final regret for $(policyname) is $(regfinal)\n")
    print("Standard deviation of the final regret for $(policyname) is $(stdevfinal))\n")
	print("Mean number of restarts for $(policyname) is $(restarts)\n\n")
	# plot regret
	plot(vec(tsave),vec(mean(Regret,dims=1)),label="$(policyname)")
end

if withAdSwitch == "on"
    nameAd = "$(fname)_AdSwitch_T_$(T)_N_$(NAd)"
    Regret = h5read(nameAd,"Regret")
    Restarts=h5read(nameAd,"Restarts")
    regfinal = mean(Regret[:,end])
    stdevfinal = std(Regret[:,end])
    restarts = mean(Restarts)
    push!(df,("AdSwitch",regfinal,stdevfinal,restarts))
    tsave=floor.(Int,collect(range(1,T,length=500)))
    print("Mean final regret for AdSwitch is $(regfinal)\n")
    print("Standard deviation of the final regret for AdSwitch is $(stdevfinal))\n")
	print("Mean number of restarts for AdSwitch is $(restarts)\n\n")
    plot(vec(tsave),vec(mean(Regret,dims=1)),label="AdSwitch")
end

legend()

if (updateCSV == "on")
    CSV.write("$(fname)T_$(T)_N_$(N).csv",df)
end


text:
module BSON

using Core: SimpleVector, TypeName
export bson

const BSONDict = Dict{Symbol,Any}
const BSONArray = Vector{Any}
const Primitive = Union{Nothing,Bool,Int32,Int64,Float64,String,Vector{UInt8},BSONDict,BSONArray}

@enum(BSONType::UInt8,
  eof, double, string, document, array, binary, undefined, objectid, boolean,
  datetime, null, regex, dbpointer, javascript, symbol, javascript_scoped,
  int32, timestamp, int64, decimal128, minkey=0xFF, maxkey=0x7F)

applychildren!(::Function, x) = x

function applychildren!(f::Function, x::BSONDict)::BSONDict
  for k in keys(x)
    x[k] = f(x[k])
  end
  return x
end

function applychildren!(f::Function, x::BSONArray)::BSONArray
  for i = 1:length(x)
    x[i] = f(x[i])
  end
  return x
end

include("write.jl")
include("read.jl")
include("extensions.jl")
include("anonymous.jl")
include("fileio_interface.jl")

using Base.Meta

macro save(file, ks...)
  ks = map(k -> k isa Symbol ? (k,k) :
                isexpr(k,:(=)) && k.args[1] isa Symbol ? (k.args[1],k.args[2]) :
                error("Unrecognised @save expression $k"), ks)
  :(bson($(esc(file)), Dict($([:($(Expr(:quote,s))=>$(esc(k))) for (s,k) in ks]...) )))
end

macro load(file, ks...)
  @assert all(k -> k isa Symbol, ks)
  ss = Expr.(:quote, ks)
  quote
    data = load($(esc(file)))
    ($(esc.(ks)...),) = ($([:(data[$k]) for k in ss]...),)
    nothing
  end
end

end # module


text:
#include("../dataset_reader.jl")
const dataset4_labels = Float64[1, 2, 3, 4, 5, 6, 7]
const dataset4_nfeatures = 12
const dataset4_filename = "drug_consumption.data"

function get_numerical_dataset4(path::String="")
    data = get_data(joinpath(path, dataset4_filename); delete_first_line=false)

    class_dict = create_numerical_map(getindex.(data, 14))

    columns_number = 13
    for line = 1:length(data) data[line] = data[line][2:columns_number+1] end

    columns_number = length(data[1])
    conversion_pattern = Vector{Union{DataType, Dict{String, <: Number}}}(undef, columns_number)
    for i = 1:12 conversion_pattern[i] = Float64 end
    for i = 13:columns_number conversion_pattern[i] = class_dict end

    convert_values(data, conversion_pattern)
end

function get_dataset4(path::String=""; shuffle::Bool=false, test_perc::Float64=0.2,
                      normalise::Bool=false, separate_test_data::Bool=true, balance_classes::Bool=false)
    data = get_numerical_dataset4(path)
    separate_dataset(data, test_percentage=test_perc, shuffle=shuffle, normalise=normalise,
                     separate_test_data=separate_test_data, balance_classes=balance_classes)
end

function get_dataset4_setup(path::String=""; normalise::Bool=false, separate::Bool=false,
                            balance::Bool=false)
    dataset = get_dataset4(path, normalise=normalise, separate_test_data=separate, balance_classes=balance)
    if separate
        SeparatedDataset(dataset..., dataset4_labels, dataset4_nfeatures, 3)
    else
        Dataset(dataset..., dataset4_labels, dataset4_nfeatures, 3)
    end
end

max_stars_repo_name: Lucas-Muniz/MNEAT

max_issues_repo_name: Lucas-Muniz/MNEAT

max_forks_repo_name: Lucas-Muniz/MNEAT


text:
using Gadfly
using CSV
using DataFrames

data = CSV.read("C:/Users/ztjsw/Desktop/data.txt", delim = ',', header = true)

data[:chromosome] = string.(data[:chromosome])
data[:neg_log_p] = -log10.(data[:pvalue])
data[:adj_pos] = copy(data[:position])

function cal_adj_pos(data)
    #calculate adjusted position
    start = 0
    for i in 2:nrow(data)
        if data[:chromosome][i] != data[:chromosome][i-1]
            start += data[:position][i-1]
        end
        data[:adj_pos][i] += start
    end
    data
end

data_new = cal_adj_pos(data)

## Plot
#calculate postion for xticks
x_chrom_pos = by(data_new, :chromosome, mid_pos = :adj_pos => d -> (maximum(d) + minimum(d))/2)[2]
x_chrom_label = unique(data_new[:chromosome])

function pos_to_chrom(iPos)
    # return the name of its chromosome for a position
    return x_chrom_label[x_chrom_pos .== iPos][1]
end

p=plot(data_new, x=:adj_pos, y=:neg_log_p, color=:chromosome, Geom.point,
        Guide.xlabel("Chromosome"),
        Guide.ylabel("-log10(p)"),
        Guide.xticks(ticks=x_chrom_pos),
        Scale.x_continuous(labels = pos_to_chrom),  #labels is a function which maps a coordinate value to a string label
        Theme(key_position = :none, grid_line_width=0mm, point_size=0.1cm))

display(p)


text:
mutable struct Point <: Shape
    position::Luxor.Point
    color::Union{Blending,Coloring}
    size::Float64
end

function Point( ; position, color, size, kwargs...)
    return Point(position, color, size)
end


text:
# just to avoid magical constants in code
const CODON_LENGTH = 3

const BASES = DNASeq("ACGT")
const BASEINTS = Dict(DNA_A => 1, DNA_C => 2, DNA_G => 3, DNA_T => 4)

const DEBUG = true

"""Only assert when DEBUG is true"""
macro myassert(exp, msg)
    if DEBUG
        return :($exp || error($msg))
    end
    return :()
end

macro includeif(test, exp)
    if test
        return exp
    end
    return :()
end

"""LogSumExp in base 10.

Borrowed from StatsFuns.jl.
"""
function logsumexp10(x::AbstractArray{T}) where {T <: Real}
    S = typeof(exp10(zero(T)))    # because of 0.4.0
    isempty(x) && return -S(Inf)
    u = maximum(x)
    abs(u) == Inf && return any(isnan, x) ? S(NaN) : u
    s = zero(S)
    for i = 1:length(x)
        @inbounds s += exp10(x[i] - u)
    end
    log10.(s) + u
end

@generated function summax(a, b)
    return quote
        result = a[1] + b[1]
        for i = 2:min(length(a), length(b))
            result = max(result, a[i] + b[i])
        end
        return result
    end
end


text:
@testset "signed FastMod" begin
    a = convert(Vector{Int32},sample(typemin(Int32):typemax(Int32),1_000_000))
    a2 = copy(a)
    b = convert(Vector{Int32},sample(typemin(Int32):typemax(Int32),1_000))
    c1 = similar(a)
    c2 = similar(a)

    @testset for m in b
        modm = SignedMod(m)
        for i in eachindex(a)
            c1[i] = mod(a[i],m)
            c2[i] = fastmod(a[i],modm)
        end
        @test all(c1.==c2)
    end
    @testset for m in b
        modm = Mod(m)
        fastmod!(a,modm)
        inplace_mod!(a2,m)
        @test all(a.==a2)
    end
end


text:
#
# This file contains a list of folders, which we have
# to add to include search path. After that we may use
# import ModuleName. This file should be included from
# root directory. e.g.: in AppXXX.jl files...
#
# @author DeadbraiN
#
push!(LOAD_PATH, "$(pwd())/src")
push!(LOAD_PATH, "$(pwd())/src/util")
push!(LOAD_PATH, "$(pwd())/src/manager")
push!(LOAD_PATH, "$(pwd())/src/manager/plugins")
push!(LOAD_PATH, "$(pwd())/src/net")
push!(LOAD_PATH, "$(pwd())/src/organism")
push!(LOAD_PATH, "$(pwd())/src/visual")
push!(LOAD_PATH, "$(pwd())/src/global")
push!(LOAD_PATH, "$(pwd())/test")

max_stars_repo_name: tmptrash/jevo

max_issues_repo_name: tmptrash/jevo

max_forks_repo_name: tmptrash/jevo


text:
using Revise, StatisticalNonlocality, Test, LinearAlgebra
using SparseArrays

include("transition_matrix_test.jl")
include("discretization_test.jl")

max_stars_repo_name: sandreza/StatisticalNonlocality

max_issues_repo_name: sandreza/StatisticalNonlocality

max_forks_repo_name: sandreza/StatisticalNonlocality


text:
#  Copyright 2016-2019 Eric S. Tellez <eric.tellez@infotec.mx>
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

export Laesa

mutable struct Laesa{T} <: Index
    db::Vector{T}
    pivots::Vector{T}
    table::Matrix{Float64} # rows: number of pivots; cols: number of objects
end

function fit(::Type{Laesa}, dist::Function, db::AbstractVector{T}, pivots::Vector{T})  where T
    @info "Creating a pivot table with $(length(pivots)) pivots and distance=$(dist)"
    table = Matrix{Float64}(undef, length(pivots), length(db))

    for j in 1:length(db)
        for i in 1:length(pivots)
            table[i, j] = dist(db[j], pivots[i])
        end
    end

    Laesa(db, pivots, table)
end

function fit(::Type{Laesa}, dist::Function, db::AbstractVector{T}, numPivots::Integer) where T
    pivots = rand(db, numPivots)
    fit(Laesa, dist, db, pivots)
end

function search(index::Laesa{T}, dist::Function, q::T, res::KnnResult) where T
    dqp = [dist(q, piv) for piv in index.pivots]
    for i in 1:length(index.db)
        dpu = @view index.table[:, i]

        evaluate = true
        for pivID in 1:length(index.pivots)
            if abs(dqp[pivID] - dpu[pivID]) > covrad(res)
                evaluate = false
                break
            end
        end

        if evaluate
            d = dist(q, index.db[i])
            push!(res, i, d)
        end
    end

    return res
end

function push!(index::Laesa{T}, dist::Function, obj::T) where T
    push!(index.db, obj)
    vec = Vector{Float64}(undef, length(index.pivots))
    for pivID in 1:length(vec)
        vec[pivID] = dist(index.pivots[pivID], obj)
    end

    index.table = hcat(index.table, vec)
    length(index.db)
end


text:
# Full positive definite matrix together with a Cholesky factorization object
immutable PDMat{T<:Real,S<:AbstractMatrix} <: AbstractPDMat{T}
    dim::Int
    mat::S
    chol::CholType{T,S}
    PDMat(d::Int,m::AbstractMatrix{T},c::CholType{T,S}) = new(d,m,c)
end

function PDMat(mat::AbstractMatrix,chol::CholType)
    d = size(mat, 1)
    size(chol, 1) == d ||
        throw(DimensionMismatch("Dimensions of mat and chol are inconsistent."))
    PDMat{eltype(mat),typeof(mat)}(d, mat, chol)
end

PDMat(mat::Matrix) = PDMat(mat,cholfact(mat))
PDMat(mat::Symmetric) = PDMat(full(mat))
PDMat(fac::CholType) = PDMat(full(fac),fac)

### Conversion
convert{T<:Real}(::Type{PDMat{T}}, a::PDMat) = PDMat(convert(Matrix{T}, a.mat))

### Basics

dim(a::PDMat) = a.dim
full(a::PDMat) = copy(a.mat)
diag(a::PDMat) = diag(a.mat)

### Arithmetics

function pdadd!(r::Matrix, a::Matrix, b::PDMat, c)
    @check_argdims size(r) == size(a) == size(b)
    _addscal!(r, a, b.mat, c)
end

*{T<:Real}(a::PDMat{T}, c::T) = PDMat(a.mat * c)
*(a::PDMat, x::StridedVecOrMat) = a.mat * x
\(a::PDMat, x::StridedVecOrMat) = a.chol \ x

### Algebra

inv(a::PDMat) = PDMat(inv(a.chol))
logdet(a::PDMat) = logdet(a.chol)
eigmax(a::PDMat) = eigmax(a.mat)
eigmin(a::PDMat) = eigmin(a.mat)

### whiten and unwhiten

function whiten!(r::StridedVecOrMat, a::PDMat, x::StridedVecOrMat)
    cf = a.chol[:UL]
    istriu(cf) ? Ac_ldiv_B!(cf, _rcopy!(r, x)) : A_ldiv_B!(cf, _rcopy!(r, x))
    return r
end

function unwhiten!(r::StridedVecOrMat, a::PDMat, x::StridedVecOrMat)
    cf = a.chol[:UL]
    istriu(cf) ? Ac_mul_B!(cf, _rcopy!(r, x)) : A_mul_B!(cf, _rcopy!(r, x))
    return r
end

### quadratic forms

quad(a::PDMat, x::StridedVector) = dot(x, a * x)
invquad(a::PDMat, x::StridedVector) = dot(x, a \ x)

quad!(r::AbstractArray, a::PDMat, x::StridedMatrix) = colwise_dot!(r, x, a.mat * x)
invquad!(r::AbstractArray, a::PDMat, x::StridedMatrix) = colwise_dot!(r, x, a.mat \ x)

### tri products

function X_A_Xt(a::PDMat, x::StridedMatrix)
    z = copy(x)
    cf = a.chol[:UL]
    istriu(cf) ? A_mul_Bc!(z, cf) : A_mul_B!(z, cf)
    A_mul_Bt(z, z)
end

function Xt_A_X(a::PDMat, x::StridedMatrix)
    z = copy(x)
    cf = a.chol[:UL]
    istriu(cf) ? A_mul_B!(cf, z) : Ac_mul_B!(cf, z)
    At_mul_B(z, z)
end

function X_invA_Xt(a::PDMat, x::StridedMatrix)
    z = copy(x)
    cf = a.chol[:UL]
    istriu(cf) ? A_rdiv_B!(z, cf) : A_rdiv_Bc!(z, cf)
    A_mul_Bt(z, z)
end

function Xt_invA_X(a::PDMat, x::StridedMatrix)
    z = copy(x)
    cf = a.chol[:UL]
    istriu(cf) ? Ac_ldiv_B!(cf, z) : A_ldiv_B!(cf, z)
    At_mul_B(z, z)
end


text:
module SCF

using LinearAlgebra
using BandedMatrices
using SparseArrays
using IntervalSets

using ArnoldiMethod
using IterativeSolvers
using Optim
using LineSearches

using SolverTraces
import SolverTraces: base_exp

using UnicodeFun
using Formatting
using Crayons

using Compat

include("quantum_systems.jl")
include("krylov_wrapper.jl")
include("fock.jl")
include("solver_trace.jl")
include("utils.jl")
include("secular_problem.jl")
include("gershgorin_discs.jl")
include("rotate.jl")
include("solve_orbital_equation.jl")
include("self_consistent_iteration.jl")
include("fock_problem.jl")
include("manifolds.jl")
include("optim.jl")

export Fock, scf!

end # module


text:
module GeneralizedMinimalResidualSolver

export GeneralizedMinimalResidual

using ..LinearSolvers
const LS = LinearSolvers
using ..MPIStateArrays: device, realview

using LinearAlgebra
using LazyArrays
using StaticArrays
using KernelAbstractions

"""
# GMRES
    GeneralizedMinimalResidual(Q; M, rtol, atol)

This is an object for solving linear systems using an iterative Krylov method.
The constructor parameter `M` is the number of steps after which the algorithm
is restarted (if it has not converged), `Q` is a reference state used only
to allocate the solver internal state, and `rtol` specifies the convergence
criterion based on the relative residual norm. The amount of memory
required for the solver state is roughly `(M + 1) * size(Q)`.
This object is intended to be passed to the [`linearsolve!`](@ref) command.

This uses the restarted Generalized Minimal Residual method of Saad and Schultz (1986).

### References
    @article{saad1986gmres,
      title={GMRES: A generalized minimal residual algorithm for solving nonsymmetric linear systems},
      author={Saad, Youcef and Schultz, Martin H},
      journal={SIAM Journal on scientific and statistical computing},
      volume={7},
      number={3},
      pages={856--869},
      year={1986},
      publisher={SIAM}
    }
"""
mutable struct GeneralizedMinimalResidual{M, MP1, MMP1, T, AT} <:
               LS.AbstractIterativeLinearSolver
    krylov_basis::NTuple{MP1, AT}
    "Hessenberg matrix"
    H::MArray{Tuple{MP1, M}, T, 2, MMP1}
    "rhs of the least squares problem"
    g0::MArray{Tuple{MP1, 1}, T, 2, MP1}
    rtol::T
    atol::T

    function GeneralizedMinimalResidual(
        Q::AT;
        M = min(20, eltype(Q)),
        rtol = √eps(eltype(AT)),
        atol = eps(eltype(AT)),
    ) where {AT}
        krylov_basis = ntuple(i -> similar(Q), M + 1)
        H = @MArray zeros(M + 1, M)
        g0 = @MArray zeros(M + 1)

        new{M, M + 1, M * (M + 1), eltype(Q), AT}(
            krylov_basis,
            H,
            g0,
            rtol,
            atol,
        )
    end
end

const weighted = false

function LS.initialize!(
    linearoperator!,
    Q,
    Qrhs,
    solver::GeneralizedMinimalResidual,
    args...,
)
    g0 = solver.g0
    krylov_basis = solver.krylov_basis
    rtol, atol = solver.rtol, solver.atol

    @assert size(Q) == size(krylov_basis[1])

    # store the initial residual in krylov_basis[1]
    linearoperator!(krylov_basis[1], Q, args...)
    @. krylov_basis[1] = Qrhs - krylov_basis[1]

    threshold = rtol * norm(krylov_basis[1], weighted)
    residual_norm = norm(krylov_basis[1], weighted)

    converged = false
    # FIXME: Should only be true for threshold zero
    if threshold < atol
        converged = true
        return converged, threshold
    end

    fill!(g0, 0)
    g0[1] = residual_norm
    krylov_basis[1] ./= residual_norm

    converged, max(threshold, atol)
end

function LS.doiteration!(
    linearoperator!,
    Q,
    Qrhs,
    solver::GeneralizedMinimalResidual{M},
    threshold,
    args...,
) where {M}

    krylov_basis = solver.krylov_basis
    H = solver.H
    g0 = solver.g0

    converged = false
    residual_norm = typemax(eltype(Q))

    Ω = LinearAlgebra.Rotation{eltype(Q)}([])
    j = 1
    for outer j in 1:M

        # Arnoldi using the Modified Gram Schmidt orthonormalization
        linearoperator!(krylov_basis[j + 1], krylov_basis[j], args...)
        for i in 1:j
            H[i, j] = dot(krylov_basis[j + 1], krylov_basis[i], weighted)
            @. krylov_basis[j + 1] -= H[i, j] * krylov_basis[i]
        end
        H[j + 1, j] = norm(krylov_basis[j + 1], weighted)
        krylov_basis[j + 1] ./= H[j + 1, j]

        # apply the previous Givens rotations to the new column of H
        @views H[1:j, j:j] .= Ω * H[1:j, j:j]

        # compute a new Givens rotation to zero out H[j + 1, j]
        G, _ = givens(H, j, j + 1, j)

        # apply the new rotation to H and the rhs
        H .= G * H
        g0 .= G * g0

        # compose the new rotation with the others
        Ω = lmul!(G, Ω)

        residual_norm = abs(g0[j + 1])

        if residual_norm < threshold
            converged = true
            break
        end
    end

    # solve the triangular system
    y = SVector{j}(@views UpperTriangular(H[1:j, 1:j]) \ g0[1:j])

    ## compose the solution
    rv_Q = realview(Q)
    rv_krylov_basis = realview.(krylov_basis)
    groupsize = 256
    event = Event(device(Q))
    event = LS.linearcombination!(device(Q), groupsize)(
        rv_Q,
        y,
        rv_krylov_basis,
        true;
        ndrange = length(rv_Q),
        dependencies = (event,),
    )
    wait(device(Q), event)

    # if not converged restart
    converged || LS.initialize!(linearoperator!, Q, Qrhs, solver, args...)

    (converged, j, residual_norm)
end

end


text:
include("../src/FiberNlse.jl")
using .FiberNlse

# Simulation dimension
Nₜ, Nₗ = (2000,2000);

# Fiber properties

D = -0.9*U.ps/U.nm/U.km;
α = 0.026/U.km;
γ =  10.1/U.W/U.km;
L = 100.0*U.m;

# Signal propoerties
T = 1000*U.ps; # Signal duration
λ = 1550*U.nm; # Wavelength
τ = 5 * U.ps; # Pulse duration
N = 1 # Soliton number
sim = U.configure(Nₜ,Nₗ,D, γ, α, L, T, λ);
Pp = 0.550;#abs((sim.β2/γ/τ^2)*N^2) # Soliton power
f0=1e10;
β = 2
φ = 1.8pi
# Input construction
t = T*0.5*range(-1, stop=1, length=Nₜ); # Time vector
noiseP = sqrt(1e-8)*randn(ComplexF64,length(t))
noisePhi = pi*1e-14*randn(ComplexF64,length(t))
Ψₒ = noiseP + sqrt(Pp)*cos(2pi*f0*t)*exp(1im*(noisePhi+β*cos(2pi*f0*t)+φ))

U.inputSignal(sim,Ψₒ);
U.simulate(sim, false);

D = -17*U.ps/U.nm/U.km;
α = 0.046/U.km;
γ =  1.1/U.W/U.km;
L = 5.0*U.km;

sim2 = U.configure(Nₜ,Nₗ,D, γ, α, L, T, λ);
U.transition(sim, sim2)

l = range(0,stop=L, length=Nₗ);

heatmap(abs.(sim2.Ψ).^2)
xlabel!("Local time [ps]")

ylabel!("Distance [km]")
fs = 1/(t[2]-t[1])
f = fftshift(fftfreq(length(t), fs))
plot(f/U.GHz,abs.(fftshift(fft(sim.Ψ[1,:].^2))), yaxis=:log10)
plot!(f/U.GHz,abs.(fftshift(fft(sim2.Ψ[end,:].^2))),yaxis=:log10)
#xlims!(-500,500)
hline!([maximum(abs.(fft(sim2.Ψ[end,:].^2)))/1000], linestyle=:dash, label="-30 dB")


text:
using StabTrussTopOpt
sto = StabTrussTopOpt
using JuMP
using MosekTools
const MOI = JuMP.MathOptInterface
using Arpack

function TO_solve_relaxed_stab(pb::sto.TOProblem)
    # here we solve the following relaxed layout opt problem w/ global stability:
    # (single load case)
    #       min_{a, q} l' * a
    # s.t.  ∑_{i ∈ [m]} q_i γ_i = f  (force equilibrium)
    #       - a σ^- ≦ q ≦ a σ^+      (stress)
    #       K(a) + τ G(q) ⪰ 0        (global stability)
    #       a ≧ 0
    #
    # JuMP modeling reference:
    #   https://github.com/JuliaOpt/JuMP.jl/blob/master/examples/robust_uncertainty.jl

    # scalar parameter
    m = pb.m # num of element
    E = pb.E
    σ_c = pb.σ_c
    σ_t = pb.σ_t
    τ = pb.τ # stability load factor
    nfree = size(pb.eK[1], 1)

    # vector & spmat data
    L = pb.l
    f = pb.f # ∈ nfree x 1
    Γ = pb.Γ # ∈ m x nfree

    # sparse matrices array indexed by element id
    eK = pb.eK
    Δ = pb.Δ

    # model = JuMP.Model(with_optimizer(SCS.Optimizer, verbose = 0))
    model = JuMP.Model(with_optimizer(Mosek.Optimizer, QUIET=true))
    @variable(model, a[1:m] >= 0)
    @variable(model, q[1:m])

    @constraint(model, Γ'*q .== f)

    K = spzeros(nfree, nfree)
    for i=1:m
        K += a[i] * eK[i]
    end
    G = spzeros(nfree, nfree)
    for i=1:m
        G += (q[i]/L[i]) * Δ[i]
    end
    # @SDconstraint(model, K + τ * G >= 0)
    @SDconstraint(model, Symmetric(K + τ * G) >= 0)

    @constraint(model, - σ_c * a - q .<= 0)
    @constraint(model, - σ_t * a + q .<= 0)

    @objective(model, Min, dot(pb.c, a))

    JuMP.optimize!(model)
    @show JuMP.objective_value(model)
    @show opt_a = JuMP.value.(a)
    @show opt_q = JuMP.value.(q)

    @show JuMP.termination_status(model) == MOI.OPTIMAL
    @show JuMP.primal_status(model) == MOI.FEASIBLE_POINT

    # check eigen value
    opt_K = JuMP.value.(K)
    opt_G = JuMP.value.(G)
    λ_min, _ = Arpack.eigs(opt_K, opt_G, which=:SM, nev=1)
    @show λ_min

    # check violation on elastic compatibility
    # solve: min_{u} || opt_K_local * u - opt_q ||^2
    opt_hk = [opt_a[i]*E/L[i] for i=1:m]
    opt_K_local = spdiagm(0 => opt_hk) * Γ
    opt_u = (opt_K_local' * opt_K_local) \ (opt_K_local' * opt_q)
    @show res_elastic_compat = norm(opt_K_local * opt_u - opt_q)^2

    opt_a = abs.(opt_a)
    return opt_a, opt_q
end


text:
# ~/.julia/bin/mpiexecjl -n 1 julia --project -O3 --check-bounds=no ./scripts-part1/diffusion3D_benchmark_multigpu.jl
# ~/.julia/bin/mpiexecjl -n 2 julia --project -O3 --check-bounds=no ./scripts-part1/diffusion3D_benchmark_multigpu.jl
# ~/.julia/bin/mpiexecjl -n 3 julia --project -O3 --check-bounds=no ./scripts-part1/diffusion3D_benchmark_multigpu.jl
# ~/.julia/bin/mpiexecjl -n 4 julia --project -O3 --check-bounds=no ./scripts-part1/diffusion3D_benchmark_multigpu.jl
using MAT

# Initialize Globals
const USE_GPU = true
const BENCHMARK = true
const VISUALIZE = false

# Include main script
include("./diffusion3D_multixpu.jl")

# Run simulation
T_eff, t_toc, nprocs, me = diffusion_3D(256)
println("Scaling Experiment MultiXPU: $(T_eff)")

# Prepare visualization
gr()
ENV["GKSwstype"] = "nul"

# Run visualization on process 0
if me == 0
    # read file
    file = matopen("$(@__DIR__)/../docs/img/scaling_mgpu.mat")
    T_effs = read(file, "T_effs")
    t_tocs = read(file, "t_tocs")
    close(file)

    # update Teff of this run
    T_effs[nprocs] = T_eff
    t_tocs[nprocs] = t_toc

    # write back to file
    file = matopen("$(@__DIR__)/../docs/img/scaling_mgpu.mat", "w")
    write(file, "T_effs", T_effs)
    write(file, "t_tocs", t_tocs)
    close(file)

    # plot
    plot([1,2,3,4], T_effs, xlabel="Number of GPUs", ylabel="T_eff (GB/s)", ylims=(0, extrema(T_effs)[2]), title="GPU Weak Scaling Experiment", label=false, lw=3)
    png("$(@__DIR__)/../docs/img/scaling_experiment_mgpu_teff")
    plot([1,2,3,4], t_tocs[1]./t_tocs, ylims=(0, 1.0), xlabel="Number of GPUs", ylabel="Parallal Efficiency", title="GPU Weak Scaling Experiment", label=false, lw=3)
    png("$(@__DIR__)/../docs/img/scaling_experiment_mgpu_pareff")
end