text:
is_constraint(::typeof(active_power_constraints)) = true

# g = [P_ref]
function active_power_constraints(polar::PolarForm, cons, buffer)
    ref_to_gen = polar.indexing.index_ref_to_gen
    # Constraint on P_ref (generator) (P_inj = P_g - P_load)
    # NB: Active power generation has been updated previously inside buffer
    copy!(cons, buffer.pg[ref_to_gen])
    return
end

# Function for AutoDiff
function active_power_constraints(polar::PolarForm, cons, vmag, vang, pinj, qinj)
    kernel! = active_power_kernel!(polar.device)
    ref = polar.indexing.index_ref
    ref_to_gen = polar.indexing.index_ref_to_gen
    ybus_re, ybus_im = get(polar.topology, PS.BusAdmittanceMatrix())

    for i_ in 1:length(ref)
        bus = ref[i_]
        i_gen = ref_to_gen[i_]
        inj = 0.0
        @inbounds for c in ybus_re.colptr[bus]:ybus_re.colptr[bus+1]-1
            to = ybus_re.rowval[c]
            aij = vang[bus] - vang[to]
            # f_re = a * cos + b * sin
            # f_im = a * sin - b * cos
            coef_cos = vmag[bus]*vmag[to]*ybus_re.nzval[c]
            coef_sin = vmag[bus]*vmag[to]*ybus_im.nzval[c]
            cos_val = cos(aij)
            sin_val = sin(aij)
            inj += coef_cos * cos_val + coef_sin * sin_val
        end
        cons[i_] = inj + polar.active_load[bus]
    end
end

function size_constraint(polar::PolarForm{T, IT, VT, MT}, ::typeof(active_power_constraints)) where {T, IT, VT, MT}
    return PS.get(polar.network, PS.NumberOfSlackBuses())
end

function bounds(polar::PolarForm{T, IT, VT, MT}, ::typeof(active_power_constraints)) where {T, IT, VT, MT}
    # Get all bounds (lengths of p_min, p_max, q_min, q_max equal to ngen)
    p_min, p_max = PS.bounds(polar.network, PS.Generators(), PS.ActivePower())
    ref_to_gen = polar.indexing.index_ref_to_gen
    pq_min = p_min[ref_to_gen]
    pq_max = p_max[ref_to_gen]
    return convert(VT, pq_min), convert(VT, pq_max)
end

function _put_active_power_injection!(fr, v_m, v_a, adj_v_m, adj_v_a, adj_P, ybus_re, ybus_im)
    @inbounds for c in ybus_re.colptr[fr]:ybus_re.colptr[fr+1]-1
        to = ybus_re.rowval[c]
        aij = v_a[fr] - v_a[to]
        cosθ = cos(aij)
        sinθ = sin(aij)
        cθ = ybus_re.nzval[c]*cosθ
        sθ = ybus_im.nzval[c]*sinθ
        adj_v_m[fr] += v_m[to] * (cθ + sθ) * adj_P
        adj_v_m[to] += v_m[fr] * (cθ + sθ) * adj_P

        adj_aij = -(v_m[fr]*v_m[to]*(ybus_re.nzval[c]*sinθ))
        adj_aij += v_m[fr]*v_m[to]*(ybus_im.nzval[c]*cosθ)
        adj_aij *= adj_P
        adj_v_a[to] += -adj_aij
        adj_v_a[fr] += adj_aij
    end
end

# Adjoint
function adjoint!(
    polar::PolarForm,
    pbm::AutoDiff.TapeMemory{F, S, I},
    pg, ∂pg,
    vm, ∂vm,
    va, ∂va,
    pinj, ∂pinj,
) where {F<:typeof(active_power_constraints), S, I}
    nbus = PS.get(polar.network, PS.NumberOfBuses())
    nref = PS.get(polar.network, PS.NumberOfSlackBuses())
    index_ref = polar.indexing.index_ref
    ref_to_gen = polar.indexing.index_ref_to_gen

    ybus_re, ybus_im = get(polar.topology, PS.BusAdmittanceMatrix())
    # Constraint on P_ref (generator) (P_inj = P_g - P_load)
    for i in 1:nref
        ibus = index_ref[i]
        igen = ref_to_gen[i]
        _put_active_power_injection!(ibus, vm, va, ∂vm, ∂va, ∂pg[i], ybus_re, ybus_im)
    end
    return
end

# MATPOWER Jacobian
function matpower_jacobian(polar::PolarForm, X::Union{State,Control}, ::typeof(active_power_constraints), V)
    nbus = get(polar, PS.NumberOfBuses())
    pf = polar.network
    ref = pf.ref ; nref = length(ref)
    pv = pf.pv ; npv = length(pv)
    pq = pf.pq ; npq = length(pq)
    gen2bus = polar.indexing.index_generators
    ngen = length(gen2bus)
    Ybus = pf.Ybus

    dSbus_dVm, dSbus_dVa = PS.matpower_residual_jacobian(V, Ybus)
    # w.r.t. state
    if isa(X, State)
        j11 = real(dSbus_dVa[ref, [pv; pq]])
        j12 = real(dSbus_dVm[ref, pq])
    # w.r.t. control
    elseif isa(X, Control)
        j11 = real(dSbus_dVm[ref, [ref; pv]])
        j12 = spzeros(length(ref), npv)
    end
    return [j11 j12]::SparseMatrixCSC{Float64, Int}
end

function matpower_hessian(polar::PolarForm, ::typeof(active_power_constraints), buffer, λ)
    pv = polar.indexing.index_pv
    pq = polar.indexing.index_pq
    ref = polar.indexing.index_ref
    # Check consistency: currently only support a single slack node
    @assert length(λ) == 1
    V = buffer.vmag .* exp.(im .* buffer.vang)
    hxx, hxu, huu = PS.active_power_hessian(V, polar.network.Ybus, pv, pq, ref)

    λₚ = λ[1]
    return FullSpaceHessian(
        λₚ .* hxx,
        λₚ .* hxu,
        λₚ .* huu,
    )
end


text:
using Revise
using NLS_Fit
using DelimitedFiles
using LinearAlgebra

rootDir = joinpath(dirname(pathof(NLS_Fit)), "..")
dataDir = joinpath(rootDir,"data")
dataFile = joinpath(dataDir,"recalibration.txt")

# Getting started example ****************
#

# Data ================
#
XY=readdlm(dataFile);
X=XY[:,1];
Y=XY[:,2];

# Model without recalibration ================
#
model = Gaussian_Peak() + Gaussian_Peak() + Gaussian_Peak()

θ1 = Float64[1.0,5,1]
θ2 = Float64[1.5,10,2]
θ3 = Float64[0.75,20,3]
θ_uncalibrated_model = vcat(θ1,θ2,θ3)

# Add an affine calibration ================
#
recalibration_map = Map_Affine(X[1]=>X[1],X[end]=>X[end])
recalibration_model = Model2Fit_Recalibration(model,recalibration_map)

θ_map = Float64[1,1]
θ_init_recalibration_model = vcat(θ_uncalibrated_model, θ_map)

# Solve a bound constrained nonlinear least squares ================
#

# Objective function ----------------
#
nls = NLS_ForwardDiff_From_Model2Fit(recalibration_model,X,Y);

# Bound constraints ----------------
#
ε = eps(Float64)
#                       h1 μ1   σ1,  h2  μ2   σ2   h3   μ3   σ3  θA   θB
lower_bound = Float64[   0, 5,   ε,   0, 10,   ε,   0,  20,   ε, 0.5, 0.5]
upper_bound = Float64[ Inf, 5, Inf, Inf, 10, Inf, Inf, 20, Inf, 1.5, 1.5]

bc = NLS_Solver.BoundConstraints(lower_bound,upper_bound)

# Solver ----------------
#
conf = NLS_Solver.LevenbergMarquardt_Conf()

# Solve the problem ----------------
#
result = NLS_Solver.solve(nls,θ_init_recalibration_model,conf)

# Fit result ================
#
@assert NLS_Solver.converged(result)

θ_fit_recalibration_model = NLS_Solver.solution(result)
Y_fit_recalibration_model = eval_y(recalibration_model,X,θ_fit_recalibration_model)

# Retrieve calibrated X ----------------
#
X_recalibrated = eval_calibrated_x(recalibration_model,X,θ_fit_recalibration_model);
X_true = (X .- 0.2)/1.1;

norm(X_recalibrated-X,Inf)
norm(X_recalibrated-X_true,Inf)

get_calibrated_model(recalibration_model)
get_calibrated_model_θ(recalibration_model,θ_fit_recalibration_model)

ext: jl


text:
@doc raw"""
    function aic(t :: Trace)

Computes the Akaike Information Criterion for a *single trace* (thus replacing the definition) with
"maximum likelihood" by one with "likelihood". The formula is

```math
\text{AIC}(t)/2 = |\text{params}(t)| - \ell(t),
```

where ``\text{params}(t)|`` is the sum of the dimensionalities of non-observed
and non-deterministic sample nodes.
"""
function aic(t :: Trace)
    ll = 0.0
    k = 0
    for v in values(t)
        if !(v.interpretation == DETERMINISTIC || v.observed)
            length(size(v.value)) == 0 ? k += 1 : k += length(v.value)
        end
        if v.observed
            ll += v.logprob_sum
        end
    end
    2.0 * (k - ll)
end

function numparams(t::Trace)
    k = 0
    for v in values(t)
        if !(v.interpretation == DETERMINISTIC || v.observed)
            length(size(v.value)) == 0 ? k += 1 : k += length(v.value)
        end
    end
    k
end

@doc raw"""
    function aic(r :: SamplingResults{I}) where I <: InferenceType

Computes an empirical estimate of the Akaike Information Criterion from a `SamplingResults`.
The formula is

```math
\text{AIC}(r)/2 = \min_{t \in \text{traces}(r)}|\text{params}(t)| - \hat\ell(t),
```

where ``\text{params}(t)|`` is the sum of the dimensionalities of non-observed
and non-deterministic sample nodes and
``\hat\ell(t)`` is the empirical maximum likelihood.
"""
function aic(r :: SamplingResults{I}) where I <: InferenceType
    min_aic = Inf
    for t in r.traces
        a = aic(t)
        if a < min_aic
            min_aic = a
        end
    end
    min_aic
end

@doc raw"""
    function hpds(r::SamplingResults{I}, pct::Float64) where I <: InferenceType

Computes the highest posterior density set (HPDS) of the `SamplingResults` object.
Let ``\mathcal T`` be the set of traces. The ``100\times Q \%``-percentile HPDS is defined as the
set that satisfies ``\sum_{t \in \mathrm{HPDS}} p(t) = Q`` and, for all ``t \in \mathrm{HPDS}``,
``p(t) > p(s)`` for every ``s \in \mathcal T  - \mathrm{HPDS}``.
It is possible to compute the HPDS using the full joint density ``p(t) \equiv p(x, z)``, where ``x`` is the
set of observed rvs and ``z`` is the set of latent rvs, since ``p(z|x) \propto p(x, z)``.

`pct` should be a float in (0.0, 1.0). E.g., `pct = 0.95` returns the 95% HPDS.
"""
function hpds(r::SamplingResults{I}, pct::Float64) where I <: InferenceType
    (pct > 0.0 && pct < 1.0) || error("pct must be in (0, 1)")
    ix = Int(floor(length(r.traces) * pct))
    for t in r.traces
        if t.logprob_sum == 0.0
            logprob!(t)
        end
    end
    # p(z|x) \propto p(x, z), so we can sort by the joint density, which is easy to
    # calculate, instead of needing to approximate the posterior density
    tr_ix = partialsortperm(r.traces, 1:ix, by=(t,) -> t.logprob_sum, rev=true)
    NonparametricSamplingResults{typeof(r.interpretation)}(
        r.interpretation,
        r.log_weights[tr_ix],
        r.return_values[tr_ix],
        r.traces[tr_ix]
    )
end

