text:
# The code to generate S5 Fig
using PyPlot
using SurNoR_2020
using Statistics
PyPlot.svg(true)
rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
rcParams["svg.fonttype"] = "none"
rcParams["pdf.fonttype"] = 42

# WARNING: put n_set = 1 for the left panel, put n_set = 2 for the middle,
#          and put n_set = 3 for the right panel
n_set = 1
# ------------------------------------------------------------------------------
# Training Algorithms
# ------------------------------------------------------------------------------
Agents_SurNoR = SurNoR_2020.Func_SurNoR_Overall_train_opt_ind_i(3)
Model = Agents_SurNoR[[2,3,4,5,6,7,8,9,10,12],:,:]

Param_Path = string("data/Fitted_parameters_simulated/Overall/nset_", string(n_set), "/")
Agents_SurNoR_nset   = SurNoR_2020.Func_SurNoR_Overall_train_opt_ind_i(1,Param_Path = Param_Path)
Model_nset = Agents_SurNoR_nset[[2,3,4,5,6,7,8,9,10,12],:,:]

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Correlation between model variables given the parameters fitted to the behavior
# of participants and the behavior of the simulated participants
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Putting all RPE, NPE, Novelty, and Surprise together
Sub_Set = 1:10

Block_Set = 1:2
Epi_Set = 1:5

Concat_Signals = []
Concat_Signals_nset = []
for Sub=Sub_Set
    for Block_ind = 1:length(Block_Set)
        Block = Block_Set[Block_ind]
        for Epi_ind = 1:length(Epi_Set)
            Epi = Epi_Set[Epi_ind]
            temp = Model[Sub,Block,Epi]
            temp_nset = Model_nset[Sub,Block,Epi]
            if (Epi_ind==1)&(Block_ind==1)
                Sub_Signal = cat(temp.Surprise,temp.Novelty_Seq,temp.δ_N_Seq,temp.δ_R_Seq,temp.Input.obs.==0,dims=2)
                Sub_Signal_nset = cat(temp_nset.Surprise,temp_nset.Novelty_Seq,temp_nset.δ_N_Seq,temp_nset.δ_R_Seq,temp_nset.Input.obs.==0,dims=2)
            else
                Sub_Signal = cat(Sub_Signal,
                                 cat(temp.Surprise,temp.Novelty_Seq,temp.δ_N_Seq,temp.δ_R_Seq,temp.Input.obs.==0,dims=2),
                                 dims=1)
                Sub_Signal_nset = cat(Sub_Signal_nset,
                                 cat(temp_nset.Surprise,temp_nset.Novelty_Seq,temp_nset.δ_N_Seq,temp_nset.δ_R_Seq,temp_nset.Input.obs.==0,dims=2),
                                 dims=1)
            end
        end
    end
    Sub_Signal = Sub_Signal[Sub_Signal[:,1] .!= -1,:]
    append!(Concat_Signals,[deepcopy(Sub_Signal)])
    Sub_Signal_nset = Sub_Signal_nset[Sub_Signal_nset[:,1] .!= -1,:]
    append!(Concat_Signals_nset,[deepcopy(Sub_Signal_nset)])
end

# Barplot
Signal_names = ["Surprise","Novelty","NPE","RPE","Goal"]
N_Sig = length(Signal_names)
N_Sub = length(Sub_Set)
Corr_Mat = zeros(N_Sub,N_Sig)
for Sub=1:N_Sub
    for Sig = 1:N_Sig
        Corr_Mat[Sub,Sig] = cor(Concat_Signals[Sub][:,Sig],Concat_Signals_nset[Sub][:,Sig])
    end
end

Colors = ["#B159E0","#3950A3","#02B0FB","#391748","#B21F24"]

fig = figure(); ax = gca()
x = 1:N_Sig

y = mean(Corr_Mat,dims=1)
dy = std(Corr_Mat,dims=1)/sqrt(N_Sub)
Sig_inds = [1,3,4]
ax.plot([x[1]-0.5,x[length(Sig_inds)]+0.5],ones(2) .* 1, "--k",alpha=0.5)
ax.plot([x[1]-0.5,x[length(Sig_inds)]+0.5],ones(2) .* 0.95, "--k",alpha=0.5)
for i = 1:length(Sig_inds)
    ax.bar(i, y[Sig_inds[i]] , yerr=dy[Sig_inds[i]], color=Colors[Sig_inds[i]])
    ax.errorbar(i, y[Sig_inds[i]] , yerr=dy[Sig_inds[i]],color="k",linewidth=1,capsize=5,linestyle ="")
end

σ = 0.1
for j = 1:N_Sub
  x_1 = (1:length(Sig_inds)) .+ 2*σ*(rand() - 0.5)
  y = Corr_Mat[j,Sig_inds]
  ax.plot(x_1, y[:],".",color="k",alpha=0.1)
end
ax.set_ylim([0.8,1.02])
ax.set_xlim([x[1]-0.5,x[length(Sig_inds)]+0.5])

ax.set_xticks(1:length(Sig_inds))
ax.set_xticklabels(Signal_names[Sig_inds])


text:
"""
output index format:
x[i][j][k][l]
i-th spin system, k-th magnetically equivalent nuclei, l-th local spin index.
"""
function getmageqcompound(g,
    H_inds_sys,
    dict_ind_to_H_ID,
    dict_H_inds_to_css,
    dict_H_IDs_to_css,
    dict_J_ID_to_val,
    J_IDs;
    atol = 1e-6)

    C_g = Graphs.maximal_cliques(g)

    N_spin_systems = length(H_inds_sys)

    mag_eq_sys_inds_local = Vector{Vector{Vector{Int}}}(undef, N_spin_systems)
    mag_eq_sys_inds_global = Vector{Vector{Vector{Int}}}(undef, N_spin_systems)
    mag_eq_sys_IDs = Vector{Vector{Vector{Int}}}(undef, N_spin_systems)

    for i = 1:N_spin_systems
        C = keeptargetintegers(C_g, H_inds_sys[i])

        mag_eq_IDs0, mag_eq_inds0 = getmageqIDs(C,
            dict_ind_to_H_ID,
            dict_H_inds_to_css,
            dict_H_IDs_to_css, dict_J_ID_to_val, J_IDs;
            atol = atol)

        mag_eq_inds = combinetransitiveeqgroups(mag_eq_inds0)
        mag_eq_IDs = combinetransitiveeqgroups(mag_eq_IDs0)

        # mag_eq_inds = mag_eq_inds0
        # mag_eq_IDs = mag_eq_IDs0

        mag_eq_sys_inds_local[i] = getmageqlocalinds(H_inds_sys[i], mag_eq_inds)
        mag_eq_sys_IDs[i] = mag_eq_IDs
        mag_eq_sys_inds_global[i] = mag_eq_inds
    end

    return mag_eq_sys_inds_local, mag_eq_sys_IDs, mag_eq_sys_inds_global
end

"""
convertlabelsglobaltolocal(H_inds_sys::Vector{Vector{Int}},
    mag_eq_sys_inds::Vector{Vector{Int}})

The return type is Vector{Vector{Int}}.

Local: Each spin system will have its own spin nucleui numbering that starts at 1.
Global: The spins systems will together have one single spin nucleui numbering that starts at 1.
"""
function getmageqlocalinds(H_inds::Vector{Int},
    mag_eq_sys_inds::Vector{Vector{Int}})

    N_eqs = length(mag_eq_sys_inds)
    labels_sys_local = Vector{Vector{Int}}(undef, N_eqs)

    for i = 1:N_eqs

        x = H_inds
        conversion_dict = Dict(x .=> collect(1:length(x)))

        z = mag_eq_sys_inds[i]
        labels_sys_local[i] = collect( conversion_dict[z[k]] for k = 1:length(z) )
    end

    return labels_sys_local
end

"""
getmageqIDs(C::Vector{Vector{Int}},
    dict_ind_to_H_ID,
    dict_H_inds_to_css;
    atol::Float64 = 1e-6)

Given a list of node indices `C` of cliques of an undirected graph, with the indices of a clique being C[i],
Returns the list of node indices for each clique that are magnetically equivalent.
"""
function getmageqIDs(C::Vector{Vector{Int}},
    dict_ind_to_H_ID,
    dict_H_inds_to_css,
    dict_H_IDs_to_css,
    dict_J_ID_to_val,
    J_IDs;
    atol::Float64 = 1e-6)

    # traverse each maximally connected cliques.
    C_IDs_mag_eq = Vector{Vector{Vector{Int}}}(undef, length(C))
    C_inds_mag_eq = Vector{Vector{Vector{Int}}}(undef, length(C))
    for i = 1:length(C)

        ### partition the nucleus by unique cs.
        cs = collect( dict_H_inds_to_css[C[i][k]] for k = 1:length(C[i]) )
        cs_IDs = collect( dict_ind_to_H_ID[C[i][k]] for k = 1:length(C[i]) )
        cs_inds = C[i]

        # unique_cs = unique(round.(cs, digits = cs_round_digits))
        # unique_cs_inds = collect( inds = findall(xx->isapprox(unique_cs[k], xx; atol = atol), cs) for k = 1:length(unique_cs) )

        unique_cs, unique_cs_inds = uniqueinds(cs; atol = atol)

        # check magnetic equivalence for the chemically equivalent entries in `unique_cs_inds`.
        pass_flags = collect( checkmageq(unique_cs_inds[k], cs_IDs, dict_J_ID_to_val,
            dict_H_IDs_to_css, J_IDs; atol = atol) for k = 1:length(unique_cs_inds) )

        # store result.
        common_cs_IDs = collect( cs_IDs[unique_cs_inds[k]] for k = 1:length(unique_cs_inds) )
        C_IDs_mag_eq[i] = common_cs_IDs[pass_flags]

        common_cs_inds = collect( cs_inds[unique_cs_inds[k]] for k = 1:length(unique_cs_inds) )
        C_inds_mag_eq[i] = common_cs_inds[pass_flags]

        if length(common_cs_IDs[pass_flags]) > 2
            println("Warning for clique $(i): more than two groups of magnetically equivalent nuclei!")
        end
    end

    # remove empty entries.
    keep_flags = trues(length(C))
    for i = 1:length(C)
        if isempty(C_IDs_mag_eq[i])
            keep_flags[i] = false
        end
    end
    C_IDs_mag_eq = C_IDs_mag_eq[keep_flags]
    C_inds_mag_eq = C_inds_mag_eq[keep_flags]

    mag_eq_IDs::Vector{Vector{Int}} = combinevectors(C_IDs_mag_eq)
    mag_eq_inds::Vector{Vector{Int}} = combinevectors(C_inds_mag_eq)

    return mag_eq_IDs, mag_eq_inds
end

### routines related to checking whether the J-coupling value of A-C and B-C
# are the same, for all C's connected to either A or B. A and B are
# chemically equivalent nuclei.
"""
matchJlabels((label_pair_list::Vector{Tuple{Int,Int}}, search_list::Vector{Int})::Vector{Int}

Returns the indices of entries of `label_pair_list` that has any of the labels in the `search_list`.
"""
function matchanyJlabels(label_pair_list::Vector{Tuple{Int,Int}}, search_list::Vector{Int})::Vector{Int}

    inds = Vector{Int}(undef, 0)
    for k = 1:length(label_pair_list)
        i, j = label_pair_list[k]

        if any(i .== search_list) || any(j .== search_list)
            push!(inds, k)
        end
    end

    return inds
end

"""
getJIDstest(J_IDs::Vector{Tuple{Int,Int}},
    common_cs_IDs::Vector{Int},
    dict_H_IDs_to_css; atol = 1e-6)

Get the list of A-C, B-C, A-D, B-D, etc pairs to test.
"""
function getJIDstest(J_IDs::Vector{Tuple{Int,Int}},
    common_cs_IDs::Vector{Int},
    dict_H_IDs_to_css; atol = 1e-6)

    inds_any = matchanyJlabels(J_IDs, common_cs_IDs)
    inds_both = matchJlabels(J_IDs, common_cs_IDs)
    inds = setdiff(inds_any, inds_both)

    #inds_both =
    J_IDs_tmp = J_IDs[inds]

    target_IDs = Vector{Int}(undef, 0)
    for l = 1:length(J_IDs_tmp)
        i, j = J_IDs_tmp[l]

        # look for nuclei pairs that have the same chemical shift.
        if !isapprox(dict_H_IDs_to_css[i], dict_H_IDs_to_css[j], atol = atol)

            # find the value between i and j that isn't in common_cs_IDs.
            target_ID = i
            if any(i .== common_cs_IDs)
                target_ID = j
            end

            push!(target_IDs, target_ID)
        end
    end

    return target_IDs
end

"""
checkmageq(test_inds::Vector{Vector{Int}},
    cs_IDs::Vector{Int},
    dict_J_ID_to_val;
    atol = 1e-6)::Bool

`test_inds` is a list of indices in cs_IDs. This function checks for magnetic equivalence for the nuclei in `test_inds`.
"""
function checkmageq(test_inds::Vector{Int},
    cs_IDs::Vector{Int},
    dict_J_ID_to_val,
    dict_H_IDs_to_css,
    J_IDs;
    atol = 1e-6)::Bool

    if length(test_inds) == 1
        # need to have at least two nuclei to be magnetically equivalent.
        return false
    end

    common_cs_IDs = cs_IDs[test_inds]
    p_IDs = getpairs(common_cs_IDs)

    ## test common J within pairs.
    J_p = collect( getJfromdict(p_IDs[l][1], p_IDs[l][2], dict_J_ID_to_val) for l = 1:length(p_IDs) )
    pass_common_J_flag = isallsame(J_p; atol = atol)

    ## test common J with other IDs.

    # find all connections to first pair.
    t_IDs = getJIDstest(J_IDs, common_cs_IDs, dict_H_IDs_to_css; atol = atol)

    # check J-values.
    pass_test_J_flag = true
    for k = 1:length(t_IDs)

        J_t = collect( getJfromdict(t_IDs[k], base_ID, dict_J_ID_to_val) for base_ID in common_cs_IDs )

        pass_test_J_flag = pass_test_J_flag & isallsame(J_t; atol = atol)
    end

    return pass_common_J_flag & pass_test_J_flag