@doc raw"""
    function hpdi(r::SamplingResults{I}, pct::Float64, addresses::AbstractArray{T}) where {I <: InferenceType, T}

Computes the highest posterior density interval(s) for a univariate variable. Does *not* check that the
data corresponding to each address in `addresses` is actually univariate; if in doubt, use `hpds` instead.
"""
function hpdi(r::SamplingResults{I}, pct::Float64, addresses::AbstractArray{T}) where {I <: InferenceType, T}
    r = hpds(r, pct)
    Dict(a => (minimum(r[a]), maximum(r[a])) for a in addresses)
end

export aic, hpds, hpdi, numparams


text:
using DataStructures
using ITensors
using QXTns

using OMEinsum

# TensorNetwork struct and public functions
export next_tensor_id!
export TensorNetwork, bonds, simple_contraction, simple_contraction!, neighbours
export decompose_tensor!, replace_with_svd!
export contract_tn!, contract_pair!, replace_tensor_symbol!
export get_hyperedges, disable_hyperindices!, find_connected_indices
export contraction_indices, contract_pair

"""Tensor network data-structure"""
mutable struct TensorNetwork
    tensor_map::OrderedDict{Symbol, QXTensor}
    bond_map::OrderedDict{Index, Vector{Symbol}}
    next_id::Int64
end

# constructors
TensorNetwork() = TensorNetwork(OrderedDict{Symbol, QXTensor}(), OrderedDict{Index, Vector{Symbol}}(), 1)

"""
    TensorNetwork(array::Vector{<: QXTensor})

Outer constructor to create a tensor network object from an array of ITensor objects.
"""
function TensorNetwork(array::Vector{<: QXTensor})
    tensor_map = OrderedDict{Symbol, QXTensor}()
    bond_map = OrderedDict{Index, Vector{Symbol}}()
    next_id = 1
    for t in array
        tensor_id = Symbol("t$(next_id)")
        next_id += 1
        tensor_map[tensor_id] = t
        for bond in inds(t)
            if haskey(bond_map, bond)
                push!(bond_map[bond], tensor_id)
            else
                bond_map[bond] = [tensor_id]
            end
        end
    end
    TensorNetwork(tensor_map, bond_map, next_id)
end

Base.copy(tn::TensorNetwork) = TensorNetwork(deepcopy(tn.tensor_map), deepcopy(tn.bond_map), tn.next_id)
Base.length(tn::TensorNetwork) = length(tn.tensor_map)
Base.values(tn::TensorNetwork) = values(tn.tensor_map)
Base.iterate(tn::TensorNetwork) = iterate(values(tn))
Base.iterate(tn::TensorNetwork, state) = iterate(values(tn), state)
Base.eltype(::TensorNetwork) = QXTensor
Base.keys(tn::TensorNetwork) = keys(tn.tensor_map)
Base.getindex(tn::TensorNetwork, i::Symbol) = tn.tensor_map[i]
Base.haskey(tn::TensorNetwork, i::Symbol) = haskey(tn.tensor_map, i)
Base.getindex(tn::TensorNetwork, i::T) where T <: Index = tn.bond_map[i]
Base.haskey(tn::TensorNetwork, i::T) where T <: Index = haskey(tn.bond_map, i)
Base.show(io::IO, ::MIME"text/plain", tn::TensorNetwork) = print(io, "TensorNetwork(tensors => $(length(tn)), bonds => $(length(bonds(tn))))")

next_tensor_id!(tn::TensorNetwork) = begin tn.next_id += 1; return Symbol("t$(tn.next_id - 1)") end
bonds(tn::TensorNetwork) = keys(tn.bond_map)
disable_hyperindices!(tn) = begin disable_hyperindices!.(tn); nothing; end

"""
    hyperindices(tn::TensorNetwork, i::Symbol; global_hyperindices=true)

Find groups of hyper indices for the given tensor. When global_hyperindices is set to true, then
indices which are identified as hyperindices because of groups of hyperindices in conneted tensors
in the network are also included.
"""
function hyperindices(tn::TensorNetwork, i::Symbol; global_hyperindices=true, all_indices=false)
    if !global_hyperindices
        return hyperindices(tn[i], all_indices=all_indices)
    else
        all_hyperindices = Vector{Vector{Index}}()
        tensor_indices = copy(inds(tn[i]))
        while length(tensor_indices) > 0
            index = tensor_indices[1]
            connected_indices = find_connected_indices(tn, index)
            push!(all_hyperindices, connected_indices)
            setdiff!(tensor_indices, connected_indices)
        end
        return filter(x -> length(x) >= (all_indices ? 1 : 2), all_hyperindices)
    end
end

"""
    tensor_data(tn::TensorNetwork, i::Symbol; consider_hyperindices=true, global_hyperindices=true)

Retrieve the tensor data for the given tensor. If the consider_hyperindices flag is true then
then the data is reshaped to take into account the local hyperindices of the tensor. If the global_hyperindices
index is also true then groups of hyperindices related via hyperindices for other tensors in the network are
also considered.
"""
function tensor_data(tn::TensorNetwork, i::Symbol; consider_hyperindices=true, global_hyperindices=true)
    if consider_hyperindices && global_hyperindices
        global_hi = hyperindices(tn, i; global_hyperindices=true)
        # filter any indices not connected to tensor
        global_hi = map(global_hi) do x
            x = filter(y -> y ∈ inds(tn[i]), x)
        end
        global_hi = filter(x -> length(x) > 1, global_hi)

        local_hi = hyperindices(tn, i; global_hyperindices=false)
        missing_hi = false
        for g in global_hi
            if findfirst(x -> Set(x) == Set(g), local_hi) === nothing
                missing_hi = true
            end
        end
        if missing_hi
            t = store(tn[i])
            if length(local_hi) > 0
                local_hi_ranks = indices2ranks(tn[i], local_hi)
                t = expand_tensor(t, local_hi_ranks)
            end
            global_hi_ranks = indices2ranks(tn[i], global_hi)
            t = reduce_tensor(t, global_hi_ranks)
            return t
        end
    end
    # can delegate to tensor specific tensor_data function
    return tensor_data(tn.tensor_map[i]; consider_hyperindices)
end

"""
    neighbours(tn::TensorNetwork, tensor::Symbol)

Function get the symbols of the neighbouring tensors
"""
function neighbours(tn::TensorNetwork, tensor::Symbol)
    tensor_indices = inds(tn[tensor])
    connected_tensors = unique(vcat([tn[x] for x in tensor_indices]...))
    setdiff(connected_tensors, [tensor])
end

"""
    merge(a::TensorNetwork, b::TensorNetwork)

Join two networks together
"""
function Base.merge(a::TensorNetwork, b::TensorNetwork)
    c = copy(a)
    for b_tensor in b
        push!(c, b_tensor)
    end
    c
end

"""
    push!(tn::TensorNetwork,
          indices::Vector{Index},
          data::Array{T, N}
          tid::Union{Nothing, Symbol}=nothing) where {T, N}

Function to add a tensor to the tensor network.

# Keywords
- `tid::Union{Nothing, Symbol}=nothing`: the id for the new tensor in `tn`. An id is
generated if one is not set.
"""
function Base.push!(tn::TensorNetwork,
                    indices::Vector{<:Index},
                    data::Array{T, N};
                    tid::Union{Nothing, Symbol}=nothing) where {T, N}
    @assert size(data) == Tuple(dim.(indices))
    tensor = QXTensor(indices, data)
    if tid === nothing tid = next_tensor_id!(tn) end
    @assert !(tid in keys(tn))
    tn.tensor_map[tid] = tensor
    for bond in indices
        if haskey(tn.bond_map, bond)
            push!(tn.bond_map[bond], tid)
        else
            tn.bond_map[bond] = [tid]
        end
    end
    tid
end

"""
    push!(tn::TensorNetwork,
          tensor::QXTensor;
          tid::Union{Nothing, Symbol}=nothing)

Function to add a tensor to the tensor network.

# Keywords
- `tid::Union{Nothing, Symbol}=nothing`: the id for the new tensor in `tn`. An id is
generated if one is not set.
"""
function Base.push!(tn::TensorNetwork,
                    tensor::QXTensor;
                    tid::Union{Nothing, Symbol}=nothing)
    if tid === nothing tid = next_tensor_id!(tn) end
    # TODO: It might be a good idea to assert tid doesn't already exist in tn.
    tn.tensor_map[tid] = tensor
    for bond in inds(tensor)
        if haskey(tn.bond_map, bond)
            push!(tn.bond_map[bond], tid)
        else
            tn.bond_map[bond] = [tid]
        end
    end
    tid
end

"""
    simple_contraction(tn::TensorNetwork)

Function to perfrom a simple contraction, contracting all tensors in order.
Only useful for very small networks for testing.
"""
function simple_contraction(tn::TensorNetwork)
    a = QXTensor(1.)
    for t in tn a = contract_tensors(a, t) end
    tensor_data(a; consider_hyperindices=false)
end

"""
    simple_contraction!(tn::TensorNetwork)

Function to perfrom a simple contraction, contracting all tensors in order.
Only useful for very small networks for testing.
"""
function simple_contraction!(tn::TensorNetwork)
    tensor_syms = collect(keys(tn))
    A = tensor_syms[1]
    for B in tensor_syms[2:end]
        A = contract_pair!(tn, A, B)
    end
    store(tn[A])
end

"""
    contract_pair(tn::TensorNetwork, a_sym::Symbol, b_sym::Symbol; mock::Bool=false)

Contract the tensors in 'tn' with ids 'a_sym' and 'b_sym'. If the mock flag is true then the
new tensor will be a mock tensor with the right dimensions but without the actual data.
"""
function contract_pair(tn::TensorNetwork, a_sym::Symbol, b_sym::Symbol; mock::Bool=false)
    r = contraction_indices(tn, a_sym, b_sym)

    # identify positions of local hyper index groups in c
    pos = 1
    c_hyper_indices = map(r.c_indices) do g
        pos += length(g)
        collect(pos-length(g):pos-1)
    end
    c_hyper_indices = convert(Vector{Vector{Int}}, c_hyper_indices)
    filter!(x -> length(x) > 1, c_hyper_indices)

    if mock || (tn[a_sym].storage isa MockTensor || tn[b_sym].storage isa MockTensor)
        # flatten c_indices to get indices
        c = QXTensor(convert(Vector{Index}, vcat(r.c_indices...)), c_hyper_indices)
    else
        c_data = EinCode((Tuple(r.a_labels), Tuple(r.b_labels)), Tuple(r.c_labels))(tensor_data(tn, a_sym), tensor_data(tn, b_sym))
        c = QXTensor(convert(Vector{Index}, vcat(r.c_indices...)), c_hyper_indices, c_data)
    end
    c
end

"""
    contract_pair!(tn::TensorNetwork, a_sym::Symbol, b_sym::Symbol, c_sym::Symbol=:_; mock::Bool=false)

Contract the tensors in 'tn' with ids 'a_sym' and 'b_sym'. If the mock flag is true then the
new tensor will be a mock tensor with the right dimensions but without the actual data.

The resulting tensor is stored in `tn` under the symbol `c_sym` if one is provided, otherwise
a new id is created for it.
"""
function contract_pair!(tn::TensorNetwork, a_sym::Symbol, b_sym::Symbol, c_sym::Symbol=:_;
                        mock::Bool=false)

    c = contract_pair(tn, a_sym, b_sym; mock=mock)

    # Get and contract the tensors a and b to create tensor c
    a = tn.tensor_map[a_sym]
    b = tn.tensor_map[b_sym]
    c_sym == :_ && (c_sym = next_tensor_id!(tn))

    # Remove the contracted indices from the bond map in tn. Also, replace all references
    # in tn to tensors a and b with a reference to tensor c
    common_indices = intersect(inds(a), inds(b))
    for ind in common_indices
        delete!(tn.bond_map, ind)
    end
    for ind in symdiff(inds(a), inds(b))
        tn.bond_map[ind] = replace(tn.bond_map[ind], a_sym => c_sym, b_sym => c_sym)
    end

    # Add tensor c to the tn and remove both a and b
    tn.tensor_map[c_sym] = c
    delete!(tn.tensor_map, a_sym); delete!(tn.tensor_map, b_sym)
    c_sym
end

"""
    contraction_indices(tn::TensorNetwork, a_sym::Symbol, b_sym::Symbol)

Function to work out the contraction indices that would be used to contract the given
tensors. Expected indices in Einstein notation using positive integers
"""
function contraction_indices(tn::TensorNetwork, a_sym::Symbol, b_sym::Symbol)
    a_indices = hyperindices(tn, a_sym; all_indices=true)
    b_indices = hyperindices(tn, b_sym; all_indices=true)
    common_indices = intersect(inds(tn[a_sym]), inds(tn[b_sym]))
    remaining_indices = symdiff(inds(tn[a_sym]), inds(tn[b_sym]))

    # label a_indices in order
    a_labels = collect(1:length(a_indices))
    num_unique = length(a_indices)
    b_labels = map(b_indices) do g
        pos = findfirst(x -> length(intersect(g, x)) > 0, a_indices)
        if pos === nothing
            pos = num_unique += 1
        end
        pos
    end

    # to work out final indices we look at the indices which survive from each tensor
    # find all index groups from a that are not fully contracted
    c_indices = Vector{Vector{Index}}()
    a_labels′ = map(zip(a_labels, a_indices)) do (i, g)
        if length(setdiff(g, common_indices)) == 0
            return nothing
        else
            push!(c_indices, intersect(g, remaining_indices))
            return i
        end
    end

    # find any index groups from b that survive
    # these are groups that are not fully contracted or also present in a
    b_labels′ = map(zip(b_labels, b_indices)) do (i, g)
        if any(map(x -> length(intersect(g, x)) > 0, a_indices))
            return nothing
        else
            push!(c_indices, intersect(g, remaining_indices))
            return i
        end
    end
    c_labels = convert(Vector{Int}, filter(x -> x !== nothing, [a_labels′..., b_labels′...]))
    (a_labels = a_labels, b_labels = b_labels, c_labels = c_labels, c_indices = c_indices)
end

"""
    contract_tn!(tn::TensorNetwork, plan)

Contract the indices of 'tn' according to 'plan'.
"""
function contract_tn!(tn::TensorNetwork, plan::Array{NTuple{3, Symbol}, 1})
    for (A_id, B_id, C_id) in plan
        @assert haskey(tn.tensor_map, A_id)
        @assert haskey(tn.tensor_map, B_id)
        contract_pair!(tn, A_id, B_id, C_id)
    end

    # Contract any disjoint tensors that may remain before returning the result.
    simple_contraction!(tn)
end

"""
    decompose_tensor!(tn::TensorNetwork,
                      tensor_id::Symbol,
                      left_indices::Array{<:Index, 1};
                      contract_S_with::Symbol=:V,
                      kwargs...)

Function to decompose a tensor in a tensor network using svd.

# Keywords
- `contract_S_with::Symbol=:V`: the maxtrix which should absorb the matrix of singular values
- `maxdim::Int`: the maximum number of singular values to keep.
- `mindim::Int`: the minimum number of singular values to keep.
- `cutoff::Float64`: set the desired truncation error of the SVD.
"""
function decompose_tensor!(tn::TensorNetwork,
                           tensor_id::Symbol,
                           left_indices::Array{<:Index, 1};
                           contract_S_with::Symbol=:V,
                           kwargs...)
    # Decompose the tensor into its svd factors.
    U_id, S_id, V_id = replace_with_svd!(tn, tensor_id, left_indices; kwargs)

    # Absorb the singular values tensor into eith U or V.
    if contract_S_with == :V_id
        return U_id, contract_pair!(tn, S_id, V_id)
    else
        return contract_pair!(tn, S_id, U_id), V_id
    end
end

# """
#     replace_with_svd!(tn::TensorNetwork,
#                       tensor_id::Symbol,
#                       left_indices::Array{<:Index, 1};
#                       kwargs...)

# Function to replace a tensor in a tensor network with its svd.

# The indices contained in 'left_indices' are considered the row indices of the tensor when
# the svd is performed.

# # Keywords
# - `maxdim::Int`: the maximum number of singular values to keep.
# - `mindim::Int`: the minimum number of singular values to keep.
# - `cutoff::Float64`: set the desired truncation error of the SVD.
# """
# function replace_with_svd!(tn::TensorNetwork,
#                            tensor_id::Symbol,
#                            left_indices::Array{<:Index, 1};
#                            kwargs...)
#     # Get the tensor and decompose it.
#     tensor = tn.tensor_map[tensor_id]
#     U, S, V = svd(convert(ITensor, tensor), left_indices; use_absolute_cutoff=true, kwargs...)
#     S_data = reshape(collect(Diagonal(store(S))), prod(size(S)))
#     S_QXTensor = QXTensor(S_data, collect(inds(S)), diagonal_check=false)

#     # Remove the original tensor and add its svd factors to the network.
#     delete!(tn, tensor_id)
#     U_id = push!(tn, convert(QXTensor, U))
#     S_id = push!(tn, S_QXTensor)
#     V_id = push!(tn, convert(QXTensor, V))
#     U_id, S_id, V_id
# end

"""
    delete!(tn::TensorNetwork, tensor_id::Symbol)

Function to remove a tensor from a tensor network.
"""
function Base.delete!(tn::TensorNetwork, tensor_id::Symbol)
    tensor = tn.tensor_map[tensor_id]
    for index in inds(tensor)
        # remove tensor_id from tn.bond_map[index]
        filter!(id -> id ≠ tensor_id, tn.bond_map[index])
    end
    delete!(tn.tensor_map, tensor_id)
end

"""
    replace_tensor_symbol!(tn::TensorNetwork, orig_sym::Symbol, new_sym::Symbol)

Replace the given symbol with the given new symbol
"""
function replace_tensor_symbol!(tn::TensorNetwork, orig_sym::Symbol, new_sym::Symbol)
    tensor = tn[orig_sym]
    for i in inds(tensor)
        replace!(tn[i], orig_sym => new_sym)
    end
    tn.tensor_map[new_sym] = tensor
    delete!(tn.tensor_map, orig_sym)
end

"""
    get_hyperedges(tn::TensorNetwork)::Array{Array{Symbol, 1}, 1}

Return an array of hyperedges in the given tensornetwork `tn`.

Hyperedges are represented as arrays of tensor symbols.
"""
function get_hyperedges(tn::TensorNetwork)
    hyperedges = Array{Array{Symbol, 1}, 1}()
    edges = collect(bonds(tn))
    while length(edges) > 0
    # for edge in edges
        edge = edges[1]
        hyper_indices = find_connected_indices(tn, edge)
        push!(hyperedges, union(map(x -> tn[x], hyper_indices)...))
        setdiff!(edges, hyper_indices)
    end
    hyperedges
end

"""
    find_connected_indices(tn::TensorNetwork, bond::Index)

Given a tensor network and an index in the network, find all indices that are related via hyper edge
relations. Involves recurisively checking bonds connected to neighbouring tensors of any newly
related edges found. Returns an array of all edges in the group including the initial edge.
"""
function find_connected_indices(tn::TensorNetwork, bond::Index)
    tensors_to_visit = Set{Symbol}()
    push!.([tensors_to_visit], tn[bond])
    related_edges = Set{Index}([bond])
    while length(tensors_to_visit) > 0
        tensor_sym = pop!(tensors_to_visit)
        for g in hyperindices(tn[tensor_sym])
            if length(intersect(related_edges, g)) > 0
                new_edges = setdiff(g, related_edges)
                for e in new_edges
                    push!(related_edges, e)
                    for t in tn[e]
                        push!(tensors_to_visit, t)
                    end
                end
            end
        end
    end
    collect(related_edges)
end

"""
    contraction_indices((a::QXTensor,b::QXTensor)

Function to work out the contraction indices that would be used to contract the given
tensors. Exptected indices in Einstein notation using positive integers
"""
function contraction_indices(a::QXTensor, b::QXTensor)
    tn = TensorNetwork([a, b])
    contraction_indices(tn, keys(tn)...)
end

"""
    contract_tensors(a::QXTensor, b::QXTensor; mock::Bool=false)

Function to contract two QXTensors and return another QXTensor. If the mock flag
is false or either of the input tensors use MockTensor then the storage for the final
tensor will be of type MockTensor.
"""
function contract_tensors(a::QXTensor, b::QXTensor; mock::Bool=false)
    tn = TensorNetwork([a, b])
    contract_pair(tn, keys(tn)...; mock=mock)
end


text:
using ISA #References: Sandoval, Steven, and Phillip L. De Leon. "The Instantaneous Spectrum: A General Framework for Time-Frequency Analysis." IEEE Transactions on Signal Processing 66.21 (2018): 5679-5693.

#----------------------------
# BASIC EXAMPLES
#----------------------------

#DEFINE 0th CANONICAL TRIPLET
a₀(t) = exp(-t^2)
ω₀(t) = 100
φ₀ = 0
𝐶₀ = Tuple([a₀,ω₀,φ₀])

#DEFINE 0th COMPONENT
ψ₀ = AMFMcomp(𝐶₀)

#DEFINE 1st CANONICAL TRIPLET
a₁(t) = 1
ω₁(t) = 10*t
φ₁ = π
𝐶₁ = Tuple([a₁,ω₁,φ₁])

#DEFINE 2nd CANONICAL TRIPLET
a₂(t) = 0.8*cos(2t)
ω₂(t) = 10 + 7.5*sin(t)
φ₂ = π
𝐶₂ = Tuple([a₂,ω₂,φ₂])

#DEFINE THE COMPONENT SET
𝑆 = [𝐶₀,𝐶₁,𝐶₂]

#DEFINE THE AMFM MODEL
z = AMFMmodel(𝑆)

#DEFINE A TIME INDEX
t = Array(0.0:0.005:2.0)

#EVALUATE THE 0th COMPONENT AT THE POINTS IN THE TIME INDEX
ψ₀(t)

#EVALUATE THE AMFM MODEL AT THE POINTS IN THE TIME INDEX
z(t)

#----------------------------
# FOURIER SERIES
#----------------------------

T = 1
aₖ(k) = 1
kInds = Vector(-10:10)

z₀ = fourierSeries(T,aₖ,kInds)

#----------------------------
# PLOTS
#----------------------------

p1 = isaPlot3d(𝐶₁, t)

p2 = isaPlot3d(ψ₀, t)

p3 = isaPlot3d(z₀, t, FreqUnits="Hz")

p4 = isaPlot3d(𝑆, t, FreqUnits="Hz")

using Plots
Plots.plot(p1,p2,p3,p4)

max_stars_repo_name: kricheso/ISA

max_issues_repo_name: kricheso/ISA

max_forks_repo_name: kricheso/ISA