end

"""
getJfromdict(i::Int, j::Int, dict::Dict{Tuple{Int,Int},T} ) where T

query `dict` for the entry `(i,j)` and `(j,i)`. Returns the entry for `(i,j)`. Returns the entry for `(j,i)` if `(i,j)` is not found. Return zero if both entries are not found.
"""
function getJfromdict(i::Int, j::Int, dict::Dict{Tuple{Int,Int},T} ) where T
    #
    if haskey(dict, (i,j))
        return dict[(i,j)]
    end

    if haskey(dict, (j,i))
        return dict[(j,i)]
    end

    return zero(T)
end

"""
combinetransitiveeqgroups(eq_inds::Vector{Vector{Int}})

`eq_inds` contains theequivalent nuclei indices.
This function combine the groups of equivalent nuclei via transitivity.

Assumes the equivalence relation cannot be true for a group containg only one nucleus.
"""
function combinetransitiveeqgroups(eq_inds::Vector{Vector{Int}})

    if isempty(eq_inds)
        return Vector{Vector{Int}}(undef, 0)
    end

    eq_inds_flat_unique = unique(combinevectors(eq_inds))
    V = 1:length(eq_inds_flat_unique)
    dict_eq_to_V = Dict(eq_inds_flat_unique .=> V)
    dict_V_to_eq = Dict(V .=> eq_inds_flat_unique)

    eq_inds_V = collect( collect( dict_eq_to_V[eq_inds[i][k]] for k = 1:length(eq_inds[i])) for i = 1:length(eq_inds) )

    # get the edges of fully connected
    E = Vector{Tuple{Int,Int}}(undef, 0)
    for k = 1:length(eq_inds)
        E_k = getconnectpath(eq_inds_V[k])

        push!(E, E_k...)
    end

    #
    h = Graphs.SimpleGraph(V[end])
    for k = 1:length(E)
        Graphs.add_edge!(h, E[k][1], E[k][2])
    end

    H = Graphs.connected_components(h)

    # convert back to the nuclei indexing of `eq_inds`.
    H_out = collect( collect( dict_V_to_eq[H[i][k]] for k = 1:length(H[i])) for i = 1:length(H) )

    return H_out
end

### final package for a compound.
function getmageqinfo(H_IDs, H_css::Vector{T}, J_IDs, J_vals;
    unique_cs_atol = 1e-6)::Tuple{Vector{Vector{Vector{Int}}},Vector{Vector{Vector{Int}}},Vector{Vector{Vector{Int}}}} where T <: Real

    # parse the spin systems and singlets.
    J_inds_sys, J_inds_sys_local, J_IDs_sys, J_vals_sys, H_inds_sys,
    cs_sys, H_inds_singlets, cs_singlets, H_inds, J_inds,
    g = setupcsJ(H_IDs, H_css, J_IDs, J_vals)

    # set up look-up.
    dict_H_IDs_to_css = Dict(H_IDs .=> H_css)
    dict_H_inds_to_css = Dict(H_inds .=> H_css)
    dict_H_ID_to_ind = Dict(H_IDs .=> collect(1:length(H_IDs)))
    dict_ind_to_H_ID = Dict( collect(1:length(H_IDs)) .=> H_IDs)

    dict_J_ID_to_val = Dict(J_IDs .=> J_vals)
    dict_J_ind_to_val = Dict(J_inds .=> J_vals)

    dict_J_ID_to_ind = Dict(J_IDs .=> J_inds)
    dict_J_ind_to_ID = Dict(J_inds .=> J_IDs)

    #
    mag_eq_sys_inds_local, mag_eq_sys_IDs,
    mag_eq_sys_inds_global = getmageqcompound(g,
    H_inds_sys, dict_ind_to_H_ID, dict_H_inds_to_css,
    dict_H_IDs_to_css, dict_J_ID_to_val, J_IDs; atol = unique_cs_atol)

    return mag_eq_sys_inds_local, mag_eq_sys_IDs,
    mag_eq_sys_inds_global
end

function getmageqinfomixture(target_names::Vector{String},
    base_path::String,
    dict_compound_to_filename;
    unique_cs_atol = 1e-6)

    N_compounds = length(target_names)
    MEs = Vector{Vector{Vector{Vector{Int}}}}(undef, N_compounds)

    for n = 1:N_compounds

        # TODO add error-handling if name is not found in the dictionary, or filename does not exist.
        load_path = joinpath(base_path, dict_compound_to_filename[target_names[n]]["file name"])
        H_IDs, H_css, J_IDs, J_vals = loadcouplinginfojson(load_path)

        MEs[n], _ = getmageqinfo(H_IDs, H_css, J_IDs, J_vals; unique_cs_atol = unique_cs_atol)
    end

    return MEs
end

ext: jl


text:
yflip(a::Int) = let D = [0,4,2,6,1,5,3,7]
    reduce(|, ((1 & (a >> (i - 1))) << D[i] for i in 1:8))
end

complement(a::Int) = reduce(|, ((1 ⊻ (1 & (a >> ((i-1) ⊻ 7)))) << (i - 1) for i in 1:8))

canonical(a::Int) = min(a, yflip(a), complement(a), complement(yflip(a)))

const EQUIVALENT_RULES = Dict(c => (yflip(c), complement(c), complement(yflip(c))) for c in 0:255)
const ECA_RULES = [Set(vcat(map(canonical, 0:255)))...]
const ECA_CLASS_I = [0,8,32,40,128,136,160,168]
const ECA_CLASS_III = [18,22,30,45,60,90,105,122,126,146,150]
const ECA_CLASS_IV = [54,106,110]
const ECA_CLASS_II = setdiff(ECA_RULES, ECA_CLASS_I, ECA_CLASS_III, ECA_CLASS_IV)

@assert ECA_CLASS_I ⊆ ECA_RULES
@assert ECA_CLASS_III ⊆ ECA_RULES
@assert ECA_CLASS_IV ⊆ ECA_RULES
@assert length(ECA_CLASS_II) == 66

function update!(state, rule)
    if length(state) < 2
        throw(DomainError(length(state), "update! requires that length(state) ≥ 2"))
    end
    left = state[end]
    right = state[1]
    temp = 2 * left + state[1]
    for i in 2:length(state)
        temp = 7 & (2 * temp + state[i])
        state[i - 1] = 1 & (rule >> temp)
    end
    temp = 7 & (2 * temp + right)
    state[end] = 1 & (rule >> temp)

    state
end

function simulate(rule, state::AbstractVector{Int}, t::Int)
    traj = Array{Int}(undef, t + 1, length(state))
    traj[1,:] = state
    for i in 1:t
        traj[i + 1, :] = update!(traj[i,:], rule)
    end
    traj
end

function simulate(rule, width::Int, t::Int)
    state = zeros(Int, width)
    state[length(state) ÷ 2 + 1] = 1
    simulate(rule, state, t)
end

max_stars_repo_name: ELIFE-ASU/dglmoore

max_issues_repo_name: ELIFE-ASU/dglmoore

max_forks_repo_name: ELIFE-ASU/dglmoore


text:
# Tree API implementations for builtin types

children(x::AbstractArray) = x

children(x::Tuple) = x

children(x::Expr) = x.args

children(p::Pair) = (p[2],)

children(dict::AbstractDict) = pairs(dict)


text:
using Base.Test
using MathProgComplex

function run_pb_orders(problem, order_to_obj)
    data = SortedDict()
    for d in keys(order_to_obj)
        logpath = joinpath(Pkg.dir("MathProgComplex"), "Mosek_runs", "Lasserreex", "d_$d")

        relax_ctx = set_relaxation(problem; hierarchykind=:Real,
                                            d = d,
                                            params = Dict(:opt_outlev=>0,
                                                          :opt_logpath=>logpath,
                                                          :opt_solver=>testsolver))

        cur_obj, dualobj = run_hierarchy(problem, relax_ctx, indentedprint=true, save_pbs=true)

        data[d] = (cur_obj, dualobj)
    end
    return data
end

@testset "GloptiPoly small real problems" begin

    ## http://homepages.laas.fr/henrion/papers/gpcocos.pdf p7
    @testset "example1" begin
        problem = Problem()
        x1 = Variable("x1", Real); add_variable!(problem, x1)
        x2 = Variable("x2", Real); add_variable!(problem, x2)
        x3 = Variable("x3", Real); add_variable!(problem, x3)

        set_objective!(problem, -2*x1+x2-x3)
        add_constraint!(problem, "ctr1", (x1*(4*x1-4*x2+4*x3-20)+x2*(2*x2-2*x3+9)+x3*(2*x3-13)+24) >> 0)
        add_constraint!(problem, "ctr2", (x1+x2+x3) << 4)
        add_constraint!(problem, "ctr3", (3*x2+x3) << 6)
        add_constraint!(problem, "def_x1", 0 << x1 << 2)
        add_constraint!(problem, "def_x2", 0 << x2)
        add_constraint!(problem, "def_x3", 0 << x3 << 3)

        ε=1e-3
        order_to_obj = SortedDict(1=>-6.0000,
                                2=>-5.6923,
                                3=>-4.0685,
                                4=>-4.0000)

        data_sol = run_pb_orders(problem, order_to_obj)

        @printf("%5s  %15s  %15s  | %15s", "order", "primal obj", "dual obj", "expected obj\n")
        for (d, (primobj, dualobj)) in data_sol
            @printf("%5i  %15f  %15f  | %15f\n", d, primobj, dualobj, order_to_obj[d])
            @test primobj ≈ order_to_obj[d] atol=ε
            @test dualobj ≈ primobj atol=mosek_optgap*min(abs(primobj), abs(dualobj))
        end
    end

    @testset "example2" begin
            ## http://homepages.laas.fr/henrion/papers/gloptipoly.pdf p14
        problem = Problem()
        x1 = Variable("x1", Real); add_variable!(problem, x1)
        x2 = Variable("x2", Real); add_variable!(problem, x2)
        x3 = Variable("x3", Real); add_variable!(problem, x3)
        x4 = Variable("x4", Real); add_variable!(problem, x4)
        x5 = Variable("x5", Real); add_variable!(problem, x5)

        set_objective!(problem, 42*x1+44*x2+45*x3+47*x4+47.5*x5-50*(x1^2+x2^2+x3^2+x4^2+x5^2))
        add_constraint!(problem, "ctr1", (20*x1+12*x2+11*x3+7*x4+4*x5) << 40)
        add_constraint!(problem, "def_x1", 0 << x1 << 1)
        add_constraint!(problem, "def_x2", 0 << x2 << 1)
        add_constraint!(problem, "def_x3", 0 << x3 << 1)
        add_constraint!(problem, "def_x4", 0 << x4 << 1)
        add_constraint!(problem, "def_x5", 0 << x5 << 1)

        ε=1e-3
        order_to_obj = SortedDict(1=>NaN,         # order 1 relaxation is primal infeasable
                                2=>-17.9189,
                                3=>-17.0000)

        data_sol = run_pb_orders(problem, order_to_obj)

        @printf("%5s  %15s  %15s  | %15s", "order", "primal obj", "dual obj", "expected obj\n")
        for (d, (primobj, dualobj)) in data_sol
            @printf("%5i  %15f  %15f  | %15f\n", d, primobj, dualobj, order_to_obj[d])
            if d==1
                @test isnan(primobj)
                @test isnan(dualobj)
            else
                @test primobj ≈ order_to_obj[d] atol=ε
                @test dualobj ≈ primobj atol = mosek_optgap*min(abs(primobj), abs(dualobj))
            end
        end
    end
end


text:
using Documenter
using LOLTools

makedocs(
    build = joinpath(@__DIR__, "local" in ARGS ? "build_local" : "build"),
    modules = [LOLTools,
               LOLTools.ChampionMasteryV4,
               LOLTools.ChampionV3,
               LOLTools.LoLStatusV3,
               LOLTools.LeagueV4,
               LOLTools.MatchV5,
               LOLTools.SpectatorV4,
               LOLTools.SummonerV4,
               LOLTools.GameClient,
              ],
    clean = false,
    format = Documenter.HTML(
        prettyurls = !("local" in ARGS),
        assets = ["assets/custom.css"],
    ),
    sitename = "LOLTools.jl",
    authors = "WooKyoung Noh",
    pages = Any[
        "Home" => "index.md",
        "API" => "API.md",
    ],
)


text:
const HTML_TEMPLATES = Dict{Symbol,BlockTemplate}(e.opening => e for e in [
   BlockTemplate(:COMMENT,     :COMMENT_OPEN,     :COMMENT_CLOSE     ),
   BlockTemplate(:SCRIPT,      :SCRIPT_OPEN,      :SCRIPT_CLOSE      ),
   BlockTemplate(:DBB,         :DBB_OPEN,         :DBB_CLOSE         ),
   BlockTemplate(:MATH_INLINE, :MATH_INLINE_OPEN, :MATH_INLINE_CLOSE ),
   BlockTemplate(:MATH_BLOCK,  :MATH_BLOCK_OPEN,  :MATH_BLOCK_CLOSE  ),
   ])


text:
#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons, Josef Kircher
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

using ChainRulesCore: ignore_derivatives

# Receives one or an array of value references in an arbitrary format (see fmi2ValueReferenceFormat) and converts it into an Array{fmi2ValueReference} (if not already).
function prepareValueReference(md::fmi2ModelDescription, vr::fmi2ValueReferenceFormat)
    tvr = typeof(vr)
    if tvr == Array{fmi2ValueReference,1}
        return vr
    elseif tvr == fmi2ValueReference
        return [vr]
    elseif tvr == String
        return [fmi2StringToValueReference(md, vr)]
    elseif tvr == Array{String,1}
        return fmi2StringToValueReference(md, vr)
    elseif tvr == Int64
        return [fmi2ValueReference(vr)]
    elseif tvr == Array{Int64,1}
        return fmi2ValueReference.(vr)
    elseif tvr == Nothing
        return Array{fmi2ValueReference,1}()
    elseif tvr == Symbol
        if vr == :states
            return md.stateValueReferences
        elseif vr == :derivatives
            return md.derivativeValueReferences
        elseif vr == :inputs
            return md.inputValueReferences
        elseif vr == :outputs
            return md.outputValueReferences
        elseif vr == :all
            return md.valueReferences
        elseif vr == :none
            return Array{fmi2ValueReference,1}()
        else
            @assert false "Unknwon symbol `$vr`, can't convert to value reference."
        end
    end

    @assert false "prepareValueReference(...): Unknown value reference structure `$tvr`."
end
function prepareValueReference(fmu::FMU2, vr::fmi2ValueReferenceFormat)
    prepareValueReference(fmu.modelDescription, vr)
end
function prepareValueReference(comp::FMU2Component, vr::fmi2ValueReferenceFormat)
    prepareValueReference(comp.fmu.modelDescription, vr)
end

"""
Returns an array of ValueReferences coresponding to the variable names.
"""
function fmi2StringToValueReference(md::fmi2ModelDescription, names::Array{String})
    vr = Array{fmi2ValueReference}(undef,0)
    for name in names
        reference = fmi2StringToValueReference(md, name)
        if reference == nothing
            @warn "Value reference for variable '$name' not found, skipping."
        else
            push!(vr, reference)
        end
    end
    vr
end

"""
Returns the model variable(s) fitting the value reference.
"""
function fmi2ModelVariablesForValueReference(md::fmi2ModelDescription, vr::fmi2ValueReference)
    ar = []
    for modelVariable in md.modelVariables
        if modelVariable.valueReference == vr
            push!(ar, modelVariable)
        end
    end
    ar
end

"""
Returns the ValueReference coresponding to the variable name.
"""
function fmi2StringToValueReference(md::fmi2ModelDescription, name::String)
    reference = nothing
    if haskey(md.stringValueReferences, name)
        reference = md.stringValueReferences[name]
    else
        @warn "No variable named '$name' found."
    end
    reference
end

function fmi2StringToValueReference(fmu::FMU2, name::Union{String, Array{String}})
    fmi2StringToValueReference(fmu.modelDescription, name)
end

"""
Returns an array of variable names matching a fmi2ValueReference.
"""
function fmi2ValueReferenceToString(md::fmi2ModelDescription, reference::fmi2ValueReference)
    [k for (k,v) in md.stringValueReferences if v == reference]
end
function fmi2ValueReferenceToString(md::fmi2ModelDescription, reference::Int64)
    fmi2ValueReferenceToString(md, fmi2ValueReference(reference))
end

function fmi2ValueReferenceToString(fmu::FMU2, reference::Union{fmi2ValueReference, Int64})
    fmi2ValueReferenceToString(fmu.modelDescription, reference)
end

function fmi2GetSolutionState(solution::FMU2Solution, vr::fmi2ValueReferenceFormat; isIndex::Bool=false)

    index = 0

    if isIndex
        index = vr
    else
        ignore_derivatives() do
            vr = prepareValueReference(solution.fmu, vr)[1]

            if solution.states !== nothing
                for i in 1:length(solution.fmu.modelDescription.stateValueReferences)
                    if solution.fmu.modelDescription.stateValueReferences[i] == vr
                        index = i
                        break
                    end
                end
            end

        end # ignore_derivatives
    end

    if index > 0
        return collect(u[index] for u in solution.states.u)
    end

    return nothing
end

function fmi2GetSolutionValue(solution::FMU2Solution, vr::fmi2ValueReferenceFormat; isIndex::Bool=false)

    index = 0

    if isIndex
        index = vr
    else
        ignore_derivatives() do
            vr = prepareValueReference(solution.fmu, vr)[1]

            if solution.states !== nothing
                for i in 1:length(solution.fmu.modelDescription.stateValueReferences)
                    if solution.fmu.modelDescription.stateValueReferences[i] == vr
                        index = i
                        break
                    end
                end
            end

            if index > 0
                return collect(u[index] for u in solution.states.u)
            end

            if solution.values !== nothing
                for i in 1:length(solution.valueReferences)
                    if solution.valueReferences[i] == vr
                        index = i
                        break
                    end
                end
            end

        end # ignore_derivatives
    end

    if index > 0
        return collect(v[index] for v in solution.values.saveval)
    end

    return nothing
end

function fmi2GetSolutionTime(solution::FMU2Solution)
    if solution.states !== nothing
        return solution.states.t
    elseif solution.values !== nothing
        return solution.values.t
    else
        return nothing
    end
end


text:
using Test

@testset "Riveros" begin
    # Mostly the testing is performed in testagainstheinrich.jmd
    # These tests just ensure stablity.
    θ = deg2rad(40.0)
    algs = zafcorrection(Riveros1993, ReedFluorescence, NullCoating, mat"NaAlSi3O8", mat"NaF", n"Na K", 15.0e3)
    @test isapprox(A(algs..., n"Na K-L3",θ,θ), 1.0401, atol=0.0001)
    @test isapprox(Z(algs...), 0.96363, atol=0.0001)
    algs = zafcorrection(Riveros1993, ReedFluorescence, NullCoating, mat"NaAlSi3O8", mat"Al2O3", n"Al K", 15.0e3)
    @test isapprox(A(algs..., n"Al K-L3",θ,θ), 0.92632, atol=0.0001)
    @test isapprox(Z(algs...), 0.9900, atol=0.0001)
    algs = zafcorrection(Riveros1993, ReedFluorescence, NullCoating, mat"NaAlSi3O8", mat"Al2O3", n"O K", 15.0e3)
    @test isapprox(A(algs..., n"O K-L3",θ,θ), 0.94999, atol=0.0001)
    @test isapprox(Z(algs...), 0.9901, atol=0.0001)
    algs = zafcorrection(Riveros1993, ReedFluorescence, NullCoating, mat"NaAlSi3O8", mat"SiO2", n"Si K", 15.0e3)
    @test isapprox(A(algs..., n"Si K-L3",θ,θ), 0.89044, atol=0.0001)
    @test isapprox(Z(algs...), 1.0067, atol=0.0001)
end

max_stars_repo_name: NicholasWMRitchie/NeXLMatrixCorrection

max_issues_repo_name: NicholasWMRitchie/NeXLMatrixCorrection

max_forks_repo_name: NicholasWMRitchie/NeXLMatrixCorrection


text:
export LiftedHRepresentation, LiftedVRepresentation

# H-Represenation

# No copy since I do not modify anything and a copy is done when building a polyhedron
mutable struct LiftedHRepresentation{T, MT<:AbstractMatrix{T}} <: MixedHRep{T}
    # Ax >= 0, it is [b -A] * [z; x] where z = 1
    A::MT
    linset::BitSet

    function LiftedHRepresentation{T, MT}(A::MT, linset::BitSet=BitSet()) where {T, MT}
        if !isempty(linset) && last(linset) > size(A, 1)
            error("The elements of linset should be between 1 and the number of rows of A")
        end
        new{T, MT}(A, linset)
    end
end
FullDim(rep::LiftedHRepresentation) = size(rep.A, 2) - 1

similar_type(::Type{LiftedHRepresentation{S, MT}}, ::FullDim, ::Type{T}) where {S, T, MT} = LiftedHRepresentation{T, similar_type(MT, T)}
hvectortype(p::Type{LiftedHRepresentation{T, MT}}) where {T, MT} = vectortype(MT)

LiftedHRepresentation{T}(A::AbstractMatrix{T}, linset::BitSet=BitSet()) where {T} = LiftedHRepresentation{T, typeof(A)}(A, linset)
LiftedHRepresentation{T}(A::AbstractMatrix, linset::BitSet=BitSet()) where {T} = LiftedHRepresentation{T}(AbstractMatrix{T}(A), linset)
LiftedHRepresentation(A::AbstractMatrix{T}, linset::BitSet=BitSet()) where T = LiftedHRepresentation{T}(A, linset)

LiftedHRepresentation(h::HRepresentation{T}) where {T} = LiftedHRepresentation{T}(h)
LiftedHRepresentation{T}(h::HRepresentation) where {T} = convert(LiftedHRepresentation{T, hmatrixtype(typeof(h), T)}, h)

function LiftedHRepresentation{T, MT}(d::FullDim, hyperplanes::HyperPlaneIt{T},
                                      halfspaces::HalfSpaceIt{T}) where {T, MT}
    nhyperplane = length(hyperplanes)
    nhrep = nhyperplane + length(halfspaces)
    A = emptymatrix(MT, nhrep, fulldim(d)+1)
    linset = BitSet(1:nhyperplane)
    for (i, h) in enumerate(hyperplanes)
        A[i,2:end] = -h.a
        A[i,1] = h.β
    end
    for (i, h) in enumerate(halfspaces)
        A[nhyperplane+i,2:end] = -h.a
        A[nhyperplane+i,1] = h.β
    end
    LiftedHRepresentation{T}(A, linset)
end

Base.copy(ine::LiftedHRepresentation{T}) where {T} = LiftedHRepresentation{T}(copy(ine.A), copy(ine.linset))

Base.isvalid(hrep::LiftedHRepresentation{T}, idx::HIndex{T}) where {T} = 0 < idx.value <= size(hrep.A, 1) && (idx.value in hrep.linset) == islin(idx)
done(idxs::HIndices{T, <:LiftedHRepresentation{T}}, idx::HIndex{T}) where {T} = idx.value > size(idxs.rep.A, 1)
Base.get(hrep::LiftedHRepresentation{T}, idx::HIndex{T}) where {T} = valuetype(idx)(-hrep.A[idx.value,2:end], hrep.A[idx.value,1])

# V-Representation

mutable struct LiftedVRepresentation{T, MT<:AbstractMatrix{T}} <: MixedVRep{T}
    R::MT # each row is a vertex if the first element is 1 and a ray otherwise
    linset::BitSet

    function LiftedVRepresentation{T, MT}(R::MT, linset::BitSet=BitSet([])) where {T, MT}
        if !isempty(linset) && last(linset) > size(R, 1)
            error("The elements of linset should be between 1 and the number of rows of R")
        end
        new{T, MT}(R, linset)
    end
end
FullDim(rep::LiftedVRepresentation) = size(rep.R, 2) - 1

similar_type(::Type{LiftedVRepresentation{S, MT}}, ::FullDim, ::Type{T}) where {S, T, MT} = LiftedVRepresentation{T, similar_type(MT, T)}
vvectortype(::Type{LiftedVRepresentation{T, MT}}) where {T, MT} = vectortype(MT)

LiftedVRepresentation{T}(R::AbstractMatrix{T}, linset::BitSet=BitSet()) where {T} = LiftedVRepresentation{T, typeof(R)}(R, linset)
LiftedVRepresentation{T}(R::AbstractMatrix, linset::BitSet=BitSet()) where {T} = LiftedVRepresentation{T}(AbstractMatrix{T}(R), linset)
LiftedVRepresentation(R::AbstractMatrix{T}, linset::BitSet=BitSet()) where T = LiftedVRepresentation{T}(R, linset)

LiftedVRepresentation(v::VRepresentation{T}) where {T} = LiftedVRepresentation{T}(v)
LiftedVRepresentation{T}(v::VRepresentation) where {T} = convert(LiftedVRepresentation{T, vmatrixtype(typeof(v), T)}, v)

function LiftedVRepresentation{T, MT}(d::FullDim,
                                      vits::VIt{T}...) where {T, MT}
    points, lines, rays = fillvits(d, vits...)
    npoint = length(points)
    nline = length(lines)
    nray = length(rays)
    nvrep = npoint + nline + nray
    R = emptymatrix(MT, nvrep, fulldim(d) + 1)
    linset = BitSet()
    function _fill(offset, z, ps)
        for (i, p) in enumerate(ps)
            R[offset + i,2:end] = coord(p)
            R[offset + i,1] = z
            if islin(p)
                push!(linset, offset + i)
            end
        end
    end
    _fill(0, one(T), points)
    _fill(npoint, zero(T), lines)
    _fill(npoint+nline, zero(T), rays)
    LiftedVRepresentation{T}(R, linset)
end

Base.copy(ext::LiftedVRepresentation{T}) where {T} = LiftedVRepresentation{T}(copy(ext.R), copy(ext.linset))

nvreps(ext::LiftedVRepresentation) = size(ext.R, 1)

function isrowpoint(ext::LiftedVRepresentation{T}, i) where {T}
    ispoint = ext.R[i,1]
    @assert ispoint == zero(T) || ispoint == one(T)
    ispoint == one(T)
end

function Base.isvalid(vrep::LiftedVRepresentation{T}, idx::VIndex{T}) where {T}
    isp = isrowpoint(vrep, idx.value)
    isl = (idx.value in vrep.linset)
    @assert !isp || !isl # if isp && isl, it is a symmetric point but it is not allowed to mix symmetric points and points
    0 < idx.value <= nvreps(vrep) && isp == ispoint(idx) && isl == islin(idx)
end
done(idxs::VIndices{T, <:LiftedVRepresentation{T}}, idx::VIndex{T}) where {T} = idx.value > size(idxs.rep.R, 1)
Base.get(vrep::LiftedVRepresentation{T}, idx::VIndex{T}) where {T} = valuetype(idx)(vrep.R[idx.value,2:end])


text:
module TestEquality
using Compat
using Compat.Test
using CategoricalArrays