text:
###########################################################################################################
###########################################################################################################
##################### OVERLAP LIKE G VİA SERRIES REPRESENTATION OF BETA FUNCTOINS #########################
###########################################################################################################
###########################################################################################################
export AuxiliaryG
###########################################################################################################
function AuxiliaryG(n1::arb, q::Int, n2::arb, n3::arb, p1::arb, p2::arb, lim::Int)

    two = parent(n1)(2)
    lowl = lowlq = lowl2q = Zero(lim)
    upl = lim
    uplq = q
    upl2q = lim + 2*q

    coef = ((p1^n1) // (Gamma(n1 + One(n1))))

    m1c1 = (n3) // (n2 + 2*q + One(n2))
    m1c2 = (n2 - n3 + 2*q + One(n2) - p2) // (n2 + 2*q + One(n2))

    m2c1 = (n2) // (n3 + 2*q + One(n3))
    m2c2 = (n3 - n2 + 2*q + One(n3) - p2) // (n3 + 2*q + One(n3))

    m1p = zeros(RF, upl2q - lowl2q + 1)
    m2p = zeros(RF, upl2q - lowl2q + 1)
    f1 = zeros(RF, upl2q - lowl2q + 1)
    f2 = zeros(RF, upl2q - lowl2q + 1)

    m11 = zeros(RF, upl2q - lowl2q + 1)
    m12 = zeros(RF, upl2q - lowl2q + 1)
    m21 = zeros(RF, upl2q - lowl2q + 1)
    m22 = zeros(RF, upl2q - lowl2q + 1)

    m1 = zeros(RF, upl2q - lowl2q + 1)
    m2 = zeros(RF, upl2q - lowl2q + 1)

    m1p[1] = Pochhammer(-n2 - 2*q,0)
    f1[1] = RF(1)
    m2p[1] = Pochhammer(-n3,0)
    f2[1] = RF(1)

    m11[1] = real(AuxiliaryGm(n2 + 2*q + two, n3 - One(n3), p2))
    m12[1] = real(AuxiliaryGm(n2 + 2*q + One(n2), n3, p2))
    m21[1] = real(AuxiliaryGm(n3 + 2*q + two, n2 - One(n2), p2))
    m22[1] = real(AuxiliaryGm(n3 + 2*q + One(n3), n2, p2))

    m1[1] = (1//4) * m1c1 * m11[1] + (1//2) * m1c2 * m12[1]
    m2[1] = (1//4) * m2c1 * m21[1] + (1//2) * m2c2 * m22[1]

    for s in lowl2q : upl2q-1
        s1 = s - lowl2q + 1

        m1p[s1+1] = (-n2 - 2*q + s1 - One(n2)) * m1p[s1]
        m2p[s1+1] = (-n3 + s1 -One(n2)) * m2p[s1]
        f1[s1+1] = f2[s1+1] = (s+1) * f1[s1]

        m1c1 = (n3 + s1) // (n2 + 2*q + One(n2) - s1)
        m1c2 = (n2 - n3 - 2*s1 + 2*q + One(n2)-p2) // (n2 + 2*q + One(n2) - s1)
        m2c1 = (n2 + s1) // (n3 + 2*q + One(n3) - s1)
        m2c2 = (n3 - n2 - 2*s1 + 2*q + One(n3)-p2) // (n3 + 2*q + One(n3) - s1)

        m1[s1+1] = (1//4) * m1c1 * m12[s1] + (1//2) * m1c2 * m1[s1]
        m2[s1+1] = (1//4) * m2c1 * m22[s1] + (1//2) * m2c2 * m2[s1]
        m12[s1+1] = m1[s1]
        m22[s1+1] = m2[s1]
    end
    #################
    #################
    #################
    B = zeros(RF, uplq - lowlq + 1)
    b = zeros(RF, uplq - lowlq + 1)

    P1 = zeros(RF, upl - lowl + 1, uplq - lowlq + 1)
    P2 = zeros(RF, upl - lowl + 1, uplq - lowlq + 1)

    M1 = zeros(RF, upl - lowl + 1, uplq - lowlq + 1)
    M2 = zeros(RF, upl - lowl + 1, uplq - lowlq + 1)

    M1P = zeros(RF, upl - lowl + 1, uplq - lowlq + 1)
    M2P = zeros(RF, upl - lowl + 1, uplq - lowlq + 1)

    for s in lowl : upl
        s1 = s - lowl + 1

        P1[s1,1] = m1p[s1] // (n3 + s + 1)
        P2[s1,1] = m2p[s1] // (n2 + 2*q + s + 1)

        M1[s1,1] = m1[s1] // f1[s1]
        M2[s1,uplq+1] = m2[s1] // f2[s1]
    end

    B[1] = Beta(n2 + 2*q + One(n2), n3 + One(n3))
    b[1] = Binomial(q,lowlq)

    for s1 in lowlq : uplq - 1
        s1p = uplq + lowlq - s1 - 1
        s11 = s1 - lowlq + 1
        s11p = s1p - lowlq + 1

        RecB = ((n3+1+2*s11-1) // (n2+2*q+1-2*s11)) * ((n3+1+2*s11-2) // (n2+2*q+1-2*s11+1))
        B[s11 + 1] = RecB * B[s11]
        b[s11 + 1] = ((q - s1) // (s1 + 1)) * b[s11]

        for s2 in 2*s1 : lim + 2*s1
            #s2p = lim + 2*s1 + 2*s1 - s2
            s22 = s2 - 2*s1 + 1
            #s22p = s2p - 2*s1 + 1

            RecP11 = ((-n2 - 2*q + 2*s11 - One(n2) + s22-1)//(-n2 - 2*q + 2*s11 - One(n2)))
            RecP12 = ((-n2 - 2*q + 2*s11 - two + s22-1)//(-n2 - 2*q + 2*s11 - two))
            RecP13 = (n3 + 2*s11 + s22 - 2) // (n3 + 2*s11 + s22)
            P1[s22, s11+1] = RecP11 * RecP12 * RecP13 * P1[s22, s11]

            RecP21 = ((-n3 - 2*s11)//(-n3 - 2*s11 + s22-1))
            RecP22 = ((-n3 - 2*s11 + One(n3))//(-n3 - 2*s11 + One(n3) + s22-1))
            RecP23 = (n2 + 2*q - 2*s11 + s22 + 2) // (n2 + 2*q - 2*s11 + s22)
            P2[s22, s11 + 1] = RecP21 * RecP22 * RecP23 * P2[s22, s11]

            M1[s22, s11 + 1] = m1[s2 + 3] // f1[s22]
            M2[s22, s11p] = m2[s2 + 3] // f2[s22]
        end
    end
    #################
    #################
    #################
    for s1 in lowlq : uplq
        s11 = s1 - lowlq + 1

        b[s11] = ((-1)^s1) * b[s11]
        B[s11] = b[s11] * B[s11]

        for s2 in lowl : upl
            s22 = s2 - lowl + 1

            M1P[s22, s11] =  b[s11] * P1[s22, s11] * M1[s22, s11]
            M2P[s22, s11] =  b[s11] * P2[s22, s11] * M2[s22, s11]
        end
    end

    res1 = Times(
        coef,
        (two^(n2 + n3 + 2*q + One(n2))),
        sum(B),
        (ExpIntegralE(-(n2 + n3 + 2*q + One(n2)), p2))
    )
    res2 = coef * sum(M1P)
    res3 = coef * sum(M2P)
    (res1 - res2 - res3) // (two^(2*q))
end

function AuxiliaryG(n1::Real, q::Int, n2::Real, n3::Real, p1::Real, p2::Real, lim::Int)
    AuxiliaryG(RF(n1), q, RF(n2), RF(n3), RF(p1), RF(p2), lim)
end


text:
"""
    modeconversion(previous_wavefields::Wavefields, wavefields::Wavefields,
        adjoint_wavefields::Wavefields; params=LMPParams())

Compute the mode conversion matrix `a` from the modes associated with `previous_wavefields`
to modes associated with `wavefields` and its `adjoint_wavefields`.

This is used in the approach known as full mode conversion [Pappert1972b].

# References

[Pappert1972b]: R. A. Pappert and R. R. Smith, “Orthogonality of VLF height gains in the
    earth ionosphere waveguide,” Radio Science, vol. 7, no. 2, pp. 275–278, 1972,
    doi: 10.1029/RS007i002p00275.
"""
function modeconversion(previous_wavefields::Wavefields{H}, wavefields::Wavefields{H},
                        adjoint_wavefields::Wavefields{H}; params=LMPParams()) where H

    @unpack wavefieldheights = params

    product = Vector{ComplexF64}(undef, numheights(wavefields))
    pproduct = similar(product)

    modes = eigenangles(wavefields)
    adjmodes = eigenangles(adjoint_wavefields)
    prevmodes = eigenangles(previous_wavefields)

    modes == adjmodes || @warn "Full mode conversion physics assumes adjoint wavefields are
        calculated for eigenangles of the original waveguide."

    nmodes = length(modes)
    nprevmodes = length(prevmodes)

    # `a` is total conversion
    a = Matrix{ComplexF64}(undef, nprevmodes, nmodes)
    for n in eachindex(modes)  # modes == adjmodes
        for m in eachindex(prevmodes)
            for i in eachindex(wavefieldheights)
                @inbounds f = wavefields[i,n][SVector(2,3,5,6)]  # Ey, Ez, Hy, Hz
                @inbounds fp = previous_wavefields[i,m][SVector(2,3,5,6)]  # Ey, Ez, Hy, Hz
                @inbounds g = SVector{4}(adjoint_wavefields[i,n][6],
                                         -adjoint_wavefields[i,n][5],
                                         -adjoint_wavefields[i,n][3],
                                         adjoint_wavefields[i,n][2])  # Hz, -Hy, -Ez, Ey

                gtranspose = transpose(g)
                product[i] = gtranspose*f
                pproduct[i] = gtranspose*fp
            end
            N, Nerr = romberg(heights(wavefields), product)  # normalization
            I, Ierr = romberg(heights(wavefields), pproduct)
            @inbounds a[m,n] = I/N
        end
    end

    return a
end


text:
using KirchMig

# Setup 2D
nR = nS = 64
dR = dS = 32
oR = oS = 0

t = 1e-2:1e-4:1e-1

dy = dx = dz = 50
x = 0:dx:oR+(nR-1)*dR
z = 0:dz:400
nx = length(x)
nz = length(z)
ny = 10
y = (0:ny-1)*dy

vel = 2700

src_x = oR:dR:oR+(nR-1)*dR
src_z = zeros(nR)
src_y = zeros(nR)
trav = eikonal_const_vel([src_x src_y src_z], x, y, z, vel)

pts = length(ARGS) == 0 ? "serial" : ARGS[1]
L = KirchMap(t, trav; parallel_threaded_serial=pts)
Random.seed!(1234)
u = rand(size(L, 2))
v = rand(size(L, 1))

v_hat = L*u;
tic()
v_hat = L*u;
fwd_t = toq()

u_hat = L'*v;
tic()
u_hat = L'*v;
adj_t = toq()

if pts == "serial"
    i = 1
elseif pts == "threaded"
    i = Threads.nthreads()
else
    i = nworkers()
end
@printf("    %02d   |%9.2f    |%9.2f \n", i, fwd_t, adj_t)


text:
using LoopVectorization

nel  = 100

# Standard loop
res1 = zeros(nel)
for iel=1:nel
    dot = sin(iel/nel*pi) - 0.5
    if dot<0.0
        res1[iel] = 1
    else
        res1[iel] = 2
    end
end

# @avx vectorized
res2 = zeros(nel)
@avx for iel=1:nel
    dot = sin(iel/nel*pi) - 0.5
    res2[iel] = (dot<0.0)*1.0 + (dot>0.0)*2.0
end

println("Sum res1 = ", sum(res1))
println("Sum res2 = ", sum(res2))


text:
using JLD
using Distributions
using Random

N = 1000 #population size
time_step = 20.0 # time step length
N_communities = 1 # how many communites there are in the population
community_probs = [1.0] # probability of being in each community
community_contact_rate = [1.0]
density_dependent_fecundity = 0.0007
average_worm_lifespan = 5.7
max_age = 100
initial_worms = 10

initial_miracidia = time_step*5000000*N/1000
init_env_cercariae = time_step*5000000*N/1000
worm_stages = 1
egg_production_distribution = "NegBin"
female_factor = 1
male_factor = 1
age_contact_rates = [0.1094983,
0.4363393,
0.4291753,
0.02498713]
age_contact_rates = age_contact_rates/sum(age_contact_rates)
c1 = age_contact_rates[1]
c2 = age_contact_rates[2]
c3 = age_contact_rates[3]
c4 = age_contact_rates[4]
ages_for_contacts = [4,9,15,100]
contact_rate_by_age_array = fill(0.00, trunc(Int,(max_age+1)))
mda_adherence = 1
mda_access = 1

last_uptake = 0
egg_multiplier = 100
sd_decrease = 1
female_factor = 1
male_factor = 1
birth_rate = 28*time_step/(1000*365)
human_cercariae_prop = 1
predis_aggregation = 0.24
cercariae_survival = 0.1
miracidia_survival = 0.1

miracidia_maturity = 24
initial_miracidia_days =  trunc(Int,miracidia_maturity/time_step)
death_prob_by_age = [0.0656, 0.0093, 0.003, 0.0023, 0.0027, 0.0038, 0.0044, 0.0048, 0.0053,
                     0.0065, 0.0088, 0.0106, 0.0144, 0.021, 0.0333, 0.0529, 0.0851, 0.1366, 0.2183, 0.2998 , 0.3698, 1]

ages_for_death = [1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60,
                   65, 70, 75, 80, 85, 90, 95, 100, 110];
r = 0.03
kato_katz_par = 0.87

vaccine_effectiveness  = 0.86
drug_effectiveness = 0.86
spec_ages = [7639, 7082, 6524, 5674, 4725, 4147, 3928, 3362,
              2636, 1970, 1468, 1166, 943, 718, 455, 244]
ages_per_index = 5
record_frequency = 1/24;

scenario = "moderate adult";
filename = "eq_runs.jld";
use_kato_katz = 0
# main parameters that we change

predis_aggregation = 0.08489151
max_fecundity = 100
max_fec_contact_rate_product = 1/4
contact_rate = max_fec_contact_rate_product/max_fecundity

M0 = 1000
rate_acquired_immunity = 0
human_larvae_maturity_time = 40
egg_sample_size = 1/100
heavy_burden_threshold = 50

#pars = Parameters(N, time_step, N_communities, community_probs, community_contact_rate,
#        density_dependent_fecundity, average_worm_lifespan,
#        max_age, initial_worms, initial_miracidia, initial_miracidia_days, init_env_cercariae,
#        worm_stages, contact_rate, max_fec_contact_rate_product, max_fecundity, age_contact_rates,
#        ages_for_contacts, contact_rate_by_age_array, mda_adherence, mda_access,  female_factor, male_factor, miracidia_maturity,
#        birth_rate, human_cercariae_prop, predis_aggregation, cercariae_survival, miracidia_survival,
#        death_prob_by_age, ages_for_death, r, vaccine_effectiveness, drug_effectiveness,
#        spec_ages, ages_per_index, record_frequency, use_kato_katz, kato_katz_par, heavy_burden_threshold,
#        rate_acquired_immunity, M0, human_larvae_maturity_time, egg_sample_size, egg_production_distribution)


text:
@testset "PositionEmbedding" begin

    pe = PositionEmbedding(10, 10)

    @test pe(randn(Float32, 3, 5)) == pe.embedding[:, 1:5]
    @test pe(randn(Float32, 3, 5, 7)) == pe.embedding[:, 1:5]
    @test pe(randn(Float32, 3, 15, 2)) == pe.embedding[:, 1:15]
end


text:
module YeoJohnsonTrans

using Optim: optimize, minimizer
using Statistics: mean, var
using StatsBase: geomean

"""
    transform(𝐱)

Transform an array using Yeo-Johnson method.  The power parameter λ is derived
from maximizing a log-likelihood estimator.
"""
function transform(𝐱; optim_args...)
    λ, details = lambda(𝐱; optim_args...)
    #@info "estimated lambda = $λ"
    transform(𝐱, λ)
end

"""
    transform(𝐱, λ)

Transform an array using Yeo-Johnson method with the provided power parameter λ.
"""
function transform(𝐱, λ)
    𝐱′ = similar(𝐱, Float64)
    for (i, x) in enumerate(𝐱)
        if x >= 0
            𝐱′[i] = λ ≈ 0 ? log(x + 1) : ((x + 1)^λ - 1)/λ
        else
            𝐱′[i] = λ ≈ 2 ? -log(-x + 1) : -((-x + 1)^(2 - λ) - 1) / (2 - λ)
        end
    end
    𝐱′
end

"""
    lambda(𝐱; interval = (-2.0, 2.0), optim_args...)

Calculate lambda from an array using a log-likelihood estimator.

Keyword arguments:
- interval: search interval
- optim_args: keyword arguments accepted by Optim.optimize function

See also: [`log_likelihood`](@ref)
"""
function lambda(𝐱; interval = (-2.0, 2.0), optim_args...)
    i1, i2 = interval
    res = optimize(λ -> -log_likelihood(𝐱, λ), i1, i2; optim_args...)
    (value=minimizer(res), details=res)
end

"""
    log_likelihood(𝐱, λ)

Return log-likelihood for the given array and lambda.
"""
function log_likelihood(𝐱, λ)
    N = length(𝐱)
    𝐲 = transform(float.(𝐱), λ)
    σ² = var(𝐲, corrected = false)
    c = sum(sign.(𝐱) .* log.(abs.(𝐱) .+ 1))
    llf = -N / 2.0 * log(σ²) + (λ - 1) * c
    #@info "λ = $λ => σ²=$σ², c=$c, llf=$llf"
    llf
end

end # module


text:
@json struct Branch
    name::String
    merged::Bool
    protected::Bool
    default::Bool
    developers_can_push::Bool
    developers_can_merge::Bool
    can_push::Bool
    web_url::String
    commit::Commit
end

endpoint(::GitLabAPI, ::typeof(get_branch), owner::AStr, repo::AStr, branch::AStr) =
    Endpoint(:GET, "/projects/$(encode(owner, repo))/repository/branches/$branch")
into(::GitLabAPI, ::typeof(get_branch)) = Branch

# Get Repository Branches
# https://docs.gitlab.com/ee/api/branches.html#list-repository-branches
function endpoint(
    ::GitLabAPI,
    ::typeof(get_branches),
    owner::AStr,
    repo::AStr,
)
    return Endpoint(:GET, "/projects/$(encode(owner, repo))/repository/branches")
end
into(::GitLabAPI, ::typeof(get_branches)) = Vector{Branch}

# Delete Branch
# https://docs.gitlab.com/ee/api/branches.html#delete-repository-branch
function endpoint(
    ::GitLabAPI,
    ::typeof(delete_branch),
    owner::AStr,
    repo::AStr,
    branch::AStr,
)
    return Endpoint(
        :DELETE,
        "/projects/$(encode(owner, repo))/repository/branches/$(encode(branch))",
    )
end
postprocessor(::GitLabAPI, ::typeof(delete_branch)) = DoNothing()


text:
using LinearAlgebra

"""
    out = outconvex2d(x, p)

Test if each of the `x` points are inside the convex polygon `p`, and return the
number of inequalities failed by each point. `p` is a 2×n counter-clockwise list
of vertices, with the first vertex duplicated.
"""
function outconvex2d(x, p)
    n = size(p, 2)

    # form A,B such that internal points satisfy Ax <= B
    A = [p[2,2:n] - p[2,1:n-1] p[1,1:n-1] - p[1,2:n]];
    B = zeros(n-1)
    for i = 1:n-1
        B[i] = dot(A[i,:], p[:,i])
    end

    return sum(A*x .> B*ones(1, size(x,2)), dims=1)
end

max_stars_repo_name: Mahmoud-Kharsa/DeltaSigma

max_issues_repo_name: Mahmoud-Kharsa/DeltaSigma

max_forks_repo_name: Mahmoud-Kharsa/DeltaSigma


text:
"""
Reference: The origin code of GPT paper from openai (https://github.com/openai/finetune-transformer-lm)
"""

using Flux
using Flux: onehotbatch, onecold, gradient, logitcrossentropy
import Flux.Optimise: update!

using BytePairEncoding

using Transformers
using Transformers.Basic
using Transformers.GenerativePreTrain
using Transformers.Pretrain

ENV["DATADEPS_ALWAYS_ACCEPT"] = true

include("./0-data.jl")

const startsym = "_start_"
const delisym = "_deli_"
const clfsym = "_clf_"
const unksym = "<unk>"
const labels = ("1", "2")

gpt_model, bpe, vocab, tokenizer = load_pretrain("GPT-OpenAIftlm";
                                                 startsym=startsym,
                                                 delisym=delisym,
                                                 clfsym=clfsym,
                                                 unksym=unksym)

set_tokenizer(tokenizer)

const gpt = gpu(Basic.set_classifier(gpt_model, Chain(Dense(768, 1), Dropout(0.1))))
const embed = gpt.embed.embeddings.tok

function acc(p, y)
    pred = onecold(collect(p))
    label = onecold(collect(y))
    sum(pred .== label) / length(label)
end

function loss(x1, x2, y, x1_mask, x2_mask, c1_index, c2_index)
    e1 = gpt.embed(tok=x1)
    e2 = gpt.embed(tok=x2)
    t1 = gpt.transformers(e1, x1_mask)
    t2 = gpt.transformers(e2, x2_mask)
    lm = lmloss(embed, onehot(vocab, x1), t1, x1_mask) + lmloss(embed, onehot(vocab, x2), t2, x2_mask)

    c1 = gather(t1, c1_index)
    c2 = gather(t2, c2_index)

    p1 = gpt.classifier(c1)::typeof(c1)
    p2 = gpt.classifier(c2)::typeof(c1)
    p = vcat(p1, p2)::typeof(c1)

    cl = logitcrossentropy(p, y)::typeof(lm)
    #unstable type will cause performance issue
    convert(Float32, 0.5) * lm + cl, p
end

const rocs = StoryCloze()
const ps = params(gpt)
const opt = ADAM(6.25e-5)

train!(args["epoch"])


text:
function mt_pop!(r::MRG32k3a, ::Type{T}) where T<:BitInteger
    r.idxI -= sizeof(T)
    i = r.idxI
    @inbounds x128 = r.ints[1 + i >> 4]
    i128 = (i >> logsizeof(T)) & idxmask(T) # 0-based "indice" in x128
    (x128 >> (i128 * (sizeof(T) << 3))) % T
end

rand(r::MRG32k3a, T::Random.SamplerUnion(Int64, UInt64, Int128, UInt128)) =
    mt_pop!(r, T[])

rand(r::MRG32k3a, T::SamplerUnion(Bool, Int8, UInt8, Int16, UInt16, Int32, UInt32)) =
    rand(r, UInt52Raw()) % T[]

##### AbstractArray

function rand!(r::MRG32k3a, A::AbstractArray{Float64},
               I::Random.SamplerTrivial{<:FloatInterval_64})
    region = LinearIndices(A)
    # what follows is equivalent to this simple loop but more efficient:
    # for i=region
    #     @inbounds A[i] = rand(r, I[])
    # end
    m = Base.checked_sub(first(region), 1)
    n = last(region)
    while m < n
        s = mt_avail(r)
        if s == 0
            gen_rand(r)
            s = mt_avail(r)
        end
        m2 = min(n, m+s)
        for i=m+1:m2
            @inbounds A[i] = rand_inbounds(r, I[])
        end
        m = m2
    end
    A
end

##### Array : internal functions

# internal array-like type to circumevent the lack of flexibility with reinterpret
struct UnsafeView{T} <: DenseArray{T,1}
    ptr::Ptr{T}
    len::Int
end

Base.length(a::UnsafeView) = a.len
Base.getindex(a::UnsafeView, i::Int) = unsafe_load(a.ptr, i)
Base.setindex!(a::UnsafeView, x, i::Int) = unsafe_store!(a.ptr, x, i)
Base.pointer(a::UnsafeView) = a.ptr
Base.size(a::UnsafeView) = (a.len,)
Base.elsize(::UnsafeView{T}) where {T} = sizeof(T)

# this is essentially equivalent to rand!(r, ::AbstractArray{Float64}, I) above, but due to
# optimizations which can't be done currently when working with pointers, we have to re-order
# manually the computation flow to get the performance
# (see https://discourse.julialang.org/t/unsafe-store-sometimes-slower-than-arrays-setindex)
function _rand_max383!(r::MRG32k3a, A::UnsafeView{Float64}, I::FloatInterval_64)
    n = length(A)
    @assert n <= dsfmt_get_min_array_size()+1 # == 383
    mt_avail(r) == 0 && gen_rand(r)
    # from now on, at most one call to gen_rand(r) will be necessary
    m = min(n, mt_avail(r))
    GC.@preserve r unsafe_copyto!(A.ptr, pointer(r.vals, r.idxF+1), m)
    if m == n
        r.idxF += m
    else # m < n
        gen_rand(r)
        GC.@preserve r unsafe_copyto!(A.ptr+m*sizeof(Float64), pointer(r.vals), n-m)
        r.idxF = n-m
    end
    if I isa CloseOpen01
        for i=1:n
            A[i] -= 1.0
        end
    end
    A
end

fill_array!(s::DSFMT_state, A::Ptr{Float64}, n::Int, ::CloseOpen01_64) =
    dsfmt_fill_array_close_open!(s, A, n)

fill_array!(s::DSFMT_state, A::Ptr{Float64}, n::Int, ::Random.CloseOpen12_64) =
    dsfmt_fill_array_close1_open2!(s, A, n)

function rand!(r::MRG32k3a, A::UnsafeView{Float64},
               I::Random.SamplerTrivial{<:FloatInterval_64})
    # depending on the alignment of A, the data written by fill_array! may have
    # to be left-shifted by up to 15 bytes (cf. unsafe_copyto! below) for
    # reproducibility purposes;
    # so, even for well aligned arrays, fill_array! is used to generate only
    # the n-2 first values (or n-3 if n is odd), and the remaining values are
    # generated by the scalar version of rand
    n = length(A)
    n2 = (n-2) ÷ 2 * 2
    n2 < dsfmt_get_min_array_size() && return _rand_max383!(r, A, I[])

    pA = A.ptr
    align = Csize_t(pA) % 16
    if align > 0
        pA2 = pA + 16 - align
        fill_array!(r.state, pA2, n2, I[]) # generate the data in-place, but shifted
        unsafe_copyto!(pA, pA2, n2) # move the data to the beginning of the array
    else
        fill_array!(r.state, pA, n2, I[])
    end
    for i=n2+1:n
        A[i] = rand(r, I[])
    end
    A
end

# fills up A reinterpreted as an array of Float64 with n64 values
function _rand!(r::MRG32k3a, A::Array{T}, n64::Int, I::FloatInterval_64) where T
    # n64 is the length in terms of `Float64` of the target
    @assert sizeof(Float64)*n64 <= sizeof(T)*length(A) && isbitstype(T)
    GC.@preserve A rand!(r, UnsafeView{Float64}(pointer(A), n64), Random.SamplerTrivial(I))
    A
end

##### Array: Float64, Float16, Float32

rand!(r::MRG32k3a, A::Array{Float64}, I::Random.SamplerTrivial{<:FloatInterval_64}) =
    _rand!(r, A, length(A), I[])

mask128(u::UInt128, ::Type{Float16}) =
    (u & 0x03ff03ff03ff03ff03ff03ff03ff03ff) | 0x3c003c003c003c003c003c003c003c00

mask128(u::UInt128, ::Type{Float32}) =
    (u & 0x007fffff007fffff007fffff007fffff) | 0x3f8000003f8000003f8000003f800000

for T in (Float16, Float32)
    @eval function rand!(r::MRG32k3a, A::Array{$T}, ::Random.SamplerTrivial{CloseOpen12{$T}})
        n = length(A)
        n128 = n * sizeof($T) ÷ 16
        _rand!(r, A, 2*n128, CloseOpen12())
        GC.@preserve A begin
            A128 = UnsafeView{UInt128}(pointer(A), n128)
            for i in 1:n128
                u = A128[i]
                u ⊻= u << 26
                # at this point, the 64 low bits of u, "k" being the k-th bit of A128[i] and "+"
                # the bit xor, are:
                # [..., 58+32,..., 53+27, 52+26, ..., 33+7, 32+6, ..., 27+1, 26, ..., 1]
                # the bits needing to be random are
                # [1:10, 17:26, 33:42, 49:58] (for Float16)
                # [1:23, 33:55] (for Float32)
                # this is obviously satisfied on the 32 low bits side, and on the high side,
                # the entropy comes from bits 33:52 of A128[i] and then from bits 27:32
                # (which are discarded on the low side)
                # this is similar for the 64 high bits of u
                A128[i] = mask128(u, $T)
            end
        end
        for i in 16*n128÷sizeof($T)+1:n
            @inbounds A[i] = rand(r, $T) + one($T)
        end
        A
    end

    @eval function rand!(r::MRG32k3a, A::Array{$T}, ::Random.SamplerTrivial{CloseOpen01{$T}})
        rand!(r, A, CloseOpen12($T))
        I32 = one(Float32)
        for i in eachindex(A)
            @inbounds A[i] = Float32(A[i])-I32 # faster than "A[i] -= one(T)" for T==Float16
        end
        A
    end
end

#### arrays of integers

function rand!(r::MRG32k3a, A::UnsafeView{UInt128}, ::SamplerType{UInt128})
    n::Int=length(A)
    i = n
    while true
        rand!(r, UnsafeView{Float64}(A.ptr, 2i), CloseOpen12())
        n < 5 && break
        i = 0
        while n-i >= 5
            u = A[i+=1]
            A[n]    ⊻= u << 48
            A[n-=1] ⊻= u << 36
            A[n-=1] ⊻= u << 24
            A[n-=1] ⊻= u << 12
            n-=1
        end
    end
    if n > 0
        u = rand(r, UInt2x52Raw())
        for i = 1:n
            A[i] ⊻= u << (12*i)
        end
    end
    A
end

for T in BitInteger_types
    @eval rand!(r::MRG32k3a, A::Array{$T}, sp::SamplerType{$T}) =
        (GC.@preserve A rand!(r, UnsafeView(pointer(A), length(A)), sp); A)

    T == UInt128 && continue

    @eval function rand!(r::MRG32k3a, A::UnsafeView{$T}, ::SamplerType{$T})
        n = length(A)
        n128 = n * sizeof($T) ÷ 16
        rand!(r, UnsafeView{UInt128}(pointer(A), n128))
        for i = 16*n128÷sizeof($T)+1:n
            @inbounds A[i] = rand(r, $T)
        end
        A
    end
end

#### from a range

for T in BitInteger_types, R=(1, Inf) # eval because of ambiguity otherwise
    @eval Sampler(::Type{MRG32k3a}, r::AbstractUnitRange{$T}, ::Val{$R}) =
        SamplerRangeFast(r)
end


text:
using Statistics,Serialization,StatsPlots

path_to_data = "../../data/working_paper_2/baseline_GER/"
filename_prefix = "../../figures/working_paper_2/"

mkpath(filename_prefix)

include("$(path_to_data)covid_par_ini.jl")
worker_results = deserialize(open("$(path_to_data)batchdata.dat"))

datapoint = fld(T,datat)

include("../../emp_traj.jl")

fac = nhh/50000
nzer= (Int(ceil(virustime/datat))-1)*datat+2*corlatent
emptotinftraj = vcat(zeros(nzer),fac*emptotinf)
nne = size(emptotinf)[1]

# Total number of infected: model vs. empirical data for Germany
emptotinftraj = vcat(zeros(nzer),fac*emptotinfnew)
nnen = size(emptotinfnew)[1]
nnen = 180
totinfmeansc = mean(hcat(detfrac*(results -> results[:totinftraj]).(worker_results)...), dims=2)
totinfstdsc = std(hcat(detfrac*(results -> results[:totinftraj]).(worker_results)...), dims=2)
pl5b = plot(collect(1:nnen),[totinfmeansc[nzer+1:nzer+nnen], totinfmeansc[nzer+1:nzer+nnen].-totinfstdsc[nzer+1:nzer+nnen],totinfmeansc[nzer+1:nzer+nnen].+totinfstdsc[nzer+1:nzer+nnen]], linestyle= [:solid :dot :dot], linewidth = [2 1 1], linecolor = [:blue :black :black],label = false, xlabel="days", ylabel = "reported tot. inf.")
plot!(pl5b, detfrac*emptotinftraj[nzer+1:nzer+nnen], linecolor = [:green], label=false, linewidth = [2])
savefig(pl5b,"$(filename_prefix)baseline_total_number_infected.pdf")

# R-value: model vs. model vs. empirical data for Germany
RKIR0mean = mean(hcat((results -> results[:RKIR0traj]).(worker_results)...), dims=2)
RKIR0std = std(hcat((results -> results[:RKIR0traj]).(worker_results)...), dims=2)
nn1 = size(RKIR0mean)[1]

nzer=(Int(ceil(virustime/datat))-1)*datat+2*corlatent
empR0smtraj = vcat(zeros(nzer),empR0new)
RKIR0smmean = mean(hcat((results -> results[:RKIR0smtraj]).(worker_results)...), dims=2)
RKIR0smstd = std(hcat((results -> results[:RKIR0smtraj]).(worker_results)...), dims=2)
nnhc = nn1 - size(RKIR0smmean)[1]
nne = size(empR0new)[1]
nne = 180
RKIR0smmean = vcat(zeros(nnhc),RKIR0smmean)
RKIR0smstd = vcat(zeros(nnhc),RKIR0smstd)
for i=1:7
    RKIR0smmean[nzer+i] = 1/0
    empR0smtraj[nzer+i] = 1/0
end

for i in 1:length(RKIR0smmean)
    if RKIR0smmean[i] == Inf
        RKIR0smmean[i] = RKIR0smmean[i-1]
        RKIR0smstd[i] = 0.0
    end
end
for i=1:7
    RKIR0smmean[nzer+i] = 1/0
    empR0smtraj[nzer+i] = 1/0
end

RKIR0s = hcat((results -> results[:RKIR0smtraj]).(worker_results)...)
x = 0
for i in 1:length(RKIR0s)
    if RKIR0s[i] == Inf
        RKIR0s[i] = NaN
        global x = i
    end
end

nanmean1(x) = mean(filter(!isnan,x))
nanstd1(x) = std(filter(!isnan,x))

RKIR0smmean= []
RKIR0smstd= []
for i in 1:length(RKIR0s[:,1])
    RKIR0s_row = RKIR0s[i,:]
    push!(RKIR0smmean,nanmean1(RKIR0s_row))
    push!(RKIR0smstd,nanstd1(RKIR0s_row))
end
nnhc = nn1 - size(RKIR0smmean)[1]
RKIR0smmean = vcat(zeros(nnhc),RKIR0smmean)
RKIR0smstd = vcat(zeros(nnhc),RKIR0smstd)
pl16b = plot(collect(1:nne),[RKIR0smmean[nzer+1:nzer+nne], RKIR0smmean[nzer+1:nzer+nne].-RKIR0smstd[nzer+1:nzer+nne],RKIR0smmean[nzer+1:nzer+nne].+RKIR0smstd[nzer+1:nzer+nne]], linestyle= [:solid :dot :dot], linewidth = [2 1 1], linecolor = [:blue :black :black],label = false, xlabel = "days", ylabel = "R0")
plot!(pl16b, empR0smtraj[nzer+1:nzer+nne], linecolor = [:green], label=false, linewidth = [2])
savefig(pl16b,"$(filename_prefix)baseline_R-value.pdf")

# Casualites: model vs. empirical data for Germany
totcasmean = mean(cat(100*(results -> results[:totcastraj]).(worker_results)...,dims=3), dims=3)
totcasstd = std(cat(100*(results -> results[:totcastraj]).(worker_results)...,dims=3), dims=3)

nzer= (Int(ceil(virustime/datat+2*corlatent/datat))-1)
empcastraj = vcat(zeros(datat*nzer),fac*empcasnew)
nne = Int(ceil(size(empcasnew)[1]/datat))
nne = Int(ceil(180 / datat))
pl7b = plot(datat*collect(1:nne),[totcasmean[nzer+1:nzer+nne,3], totcasmean[nzer+1:nzer+nne,3].-totcasstd[nzer+1:nzer+nne,3],totcasmean[nzer+1:nzer+nne,3].+totcasstd[nzer+1:nzer+nne,3]], linestyle = [:solid :dot :dot], linewidth = [2 1 1], linecolor = [:blue :black :black] )
plot!(pl7b, empcastraj[datat*nzer+1:180+datat*nzer], linecolor = [:green], linewidth = [2], legend = false, xlabel = "days", ylabel = "casualties [%]")
savefig(pl7b,"$(filename_prefix)baseline_casualties.pdf")

# Workers on short time: model vs. empirical data for Germany
shorttimemean = mean(hcat((results -> results[:shorttimetraj]).(worker_results)...), dims=2)
shorttimestd = std(hcat((results -> results[:shorttimetraj]).(worker_results)...), dims=2)

nzerw= (Int(ceil(virustime/datat+2*corlatent/datat))-1)
nnew = Int(ceil(size(empcasnew)[1]/datat))
nnew = Int(ceil(180/datat))

pl26b = plot(datat*collect(1:nnew),[shorttimemean[nzerw+1:nzerw+nnew], shorttimemean[nzerw+1:nzerw+nnew].-shorttimestd[nzerw+1:nzerw+nnew],shorttimemean[nzerw+1:nzerw+nnew].+shorttimestd[nzerw+1:nzerw+nnew]], linestyle= [:solid :dot :dot], linewidth = [2 1 1], linecolor = [:blue :black :black], label=false, xlabel = "days", ylabel = "short time [%]")
# tatsächlich ausgezahlte kurzarbeit from april to july
scatter!(pl26b, [23], [0.078], markersize = 5, markercolor = [:green], label = false)
scatter!(pl26b, [53], [0.180], markersize = 5, markercolor = [:green], label = false)
scatter!(pl26b, [83], [0.179], markersize = 5, markercolor = [:green], label = false)
scatter!(pl26b, [114], [0.140], markersize = 5, markercolor = [:green], label = false)
scatter!(pl26b, [145], [0.128], markersize = 5, markercolor = [:green], label = false)
savefig(pl26b,"$(filename_prefix)baseline_shorttime.pdf")


text:
#=
   -x, abs(x), inv(x),
   sqrt(x), invsqrt(x)
   x+y, x-y, x*y, x/y, hypot(x,y)
   muladd(a,b,c), fma(a,b,c)
=#

function signbit(x::ArbFloat{P}) where {P}
    0 != ccall(@libarb(arb_is_negative), Int, (Ref{ArbFloat{P}},), x)
end

function abs(x::ArbFloat{P}) where {P}
    z  = initializer(ArbFloat{P})
    lo = initializer(ArfFloat{P})
    hi = initializer(ArfFloat{P})
    ccall(@libarb(arb_get_abs_lbound_arf), Cvoid, (Ref{ArfFloat{P}}, Ref{ArbFloat{P}}, Int), lo, x, P)
    ccall(@libarb(arb_get_abs_ubound_arf), Cvoid, (Ref{ArfFloat{P}}, Ref{ArbFloat{P}}, Int), hi, x, P)
    ccall(@libarb(arb_set_interval_arf), Cvoid, (Ref{ArbFloat{P}}, Ref{ArfFloat{P}}, Ref{ArfFloat{P}}, Int), z, lo, hi, P)
    return z
end

function abs2(x::T) where {T <: ArbFloat}
    a = abs(x)
    return a*a
end

for (op,cfunc) in ((:-,:arb_neg), (:sign, :arb_sgn), (:absz, :arb_abs))
  @eval begin
    function ($op)(x::ArbFloat{P}) where {P}
      z = initializer(ArbFloat{P})
      ccall(@libarb($cfunc), Cvoid, (Ref{ArbFloat{P}}, Ref{ArbFloat{P}}), z, x)
      return z
    end
  end
end

for (op,cfunc) in ((:inv, :arb_inv), (:sqrt, :arb_sqrt), (:invsqrt, :arb_rsqrt))
  @eval begin
    function ($op)(x::ArbFloat{P}) where {P}
      z = initializer(ArbFloat{P})
      ccall(@libarb($cfunc), Cvoid, (Ref{ArbFloat{P}}, Ref{ArbFloat{P}}, Int), z, x, P)
      return z
    end
  end
end

for (op,cfunc) in ((:+,:arb_add), (:-, :arb_sub), (:/, :arb_div), (:hypot, :arb_hypot))
  @eval begin
    function ($op)(x::ArbFloat{P}, y::ArbFloat{P}) where {P}
      z = initializer(ArbFloat{P})
      ccall(@libarb($cfunc), Cvoid, (Ref{ArbFloat{P}}, Ref{ArbFloat{P}}, Ref{ArbFloat{P}}, Int), z, x, y, P)
      return z
    end
    ($op)(x::ArbFloat{P}, y::ArbFloat{Q}) where {P,Q} = ($op)(promote(x,y)...)
    ($op)(x::T, y::F) where {T <: ArbFloat,F <: AbstractFloat} = ($op)(x, convert(T, y))
    ($op)(x::F, y::T) where {T <: ArbFloat,F <: AbstractFloat} = ($op)(convert(T, x), y)
    ($op)(x::T, y::R) where {T <: ArbFloat,R <: Rational} = ($op)(x, convert(T, y))
    ($op)(x::R, y::T) where {T <: ArbFloat,R <: Rational} = ($op)(convert(T, x), y)
    ($op)(x::T, y::R) where {T <: ArbFloat,R <: Real} = ($op)(x, convert(T, y))
    ($op)(x::R, y::T) where {T <: ArbFloat,R <: Real} = ($op)(convert(T, x), y)
  end
end

function (*)(x::ArbFloat{P}, y::ArbFloat{P}) where {P}
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_mul), Cvoid, (Ref{ArbFloat{P}}, Ref{ArbFloat{P}}, Ref{ArbFloat{P}}, Int), z, x, y, P)
    return z
end
(*)(x::T, y::F) where {T <: ArbFloat,F <: AbstractFloat} = (*)(x, convert(T, y))
(*)(x::F, y::T) where {T <: ArbFloat,F <: AbstractFloat} = (*)(convert(T, x), y)
(*)(x::T, y::R) where {T <: ArbFloat,R <: Rational} = (*)(x, convert(T, y))
(*)(x::R, y::T) where {T <: ArbFloat,R <: Rational} = (*)(convert(T, x), y)
(*)(x::T, y::R) where {T <: ArbFloat,R <: Real} = (*)(x, convert(T, y))
(*)(x::R, y::T) where {T <: ArbFloat,R <: Real} = (*)(convert(T, x), y)

(+)(x::T, y::T, z::T) where {T <: ArbFloat} = (x + y) + z

(+)(x::I, y::T, z::T) where {T <: ArbFloat,I <: Integer} = x + (y + z)
(+)(x::T, y::T, z::I) where {T <: ArbFloat,I <: Integer} = (x + y) + z
(+)(x::T, y::T, z::R) where {T <: ArbFloat,R <: Real} = (x + y) + z
(+)(x::R, y::T, z::T) where {T <: ArbFloat,R <: Real} =  x + (y + z)

(*)(x::T, y::T, z::T) where {T <: ArbFloat} = (x * y) * z

(*)(x::I, y::T, z::T) where {T <: ArbFloat,I <: Integer} = x * (y * z)
(*)(x::T, y::T, z::I) where {T <: ArbFloat,I <: Integer} = (x * y) * z
(*)(x::T, y::T, z::R) where {T <: ArbFloat,R <: Real} = (x * y) * z
(*)(x::R, y::T, z::T) where {T <: ArbFloat,R <: Real} =  x * (y * z)

(-)(x::T, y::T, z::T) where {T <: ArbFloat} = (x - y) - z
(/)(x::T, y::T, z::T) where {T <: ArbFloat} = (x / y) / z

for (op,cfunc) in ((:+,:arb_add_si), (:-, :arb_sub_si), (:*, :arb_mul_si), (:/, :arb_div_si))
  @eval begin
    function ($op)(x::ArbFloat{P}, y::Int) where {P}
      z = initializer(ArbFloat{P})
      ccall(@libarb($cfunc), Cvoid, (Ref{ArbFloat{P}}, Ref{ArbFloat{P}}, Int, Int), z, x, y, P)
      return z
    end
  end
end

(+)(x::Int, y::T) where {T <: ArbFloat} = (+)(y,x)
(-)(x::Int, y::T) where {T <: ArbFloat} = -((-)(y,x))
(*)(x::Int, y::T) where {T <: ArbFloat} = (*)(y,x)
(/)(x::Int, y::T) where {T <: ArbFloat} = (/)(T(x),y)

(+)(x::T, y::Integer) where {T <: ArbFloat} = (+)(x, convert(T, y))
(-)(x::T, y::Integer) where {T <: ArbFloat} = (-)(x, convert(T, y))
(*)(x::T, y::Integer) where {T <: ArbFloat} = (*)(x, convert(T, y))
(/)(x::T, y::Integer) where {T <: ArbFloat} = (/)(x, convert(T, y))

(+)(x::Integer, y::T) where {T <: ArbFloat} = (+)(convert(T, x), y)
(-)(x::Integer, y::T) where {T <: ArbFloat} = -((-)(y, convert(T, x)))
(*)(x::Integer, y::T) where {T <: ArbFloat} = (*)(convert(T,x), y)
(/)(x::Integer, y::T) where {T <: ArbFloat} = (/)(convert(T,x), y)

for (op,cfunc) in ((:addmul,:arb_addmul), (:submul, :arb_submul))
  @eval begin
    function ($op)(w::ArbFloat{P}, x::ArbFloat{P}, y::ArbFloat{P}) where {P}
      z = initializer(ArbFloat{P})
      ccall(@libarb($cfunc), Cvoid, (Ref{ArbFloat{P}}, Ref{ArbFloat{P}}, Ref{ArbFloat{P}}, Ref{ArbFloat{P}}, Int), z, w, x, y, P)
      z
    end
  end
end

muladd(a::T, b::T, c::T) where {T <: ArbFloat} = addmul(c,a,b)
mulsub(a::T, b::T, c::T) where {T <: ArbFloat} = addmul(-c,a,b)


text:
## @lc code=start
using LeetCode

using DataStructures
function repeat_limited_string(s::String, repeat_limit::Int)::String
    table = counter(s)
    chars = sort!([i for i in keys(table)]; rev=true)
    res, i, n, char_len = Char[], 1, length(chars), table[chars[1]]
    while true
        if char_len <= repeat_limit
            append!(res, fill(chars[i], char_len))
            table[chars[i]] = 0
            i == n && return join(res)
            char_len = table[chars[i += 1]]
        else
            append!(res, fill(chars[i], repeat_limit))
            char_len -= repeat_limit
            j = findfirst(x -> table[x] != 0, @view chars[i + 1:end])
            j === nothing && return join(res)
            push!(res, chars[i + j])
            table[chars[i + j]] -= 1
        end
    end
end

## @lc code=end
## @lc test=start
@testset "2182.construct-string-with-repeat-limit.jl" begin
    @test repeat_limited_string("cczazcc", 3) == "zzcccac"
    @test repeat_limited_string("aababab", 2) == "bbabaa"
end
## @lc test=end

## @lc info=start
# link: [solution to 2182]()
## @lc info=end


text:
function wrap!(w_api, fdef::FDefinition)
    if is_command_type(fdef.name, ENUMERATE)
        new_fdef = wrap_enumeration_command(typed_fdef(fdef))
    elseif startswith(fdef.name, "vkDestroy")
        return
    elseif !is_command_type(fdef.name, CREATE)
        # name = name_transform(fdef)
        # body = statements(patterns(fdef))
        # new_fdef = FDefinition(name, Signature(name, args, kwargs), fdef.short, body)
        new_fdef = wrap_generic(typed_fdef(fdef))
    else
        return
    end
    w_api.funcs[new_fdef.name] = new_fdef
end

function typed_fdef(fdef::FDefinition)
    @assert startswith(fdef.body[1].body, "ccall")
    called_fun, ret_type, types, names = deconstruct_ccall(fdef.body[1].body)
    FDefinition(called_fun, Signature(called_fun, PositionalArgument.(names, types), KeywordArgument[]), fdef.short, fdef.body, "", ret_type)
end

function arg_transform(arg::PositionalArgument)
    PositionalArgument(argname_transform(arg), argtype_transform(arg))
end

argname_transform(arg::PositionalArgument) = fieldname_transform(arg.name, arg.type)
argtype_transform(arg::PositionalArgument) = fieldtype_transform(arg.name, arg.type)

function name_transform(str, ::Type{FDefinition})
    nc_convert(SnakeCaseLower, str[3:end])
end

pass_new_nametype(::Type{FDefinition}) = (x, y, z) -> (arg = arg_transform(PositionalArgument(x, y)) ; (arg.name, arg.type))

"""
Get initialization arguments from the results of `accumulate_passes`.
"""
function init_args(pass_results, body, fdef::FDefinition; use_all_args=true, name_depth=2)
    args = String[]
    for (name, pr_list) ∈ pass_results
        pass_args = last(pr_list).pass_args
        @unpack tmp_name, new_name, type, new_type, sname = pass_args
        # if is_ptr(type) && inner_type(type) == "Cstring"
        #     push!(args, new_name * "_ptrarray")
        if (!any(values(pass_args.passes)) || use_all_args)
            last_name = last_argname(body, name_hierarchy(name_depth, name, tmp_name, new_name)...)
            push!(args, last_name)
        end
    end
    args
end

function wrap_enumeration_command(fdef)
    sig = fdef.signature
    fname = fdef.name
    !has_count_to_be_filled(fname) && return wrap_generic(fdef)
    arr_row = enumeration_command_array_variable(fname)
    count_row = enumeration_command_count_variable(fname)
    count_new_name = fieldname_transform(count_row.name, count_row.type)
    arr_new_name = fieldname_transform(arr_row.name, arr_row.type)
    enumerated_type = inner_type(arr_row.type)
    args = arguments(sig)
    kwargs = keyword_arguments(sig)
    new_fname = name_transform(fdef)
    new_sig = Signature(new_fname, args, kwargs)
    # body, pass_results = accumulate_passes(fname, sig.args, pass_new_nametype(FDefinition), [ComputeLengthArgument(), InitializePointers()])
    body = Statement[]
    command_args = []
    call_args = map(x -> fieldname_transform(x.name, x.type), sig.args)
    first_call_args = replace(join_args(call_args), arr_new_name => "C_NULL")
    second_call_args = join_args(call_args)
    body = [
        body...,
        Statement("$(count_new_name) = Ref{UInt32}(0)", count_new_name),
        Statement("$(checked_function_call(fdef))($first_call_args)"),
        Statement("$(arr_new_name) = Array{$(enumerated_type)}(undef, $(count_new_name)[])", arr_new_name),
        Statement("$(checked_function_call(fdef))($second_call_args)"),
    ]
    enumerated_type_new = name_transform(enumerated_type, SDefinition)
    convert_statement = Statement(enumerated_type == "void" ? arr_new_name : "$enumerated_type_new.($(arr_new_name))")
    push!(body, convert_statement)
    FDefinition(new_fname, new_sig, false, body)
end

deconstruct_ccall(fdef::FDefinition) = deconstruct_ccall(fdef.body[1].body)
function deconstruct_ccall(expr)
    matches = match(r"ccall\(\(:(.*?),.*?\), (\w+), \((.*?)\), (.*?)\)", expr)
    called_fun, ret_type, types_unsplit, names_unsplit = matches.captures
    called_fun, ret_type, split(types_unsplit, ", "), split(names_unsplit, ", ")
end

function skip_wrap(fdef; warn=true)
    warn ? @warn("Skipping creation command $(fdef.name)") : nothing
    name = name_transform(fdef)
    FDefinition(name, Signature(name, fdef.signature.args, fdef.signature.kwargs), fdef.short, fdef.body)
end

checked_function_call(fdef) = (fdef.return_type == "VkResult" ? "@check " : "") * fdef.name

function wrap_generic(fdef)
    sig = fdef.signature
    args, kwargs = arguments(sig, transform_name=true), keyword_arguments(sig, transform_name=true)
    new_fname = name_transform(fdef.name, FDefinition)
    # body = [Statement("$(sig.name)($(join_args(argnames(new_sig))))")]
    # kept_args = arguments(sig, transform_name=false, drop_type=false)
    # fname = name_transform(fdef.name, FDefinition)
    body, pass_results = accumulate_passes(fdef.name, fdef.signature.args, pass_new_nametype(FDefinition), [ComputeLengthArgument(), InitializePointers()])
    # new_sig = Signature(fname, remove_type.(args), kwargs)
    last_args_used = init_args(pass_results, body, fdef)
    triggered_vars_ip = [x.first for x ∈ pass_results if x.second[2].is_triggered]
    triggered_vars_ip_new = fieldname_transform.(triggered_vars_ip, nothing)
    new_sig = Signature(new_fname, filter(x -> x.name ∉ triggered_vars_ip_new, args), kwargs)
    # for (i, arg) ∈ enumerate(sig.args)
    #     tmp_argname(arg.name, arg.type) ∉ last_args_used && arg.name ∉ last_args_used && insert!(last_args_used, i, arg.name)
    # end
    _m = fdef.return_type == "VkResult" ? "@check " : ""
    push!(body, Statement("$(checked_function_call(fdef))($(join_args(last_args_used)))"))
    @assert length(triggered_vars_ip) <= 1
    if !isempty(triggered_vars_ip)
        var = first(triggered_vars_ip)
        type = inner_type(last(pass_results[var]).pass_args.type)
        push!(body, Statement(is_vulkan_struct(type) ? "$(name_transform(type, SDefinition))($var[])" : "$var[]", var))
    end
    FDefinition(new_fname, new_sig, false, body, "Generic definition")
end

function type_dependencies(fdef::FDefinition)
    types = filter(!isnothing, argtypes(fdef.signature))
    isempty(types) ? String[] : unique(vcat(type_dependencies.(types))...)
end


text:
function unzip_file(zipfile::String, out_path::String)::Nothing
  r = Reader(zipfile)
  for f in r.files
      path = joinpath(out_path,f.name)
      println(path)
      try
          write(path,read(f))
      catch e
          println(e)
      finally
          close(f)
      end
  end
  close(r)
  nothing
end

# function HTTP response to JSON
function r2j(response)
    JSON.Parser.parse(String(response))
end

# results
function results(success::Bool,msg::String,data::Any)
    return Dict("success"=>success,"msg"=>msg,"data"=>data)
  end

  function results(success::Bool,msg::String)
    return Dict("success"=>success,"msg"=>msg,"data"=>nothing)
  end

  function results(err::Exception)
    msg = ""
    if typeof(err) == HTTP.ExceptionRequest.StatusError
      try
        d = r2j(err.response.body)
        msg = string("[HTTP_STATUS:",err.status,", ERR_CODE:",d["code"]
          ,"] ",d["msg"])
      catch e
        msg = string("[HTTP_STATUS:",err.status,"] ",string(err.response))
      end
    else
      msg = string(err)
    end
    return Dict("success"=>false,"msg"=>msg,"data"=>nothing)
  end

  # timestamp to datetime
function ts2dt(timestamp::Int64)
    unix2datetime(timestamp/1_000)
end

# datetime to timestamp
function dt2ts(dt::DateTime)
    Int64(floor(datetime2unix(dt)*1_000))
end

# yyyy-mm-dd hh:mm:ss (UTC) string to timestamp
function st2ts(st::String)
    dt = DateTime(st,dateformat"y-m-d H:M:S")
    dt2ts(dt)
end

# signing with apikey and apisecret
function timestamp()
    Int64(floor(datetime2unix(now(Dates.UTC))*1_000))
end

#local time to UTC
# yyyy-mm-dd hh:mm:ss string to UTC datetime
function local_st2utc_dt(st::String)
    zdt= ZonedDateTime(
            DateTime(st,dateformat"y-m-d H:M:S"),
            tz"Asia/Seoul"
          )
    dt = DateTime(zdt,Dates.UTC)
end

# UTC time stamp to local datetime
# timestamp (UTC) string to datetime
function utc_ts2local_dt(ts::Int64)
    zdt= ZonedDateTime(
            ts2dt(ts),
            tz"UTC"
          )
    DateTime(
      astimezone(zdt, tz"Asia/Seoul"))
end

# dict to params
function d2p(dict::Dict)
    #?a=1&b=2 ==> a=1&b=2
    string(URIs.URI(query=dict))[2:end]
end

function sign(secret,msg)
    bytes2hex(SHA.hmac_sha256(Vector{UInt8}(secret),Vector{UInt8}(msg)))
    #Nettle.hexdigest("sha256",secret,msg)
end

function monitor_input()
  # Put STDIN in 'raw mode'
  # key input
  """
  inputBuffer = monitorInput()
  if isready(inputBuffer) && take!(inputBuffer) == 'q'
    println("==================END!====================")
    return
  end
  """
  ccall(:jl_tty_set_mode, Int32, (Ptr{Nothing}, Int32), stdin.handle, true) == 0 ||
        throw("FATAL: Terminal unable to enter raw mode.")

  inputBuffer = Channel{Char}(100)

  @async begin
      while true
          c = read(stdin, Char)
          put!(inputBuffer, c)
      end
  end
  return inputBuffer
end
# END

max_stars_repo_name: mrchaos/fxrobot

max_issues_repo_name: mrchaos/fxrobot

max_forks_repo_name: mrchaos/fxrobot


text:
module test_julia_boolean

using Test

@test true isa Bool
@test true isa Integer
@test true isa Number
@test 2true + true == 3
@test 2false == 0
@test false - true == -1

@test true == 1
@test false == 0

end # module test_julia_boolean


text:
function main(x, y, dir, moves)
    found_place = false
    known_places = Set([[x, y]])

    for move in moves
        turn = move[1]
        steps = parse(Int32, move[2:end])

        if     dir == 'N' && turn == 'L'
            dir = 'W'
        elseif dir == 'N' && turn == 'R'
            dir = 'E'
        elseif dir == 'W' && turn == 'L'
            dir = 'S'
        elseif dir == 'W' && turn == 'R'
            dir = 'N'
        elseif dir == 'S' && turn == 'L'
            dir = 'E'
        elseif dir == 'S' && turn == 'R'
            dir = 'W'
        elseif dir == 'E' && turn == 'L'
            dir = 'N'
        elseif dir == 'E' && turn == 'R'
            dir = 'S'
        end

        for _ in 1:steps
            if     (dir == 'N')
                y += 1
            elseif (dir == 'W')
                x -= 1
            elseif (dir == 'S')
                y -= 1
            elseif (dir == 'E')
                x += 1
            end

            current_place = [x, y]
            if !found_place && current_place in known_places
                found_place = true
                println("qb: ", abs(x) + abs(y))
            end

            known_places = union(known_places, Set([current_place]))
        end
    end

    println("qa: ", abs(x) + abs(y))
end

content = strip(read(stdin, String))
moves = split(content, ", ")
main(0, 0, 'N', moves)


text:
module SUNRepresentations

using TensorOperations
using SparseArrayKit
using RationalRoots
using LinearAlgebra
using TensorKit;
using TensorKit: fusiontensor, Nsymbol

export SUNIrrep, basis, weight, Zweight, creation, annihilation, highest_weight,dim
export directproduct, CGC

struct SUNIrrep{N} <: TensorKit.AbstractIrrep{TensorKit.SU{N}}
    I::NTuple{N,Int}
end
SUNIrrep(args::Vararg{Int}) = SUNIrrep(args)
Base.isless(s1::SUNIrrep{N}, s2::SUNIrrep{N}) where N = isless(s1.I, s2.I)

_normalize(s::SUNIrrep) = (I = weight(s); return SUNIrrep(I .- I[end]))

Base.getproperty(s::SUNIrrep{N}, f::Symbol) where {N} = f == :N ? N : getfield(s, f)
weight(s::SUNIrrep) = getfield(s, :I)

function TensorKit.dim(s::SUNIrrep)
    N = s.N
    I = weight(s)
    dim = 1//1
    for k2 = 2:N, k1 = 1:k2-1
        dim *= (k2 - k1 + I[k1] - I[k2]) // (k2 - k1)
    end
    @assert denominator(dim) == 1
    return numerator(dim)
end

include("gtpatterns.jl")

basis(s::SUNIrrep) = GTPatternIterator(s)

# direct product: return dictionary with new irreps as keys, outer multiplicities as value
function directproduct(s1::SUNIrrep{N}, s2::SUNIrrep{N}) where {N}
    dim(s1) > dim(s2) && return directproduct(s2, s1)
    result = Dict{SUNIrrep{N}, Int}()
    for m in basis(s1)
        t = weight(s2)
        bad_pattern = false
        for k in 1:N, l=N:-1:k
            bad_pattern && break
            bkl = m[k,l]
            if checkbounds(m, k, l-1)
                bkl -= m[k,l-1]
            end
            t = Base.setindex(t, t[l] + bkl, l)
            if l>1
                bad_pattern = t[l-1]< t[l]
            end
        end
        if !(bad_pattern)
            s = _normalize(SUNIrrep(t))
            result[s] = get(result, s, 0) + 1
        end
    end
    return result
end

include("clebschgordan.jl")
include("sector.jl")

end