@testset "== and isequal() for CategoricalPool{Int} and CategoricalPool{Float64}" begin
    pool1 = CategoricalPool([1, 2, 3])
    pool2 = CategoricalPool([2.0, 1.0, 3.0])

    @test isequal(pool1, pool1) === true
    @test isequal(pool1, pool2) === false
    @test isequal(pool2, pool2) === true

    @test (pool1 == pool1) === true
    @test (pool1 == pool2) === false
    @test (pool2 == pool2) === true

    @test (pool1 === pool1) === true
    @test (pool1 === pool2) === false
    @test (pool2 === pool2) === true

    opool1 = CategoricalPool([1, 2, 3], true)
    opool2 = CategoricalPool([2.0, 1.0, 3.0], true)

    @test isequal(opool1, opool1) === true
    @test isequal(opool1, opool2) === false
    @test isequal(opool2, opool2) === true

    @test (opool1 == opool1) === true
    @test (opool1 == opool2) === false
    @test (opool2 == opool2) === true

    @test (opool1 === opool1) === true
    @test (opool1 === opool2) === false
    @test (opool2 === opool2) === true

    nv1a = CategoricalArrays.catvalue(1, pool1)
    nv2a = CategoricalArrays.catvalue(1, pool2)
    nv1b = CategoricalArrays.catvalue(2, pool1)
    nv2b = CategoricalArrays.catvalue(2, pool2)

    @test isequal(nv1a, nv1a) == true
    @test isequal(nv1a, nv2a) == false
    @test isequal(nv1a, nv1b) == false
    @test isequal(nv1a, nv2b) == true

    @test isequal(nv1b, nv1a) == false
    @test isequal(nv1b, nv2a) == true
    @test isequal(nv1b, nv1b) == true
    @test isequal(nv1b, nv2b) == false

    @test isequal(nv2a, nv1a) == false
    @test isequal(nv2a, nv2a) == true
    @test isequal(nv2a, nv1b) == true
    @test isequal(nv2a, nv2b) == false

    @test isequal(nv2b, nv1a) == true
    @test isequal(nv2b, nv2a) == false
    @test isequal(nv2b, nv1b) == false
    @test isequal(nv2b, nv2b) == true

    @test isequal(1, nv1a) == true
    @test isequal(1, nv2a) == false
    @test isequal(1, nv1b) == false
    @test isequal(1, nv2b) == true

    @test isequal(nv1a, 2) == false
    @test isequal(nv2a, 2) == true
    @test isequal(nv1b, 2) == true
    @test isequal(nv2b, 2) == false

    ov1a = CategoricalArrays.catvalue(1, opool1)
    ov2a = CategoricalArrays.catvalue(1, opool2)
    ov1b = CategoricalArrays.catvalue(2, opool1)
    ov2b = CategoricalArrays.catvalue(2, opool2)

    @test isequal(ov1a, ov1a) == true
    @test isequal(ov1a, ov2a) == false
    @test isequal(ov1a, ov1b) == false
    @test isequal(ov1a, ov2b) == true

    @test isequal(ov1b, ov1a) == false
    @test isequal(ov1b, ov2a) == true
    @test isequal(ov1b, ov1b) == true
    @test isequal(ov1b, ov2b) == false

    @test isequal(ov2a, ov1a) == false
    @test isequal(ov2a, ov2a) == true
    @test isequal(ov2a, ov1b) == true
    @test isequal(ov2a, ov2b) == false

    @test isequal(ov2b, ov1a) == true
    @test isequal(ov2b, ov2a) == false
    @test isequal(ov2b, ov1b) == false
    @test isequal(ov2b, ov2b) == true

    @test isequal(1, ov1a) == true
    @test isequal(1, ov2a) == false
    @test isequal(1, ov1b) == false
    @test isequal(1, ov2b) == true

    @test isequal(ov1a, 2) == false
    @test isequal(ov2a, 2) == true
    @test isequal(ov1b, 2) == true
    @test isequal(ov2b, 2) == false

    @test (ov1a == ov1a) == true
    @test (ov1a == ov2a) == false
    @test (ov1a == ov1b) == false
    @test (ov1a == ov2b) == true

    @test (ov1b == ov1a) == false
    @test (ov1b == ov2a) == true
    @test (ov1b == ov1b) == true
    @test (ov1b == ov2b) == false

    @test (ov2a == ov1a) == false
    @test (ov2a == ov2a) == true
    @test (ov2a == ov1b) == true
    @test (ov2a == ov2b) == false

    @test (ov2b == ov1a) == true
    @test (ov2b == ov2a) == false
    @test (ov2b == ov1b) == false
    @test (ov2b == ov2b) == true

    @testset "ordered and non-ordered values are equal" begin
    @test (ov1a == nv1a) === true
    @test (ov1a == nv2a) === false
    @test (ov1a == nv1b) === false
    @test (ov1a == nv2b) === true

    @test (ov1b == nv1a) === false
    @test (ov1b == nv2a) === true
    @test (ov1b == nv1b) === true
    @test (ov1b == nv2b) === false

    @test (ov2a == nv1a) === false
    @test (ov2a == nv2a) === true
    @test (ov2a == nv1b) === true
    @test (ov2a == nv2b) === false

    @test (ov2b == nv1a) === true
    @test (ov2b == nv2a) === false
    @test (ov2b == nv1b) === false
    @test (ov2b == nv2b) === true
    end

    @testset "non-equality with missing" begin
    @test ismissing(nv1a == missing)
    @test ismissing(ov1a == missing)
    @test ismissing(missing == nv1a)
    @test ismissing(missing == ov1a)

    @test isequal(nv1a, missing) == false
    @test isequal(ov1a, missing) == false
    @test isequal(missing, nv1a) == false
    @test isequal(missing, ov1a) == false
    end
end

@testset "in()" begin
    pool = CategoricalPool([5, 1, 3])
    nv = CategoricalArrays.catvalue(2, pool)

    @test (nv in 1:3) === true
    @test (nv in [1, 2, 3]) === true
    @test (nv in 2:3) === false
    @test (nv in [2, 3]) === false
end

end


text:
"""
    struct Straight{T} <: ContinuousSegment{T}
        l::T
        α0::typeof(0.0°)
        axis::Tensors.Vec{3, Float64}
    end
A straight line segment is parameterized by its length.
It begins at a point `p0` with initial angle `α0`.
The parametric function describing the line segment is given by
`t -> p0 + Point(t*cos(α),t*sin(α))` where `t` is a length from 0 to `l`.
"""
struct Straight{T} <: Segment{T}
    l::T
    p0::T
    axis::Tensors.Vec{3, Float64}
    α0::typeof(1.0°)
end

"""
    Straight{T<:Coordinate}(l::T; p0::Point=Point(zero(T),zero(T)), α0=0.0°)
Outer constructor for `Straight` segments.
"""
Straight(l::T;
         p0::T=zero(T),
         axis=Tensors.Vec{3}((0.0, 0.0, 1.0)),
         α0=0.0°) where {T <: Unitful.Length} =
    Straight{T}(l, p0, axis, α0)

pathlength(s::Straight) = s.l
orientation(s::Straight, z) = (s.axis, s.α0)

"""
    add_straight!(w::Waveguide{T}, l::Unitful.Length,
        med::Medium=lastmedium(w)) where {T <: Unitful.Length}
Extend a waveguide `w` straight by length `l` in the current direction.
By default, we take the last continuous medium in the waveguide
"""
function add_straight!(w::Waveguide{T}, l::Unitful.Length,
        med::Medium=lastmedium(w)) where {T <: Unitful.Length}
    Unitful.dimension(T) != Unitful.dimension(typeof(l)) && throw(Unitful.DimensionError(T(1),l))
    @assert l > zero(l) "tried to go straight by zero or a negative amount."
    p0 = pathlength(w)
    (axis, α) = orientation(w, p0)
    s = Straight{T}(l, p0, axis, α)
    push!(w.structures, Structure(s, med))
    return nothing
end


text:
# ----------------------------------------------------------------------------------- #
# Copyright (c) 2020 Varnerlab
# Robert Frederick School of Chemical and Biomolecular Engineering
# Cornell University, Ithaca NY 14850

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ----------------------------------------------------------------------------------- #

# where are we installed?
_PATH_TO_ROOT = pwd()
_PATH_TO_SRC = joinpath(_PATH_TO_ROOT,"src")

# import pkg -
import Pkg
Pkg.activate(_PATH_TO_ROOT)
Pkg.instantiate()

# system packages - these are required to be installed to solve the modeling problem
using LinearAlgebra # pre-installed w/Julia
using Statistics    # pre-installed w/Julia
using GLPK
using DelimitedFiles
using TOML
using CSV
using DataFrames
using Logging

# load my model files -
my_model_files_array = [
    "Checks.jl"         ;
    "Types.jl"          ;
    "Data.jl"           ;
    "Kinetics.jl"       ;
    "Control.jl"        ;
    "Solver.jl"         ;
    "Utility.jl"        ;
    "Flux.jl"           ;
    "Constraints.jl"    ;
]

# process my model files -
for file_name in my_model_files_array
    path_to_my_file = joinpath(_PATH_TO_SRC,file_name)
    include("$(path_to_my_file)")
end

max_stars_repo_name: varnerlab/CellFreeModelGenerationKit

max_issues_repo_name: varnerlab/CellFreeModelGenerationKit

max_forks_repo_name: varnerlab/CellFreeModelGenerationKit


text:
import LinearAlgebra: Transpose, opnorm
import LinearAlgebra
import Random: seed!
using SparseArrays
export fill_diag!, diag!

# temporary function. will actually add CUDA barrier function.
function sync()
    nothing
end

"""
    Allgather_opt(B)
Returns a tuple of proper function and argument for Allgather.
"""
function Allgather_opt(B::MPIVecOrMat{T,AT}) where {T,AT}
    if allsame(B.local_lengths)
        (Allgather!, B.localarray, B.local_lengths[1])
    else
        (Allgatherv!, B.localarray, B.local_lengths)
    end
end

function Allgather_opt(B::Transpose{T,ATT} where ATT <: MPIVecOrMat{T,AT}) where {T,AT}
    B_buf = transpose(B).localarray
    tB_local_lengths = transpose(B).local_lengths
    if allsame(tB_local_lengths)
        (Allgather!, B_buf, tB_local_lengths[1])
    else
        (Allgatherv!, B_buf, tB_local_lengths)
    end
end

@inline get_local(A::MPIVecOrMat{T,AT}) where {T,AT} = A.localarray
@inline get_local(A::Transpose{T,ATT} where ATT <: MPIVecOrMat{T,AT}) where {T,AT} = transpose(transpose(A).localarray)

"""
Scenario a: short and fat matrix multiplied by a fat matrix. Temporary space required
A: r x [p], B: p x [q], C: r x [q]
tmp: r x p: A is Allgather!-ed.
"""
function LinearAlgebra.mul!(C::MPIMatrix{T,AT}, A::MPIMatrix{T,AT}, B::MPIMatrix{T,AT};
                            tmp::AbstractArray{T,2}=AT{T}(undef, size(A,1), size(A,2))) where {T,AT}
    @assert size(C,1) == size(A,1) && size(C,2) == size(B,2)
    @assert size(A,1) == size(tmp,1) && size(A, 2) == size(tmp, 2)
    localA = get_local(A)
    localB = get_local(B)
    localC = get_local(C)

    #Allgather
    sync()
    (Allgather_ftn!, sendbuf, count_arg) = Allgather_opt(A)
    Allgather_ftn!(sendbuf, tmp, count_arg)

    LinearAlgebra.mul!(localC, tmp, localB)
    C
end

"""
Scenario h: Scenario 1 transposed.
"""
function LinearAlgebra.mul!(C::Transpose{T,MPIMatrix{T,AT}}, A::Transpose{T,MPIMatrix{T,AT}}, B::Transpose{T,MPIMatrix{T,AT}};
                            tmp::AbstractArray{T,2}=AT{T}(undef, size(B,2), size(B,1))) where {T,AT}
    LinearAlgebra.mul!(transpose(C), transpose(B), transpose(A); tmp=tmp)
    C
end

"""
Scenario b: inner product, result distributed, temporary space required
A: r x [p], B: [p] x q, C: r x [q]
tmp: r x q, for local computation
"""
function LinearAlgebra.mul!(C::MPIMatrix{T,AT}, A::MPIMatrix{T,AT}, B::Transpose{T, MPIMatrix{T,AT}};
                            tmp::AbstractArray{T,2}=AT{T}(undef, size(A,1), size(B,2))) where {T,AT}
    @assert size(C,1) == size(A,1) && size(C,2) == size(B,2)
    @assert size(A,1) == size(tmp,1) && size(B,2) == size(tmp,2)
    localA = get_local(A)
    localB = get_local(B)
    LinearAlgebra.mul!(tmp, localA, localB)
    # Ireduce would fit well here... but is not supported on both PyTorch and MPI.jl.
    sync()
    for i = 0:Size()-1
        # NOTE: an array is set to be contiguous only if the first indices are colon.
        Reduce!(@view(tmp[:, C.partitioning[i+1][2]]), C.localarray; root=i)
    end
    C
end

"""
Scenario c: Scenario 2 transposed.
"""
function LinearAlgebra.mul!(C::Transpose{T,MPIMatrix{T,AT}}, A::MPIMatrix{T,AT},B::Transpose{T,MPIMatrix{T,AT}};
                            tmp::AbstractArray{T,2}=AT{T}(undef, size(B,2), size(A,1))) where {T,AT}
    LinearAlgebra.mul!(transpose(C), transpose(B), transpose(A);tmp=tmp)
end

"""
Scenario d: inner product, result broadcasted
"""
function LinearAlgebra.mul!(C::AbstractMatrix{T}, A::MPIMatrix{T,AT}, B::Transpose{T,MPIMatrix{T,AT}}) where {T,AT}
    @assert size(C, 1) == size(A, 1) && size(C, 2) == size(B, 2)
    localA = get_local(A)
    localB = get_local(B)
    LinearAlgebra.mul!(C, localA, localB)
    sync()
    Allreduce!(C)
    C
end

"""
Scenario e: A::MPIMatrix, B, C: broadcast matrix.
"""
function LinearAlgebra.mul!(C::AbstractMatrix{T}, A::MPIMatrix{T,AT}, B::AbstractMatrix{T}) where {T, AT}
    localA = get_local(A)
    LinearAlgebra.mul!(C, localA, B[A.partitioning[Rank() + 1][2], :])
    sync()
    Allreduce!(C)
    C
end

"""
Scenario f: outer product, temporary space required
transpose(A) is Allgather!-ed.
"""
function LinearAlgebra.mul!(C::MPIMatrix{T,AT}, A::Transpose{T,MPIMatrix{T,AT}}, B::MPIMatrix{T,AT};
                            tmp::AbstractArray{T,2}=AT{T}(undef,size(A,2),size(A,1))) where {T,AT}
    @assert size(C,1) == size(A,1) && size(C,2) == size(B,2)
    @assert size(A,2) == size(tmp,1) && size(A,1) == size(tmp,2)
    localA = get_local(A)
    localB = get_local(B)
    localC = get_local(C)

    # Allgather
    sync()
    (Allgather_ftn!, sendbuf, count_arg) = Allgather_opt(A)
    Allgather_ftn!(sendbuf, tmp, count_arg)

    LinearAlgebra.mul!(localC, transpose(tmp), localB)
    C
end

"""
Scenario g: Scenario 3 transposed.
"""
function LinearAlgebra.mul!(C::Transpose{T,MPIMatrix{T,AT}}, A::Transpose{T,MPIMatrix{T,AT}}, B::MPIMatrix{T,AT};
                            tmp::AbstractArray{T,2}=AT{T}(undef,size(B,1), size(B,2))) where {T,AT}
    LinearAlgebra.mul!(transpose(C), transpose(B), transpose(A);tmp=tmp)
    C
end

"""
Scenario j: Small, broadcasted matrix multiplied by a distributed matrix.
A: s x r, B: r x [q], C: s x [q].
"""
function LinearAlgebra.mul!(C::MPIMatrix{T,AT}, A::AbstractMatrix{T}, B::MPIMatrix{T,AT}) where {T,AT}
    localB = get_local(B)
    localC = get_local(C)
    LinearAlgebra.mul!(localC, A, localB)
    C
end

"""
Scenario i: Scenario 4 transposed.
"""
function LinearAlgebra.mul!(C::Transpose{T,MPIMatrix{T,AT}}, A::Transpose{T,MPIMatrix{T,AT}},
    B::AbstractMatrix{T}) where {T,AT}
LinearAlgebra.mul!(transpose(C), transpose(B), transpose(A))
C
end

"""
Scenario k: B transposed MPIMatrix, A, C: broadcast matrix.
"""
function LinearAlgebra.mul!(C::AbstractMatrix{T}, A::AbstractMatrix{T},
    B::Transpose{T, MPIMatrix{T,AT}}) where {T, AT}
    LinearAlgebra.mul!(transpose(C), transpose(B), transpose(A))
    C
end

"""
Distributed matrix x broadcasted vector multiplications
"""
const MPIColVector{T,AT} = Union{MPIVector{T,AT},Transpose{T,MPIMatrix{T,AT}}}
get_MPIArray(m::MPIArray) = m
get_MPIArray(m::Transpose{T,MPIMatrix{T,A}}) where {T,A} = transpose(m)

"""
Scenario l: Col-dist matrix x distributed vector = distributed vector
"""
function LinearAlgebra.mul!(C::MPIColVector{T,AT}, A::MPIMatrix{T,AT}, B::MPIColVector{T,AT};
                            tmp::AbstractArray{T}=AT{T}(undef, size(C, 1))) where {T,AT}
    LinearAlgebra.mul!(tmp, A, B)
    localC = get_local(C)
    C_mpi = get_MPIArray(C)
    localC .= tmp[C_mpi.partitioning[Rank()+1][1]]
end

"""
Scenario m: Inner product
"""
function LinearAlgebra.mul!(C::AbstractVector{T}, A::MPIMatrix{T,AT}, B::MPIColVector{T,AT}) where {T,AT}
    localA = get_local(A)
    localB = get_local(B)
    LinearAlgebra.mul!(C, localA, localB)
    sync()
    Allreduce!(C)
    C
end

"""
Scenario n: no vector distributed, A MPIMatrix
"""
function LinearAlgebra.mul!(C::AbstractVector{T}, A::MPIMatrix{T,AT}, B::AbstractVector{T}) where {T,AT}
    localA = get_local(A)
    LinearAlgebra.mul!(C, localA, B[A.partitioning[Rank()+1][2]])
    sync()
    Allreduce!(C)
    C
end

"""
Scenario o: Row-dist matrix x distributed vector = distributed vector
"""
function LinearAlgebra.mul!(C::MPIColVector{T,AT}, A::Transpose{T,MPIMatrix{T,AT}},B::MPIColVector{T,AT};
                            tmp::AbstractArray{T}=AT{T}(undef, size(B,1))) where {T,AT}
    localA = get_local(A)
    localC = get_local(C)

    # Allgather
    sync()
    (Allgather_ftn!, sendbuf, count_arg) = Allgather_opt(B)
    Allgather_ftn!(sendbuf, tmp, count_arg)
    LinearAlgebra.mul!(localC, localA, tmp)
    C
end

"""
Scenario p: Row-dist matrix x broadcast vector = dist vector
"""
function LinearAlgebra.mul!(C::MPIColVector{T,AT}, A::Transpose{T,MPIMatrix{T,AT}}, B::AbstractVector{T}) where {T,AT}
    localA = get_local(A)
    localC = get_local(C)
    LinearAlgebra.mul!(localC, localA, B)
    # undefined in CuArrays: mul!(::Transpose{Float64,CuArray{Float64,2}}, ::Transpose{Float64,CuArray{Float64,2}}, ::CuArray{Float64,1})
    C
end

"""
Scenario q: no vector distributed, A transposed MPIMatrix
"""
function LinearAlgebra.mul!(C::AbstractVector{T}, A::Transpose{T, MPIMatrix{T,AT}}, B::AbstractVector{T}) where {T,AT}
    localA = get_local(A)
    fill!(C, zero(T))
    LinearAlgebra.mul!(@view(C[transpose(A).partitioning[Rank()+1][2]]), localA, B[transpose(A).partitioning[Rank()+1][1]])
    sync()
    Allreduce!(C)
    C
end

const AbstractSparseOrTranspose{T} = Union{AbstractSparseMatrix{T, <:Integer},Transpose{T,<:AbstractSparseMatrix}}

"""
Sparse matrix-vector multiplications
"""
#=
function LinearAlgebra.mul!(C::AbstractVector{T}, A::AbstractSparseOrTranspose{T}, B::MPIColVector{T,AT};
                            tmp::AbstractArray{T}=AT{T}(undef, size(B,1))) where {T,AT}
    @assert length(C) == size(A,1) && length(B) == size(A,2)
    @assert length(tmp) == length(B)

    # Allgather
    sync()
    (Allgather_ftn!, sendbuf, count_arg) = Allgather_opt(B)
    Allgather_ftn!(sendbuf, tmp, count_arg)

    LinearAlgebra.mul!(C, A, tmp)
    C
end

function LinearAlgebra.mul!(C::MPIColVector{T,AT}, A::AbstractSparseOrTranspose{T}, B::AbstractVector{T};
                            tmp::AbstractArray{T}=AT{T}(undef, size(A,1))) where {T,AT}
    @assert length(C) == size(A,1) && length(B) == size(A,2)
    @assert length(tmp) == length(C)
    localC = get_local(C)
    LinearAlgebra.mul!(tmp, A, B)
    C_mpi = get_MPIArray(C)
    localC .= tmp[C_mpi.partitioning[Rank()+1][1]]
    C
end

function LinearAlgebra.mul!(C::MPIColVector{T,AT}, A::AbstractSparseOrTranspose{T}, B::MPIColVector{T,AT};
                            tmp_m::AbstractArray{T}=AT{T}(undef, size(A,1)),
                            tmp_n::AbstractArray{T}=AT{T}(undef, size(B,1))) where {T,AT}
    @assert length(C) == size(A,1) && length(B) == size(A,2)
    @assert length(tmp_m) == size(A,1) && length(tmp_n) == size(A,2)
    LinearAlgebra.mul!(tmp_m, A, B)
    localC = get_local(C)
    C_mpi = get_MPIArray(C)
    localC .= tmp_m[C_mpi.partitioning[Rank()+1][1]]
    C
end
=#
"""
    LinearAlgebra.dot(A::MPIArray, B::MPIArray)

dot product of two MPI vectors.
"""
@inline function LinearAlgebra.dot(A::MPIArray, B::MPIArray)
   c = LinearAlgebra.dot(A.localarray, B.localarray)
   MPI.Allreduce(c, MPI.SUM, MPI.COMM_WORLD)
end

"""
    LinearAlgebra.diagind(M::MPIMatrix, k)

returns indices of the diagonal with respect to `M.localarray`.
"""
@inline function LinearAlgebra.diagind(M::MPIMatrix, k::Integer=0)
    offset = - (M.partitioning[Rank()+1][2][1]-1) + k
    return LinearAlgebra.diagind(size(M,1), size(M.localarray,2), offset)
end

"""
    LinearAlgebra.diag(M::MPIMatrix{T,A}, k; dist::Bool=false) where {T,A}

returns the diagonal of M. If dist is false, the result is a broadcasted A{T,1}. otherwise, the result is
distributed 1 x n "row vector". Distributed return is only valid for a square matrix (due to current way of distribution).
"""
function LinearAlgebra.diag(M::MPIMatrix{T,A}, k::Integer=0; dist::Bool=false) where {T,A}
    # TODO
end

"""
    diag!(d::MPIMatrix{T,A}, M::MPIMatrix{T,A})
returns a distributed 1-row matrix.
"""
function diag!(d::MPIMatrix{T,A}, M::MPIMatrix{T,A}) where {T,A}
    @assert size(d,1) == 1
    @assert size(d,2) == size(M, 2)
    @assert size(M,1) == size(M,2)
    d.localarray .= reshape(M.localarray[LinearAlgebra.diagind(M)], 1, :)
end

"""
    diag!(d::AbstractVector, M::MPIMatrix{T,A})
returns a local col vector.
"""
function diag!(d::AbstractVector, M::MPIMatrix{T,A}) where {T,A}
    @assert size(d,1) == size(M, 1)
    @assert size(M,1) == size(M,2)
    d[M.partitioning[Rank()+1][2]] .= M.localarray[LinearAlgebra.diagind(M)]
    counts = reshape(map(x->convert(Cint, length(x[2])), M.partitioning), :)
    sync()

    MPI.Allgatherv!(MPI.IN_PLACE, d, counts, MPI.COMM_WORLD)
end

"""
    fill_diag!(M, x, k=0)

fills the diagonal of M with x.
"""
@inline function fill_diag!(M::MPIMatrix, x, k::Integer=0)
    M.localarray[LinearAlgebra.diagind(M, k)] .= x
end

"""
    opnorm(A::MPIMatrix; method="power", tol=1e-8)

Estimates operator l2 norm of MPIMatrix A. If `method=="power", it uses the power iteration.
If `method=="quick"`, it computes the product of matrix l1 norm and the matrix l-infty norm.
"""
function opnorm(A::MPIMatrix; method::String="power", tol=1e-6, maxiter=1000, seed=777, verbose=false)
    if method == "power"
        _opnorm_power(A; tol=tol, maxiter=maxiter, seed=seed, verbose=verbose)
    elseif method == "quick"
        _opnorm_quick(A)
    else
        @error("Invalid method. Valid values are: \"power\" and \"quick\".")
    end
end

function opnorm(A::MPIMatrix, p::Real; method::String="power", tol=1e-6, maxiter=1000, seed=777, verbose=false)
    if p == 1
        _opnorm_l1(A)
    elseif p == 2
        opnorm(A; method=method, tol=tol, maxiter=maxiter, seed=seed, verbose=verbose)
    elseif p == Inf
        _opnorm_linfty(A)
    else
        @error("Invalid p for opnorm. Valid values are: 1, 2, and Inf.")
    end
end

function _opnorm_l1(A::MPIMatrix)
    maximum(sum(abs.(A); dims=1))
end

function _opnorm_linfty(A::MPIMatrix)
    maximum(sum(abs.(A); dims=2))
end

function _opnorm_power(A::MPIMatrix{T,AT}; tol=1e-6, maxiter=1000, seed=777, verbose=false) where {T, AT}
    verbose && Rank() == 0 && println("computing max singular value...")
    m, n = size(A)
    seed!(seed)
    v = MPIVector{T,AT}(undef, n)
    randn!(v)
    Av = MPIVector{T,AT}(undef, m)
    LinearAlgebra.mul!(Av, A, v)
    s, s_prev = -Inf, -Inf
    for i in 1:maxiter
        LinearAlgebra.mul!(v, transpose(A), Av)
        v ./= sqrt(sum(v.^2))
        LinearAlgebra.mul!(Av, A, v)
        s = sqrt(sum(Av.^2))
        if abs((s_prev - s)/s) < tol
            break
        end
        s_prev = s
        verbose && Rank() == 0 && i%100 == 0 && println("iteration $i: $s")
    end
    verbose && Rank() == 0 && println("done computing max singular value: $s")
    s
end

function _opnorm_quick(A::MPIMatrix)
    l1 = _opnorm_l1(A)
    linf = _opnorm_linfty(A)
    l1*linf
end


text:
using Pkg.Artifacts

function __init__()
    prana_path = joinpath(artifact"Prana", "prana-2.8.0.beta.r2021.01.19")
    if !isCI
        cd(prana_path) do
            mat"addpath($(prana_path))"
        end
    end
end


text:
# This code is part of QuantumCircuits.
#
# (C) Copyright Rafał Pracht 2022.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

module Gates

using QuantumCircuits.QCircuits.QBase
using QuantumCircuits.QCircuits.Math
using QuantumCircuits.QCircuits.Registers

import Base: show, length, inv
import QuantumCircuits.QCircuits.QBase: tomatrix, setparameters!, simplify,
       standardGateError, decompose, bindparameters!

export X, Y, Z, S, Sd, T, Td, H, CX, U3, Rx, Ry, Rz, U, Sx,
      getqubits, getqubitsids, toU3,
      ChromosomeGate, DNAGate, Parameters, getArgs, decompose,
      ParameterT, getvalue

"Parameter type"
const ParamT = Number

mutable struct ParameterT
    value::Union{ParamT, Nothing}
end
ParameterT() = ParameterT(nothing)
ParameterT(param::ParameterT) = ParameterT(param.value)
Base.show(io::IO, p::ParameterT) = print(io, "$(p.value)")

const Parameter = Union{ParamT, ParameterT}

getvalue(p::ParamT) = p
function getvalue(p::ParameterT)
    @assert p.value != nothing "Unable get unset parameter."
    return p.value
end

#const ParamIdx = Int
# "The gates parameters structure, one plece for all parameters in cirquit."
# struct Parameters
#     params::Vector{ParamT}
#
#     Parameters() = new(ParamT[])
# end
# "Function return the next free parameter."
# function getnext(p::Parameters)
#     push!(p.params, zero(ParamT))
#
#     return length(p.params)
# end
# "Get/Set index"
# Base.getindex(p::Parameters, idx::ParamIdx) = p.params[idx]
# Base.setindex!(p::Parameters, v::ParamT, idx::ParamIdx) = p.params[idx] = v

"Gate which can be use in chromosome"
#abstract type ChromosomeGate <: QuantumGate end
#abstract type DNAGate <: ChromosomeGate end

#"Single Qubit Gate"
#abstract type SingleQubitGate <: QuantumGate end

"Quantum Universal Gate"
abstract type UniversalGate <: QuantumGate end

"Comparator"
Base.:(==)(c1::QuantumGate, c2::QuantumGate) = false

"Single qubit create macro"
macro singleQubitGate(name)
    eval(quote
        struct $name <: QuantumGate qubit::Qubit end

        Base.:(==)(g1::$name, g2::$name) = g1.qubit == g2.qubit
        Base.hash(g::$name, h::UInt) = hash(g.qubit, h)

        getArgs(g::$name) = (getid(g.qubit), )
    end)
end

# Single qubit quantum gates
@singleQubitGate(X)
@singleQubitGate(Y)
@singleQubitGate(Z)
@singleQubitGate(S)
@singleQubitGate(Sd)
@singleQubitGate(H)
@singleQubitGate(T)
@singleQubitGate(Td)
@singleQubitGate(Sx)

"CX Quantum Gate"
struct CX <: QuantumGate
    control::Qubit
    target::Qubit
end
Base.:(==)(g1::CX, g2::CX) = g1.control == g2.control && g1.target == g2.target
Base.hash(g::CX, h::UInt) = hash((g.control, g.target), h)
function Base.:(>)(x::CX, y::CX)
    if x.control == y.control
        return x.target > y.target
    else
        return x.control > y.control
    end
end
getArgs(g::CX) = (getid(g.control), getid(g.target))

"Rotation gates"
abstract type RotationGate <: QuantumGate end

"Rotation macro"
macro rotationGate(name)
    eval(quote
        mutable struct $name <: RotationGate
            qubit::Qubit
            θ::Parameter
        end
        $name(qubit::Integer) = $name(qubit, ParameterT(rand() * 2π))

        Base.:(==)(g1::$name, g2::$name) = g1.qubit == g2.qubit && g1.θ == g2.θ
        Base.hash(g::$name, h::UInt) = hash((g.qubit, g.θ), h)

        $name(qubis::Vector{T}, θ::ParamT) where T <: Integer = [$name(q, θ) for q in qubis]
        getArgs(g::$name) = (getid(g.qubit), g.θ)
    end)
end

@rotationGate(Rx)
@rotationGate(Ry)
@rotationGate(Rz)

appendparams!(param, gate::RotationGate) = appendparams!(param, gate.θ)
appendRandParams!(param, gate::RotationGate) = appendRandParams!(param, gate.θ)
Base.length(gate::RotationGate) = getlength(gate.θ)
function setparameters!(gate::RotationGate, params)
    @assert length(params) == length(gate) "RotationGate has $(length(gate)) parameters but $(length(param)) was provided."
    setparameter!(gate.θ, params, 1)
end

"The standard U3 gate"
mutable struct U3 <: UniversalGate
    qubit::Qubit
    θ::Parameter
    ϕ::Parameter
    λ::Parameter
end
U3(qubit::Qubit) = U3(qubit, ([ParameterT(p) for p in rand(3) * 2π])...)
getArgs(g::U3) = (getid(g.qubit), g.θ, g.ϕ, g.λ)

Base.show(io::IO, g::U3) = print(io, "U3($(g.qubit), $(g.θ), $(g.ϕ), $(g.λ))")

Base.:(==)(g1::U3, g2::U3) = g1.qubit == g2.qubit && g1.θ == g2.θ && g1.ϕ == g2.ϕ && g1.λ == g2.λ
Base.hash(g::U3, h::UInt) = hash((g.qubit, g.θ, g.ϕ, g.λ), h)

bindparameters!(gate::QuantumGate) = nothing
bindparameters!(gate::RotationGate) = gate.θ = getvalue(gate.θ)
function bindparameters!(gate::U3)
    gate.θ = getvalue(gate.θ)
    gate.ϕ = getvalue(gate.ϕ)
    gate.λ = getvalue(gate.λ)
end

# "The universar single qubit unitary gate macro"
# macro universalGate(name)
#     eval(quote
#         mutable struct $name <: UniversalGate
#             qubit::Qubit
#             θ::ParamT
#             ϕ::ParamT
#             λ::ParamT
#         end
#         $name(qubit::Integer) = $name(qubit, (rand(3) * 2π)...)
#     end)
# end

"The universar single qubit unitary gate"
mutable struct U <: UniversalGate
    qubit::Qubit
    β::Union{ParamT, Nothing}
    γ::Union{ParamT, Nothing}
    δ::Union{ParamT, Nothing}
end
U(qubit::Integer) = U(qubit, nothing, nothing, nothing)
getArgs(g::U) = (getid(g.qubit), g.β, g.γ, g.δ)

function setparameter!(param::ParameterT, params, i)
    param.value = params[i]
    return i + 1
end
setparameter!(param::ParamT, params, i) = i

function setparameters!(gate::U3, params)
    @assert length(params) == length(gate) "U3 has $(length(gate)) parameters but $(length(params)) was provided."
        i = setparameter!(gate.θ, params, 1)
        i = setparameter!(gate.ϕ, params, i)
        i = setparameter!(gate.λ, params, i)
end

function setparameters!(gate::U, params)
    @assert length(params) == 3 "U has 3 parameters but $(length(params)) was provided."
    gate.β = params[1]
    gate.γ = params[2]
    gate.δ = params[3]
end

"Return the qubits on which operate the gate"
getqubits(gate::QuantumGate) = (gate.qubit,) # All Single Qubit Gates has qubit field
getqubits(gate::CX) = (gate.control, gate.target)

"Return the qubits ids on which operate the gate"
getqubitsids(gate::QuantumGate) = (getid(gate.qubit),) # All Single Qubit Gates has qubit field
getqubitsids(gate::CX) = (getid(gate.control), getid(gate.target))

"Show method"
Base.show(io::IO, gate::QuantumGate) = print(io, "$(typeof(gate))($(gate.qubit))")
Base.show(io::IO, gate::CX) = print(io, "CX($(gate.control), $(gate.target))")

"Convert gate to U3 gate"
toU3(gate::X) = U3(gate.qubit, π, 0, π)

"Return the gate matrix"
u3matrix(θ, ϕ, λ) = [cos(θ/2)+0im              -exp(1im * λ) * sin(θ/2);
                     exp(1im * ϕ) * sin(θ/2)   exp(1im * λ + 1im * ϕ) * cos(θ/2)]
u2matrix(ϕ, λ) = u3matrix(π/2, ϕ, λ)
u1matrix(λ) = u3matrix(0, 0, λ)

# u3matrix(θ, -π/2, π/2)
Rxmatrix(θ) = [cos(θ/2)+0im -1im * sin(θ/2); -1im * sin(θ/2) cos(θ/2)]
# u3matrix(θ, 0, 0)
Rymatrix(θ) = [cos(θ/2)+0im -sin(θ/2); sin(θ/2) cos(θ/2)]
# exp(-1im * θ/2) * u1matrix(θ)
Rzmatrix(θ) = [exp(-1im * θ/2) 0; 0 exp(1im * θ/2)]
# Universary matrix
umatrix(β, γ, δ) = Rzmatrix(β) * Rymatrix(γ) * Rzmatrix(δ)

# Initial stata
const _0 = [1+0im; 0]
const _1 = [0im; 1]

const I0X0I = _0 * dagger(_0)
const I1X1I = _1 * dagger(_1)

# Gate matrixs
const Xmatrix = round.(u3matrix(π, 0, π), digits=2)
const Ymatrix = round.(u3matrix(π, π/2, π/2), digits=2)
const Zmatrix = round.(u1matrix(π), digits=2)
const Imatrix = [1.0 + 0im 0; 0 1]
const Hmatrix = u2matrix(0, π)
const Smatrix = round.(u1matrix(π/2), digits=2)
const Sdmatrix = round.(u1matrix(-π/2), digits=2)
const Tmatrix = u1matrix(π/4)
const Tdmatrix = u1matrix(-π/4)
const CX0matrix = [1.0 + 0im 0 0 0; 0 0 0 1; 0 0 1 0; 0 1 0 0]
const CX1matrix = [1.0 + 0im 0 0 0; 0 1 0 0; 0 0 0 1; 0 0 1 0]

tomatrix(gate::X) = Xmatrix
tomatrix(gate::Y) = Ymatrix
tomatrix(gate::Z) = Zmatrix
tomatrix(gate::S) = Smatrix
tomatrix(gate::Sd) = Sdmatrix
tomatrix(gate::T) = Tmatrix
tomatrix(gate::Td) = Tdmatrix
tomatrix(gate::H) = Hmatrix
tomatrix(gate::U3) = u3matrix(getvalue(gate.θ), getvalue(gate.ϕ), getvalue(gate.λ))
function tomatrix(gate::U3, param)
    @assert length(param) == 3 "U3 has 3 parameters but $(length(param)) was provided."
    u3matrix(param[1], param[2], param[3])
end

tomatrix(gate::Rx) = Rxmatrix(getvalue(gate.θ))
tomatrix(gate::Ry) = Rymatrix(getvalue(gate.θ))
tomatrix(gate::Rz) = Rzmatrix(getvalue(gate.θ))
function tomatrix(gate::Rx, param)
    @assert length(param) == 1 "Rx has 1 parameters but $(length(param)) was provided."
    Rxmatrix(param[1])
end
function tomatrix(gate::Ry, param)
    @assert length(param) == 1 "Ry has 1 parameters but $(length(param)) was provided."
    Rymatrix(param[1])
end
function tomatrix(gate::Rz, param)
    @assert length(param) == 1 "Rz has 1 parameters but $(length(param)) was provided."
    Rzmatrix(param[1])
end

tomatrix(gate::U) = umatrix(gate.β, gate.γ, gate.δ)
function tomatrix(gate::U, param)
    @assert length(param) == 3 "U has 3 parameters but $(length(param)) was provided."
    umatrix(param[1], param[2], param[3])
end

function tomatrix(gate::CX)
    control = getid(gate.control)
    target = getid(gate.target)
    if abs(control - target) == 1
        if control < target
            return CX0matrix
        else
            return CX1matrix
        end
    else
        # https://quantumcomputing.stackexchange.com/questions/4252/how-to-derive-the-cnot-matrix-for-a-3-qubit-system-where-the-control-target-qu
        if control < target
            s0 = kron(Imatrix, I0X0I)
            s1 = kron(Imatrix, I1X1I)
            e0 = Imatrix
            e1 = Xmatrix
        else
            s0 = kron(Imatrix, Imatrix)
            s1 = kron(Imatrix, Xmatrix)
            e0 = I0X0I
            e1 = I1X1I
        end

        additional_update = abs(target - control) - 2
        for i in 1:additional_update
            s0 = kron(Imatrix, s0)
            s1 = kron(Imatrix, s1)
        end

        return kron(e0, s0) + kron(e1, s1)
    end
end

function tomatrix(qubits::Integer, gate::QuantumGate, params=nothing)
    # get gate unitary
    if params != nothing
        gate_unitary = tomatrix(gate, params)
    else
        gate_unitary = tomatrix(gate)
    end

    # create kroneker product
    if qubits == 1
        return gate_unitary
    end

    # Get rest
    idxs = getqubitsids(gate)
    minidx = minimum(idxs)
    maxidx = qubits - maximum(idxs) - 1
    #println("idxs: $idxs, minidx: $minidx, maxidx: $maxidx")

    return kron(eye(2^maxidx), gate_unitary, eye(2^minidx))
end

"Method return the error of standard gate"
standardGateError(gate::QuantumGate, params=nothing) = 0.0

function standardGateError(gate::UniversalGate, params=nothing)
    # get gate unitary
    if params != nothing
        m = tomatrix(gate, params)
    else
        m = tomatrix(gate)
    end

    # err = 1.0
    # err *= unitary_error(m, Xmatrix)
    # err *= unitary_error(m, Ymatrix)
    # err *= unitary_error(m, Zmatrix)
    # err *= unitary_error(m, Imatrix)
    # err *= unitary_error(m, Hmatrix)
    # err *= unitary_error(m, Smatrix)
    # err *= unitary_error(m, Sdmatrix)
    # err *= unitary_error(m, Tmatrix)
    # err *= unitary_error(m, Tdmatrix)
    #
    # return log(err + 1.0)

    return minimum([unitary_error(m, Xmatrix),
                    unitary_error(m, Ymatrix),
                    unitary_error(m, Zmatrix),
                    unitary_error(m, Imatrix),
                    unitary_error(m, Hmatrix),
                    unitary_error(m, Smatrix),
                    unitary_error(m, Sdmatrix),
                    unitary_error(m, Tmatrix),
                    unitary_error(m, Tdmatrix)])
end

appendparams!(param, gate) = nothing
appendparams!(param, value::ParameterT) = push!(param, value.value)
appendparams!(param, value::ParamT) = nothing
function appendparams!(param, gate::U3)
    appendparams!(param, gate.θ)
    appendparams!(param, gate.ϕ)
    appendparams!(param, gate.λ)
    nothing
end
function appendparams!(param, gate::U)
    push!(param, gate.β)
    push!(param, gate.γ)
    push!(param, gate.δ)
    nothing
end

appendRandParams!(param, value::ParameterT) = append!(param, rand(1) * 2π)
appendRandParams!(param, value::ParamT) = nothing
appendRandParams!(param, gate) = nothing
function appendRandParams!(param, gate::U3)
    appendRandParams!(param, gate.θ)
    appendRandParams!(param, gate.ϕ)
    appendRandParams!(param, gate.λ)
    nothing
end
appendRandParams!(param, gate::U) = append!(param, rand(3) * 2π)

getlength(value::ParameterT) = 1
getlength(value::ParamT) = 0
Base.length(gate) = 0
Base.length(gate::U3) = getlength(gate.θ) + getlength(gate.ϕ) + getlength(gate.λ)
Base.length(gate::U) = 3

################################################################################
const maxerr = 5e-5

"The library of single qubit gates"
const singlequbitsgates = Dict(Xmatrix => X,
                               Ymatrix => Y,
                               Zmatrix => Z,
                               Smatrix => S,
                               Sdmatrix => Sd,
                               Tmatrix => T,
                               Tdmatrix => Td,
                               Hmatrix => H)
# TODO  + I !!!!!

decompose(gate::QuantumGate) = [gate]

"The simplify the gates."
simplify(gate::QuantumGate) = [gate]

function simplify(gate::UniversalGate)
    exp = tomatrix(gate)

    # check for identify
    if unitary_error(Imatrix, exp) < maxerr
        return []
    end

    # single qubits gate
    for (matrix, mgate) in singlequbitsgates
        if unitary_error(matrix, exp) < maxerr
            # println("found: $(mgate(gate.qubit))")
            return [mgate(gate.qubit)]
        end
    end

    # two qubits gate
    for (matrix, mgate) in singlequbitsgates
        for (matrix2, mgate2) in singlequbitsgates
            # println("$mgate -> err: $(unitary_error(matrix, exp)) == $mgate")
            # println(matrix)
            # println(exp)
            if unitary_error(matrix2 * matrix, exp) < maxerr
                # println("found: $(mgate(gate.qubit))")
                return [mgate(gate.qubit), mgate2(gate.qubit)]
            end
        end
    end

    return [gate]
end

inv(gate::X) = gate
inv(gate::H) = gate
inv(gate::Z) = gate
inv(gate::Y) = gate
inv(gate::Rx) = Rx(gate.qubit, neg(gate.θ))
inv(gate::Ry) = Ry(gate.qubit, neg(gate.θ))
inv(gate::Rz) = Rz(gate.qubit, neg(gate.θ))
inv(gate::CX) = gate
inv(gate::T) = Td(gate.qubit)
inv(gate::S) = Sd(gate.qubit)
inv(gate::Td) = T(gate.qubit)
inv(gate::Sd) = S(gate.qubit)
inv(gate::U3) = U3(gate.qubit, neg(gate.θ), neg(gate.λ), neg(gate.ϕ))

function neg(p::ParameterT)
    @assert p.value != nothing "Unable inv unset parameter."
    return ParameterT(-p.value)
end
neg(value::ParamT) = -value

end  # module Gates


text:
using QuasiMonteCarlo
using QuasiMonteCarlo: SamplingAlgorithm

# We need to convert the matrix that QuasiMonteCarlo produces into a vector of Tuples like Surrogates expects
# This will eventually be removed once we refactor the rest of the code to work with d x n matrices instead
# of vectors of Tuples
function sample(args...; kwargs...)
    s = QuasiMonteCarlo.sample(args...; kwargs...)
    if s isa Vector
        # 1D case: s is a Vector
        return s
    else
        # ND case: s is a d x n matrix, where d is the dimension and n is the number of samples
        return collect(reinterpret(reshape, NTuple{size(s, 1), eltype(s)}, s))
    end
end

max_stars_repo_name: ludoro/Surrogates

max_issues_repo_name: ludoro/Surrogates

max_forks_repo_name: ludoro/Surrogates


text:
# types

abstract type AbstractCell end

struct DataCell <: AbstractCell
    data::AbstractArray
    width::Int
    height::Int
    style::Union{Nothing,Dict}
end

struct ColoredString
    data::String
    style::Union{Nothing,Dict}
end

struct Vertical <: AbstractCell
    data::String
    width::Int
    height::Int
    Vertical(height::Int) = new("|", 1, height)
end

struct Dash <: AbstractCell
    data::String
    repeat::Int
    Dash(dash::String, n::Int) = new(dash, n)
end

struct Connector <: AbstractCell
    data::String
    Connector(data::String) = new(data)
end

mutable struct Cell <: AbstractCell
    data::AbstractArray
    width::Int
    height::Int
end

mutable struct Margin
    leftside::Int
    rightside::Int
end

struct Mill
    board::AbstractArray
    option::Dict{Symbol, Any}
    Mill(board; options...) = new(board, Dict{Symbol, Any}(options))
end

const Linear{T} = AbstractVector{T} where T <: Union{Cell, Vertical}
const Horizontal{T} = AbstractVector{T} where T <: Union{Dash, Connector}
const PlateVector{T} = AbstractVector{T} where T <: Union{Linear, Horizontal}

struct TableMode
    style::Symbol
    corner_corner::Char
    header_fillchar::Char
    TableMode(style::Symbol; corner_corner::Char = '|', header_fillchar::Char = '-') = new(style, corner_corner, header_fillchar)
end


text:
import ForwardDiff

"Parameters for right-hand side function of degenerate variational integrator."
mutable struct ParametersCMDVI{DT, TT, D, ET <: NamedTuple} <: Parameters{DT,TT}
    equs::ET
    Δt::TT

    t::TT
    q::Vector{DT}
    v::Vector{DT}
    θ::Vector{DT}

    function ParametersCMDVI{DT,D}(equs::ET, Δt::TT) where {DT, TT, D, ET <: NamedTuple}
        new{DT, TT, D, ET}(equs, Δt, zero(TT), zeros(DT,D), zeros(DT,D), zeros(DT,D))
    end
end

@doc raw"""
Degenerate variational integrator cache.

### Fields

* `q`: internal stages of solution
* `v`: internal stages of vector field
* `Θ`: implicit function evaluated on solution
* `f`: vector field of implicit function
"""
struct IntegratorCacheCMDVI{DT,D,S} <: ODEIntegratorCache{DT,D}
    q::Vector{DT}
    v::Vector{DT}
    θ::Vector{DT}

    q̄::Vector{DT}
    θ̄::Vector{DT}
    f̄::Vector{DT}

    function IntegratorCacheCMDVI{DT,D,S}() where {DT,D,S}
        new(zeros(DT,D), zeros(DT,D), zeros(DT,D),
            zeros(DT,D), zeros(DT,D), zeros(DT,D))
    end
end

function IntegratorCache{ST}(params::ParametersCMDVI{DT,TT,D,S}; kwargs...) where {ST,DT,TT,D,S}
    IntegratorCacheCMDVI{ST,D,S}(; kwargs...)
end

@inline CacheType(ST, ::ParametersCMDVI{DT,TT,D,S}) where {DT,TT,D,S} = IntegratorCacheCMDVI{ST,D,S}

"""
Implicit midpoint Runge-Kutta integrator.
"""
struct IntegratorCMDVI{DT, TT, D, PT <: ParametersCMDVI{DT,TT},
                                   ST <: NonlinearSolver{DT},
                                   IT <: InitialGuessIODE{TT}} <: AbstractIntegratorRK{DT,TT}
    params::PT
    solver::ST
    iguess::IT
    caches::CacheDict{PT}

    function IntegratorCMDVI(params::ParametersCMDVI{DT,TT,D}, solver::ST, iguess::IT, caches) where {DT,TT,D,ST,IT}
        new{DT, TT, D, typeof(params), ST, IT}(params, solver, iguess, caches)
    end

    function IntegratorCMDVI{DT,D}(equations::NamedTuple, Δt::TT) where {DT,TT,D}
        # create params
        params = ParametersCMDVI{DT,D}(equations, Δt)

        # create cache dict
        caches = CacheDict(params)

        # create solver
        solver = create_nonlinear_solver(DT, 2D, params, caches)

        # create initial guess
        iguess = InitialGuessIODE(get_config(:ig_extrapolation), equations[:v̄], equations[:f̄], Δt)

        # create integrator
        IntegratorCMDVI(params, solver, iguess, caches)
    end

    function IntegratorCMDVI(equation::Union{IODE{DT}, LODE{DT}}, Δt::TT; kwargs...) where {DT,TT}
        IntegratorCMDVI{DT, ndims(equation)}(get_functions(equation), Δt; kwargs...)
    end
end

@inline Base.ndims(::IntegratorCMDVI{DT,TT,D}) where {DT,TT,D} = D

function initialize!(int::IntegratorCMDVI, sol::AtomicSolutionPODE)
    sol.t̄ = sol.t - timestep(int)

    equation(int, :v̄)(sol.t, sol.q, sol.v)
    equation(int, :f̄)(sol.t, sol.q, sol.v, sol.f)
    equation(int, :ϑ)(sol.t, sol.q, sol.v, sol.p)

    initialize!(int.iguess, sol.t, sol.q, sol.p, sol.v, sol.f,
                            sol.t̄, sol.q̄, sol.p̄, sol.v̄, sol.f̄)
end

function update_params!(int::IntegratorCMDVI, sol::AtomicSolutionPODE)
    # set time for nonlinear solver and copy previous solution
    int.params.t  = sol.t
    int.params.q .= sol.q
    int.params.v .= sol.v
    equations(int)[:ϑ](sol.t, sol.q, sol.v, int.params.θ)
end

function initial_guess!(int::IntegratorCMDVI{DT,TT,D}, sol::AtomicSolutionPODE{DT,TT},
                        cache::IntegratorCacheCMDVI{DT}=int.caches[DT]) where {DT,TT,D}

    # compute initial guess for solution
    evaluate!(int.iguess, sol.q̄, sol.p̄, sol.v̄, sol.f̄,
                          sol.q, sol.p, sol.v, sol.f,
                          cache.q, cache.v, one(TT))

    int.solver.x[1:D] .= cache.q                           # q(n+1)

    # compute initial guess for solution
    evaluate!(int.iguess, sol.q̄, sol.p̄, sol.v̄, sol.f̄,
                          sol.q, sol.p, sol.v, sol.f,
                          cache.q, cache.v, one(TT)/2)

    offset_v = D
    offset_x = D + div(D,2)
    for k in 1:div(D,2)
        int.solver.x[offset_v+k] = cache.v[k]              # v¹(n+1/2)
        int.solver.x[offset_x+k] = cache.q[div(D,2)+k]     # q²(n+1/2)
    end
end

function compute_stages!(x::Vector{ST}, q::Vector{ST}, v::Vector{ST}, θ::Vector{ST},
                                        q̄::Vector{ST}, θ̄::Vector{ST}, f̄::Vector{ST},
                                        params::ParametersCMDVI{DT,TT,D}) where {ST,DT,TT,D}

    # set some local variables for convenience and clarity
    local t⁻ = params.t
    local t̄  = params.t + params.Δt / 2
    local t⁺ = params.t + params.Δt

    # copy x to q
    q .= x[1:D]

    # copy x to q⁻, q⁺ and v
    for k in 1:div(D,2)
        q̄[k] = (params.q[k] + q[k]) / 2
        q̄[div(D,2)+k] = x[D+div(D,2)+k]

        v[k]          = x[D+k]
        v[div(D,2)+k] = 0
    end

    # compute f = f(q,v)
    params.equs[:f](t̄, q̄, v, f̄)

    # compute Θ = ϑ(q,v)
    params.equs[:ϑ](t̄, q̄, v, θ̄)
    params.equs[:ϑ](t⁺, q, v, θ)
end

function function_stages!(x::Vector{ST}, b::Vector{ST}, params::ParametersCMDVI{DT,TT,D},
                          caches::CacheDict) where {ST,DT,TT,D}
    # get cache for internal stages
    cache = caches[ST]

    # compute stages from nonlinear solver solution x
    compute_stages!(x, cache.q, cache.v, cache.θ, cache.q̄, cache.θ̄, cache.f̄, params)

    # compute b
    b[1:D] .= cache.θ̄  .- params.θ .- params.Δt .* cache.f̄ ./ 2

    for k in 1:div(D,2)
        b[D+k] = cache.q[k] - params.q[k] - params.Δt * cache.v[k]
        b[D+div(D,2)+k] = cache.θ[k] - params.θ[k] - params.Δt * cache.f̄[k]
    end
end

function integrate_step!(int::IntegratorCMDVI{DT,TT}, sol::AtomicSolutionPODE{DT,TT},
                         cache::IntegratorCacheCMDVI{DT}=int.caches[DT]) where {DT,TT}

    # update nonlinear solver parameters from atomic solution
    update_params!(int, sol)

    # compute initial guess
    initial_guess!(int, sol, cache)

    # reset atomic solution
    reset!(sol, timestep(int))

    # call nonlinear solver
    solve!(int.solver)

    # print solver status
    print_solver_status(int.solver.status, int.solver.params)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.solver.status, int.solver.params)

    # compute vector field at internal stages
    compute_stages!(int.solver.x, cache.q, cache.v, cache.θ, cache.q̄, cache.θ̄, cache.f̄, int.params)

    # compute final update
    sol.q .= cache.q
    sol.p .= cache.θ

    # copy solution to initial guess
    update_vector_fields!(int.iguess, sol.t, sol.q, sol.p, sol.v, sol.f)
end


text:
@testset "Gauss–Chebyshev" begin

    n = 10

    @testset "x.^2" begin
        x, w = gausschebyshev(n,1)
        @test dot(x.^2,w) ≈ pi/2
        x, w = gausschebyshev(n,2)
        @test dot(x.^2,w) ≈ pi/8
        x, w = gausschebyshev(n,3)
        @test dot(x.^2,w) ≈ pi/2
        x, w = gausschebyshev(n,4)
        @test dot(x.^2,w) ≈ pi/2
    end

    @testset "x^3" begin
        x, w = gausschebyshev(n,1)
        @test abs(dot(x.^3,w))<1e-15
        x, w = gausschebyshev(n,2)
        @test abs(dot(x.^3,w))<1e-15
        x, w = gausschebyshev(n,3)
        @test dot(x.^3,w) ≈ 3*pi/8
        x, w = gausschebyshev(n,4)
        @test dot(x.^3,w) ≈ -3*pi/8
    end
end


text:
"""
---
WOLA waveform structure
# --- Syntax
- nFFT	  : FFT size [Int]
- nCP	  : CP size [Int]
- allocatedSubcarriers	: Vector of allocated subcarriers [Array{Int}]
- windowTx	: Window used @ transmitter side
- windowRx	: Window used @ recevier side
# ---
# v 1.0 - Robin Gerzaguet.
"""
struct StrucWOLA<: Waveform
	nFFT::Int;
	nCP::Int;
	allocatedSubcarriers::Array{Int};
	windowTx::Window;
	windowRx::Window;
end

"""
---
Create and initiate a WOLA waveform structure
# --- Syntax
 wola = initWOLA(nFFT,nCP,allocatedSubcarriers,winFuncTx,winLengthTx,winFuncRx,winLengthRx;windowTx=[],windowRx=[])
# --- Input parameters
- nFFT	  : FFT size [Int]
- nCP	  : CP size [Int]
- allocatedSubcarriers	: Vector of allocated subcarriers [Array{Int}]
- winFuncTx	  : Name of window used @Tx
- winLengthTx : Size of window @Tx
- winFuncRx	  : Name of window used @Rx
- winLengthRx : Size of window @Rx
- windowTx	  : Coefficient of Tx window. By default it is empty. To force a given window, populate this vector. If let empty, the window will be created based on winLengthTx and winFuncTx
- windowRx	  : Coefficient of Rx window. By default it is empty. To force a given window, populate this vector. If let empty, the window will be created based on winLengthRx and winFuncRx
# --- Output parameters
- wola		  : Waveform structure
# ---
# v 1.0 - Robin Gerzaguet.
"""
function initWOLA(nFFT,nCP,allocatedSubcarriers,winFuncTx,winLengthTx,winFuncRx,winLengthRx;windowTx=[],windowRx=[])
	# ----------------------------------------------------
	# --- Tx window
	# ----------------------------------------------------
	# --- Creating Wola Window
	if isempty(windowTx)
		windowTx	= getWolaWindow(winFuncTx,nFFT,nCP,winLengthTx);
	else
		if !isa(windowTx,Window)
			# --- Proposed window is not Window
			# Convert taps to a Window structure
			windowTx = Window("custom",length(windowTx),windowTx);
		end
	end
	# ----------------------------------------------------
	# --- Rx window
	# ----------------------------------------------------
	# --- Creating Wola Window
	# For receiver, no CP is considered
	if isempty(windowRx)
		windowRx	= getWolaWindow(winFuncRx,nFFT,0,winLengthRx);
	else
		if !isa(windowRx,Window)
			# --- Proposed window is not Window
			# Convert taps to a Window structure
			windowRx = Window("custom",length(windowRx),windowRx);
		end
	end
	# ----------------------------------------------------
	# --- Create objects
	# ----------------------------------------------------
	# --- Create structure
	return StrucWOLA(nFFT,nCP,allocatedSubcarriers,windowTx,windowRx);
end

"""
---
Apply Weighted Overlap and Add Orthogonal Frequency Division Multiplexing (WOLA-OFDM) to the time frequency matrix qamMat and returns a time domain OFDM signal
OFDM is parametrized by its FFT size, its cyclic prefix length (in samples) and the allocSubcarriers vector.
The WOLA part is parametrized by the window applied at each beginning and ending of symbols. The window size can be (and is likely to be) higher than the length of the CP and pure OFDM compatibility is ensured by overlapping the symbols. The interested reader can refer to [1] [2] and [3] for WOLA principle description.
# ---
# Syntax
sigId	= wolaSigGen(qamMat,nFFT,nCP,allocSubcarriers,winFunc,winLength=0)
# ---  Input parameters
- qamMat  : Time frequency matrix : [Array{Complex{Float64},nbSubcarriers,nbSymb}]
- nbSymb			: Number of OFDM symbol tro be transmitted
- nbSubcarriers	: Number of allocated subcarriers (shall be < nFFT)
- nFFT	  : FFT size [Int]
- nCP	  : Cyclic prefix size (in samples) [Int]
- allocSubcarriers : Vector of index of allocated subcarriers [Array{Int,nbSubcarriers}]
- winFunc : Type of window. Can be a string of supported window or directly the window taps.
- winLength : Length of window. This parameter is not used if winFunc is an array of the window tap.
# ---  Output parameters
- sigId	  : WOLA-OFDM signal in time domain [Array{Complex{Float64},nbEch}]. nbEch	: Number of samples: nbSymb*(nFFT+nCP)
# ---
# Supported window
-	"Triangle"		: Triangle window
-	"srrc"			: Square Root Raised Cosine
-	"Meyer"			: Meyer window (See [1])
# ---
# References
- [1] R. Zayani, Y. Medjahdi, H. Shaiek and D. Roviras, "WOLA-OFDM: A Potential Candidate for Asynchronous 5G," 2016.
- [2] Y. Medjahdi and al, "On the road to 5G: Comparative study of Physical layer in MTC context", 2017.
- [3] R. Gerzaguet and al, "Comparison of Promising Candidate Waveforms for 5G: WOLA-OFDM Versus BF-OFDM", 2017.
# ---
# v 1.0 - Robin Gerzaguet.
"""
function wolaSigGen(qamMat,nFFT::Union{Int,Float64},nCP::Union{Int,Float64},allocatedSubcarriers,winFunc,winLength=0)
	# Set 4 core for FFT computation
	#FFTW.set_num_threads(4)
	# ----------------------------------------------------
	# --- Overall parameters
	# ----------------------------------------------------
	# --- Getting parameters
	sizeSymb		= nFFT + nCP;
	# --- Extended size
	sizeExt			= sizeSymb +  winLength;
	# --- WOLA signal parameters
	nbSymb			= size(qamMat,2);
	sigWOLA			= zeros(Complex{Float64},nbSymb*sizeSymb+winLength);
	# ----------------------------------------------------
	# --- Getting window
	# ----------------------------------------------------
	if isa(winFunc,String)
		# --- We create a window with a dedicated function
		# No control on window and parameters, it will be done in the function
		window = getWolaWindow(winFunc,nFFT,nCP,winLength);
	else
		# --- We have a array with the filter taps, no processing here
		window = winFunc;
	end
	# ----------------------------------------------------
	# --- Modulator stage
	# ----------------------------------------------------
	# Set 4 core for FFT computation
	#FFTW.set_num_threads(4)
	# --- Getting parameters
	nbSymb		= size(qamMat,2); 		# --- Applied on x symbols
	sizeSymb	= nFFT + nCP; 			# --- Symbol size
	nSize 		= sizeSymb*nbSymb; 		# --- Total burst length
	# --- Mapping to nFFT elements
	qamCurr		= zeros(Complex{Float64},nFFT,nbSymb);
	qamCurr[allocatedSubcarriers,:] = qamMat;
	# --- Switch to time domain
	sTmp 	= ifft(qamCurr,1);
	# --- Inserting Cyclic prefix and cyclic suffix
	sigCP	  = [sTmp[end+1-nCP-winLength÷2:end,:]; sTmp; sTmp[1:winLength÷2,:] ];
	# --- Apodisation
	sigApod	  = mapslices(x -> x .* window, sigCP,dims=1);
	# --- Overlap and add
	y = zeros(Complex{Float64},sizeSymb*nbSymb + winLength)
	for iB = 1 : 1 : nbSymb
		y[ 1 + (iB-1)*sizeSymb : iB*sizeSymb+ winLength ] = y[ 1 + (iB-1)*sizeSymb : iB*sizeSymb+ winLength ] + sigApod[:,iB];
	end
	return y;
end

# MD Tx
function wolaSigGen(qamMat::Array{Complex{Float64}},wola::StrucWOLA)
	return wolaSigGen(qamMat,wola.nFFT,wola.nCP,wola.allocatedSubcarriers,wola.windowTx.window,wola.windowTx.winLength);
end


text:
## MIT license (c) 2019 by Andrew Lyasoff
##
## note the formula
## $\int_a^bf(t)dt={1\over2}(b-a)\int_{-1}^{+1}f\bigl({1\over2}(b-a)s+{1\over2}(a+b)\bigr)d s$
##(Used by FastGaussQuadrature)
##

begin
    @everywhere using SpecialFunctions
    @everywhere using Interpolations
    @everywhere using Roots
    @everywhere using FastGaussQuadrature

    @everywhere function EUcall(S::Float64,t::Float64,K::Float64,σ::Float64,r::Float64,δ::Float64)
        local v1,v2
        v1=(δ-r+σ^2/2.0)*t+log(K/S)
        v2=(-δ + r + σ^2/2.0)*t - log(K/S)
        v3=σ*sqrt(2.0*t)
        return -K*(exp(-t*r)/2)+K*(exp(-t*r)/2)*erf(v1/v3)+S*(exp(-t*δ)/2)*erf(v2/v3)+S*(exp(-t*δ)/2)
    end

    @everywhere function F(ϵ::Int64,t::Float64,u::Float64,v::Float64,r::Float64,δ::Float64,σ::Float64)
        v1=(r-δ+ϵ*σ^2/2)*(u-t)-log(v)
        v2=σ*sqrt(2*(u-t))
        return 1.0+erf(v1/v2)
    end

    @everywhere function ah(t::Float64,z::Float64,r::Float64,δ::Float64,σ::Float64,f)
        return (exp(-δ*(z-t))*(δ/2)*F(1,t,z,f(z)/f(t),r,δ,σ))
    end

    @everywhere function bh(t::Float64,z::Float64,r::Float64,δ::Float64,σ::Float64,K::Float64,f)
        return (exp(-r*(z-t))*(r*K/2)*F(-1,t,z,f(z)/f(t),r,δ,σ))
    end

    @everywhere function mke_grd(bgn::Float64,step::Float64,size::Int64)
        return bgn:step:(bgn+step+(size-1)*step)
    end

    @everywhere function workerF(loc_range::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},r::Float64,δ::Float64,σ::Float64,T::Float64,K::Float64,nds::Array{Float64,1},wghts::Array{Float64,1},f)
        local loc,LRT,aaa,bbb,an,bn,lnlr
        lnlr=length(loc_range)
        loc=zeros(lnlr)
        for i=lnlr:-1:1
            t=loc_range[i]
            an=[(1/2)*(T-t)*ah(t,(1/2)*(T-t)*s+(1/2)*(t+T),r,δ,σ,f) for s in nds]
            bn=[(1/2)*(T-t)*bh(t,(1/2)*(T-t)*s+(1/2)*(t+T),r,δ,σ,K,f) for s in nds]
            aaa=wghts'*an
            bbb=wghts'*bn;
            LRT=find_zero(x->x-K-EUcall(x,T-t,K,σ,r,δ)-aaa*x+bbb,(K-10,K+20));
            loc[i]=LRT
        end
        return loc
    end

    function cmbn(varg::Array{Array{Float64,1},1})
        if length(varg)>1
            loc=varg[1]
            for i=2:length(varg)
                loc=vcat(loc,varg[i])
            end
            return loc
        else
            return varg
        end
    end

end

function masterF(prnmb::Int64, start_iter::Int64,nmb_of_iter::Int64,conv_tol::Float64,K::Float64,σ::Float64,δ::Float64,r::Float64,T::Float64,Δ::Float64,no_grd_pnts::Int64,vls::Array{Float64,1},nds::Array{Float64,1},wghts::Array{Float64,1})
    local no_iter,conv_check,absc,absc0,val,valPrev,loc,f
    absc=range(0.0,length=pnmb*no_grd_pnts+1,stop=T)
    absc0=range(0.0,length=pnmb*no_grd_pnts,stop=(T-Δ))
    all_grid=[range((T/prnmb)*i,length=no_grd_pnts,stop=((T/prnmb)*(i+1)-Δ)) for i=0:(pnmb-1)]
    val=vls;
    f=CubicSplineInterpolation(absc,val,extrapolation_bc = Interpolations.Line())
    for i=2:(prnmb+1)
        remotecall_fetch(()->K, i)
        remotecall_fetch(()->σ, i)
        remotecall_fetch(()->δ, i)
        remotecall_fetch(()->r, i)
        remotecall_fetch(()->T, i)
        remotecall_fetch(()->all_grid, i)
        remotecall_fetch(()->f, i)
        remotecall_fetch(()->nds, i)
        remotecall_fetch(()->wghts, i)
    end
    no_iter=start_iter;
    conv_check=1000.0
    while no_iter<nmb_of_iter&&conv_check>conv_tol
        valPrev=val;
        no_iter+=1
        lst=[@spawnat i workerF(all_grid[myid()-1],r,δ,σ,T,K,nds,wghts,f) for i=2:(prnmb+1)]
        val0=[fetch(lst[i]) for i=1:prnmb];
        val=cmbn(val0)
        push!(val,max(K,K*(r/δ)))
        f=CubicSplineInterpolation(absc,val,extrapolation_bc = Interpolations.Line())
        for i=2:(prnmb+1)
            remotecall_fetch(()->f, i)
        end
        conv_check=maximum(abs.(valPrev-val))
    end
    return absc,val,conv_check,no_iter
end

max_stars_repo_name: AndrewLyasoff/SMAP

max_issues_repo_name: lhyzh/SMAP

max_forks_repo_name: lhyzh/SMAP