text:
module SingleStackUtils

export get_vars_from_nodal_stack,
    get_vars_from_element_stack,
    get_horizontal_variance,
    get_horizontal_mean,
    reduce_nodal_stack,
    reduce_element_stack,
    horizontally_average!,
    dict_of_nodal_states,
    NodalStack,
    single_stack_diagnostics

using OrderedCollections
using UnPack
using StaticArrays
import KernelAbstractions: CPU

using ..BalanceLaws
using ..DGMethods
using ..DGMethods.Grids
using ..MPIStateArrays
using ..VariableTemplates

"""
    get_vars_from_nodal_stack(
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars;
        vrange::UnitRange = 1:size(Q, 3),
        i::Int = 1,
        j::Int = 1,
        exclude::Vector{String} = String[],
        interp = false,
    ) where {T, dim, N}

Return a dictionary whose keys are the `flattenednames()` of the variables
specified in `vars` (as returned by e.g. `vars_state`), and
whose values are arrays of the values for that variable along the vertical
dimension in `Q`. Only a single element is expected in the horizontal as
this is intended for the single stack configuration and `i` and `j` identify
the horizontal nodal coordinates.

Variables listed in `exclude` are skipped.
"""
function get_vars_from_nodal_stack(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars;
    vrange::UnitRange = 1:size(Q, 3),
    i::Int = 1,
    j::Int = 1,
    exclude::Vector{String} = String[],
    interp = false,
) where {T, dim, N}

    # extract grid information and bring `Q` to the host if needed
    FT = eltype(Q)
    Nq = N .+ 1
    # Code assumes the same polynomial order in all horizontal directions
    @inbounds begin
        Nq1 = Nq[1]
        Nq2 = Nq[2]
        Nqk = dim == 2 ? 1 : Nq[dim]
    end
    Np = dofs_per_element(grid)
    state_data = array_device(Q) isa CPU ? Q.realdata : Array(Q.realdata)

    # set up the dictionary to be returned
    var_names = flattenednames(vars)
    stack_vals = OrderedDict()
    num_vars = varsize(vars)
    vars_wanted = Int[]
    @inbounds for vi in 1:num_vars
        if !(var_names[vi] in exclude)
            stack_vals[var_names[vi]] = FT[]
            push!(vars_wanted, vi)
        end
    end
    elemtobndy = convert(Array, grid.elemtobndy)
    vmap⁻ = convert(Array, grid.vmap⁻)
    vmap⁺ = convert(Array, grid.vmap⁺)
    vgeo = convert(Array, grid.vgeo)
    # extract values from `state_data`
    @inbounds for ev in vrange, k in 1:Nqk, v in vars_wanted
        if interp && k == 1 && elemtobndy[5, ev] == 0
            # Get face degree of freedom number
            n = i + Nq1 * ((j - 1))
            # get the element numbers
            ev⁻ = ev
            # Get neighboring id data
            id⁻, id⁺ = vmap⁻[n, 5, ev⁻], vmap⁺[n, 5, ev⁻]
            ev⁺ = ((id⁺ - 1) ÷ Np) + 1
            # get the volume degree of freedom numbers
            vid⁻, vid⁺ = ((id⁻ - 1) % Np) + 1, ((id⁺ - 1) % Np) + 1

            J⁻, J⁺ = vgeo[vid⁻, Grids._M, ev⁻], vgeo[vid⁺, Grids._M, ev⁺]
            state_local = J⁻ * state_data[vid⁻, v, ev⁻]
            state_local += J⁺ * state_data[vid⁺, v, ev⁺]
            state_local /= (J⁻ + J⁺)
            push!(stack_vals[var_names[v]], state_local)
        elseif interp && k == Nqk && elemtobndy[6, ev] == 0
            # Get face degree of freedom number
            n = i + Nq1 * ((j - 1))
            # get the element numbers
            ev⁻ = ev
            # Get neighboring id data
            id⁻, id⁺ = vmap⁻[n, 6, ev⁻], vmap⁺[n, 6, ev⁻]
            # periodic and need to handle this point (otherwise handled above)
            if id⁺ == id⁻
                vid⁻ = ((id⁻ - 1) % Np) + 1

                state_local = state_data[vid⁻, v, ev⁻]
                push!(stack_vals[var_names[v]], state_local)
            end
        else
            ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))
            state_local = state_data[ijk, v, ev]
            push!(stack_vals[var_names[v]], state_local)
        end
    end

    return stack_vals
end

"""
    get_vars_from_element_stack(
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars;
        vrange::UnitRange = 1:size(Q, 3),
        exclude::Vector{String} = String[],
        interp = false,
    ) where {T, dim, N}

Return an array of [`get_vars_from_nodal_stack()`](@ref)s whose dimensions
are the number of nodal points per element in the horizontal plane.

Variables listed in `exclude` are skipped.
"""
function get_vars_from_element_stack(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars;
    vrange::UnitRange = 1:size(Q, 3),
    exclude::Vector{String} = String[],
    interp = false,
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds Nq1 = Nq[1]
    @inbounds Nq2 = Nq[2]

    return [
        get_vars_from_nodal_stack(
            grid,
            Q,
            vars,
            vrange = vrange,
            i = i,
            j = j,
            exclude = exclude,
            interp = interp,
        ) for i in 1:Nq1, j in 1:Nq2
    ]
end

"""
    get_horizontal_mean(
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars;
        vrange::UnitRange = 1:size(Q, 3),
        exclude::Vector{String} = String[],
        interp = false,
    ) where {T, dim, N}

Return a dictionary whose keys are the `flattenednames()` of the variables
specified in `vars` (as returned by e.g. `vars_state`), and
whose values are arrays of the horizontal averages for that variable along
the vertical dimension in `Q`. Only a single element is expected in the
horizontal as this is intended for the single stack configuration.

Variables listed in `exclude` are skipped.
"""
function get_horizontal_mean(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars;
    vrange::UnitRange = 1:size(Q, 3),
    exclude::Vector{String} = String[],
    interp = false,
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds Nq1 = Nq[1]
    @inbounds Nq2 = Nq[2]

    vars_avg = OrderedDict()
    vars_sq = OrderedDict()
    for i in 1:Nq1
        for j in 1:Nq2
            vars_nodal = get_vars_from_nodal_stack(
                grid,
                Q,
                vars,
                vrange = vrange,
                i = i,
                j = j,
                exclude = exclude,
                interp = interp,
            )
            vars_avg = merge(+, vars_avg, vars_nodal)
        end
    end
    map!(x -> x ./ Nq1 / Nq1, values(vars_avg))
    return vars_avg
end

"""
    get_horizontal_variance(
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars;
        vrange::UnitRange = 1:size(Q, 3),
        exclude::Vector{String} = String[],
        interp = false,
    ) where {T, dim, N}

Return a dictionary whose keys are the `flattenednames()` of the variables
specified in `vars` (as returned by e.g. `vars_state`), and
whose values are arrays of the horizontal variance for that variable along
the vertical dimension in `Q`. Only a single element is expected in the
horizontal as this is intended for the single stack configuration.

Variables listed in `exclude` are skipped.
"""
function get_horizontal_variance(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars;
    vrange::UnitRange = 1:size(Q, 3),
    exclude::Vector{String} = String[],
    interp = false,
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds Nq1 = Nq[1]
    @inbounds Nq2 = Nq[2]

    vars_avg = OrderedDict()
    vars_sq = OrderedDict()
    for i in 1:Nq1
        for j in 1:Nq2
            vars_nodal = get_vars_from_nodal_stack(
                grid,
                Q,
                vars,
                vrange = vrange,
                i = i,
                j = j,
                exclude = exclude,
                interp = interp,
            )
            vars_nodal_sq = OrderedDict(vars_nodal)
            map!(x -> x .^ 2, values(vars_nodal_sq))
            vars_avg = merge(+, vars_avg, vars_nodal)
            vars_sq = merge(+, vars_sq, vars_nodal_sq)
        end
    end
    map!(x -> (x ./ Nq1 / Nq1) .^ 2, values(vars_avg))
    map!(x -> x ./ Nq1 / Nq1, values(vars_sq))
    vars_var = merge(-, vars_sq, vars_avg)
    return vars_var
end

"""
    reduce_nodal_stack(
        op::Function,
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars::NamedTuple,
        var::String;
        vrange::UnitRange = 1:size(Q, 3),
    ) where {T, dim, N}

Reduce `var` from `vars` within `Q` over all nodal points in the specified
`vrange` of elements with `op`. Return a tuple `(result, z)` where `result` is
the final value returned by `op` and `z` is the index within `vrange` where the
`result` was determined.
"""
function reduce_nodal_stack(
    op::Function,
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars::Type,
    var::String;
    vrange::UnitRange = 1:size(Q, 3),
    i::Int = 1,
    j::Int = 1,
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds begin
        Nq1 = Nq[1]
        Nq2 = Nq[2]
        Nqk = dim == 2 ? 1 : Nq[dim]
    end

    var_names = flattenednames(vars)
    var_ind = findfirst(s -> s == var, var_names)
    if var_ind === nothing
        return
    end

    state_data = array_device(Q) isa CPU ? Q.realdata : Array(Q.realdata)
    z = vrange.start
    FT = eltype(state_data)

    # Initialize result with identity operation for operator
    result = if op isa typeof(+)
        zero(FT)
    elseif op isa typeof(*)
        one(FT)
    elseif op isa typeof(min)
        floatmax(FT)
    elseif op isa typeof(max)
        -floatmax(FT)
    else
        error("unknown operator: $op")
    end

    for ev in vrange
        for k in 1:Nqk
            ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))
            new_result = op(result, state_data[ijk, var_ind, ev])
            if !isequal(new_result, result)
                result = new_result
                z = ev
            end
        end
    end

    return (result, z)
end

"""
    reduce_element_stack(
        op::Function,
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        vars::NamedTuple,
        var::String;
        vrange::UnitRange = 1:size(Q, 3),
    ) where {T, dim, N}

Reduce `var` from `vars` within `Q` over all nodal points in the specified
`vrange` of elements with `op`. Return a tuple `(result, z)` where `result` is
the final value returned by `op` and `z` is the index within `vrange` where the
`result` was determined.
"""
function reduce_element_stack(
    op::Function,
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    vars::Type,
    var::String;
    vrange::UnitRange = 1:size(Q, 3),
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds Nq1 = Nq[1]
    @inbounds Nq2 = Nq[2]

    return [
        reduce_nodal_stack(
            op,
            grid,
            Q,
            vars,
            var,
            vrange = vrange,
            i = i,
            j = j,
        ) for i in 1:Nq1, j in 1:Nq2
    ]
end

"""
    horizontally_average!(
        grid::DiscontinuousSpectralElementGrid{T, dim, N},
        Q::MPIStateArray,
        i_vars,
    ) where {T, dim, N}

Horizontally average variables, from variable
indexes `i_vars`, in `MPIStateArray` `Q`.

!!! note
    These are not proper horizontal averages-- the main
    purpose of this method is to ensure that there are
    no horizontal fluxes for a single stack configuration.
"""
function horizontally_average!(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    Q::MPIStateArray,
    i_vars,
) where {T, dim, N}

    Nq = N .+ 1
    @inbounds begin
        Nq1 = Nq[1]
        Nq2 = Nq[2]
        Nqk = dim == 2 ? 1 : Nq[dim]
    end

    ArrType = typeof(Q.data)
    state_data = array_device(Q) isa CPU ? Q.realdata : Array(Q.realdata)

    for ev in 1:size(state_data, 3), k in 1:Nqk, i_v in i_vars
        Q_sum = 0
        for i in 1:Nq1, j in 1:Nq2
            Q_sum += state_data[i + Nq1 * ((j - 1) + Nq2 * (k - 1)), i_v, ev]
        end
        Q_ave = Q_sum / (Nq1 * Nq2)
        for i in 1:Nq1, j in 1:Nq2
            ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))
            state_data[ijk, i_v, ev] = Q_ave
        end
    end
    Q.realdata .= ArrType(state_data)
end

get_data(solver_config, ::Prognostic) = solver_config.Q
get_data(solver_config, ::Auxiliary) = solver_config.dg.state_auxiliary
get_data(solver_config, ::GradientFlux) = solver_config.dg.state_gradient_flux

"""
    dict_of_nodal_states(
        solver_config,
        state_types = (Prognostic(), Auxiliary());
        aux_excludes = [],
        interp = false,
        )

A dictionary of single stack prognostic and auxiliary
variables at the `i=1`,`j=1` node given
 - `solver_config` a `SolverConfiguration`
 - `aux_excludes` a vector of strings containing the
    variables to exclude from the auxiliary state.
"""
function dict_of_nodal_states(
    solver_config,
    state_types = (Prognostic(), Auxiliary());
    aux_excludes = String[],
    interp = false,
)
    FT = eltype(solver_config.Q)
    all_state_vars = []
    for st in state_types
        state_vars = get_vars_from_nodal_stack(
            solver_config.dg.grid,
            get_data(solver_config, st),
            vars_state(solver_config.dg.balance_law, st, FT),
            exclude = st isa Auxiliary ? aux_excludes : String[],
            interp = interp,
        )
        push!(all_state_vars, state_vars...)
    end
    return OrderedDict(all_state_vars...)
end

# A container for holding various
# global or point-wise states:
struct States{P, A, D, HD}
    prog::P
    aux::A
    diffusive::D
    hyperdiffusive::HD
end

"""
    NodalStack(
            bl::BalanceLaw,
            grid::DiscontinuousSpectralElementGrid,
            prognostic,
            auxiliary,
            diffusive,
            hyperdiffusive;
            i = 1,
            j = 1,
            interp = true,
        )

A struct whose `iterate(::NodalStack)` traverses
the nodal stack and returns a NamedTuple of
point-wise fields (`Vars`).

# Example
```julia
for state_local in NodalStack(
        bl,
        grid,
        prognostic, # global field along nodal stack
        auxiliary,
        diffusive,
        hyperdiffusive
    )
prog = state_local.prog # point-wise field along nodal stack
end
```

## TODO: Make `prognostic`, `auxiliary`, `diffusive`, `hyperdiffusive` optional

# Arguments
 - `bl` the balance law
 - `grid` the discontinuous spectral element grid
 - `prognostic` the global prognostic state
 - `auxiliary` the global auxiliary state
 - `diffusive` the global diffusive state (gradient-flux)
 - `hyperdiffusive` the global hyperdiffusive state
 - `i,j` the `i,j`'th nodal stack (in the horizontal directions)
 - `interp` a bool indicating whether to
    interpolate the duplicate Gauss-Lebotto
    points at the element faces.

!!! warn
    Before iterating, the data is transferred from the
    device (GPU) to the host (CPU), as this is intended
    for debugging / diagnostics usage.
"""
struct NodalStack{N, BL, G, S, VR, TI, TJ, CI, IN}
    bl::BL
    grid::G
    states::S
    vrange::VR
    i::TI
    j::TJ
    cart_ind::CI
    interp::IN
    function NodalStack(
        bl::BalanceLaw,
        grid::DiscontinuousSpectralElementGrid;
        prognostic,
        auxiliary,
        diffusive,
        hyperdiffusive,
        i = 1,
        j = 1,
        interp = true,
    )
        states = States(prognostic, auxiliary, diffusive, hyperdiffusive)
        vrange = 1:size(prognostic, 3)
        grid_info = basic_grid_info(grid)
        @unpack Nqk = grid_info
        if last(polynomialorders(grid)) == 0
            interp = false
        end
        # Store cartesian indices, so we can map the iter_state
        # to the cartesian space `Q[i, var, j]`
        if interp
            cart_ind = CartesianIndices(((Nqk - 1), size(prognostic, 3)))
        else
            cart_ind = CartesianIndices((Nqk, size(prognostic, 3)))
        end
        args = (bl, grid, states, vrange, i, j, cart_ind, interp)
        BL, G, S, VR, TI, TJ, CI, IN = typeof.(args)
        if interp
            len = size(prognostic, 3) * (Nqk - 1) + 1
        else
            len = size(prognostic, 3) * Nqk
        end
        new{len, BL, G, S, VR, TI, TJ, CI, IN}(args...)
    end
end

Base.length(gs::NodalStack{N}) where {N} = N

# Helper function
get_state(v, state, vid⁻, vid⁺, ev⁻, ev⁺, J⁻, J⁺) =
    (J⁻ * state[vid⁻, v, ev⁻] + J⁺ * state[vid⁺, v, ev⁺]) / (J⁻ + J⁺)

to_cpu(state) =
    array_device(state) isa CPU ? state.realdata : Array(state.realdata)

function interp_top(state, args, n_vars, bl, st)
    vs = Vars{vars_state(bl, st, eltype(state))}
    if n_vars ≠ 0
        return vs(map(v -> get_state(v, state, args...), 1:n_vars))
    else
        return nothing
    end
end

function interp_bot(state, vid⁻, ev⁻, n_vars, bl, st)
    vs = Vars{vars_state(bl, st, eltype(state))}
    if n_vars ≠ 0
        return vs(map(v -> state[vid⁻, v, ev⁻], 1:n_vars))
    else
        return nothing
    end
end

function no_interp(state, ijk, ev, n_vars, bl, st)
    vs = Vars{vars_state(bl, st, eltype(state))}
    if n_vars ≠ 0
        return vs(map(v -> state[ijk, v, ev], 1:n_vars))
    else
        return nothing
    end
end

function Base.iterate(gs::NodalStack, iter_state = 1)
    iter_state > length(gs) && return nothing

    # extract grid information
    grid = gs.grid
    FT = eltype(grid)
    grid_info = basic_grid_info(grid)
    @unpack N, Nq, Np, Nqk = grid_info
    @inbounds Nq1, Nq2 = Nq[1], Nq[2]
    states = gs.states
    bl = gs.bl
    interp = gs.interp

    # bring `Q` to the host if needed

    prognostic = to_cpu(states.prog)
    auxiliary = to_cpu(states.aux)
    diffusive = to_cpu(states.diffusive)
    hyperdiffusive = to_cpu(states.hyperdiffusive)

    n_vars_prog = size(prognostic, 2)
    n_vars_aux = size(auxiliary, 2)
    n_vars_diff = size(diffusive, 2)
    n_vars_hd = size(hyperdiffusive, 2)

    elemtobndy = convert(Array, grid.elemtobndy)
    vmap⁻ = convert(Array, grid.vmap⁻)
    vmap⁺ = convert(Array, grid.vmap⁺)
    vgeo = convert(Array, grid.vgeo)
    i, j = gs.i, gs.j
    if iter_state == length(gs)
        ijk_cart = (Nqk, size(states.prog, 3))
    else
        ijk_cart = Tuple(gs.cart_ind[iter_state])
    end
    ev = ijk_cart[2]
    k = ijk_cart[1]
    iter_state⁺ = iter_state + 1
    if interp && k == 1 && elemtobndy[5, ev] == 0
        # Get face degree of freedom number
        n = i + Nq1 * ((j - 1))
        # get the element numbers
        ev⁻ = ev
        # Get neighboring id data
        id⁻, id⁺ = vmap⁻[n, 5, ev⁻], vmap⁺[n, 5, ev⁻]
        ev⁺ = ((id⁺ - 1) ÷ Np) + 1
        # get the volume degree of freedom numbers
        vid⁻, vid⁺ = ((id⁻ - 1) % Np) + 1, ((id⁺ - 1) % Np) + 1
        J⁻, J⁺ = vgeo[vid⁻, Grids._M, ev⁻], vgeo[vid⁺, Grids._M, ev⁺]

        args = (vid⁻, vid⁺, ev⁻, ev⁺, J⁻, J⁺)
#! format: off
        prog = interp_top(prognostic, args, n_vars_prog, bl, Prognostic())
        aux = interp_top(auxiliary, args, n_vars_aux, bl, Auxiliary())
        ∇flux = interp_top(diffusive, args, n_vars_diff, bl, GradientFlux())
        hyperdiff = interp_top(hyperdiffusive, args, n_vars_hd, bl, Hyperdiffusive())
#! format: on

        return ((; prog, aux, ∇flux, hyperdiff), iter_state⁺)

    elseif interp && k == Nqk && elemtobndy[6, ev] == 0
        # Get face degree of freedom number
        n = i + Nq1 * ((j - 1))
        # get the element numbers
        ev⁻ = ev
        # Get neighboring id data
        id⁻, id⁺ = vmap⁻[n, 6, ev⁻], vmap⁺[n, 6, ev⁻]
        # periodic and need to handle this point (otherwise handled above)
        if id⁺ == id⁻
            vid⁻ = ((id⁻ - 1) % Np) + 1

#! format: off
            prog = interp_bot(prognostic, vid⁻, ev⁻, n_vars_prog, bl, Prognostic())
            aux = interp_bot(auxiliary, vid⁻, ev⁻, n_vars_aux, bl, Auxiliary())
            ∇flux = interp_bot(diffusive, vid⁻, ev⁻, n_vars_diff, bl, GradientFlux())
            hyperdiff = interp_bot( hyperdiffusive, vid⁻, ev⁻, n_vars_hd, bl, Hyperdiffusive())
#! format: on

            return ((; prog, aux, ∇flux, hyperdiff), iter_state⁺)
        else
            error("uncaught case in iterate(::NodalStack)")
        end
    else
        ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))

#! format: off
        prog = no_interp(prognostic, ijk, ev, n_vars_prog, bl, Prognostic())
        aux = no_interp(auxiliary, ijk, ev, n_vars_aux, bl, Auxiliary())
        ∇flux = no_interp(diffusive, ijk, ev, n_vars_diff, bl, GradientFlux())
        hyperdiff = no_interp(hyperdiffusive, ijk, ev, n_vars_hd, bl, Hyperdiffusive())
#! format: on

        return ((; prog, aux, ∇flux, hyperdiff), iter_state⁺)
    end
end

include("single_stack_diagnostics.jl")

end # module


text:
# Test 1 x 1 convolution module using Householder matrices
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

using InvertibleNetworks, LinearAlgebra, Test

###################################################################################################
# Initialize parameters

# Dimensions
nx = 28
ny = 28
k = 4
batchsize = 2

# Variables
v1 = randn(Float32, k)
v10 = randn(Float32, k)
dv1 = v1 - v10

v2 = randn(Float32, k)
v20 = randn(Float32, k)
dv2 = v2 - v20

v3 = randn(Float32, k)
v30 = randn(Float32, k)
dv3 = v3 - v30

# Input
X = randn(Float32, nx, ny, k, batchsize)
X0 = randn(Float32, nx, ny, k, batchsize)
dX = X - X0

# Operators
C = Conv1x1(v1, v2, v3)
C0 = Conv1x1(v10, v20, v30)

###################################################################################################
# Test invertibility

X_ = C.inverse(C.forward(X))
err1 = norm(X - X_)/norm(X)

@test isapprox(err1, 0f0; atol=1f-6)

X_ = C.forward(C.inverse(X))
err2 = norm(X - X_)/norm(X)

@test isapprox(err2, 0f0; atol=1f-6)

Y = C.forward(X)
ΔY = randn(Float32, nx, ny, k, batchsize)
ΔX_, X_ = C.inverse((ΔY, Y))
err3 = norm(X - X_)/norm(X)

@test isapprox(err3, 0f0; atol=1f-6)

###################################################################################################
# Test gradients are set in inverse pass

# Predicted data and misfit
C0.v1.grad = nothing
Y_ = C0.forward(X)
@test isnothing(C0.v1.grad)

ΔY = Y_ - Y

# Compute gradients w.r.t. v
ΔX, X_ = C0.inverse((ΔY, Y_))
@test ~isnothing(C0.v1.grad)

###################################################################################################
# Gradient test

loss(ΔY) = .5f0*norm(ΔY)^2

function objective(C, X, Y)
    Y0 = C.forward(X)
    ΔY = Y0 - Y
    f = loss(ΔY)
    ΔX, X_ = C.inverse((ΔY, Y0))
    @test isapprox(norm(X - X_)/norm(X), 0f0, atol=1f-6)
    return f, ΔX, C.v1.grad, C.v2.grad, C.v3.grad
end

# Observed data
Y = C.forward(X)

# Gradient test for X
maxiter = 5
print("Gradient test ΔX\n")
clear_grad!(C)
f0, ΔX = objective(C, X0, Y)[1:2]
h = .01f0
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)
for j=1:maxiter
    f = loss(C.forward(X0 + h*dX) - Y)
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f0)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f0)

# Gradient test for v
print("\nGradient test Δv1\n")
clear_grad!(C0)
f0, ΔX, Δv1, Δv2, Δv3 = objective(C0, X, Y)
h = .01f0
err3 = zeros(Float32, maxiter)
err4 = zeros(Float32, maxiter)
for j=1:maxiter
    C0.v1.data = v10 + h*dv1
    C0.v2.data = v20 + h*dv2
    C0.v3.data = v30 + h*dv3
    f = loss(C0.forward(X) - Y)
    err3[j] = abs(f - f0)
    err4[j] = abs(f - f0 - h*dot(dv1, Δv1) - h*dot(dv2, Δv2) - h*dot(dv3, Δv3))
    print(err3[j], "; ", err4[j], "\n")
    global h = h/2f0
end

@test isapprox(err3[end] / (err3[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err4[end] / (err4[1]/4^(maxiter-1)), 1f0; atol=1f1)

###################################################################################################
# Gradient test: forward-inverse in reverse order

loss(ΔY) = .5f0*norm(ΔY)^2

function objectiveT(C, X, Y)
    Y0 = C.inverse(X)
    ΔY = Y0 - Y
    f = loss(ΔY)
    ΔX, X_ = C.forward((ΔY, Y0))
    @test isapprox(norm(X - X_)/norm(X), 0f0, atol=1f-6)
    return f, ΔX, C.v1.grad, C.v2.grad, C.v3.grad
end

# Observed data
Y = C.forward(X)

# Gradient test for X
maxiter = 5
print("\nGradient test ΔX\n")
C = Conv1x1(v1, v2, v3)
f0, ΔX = objectiveT(C, X0, Y)[1:2]
h = .01f0
err5 = zeros(Float32, maxiter)
err6 = zeros(Float32, maxiter)
for j=1:maxiter
    f = loss(C.inverse(X0 + h*dX) - Y)
    err5[j] = abs(f - f0)
    err6[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err5[j], "; ", err6[j], "\n")
    global h = h/2f0
end

@test isapprox(err5[end] / (err5[1]/2^(maxiter-1)), 1f0; atol=1f0)
@test isapprox(err6[end] / (err6[1]/4^(maxiter-1)), 1f0; atol=1f0)

# Gradient test for v
print("\nGradient test Δv1\n")
C0 = Conv1x1(v10, v20, v30)
f0, ΔX, Δv1, Δv2, Δv3 = objectiveT(C0, X, Y)
h = .01f0
err7 = zeros(Float32, maxiter)
err8 = zeros(Float32, maxiter)
for j=1:maxiter
    C0.v1.data = v10 + h*dv1
    C0.v2.data = v20 + h*dv2
    C0.v3.data = v30 + h*dv3
    f = loss(C0.inverse(X) - Y)
    err7[j] = abs(f - f0)
    err8[j] = abs(f - f0 - h*dot(dv1, Δv1) - h*dot(dv2, Δv2) - h*dot(dv3, Δv3))
    print(err7[j], "; ", err8[j], "\n")
    global h = h/2f0
end

@test isapprox(err7[end] / (err7[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err8[end] / (err8[1]/4^(maxiter-1)), 1f0; atol=1f1)

###################################################################################################
# Jacobian-related tests

# Gradient test

# Initialization
batchsize=10
v10 = randn(Float32, k)
v20 = randn(Float32, k)
v30 = randn(Float32, k)
C0 = Conv1x1(v10, v20, v30; logdet=true)
θ0 = deepcopy(get_params(C0))
v1 = randn(Float32, k)
v2 = randn(Float32, k)
v3 = randn(Float32, k)
C = Conv1x1(v1, v2, v3; logdet=true)
θ = deepcopy(get_params(C))
X = randn(Float32, nx, ny, k, batchsize)

# Perturbation (normalized)
dθ = θ-θ0
for i = 1:length(θ)
    dθ[i] = norm(θ0[i])*dθ[i]/(norm(dθ[i]).+1f-10)
end
dX = randn(Float32, nx, ny, k, batchsize); dX = norm(X)*dX/norm(dX)

# Jacobian eval
dY, Y = C.jacobian(dX, dθ, X)

# Test
print("\nJacobian test\n")
h = 0.1f0
maxiter = 7
err9 = zeros(Float32, maxiter)
err10 = zeros(Float32, maxiter)
for j=1:maxiter
    set_params!(C, θ+h*dθ)
    Y_loc, _ = C.forward(X+h*dX)
    err9[j] = norm(Y_loc - Y)
    err10[j] = norm(Y_loc - Y - h*dY)
    print(err9[j], "; ", err10[j], "\n")
    global h = h/2f0
end

@test isapprox(err9[end] / (err9[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err10[end] / (err10[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Adjoint test

set_params!(C, θ)
dY, Y = C.jacobian(dX, 0f0*dθ, X)
dY_ = randn(Float32, size(dY))
dX_, dθ_, _ = C.adjointJacobian(dY_, Y)
a = dot(dY, dY_)
b = dot(dX, dX_)+dot(0f0*dθ, dθ_)
@test isapprox(a, b; rtol=1f-3)

# Gradient test (inverse)

# Perturbation (normalized)
dY = randn(Float32, nx, ny, k, batchsize); dY *= norm(Y)/norm(dY)

# Jacobian (inverse) eval
dX, X = C.jacobianInverse(dY, dθ, Y)

# Test
print("\nJacobian (inverse) test\n")
h = 0.1f0
maxiter = 7
err11 = zeros(Float32, maxiter)
err12 = zeros(Float32, maxiter)
for j=1:maxiter
    set_params!(C, θ+h*dθ)
    X_loc, _ = C.inverse(Y+h*dY)
    err11[j] = norm(X_loc - X)
    err12[j] = norm(X_loc - X - h*dX)
    print(err11[j], "; ", err12[j], "\n")
    global h = h/2f0
end

@test isapprox(err11[end] / (err11[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err12[end] / (err12[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Inverse test

dY = randn(Float32, nx, ny, k, batchsize)
Y = randn(Float32, nx, ny, k, batchsize)
dX = randn(Float32, nx, ny, k, batchsize)
X = randn(Float32, nx, ny, k, batchsize)
dX_, X_ = C.jacobianInverse(dY, 0f0*dθ, Y)
dY_, Y_ = C.jacobian(dX_, 0f0*dθ, X_)
@test isapprox(dY_, dY; rtol=1f-3)
@test isapprox(Y_, Y; rtol=1f-3)
dY_, Y_ = C.jacobian(dX, 0f0*dθ, X)
dX_, X_ = C.jacobianInverse(dY_, 0f0*dθ, Y_)
@test isapprox(dX_, dX; rtol=1f-3)
@test isapprox(X_, X; rtol=1f-3)

# Adjoint test (inverse)

set_params!(C, θ)
dY, Y = C.jacobianInverse(dX, dθ, X)
dY_ = randn(Float32, size(dY))
dX_, dθ_, _ = C.adjointJacobianInverse(dY_, Y)
a = dot(dY, dY_)
b = dot(dX, dX_)+dot(dθ, dθ_)
@test isapprox(a, b; rtol=1f-3)

ext: jl


text:
@testset "Loggers" begin
    FMT_STR = "[{level}]:{name} - {msg}"

    LEVELS = Dict(
        "not_set" => 0,
        "debug" => 10,
        "info" => 20,
        "warn" => 30,
        "error" => 40,
    )

    @testset "Simple" begin
        io = IOBuffer()

        try
            handler = DefaultHandler(io, DefaultFormatter(FMT_STR))

            logger = Logger(
                "Logger.example",
                Dict("Buffer" => handler),
                "error",
                LEVELS,
                DefaultRecord,
                true
            )

            @test logger.name == "Logger.example"
            @test logger.level == "error"
            @test length(gethandlers(logger)) == 1

            push!(logger, DefaultHandler(tempname()))
            @test length(gethandlers(logger)) == 2

            @test ispropagating(logger)
            @test setpropagating!(logger, true)

            setlevel!(logger, "info")
            @test logger.level == "info"

            push!(logger, Memento.Filter(logger))
            @test length(getfilters(logger)) == 3

            setlevel!(logger, "error") do
                @test getlevel(logger) == "error"
                Memento.warn(logger, "silenced message should not be displayed")
            end
            @test getlevel(logger) == "info"

            setrecord!(logger, DefaultRecord)

            addlevel!(logger, "fubar", 50)

            show(io, logger)
            @test occursin("Logger(Logger.example)", String(take!(io)))

            msg = "It works!"
            Memento.info(logger, msg)
            @test occursin("[info]:Logger.example - $msg", String(take!(io)))

            Memento.debug(logger, "This shouldn't get logged")
            @test isempty(String(take!(io)))

            @test_throws TestError Memento.error(logger, TestError("I failed."))
            @test occursin("I failed", String(take!(io)))

            # CompositeException
            comp = CompositeException([
                ErrorException("I am the first error"),
                ArgumentError("I am the second error"),
            ])
            @test_throws CompositeException Memento.error(logger, comp)
            output = String(take!(io))
            @test occursin("I am the first error", output)
            @test occursin("I am the second error", output)

            comp = CompositeException([
                ArgumentError("Error numero uno"),
                ErrorException("Error numero dos"),
            ])
            Memento.warn(logger, comp)
            output = String(take!(io))
            @test occursin("Error numero uno", output)
            @test occursin("Error numero dos", output)

            msg = "Something went very wrong"
            log(logger, "fubar", msg)
            @test occursin("[fubar]:Logger.example - $msg", String(take!(io)))

            new_logger = Logger("new_logger")
        finally
            close(io)
        end
    end
    @testset "Lazy Messages" begin
        # A test utility function that gives
        # us a function to pass to the log method
        # which will execute only if the message is
        # evaluated.
        function msg_func(msg)
            inner() = msg
            return inner
        end

        io = IOBuffer()

        try
            handler = DefaultHandler(io, DefaultFormatter(FMT_STR))

            logger = Logger(
                "Logger.example",
                Dict("Buffer" => handler),
                "error",
                LEVELS,
                DefaultRecord,
                true
            )

            @test logger.name == "Logger.example"
            @test logger.level == "error"
            @test length(gethandlers(logger)) == 1

            push!(logger, DefaultHandler(tempname()))
            @test length(gethandlers(logger)) == 2

            setlevel!(logger, "info")
            @test logger.level == "info"

            setrecord!(logger, DefaultRecord)
            addlevel!(logger, "fubar", 50)

            show(io, logger)
            @test occursin("Logger(Logger.example)", String(take!(io)))

            msg = "It works!"
            Memento.info(msg_func(msg), logger)
            @test occursin("[info]:Logger.example - $msg", String(take!(io)))

            Memento.debug(msg_func("This shouldn't get logged"), logger)
            @test isempty(String(take!(io)))

            msg = "Something went very wrong"
            @test_throws ErrorException error(msg_func(msg), logger)
            @test occursin("[error]:Logger.example - $msg", String(take!(io)))

            new_logger = Logger("new_logger")
        finally
            close(io)
        end
    end

    @testset "Recursive setlevel" begin
        getlogger.(("foo.bar.baz", "a.b.c"))
        foo_children = Memento.getchildren("foo")
        a_children = Memento.getchildren("a")

        @test all(==("info"), getlevel.(foo_children))
        setlevel!(getlogger("foo"), "debug"; recursive=true)

        @test all(==("debug"), getlevel.(foo_children))
        @test all(==("info"), getlevel.(a_children))

        setlevel!(getlogger("foo"), "notice"; recursive=true) do
            Memento.info(getlogger("foo"), "I shouldn't be getting printed.")
        end

        # Check that our child logging levels remain the same after
        # the above call.
        @test all(==("debug"), getlevel.(foo_children))
        @test all(==("info"), getlevel.(a_children))
    end

    @testset "Memento.config" begin
        root_logger = Memento.config!(
            "info";
            fmt="[{date} | {level} | {name}]: {msg}", colorized=false
        )

        my_logger = Memento.config!(
            getlogger("MyLogger"), "info";
            fmt="[{date} | {level} | {name}]: {msg}", colorized=false
        )

        str_logger = Memento.config!(
            "StrLogger", "info";
            fmt="[{date} | {level} | {name}]: {msg}", colorized=false
        )

        @test length(gethandlers(root_logger)) == 1
        @test length(gethandlers(my_logger)) == 1
        @test length(gethandlers(str_logger)) == 1

        @test ispropagating(my_logger)
        @test ispropagating(str_logger)

        @test getlevel(root_logger) == getlevel(my_logger) == getlevel(str_logger)

        non_prop_logger = Memento.config!(
            getlogger("NonPropLogger"), "info"; propagate=false
        )
        @test !ispropagating(non_prop_logger)

        Memento.config("notice"; recursive=true)
        @test all(l -> getlevel(l) == "notice", values(Memento._loggers))

        # reconfig
        logger = getlogger("NoReplace")
        new_logger = Memento.config("NoReplace", "info")
        @test new_logger === getlogger("NoReplace")
        @test new_logger === logger
    end

    @testset "Stacktraces" begin
        io = IOBuffer()

        try
            handler = DefaultHandler(
                io,
                DefaultFormatter(string(FMT_STR, " | {stacktrace}"))
            )

            logger = Logger(
                "Logger.example",
                Dict("Buffer" => handler),
                "info",
                LEVELS,
                DefaultRecord,
                true
            )

            # Define a test function that logs a message
            test_func() = Memento.info(logger, "Hello")
            test_func()

            msg = String(take!(io))
            @test !isempty(msg)
            @test occursin("test_func", msg)
        finally
            close(io)
        end
    end

    @testset "Don't overwrite registered logger" begin
        original_logger = Logger("new_logger")
        Memento.register(original_logger)

        logger = getlogger("new_logger")
        @test logger == original_logger

        # Makre sure registration doesn't overwrite the original registered logger
        Memento.register(Logger("new_logger"))
        logger = getlogger("new_logger")
        @test logger == original_logger
    end
end


text:
using DrWatson
@quickactivate
using ArgParse
using BSON
using Flux
using GenerativeAD
using StatsBase: fit!, predict, sample

s = ArgParseSettings()
@add_arg_table! s begin
   "max_seed"
		default = 1
		arg_type = Int
		help = "seed"
	"dataset"
		default = "iris"
		arg_type = String
		help = "dataset"
    "contamination"
    	arg_type = Float64
    	help = "contamination rate of training data"
    	default = 0.0
end
parsed_args = parse_args(ARGS, s)
@unpack dataset, max_seed, contamination = parsed_args

modelname = "sptn_005_010"

function sample_params()
	parameter_rng = (
		firstdense = [true, false],
		batchsize = 2 .^ (5:7),
		ncomp = 2 .^ (1:4),
		nlayers = 1:3,
		activation = [identity],
		unitary = [:butterfly],
		sharing = [:dense, :all, :none],
		init_seed = 1:Int(1e8),
	)

	(;zip(keys(parameter_rng), map(x->sample(x, 1)[1], parameter_rng))...)
end

function fit(data, parameters)
	model = GenerativeAD.Models.SPTN(;idim=size(data[1][1],1), parameters...)

	try
		global info, fit_t, _, _, _ = @timed fit!(model, data; max_train_time=82800/max_seed,
			patience=20, check_interval=10, quantile=(0.05, 0.1), parameters...)
	catch e
		# return an empty array if fit fails so nothing is computed
		@info "Failed training due to \n$e"
		return (fit_t = NaN, history=nothing, npars=nothing, model=nothing), []
	end

	# construct return information - put e.g. the model structure here for generative models
	training_info = (
		fit_t = fit_t,
		history = info.history,
		npars = info.npars,
		model = info.model
		)

	# now return the different scoring functions
	training_info, [(x -> predict(info.model, x), parameters)]
end

if abspath(PROGRAM_FILE) == @__FILE__
	# set a maximum for parameter sampling retries
	try_counter = 0
	max_tries = 10*max_seed
	cont_string = (contamination == 0.0) ? "" : "_contamination-$contamination"
	while try_counter < max_tries
	    parameters = sample_params()

	    for seed in 1:max_seed
			savepath = datadir("experiments/tabular$cont_string/$(modelname)/$(dataset)/seed=$(seed)")
			mkpath(savepath)

			# get data
			data = GenerativeAD.load_data(dataset, seed=seed, contamination=contamination)
			# edit parameters
			edited_parameters = GenerativeAD.edit_params(data, parameters)

			@info "Trying to fit $modelname on $dataset with parameters $(edited_parameters)..."
			@info "Train/valdiation/test splits: $(size(data[1][1], 2)) | $(size(data[2][1], 2)) | $(size(data[3][1], 2))"
			@info "Number of features: $(size(data[1][1], 1))"

			# check if a combination of parameters and seed alread exists
			if GenerativeAD.check_params(savepath, edited_parameters)
				# fit
				training_info, results = fit(data, edited_parameters)

				# save the model separately
				if training_info.model != nothing
					tagsave(joinpath(savepath, savename("model", edited_parameters, "bson", digits=5)),
						Dict("model"=>training_info.model,
							"fit_t"=>training_info.fit_t,
							"history"=>training_info.history,
							"parameters"=>edited_parameters
							), safe = true)
					training_info = merge(training_info, (model = nothing,))
				end

				# here define what additional info should be saved together with parameters, scores, labels and predict times
				save_entries = merge(training_info, (modelname = modelname, seed = seed, dataset = dataset, contamination = contamination))

				# now loop over all anomaly score funs
				for result in results
					GenerativeAD.experiment(result..., data, savepath; save_entries...)
				end
				global try_counter = max_tries + 1
			else
				@info "Model already present, trying new hyperparameters..."
				global try_counter += 1
			end
		end
	end
	(try_counter == max_tries) ? (@info "Reached $(max_tries) tries, giving up.") : nothing
end


text:
function sample_HG(rng_base::StableRNG, g::Grid, nsamples::Int)
    seeds = rand(rng_base, 1:2^32, nsamples)
    hg = Vector{Float64}(undef, nsamples)
    vs = vertices(g)
    Threads.@threads for i in 1:nsamples
        rng = StableRNG(seeds[i])
        v = rand(rng, vs)
        dv = sps(g, v)
        hg[i] = HG(g, v, dv)
    end
    hg
end

function sample_AG_NG_TG(rng_base::StableRNG, g::Grid, dxy::Int, nsamples::Int;
                         ng_only=nothing)
    seeds = rand(rng_base, 1:2^32, nsamples)
    isnothing(ng_only) && (ng_only = dxy < 2*div(7*g.n, 200))
    ag = Vector{Float64}(undef, ng_only ? 0 : nsamples)
    ng = Vector{Int64}(undef, nsamples)
    tg = Vector{Float64}(undef, ng_only ? 0 : nsamples)
    Threads.@threads for i in 1:nsamples
        rng = StableRNG(seeds[i])

        x = Vertex(rand(rng, 1+dxy:g.n-dxy), rand(rng, 1+dxy:g.n-dxy))
        dx = sps(g, x, dxy + 1)

        y = rand(rng, circle_points(g, x, dxy, dx))
        @assert y != x

        dy = sps(g, y, dxy + 1)
        ng[i] = NG(g, x, y, dx, dy)

        ng_only && continue

        ag[i] = AG(rng, g, dxy, x, dx)
        tg[i] = TG(rng, g, x, y, dxy, dx, dy)
    end
    ag, ng, tg
end

struct Sampling
    hg::Dict{Int, Vector{Float64}}
    ag::Dict{Int, Dict{Int, Vector{Float64}}}
    ng::Dict{Int, Dict{Int, Vector{Int}}}
    tg::Dict{Int, Dict{Int, Vector{Float64}}}
end

Sampling() = Sampling(Dict(), Dict(), Dict(), Dict())

function geodesic_model(ngs::Dict{Int, Vector{Int}})
    xs = Float64[]
    ys = Float64[]
    for (d, ng_d) in ngs
        for ng in ng_d
            push!(xs, d)
            push!(ys, ng)
        end
    end
    GLM.lm(@GLM.formula(log2(Y) ~ 1 + log2(X)), DataFrames.DataFrame(X=xs, Y=ys))
end

function sample!(rng::StableRNG, s::Sampling, g::Grid, nsamples1::Int, nsamples2::Int;
                 ng_only=nothing)
    println(g.n, " hg")
    hg = sample_HG(rng, g, nsamples1)
    append!(get!(()->Float64[], s.hg, g.n), hg)

    ags = get!(()->Dict{Int, Vector{Float64}}(), s.ag, g.n)
    ngs = get!(()->Dict{Int, Vector{Int}}(), s.ng, g.n)
    tgs = get!(()->Dict{Int, Vector{Float64}}(), s.tg, g.n)

    mingeo = 4
    maxgeo = 2*div(7*g.n, 40)

    for d in mingeo:2:maxgeo
        println(g.n, " ag_ng_tg: ", d, " / ", maxgeo)

        ag, ng, tg = sample_AG_NG_TG(rng, g, d, nsamples2, ng_only=ng_only)
        append!(get!(()->Float64[], ags, d), ag)
        append!(get!(()->Int[], ngs, d), ng)
        append!(get!(()->Float64[], tgs, d), tg)
    end
end

function diags_sampling(diags, nsamples1=100000, nsamples2=1000, ns=[2^i+1 for i in 6:13];
                        seed=1)
    rng = StableRNG(seed)
    s = Sampling()
    for n in ns
        g = Grid(diags, n)
        sample!(rng, s, g, nsamples1, nsamples2)
    end
    s
end

randgrid(rng::StableRNG, n::Int, p::Float64) =
    Grid(rand(rng, Float64, (n-1, n-1)) .< p, rand(rng, Float64, (n-1, n-1)) .< p)

function randmin(rng_base::StableRNG, n::Int)
    seed = rand(rng_base, 1:2^32)
    Optim.optimize(0.0, 1.0) do p
        rng = StableRNG(seed)
        g = randgrid(rng, n, p)
        hg = sample_HG(rng, g, 128)
        Statistics.mean(hg .^ 2)
    end
end

function rand_sampling(nreplicates=25, nsamples1=10000, nsamples2=100,
                       ns=[2^i+1 for i in 6:13]; seed=2)
    rng = StableRNG(seed)
    s = Sampling()
    for n in ns
        res = randmin(rng, n)
        println(res)
        p = res.minimizer
        for _ in 1:nreplicates
            g = randgrid(rng, n, p)
            sample!(rng, s, g, nsamples1, nsamples2)
        end
    end
    s
end

function disordered_NG_sampling(diags, nreplicates=25, nsamples=100,
                                ns=[2^i+1 for i in 6:13]; seed=3)
    rng = StableRNG(seed)
    s = Sampling()
    for n in ns
        g = Grid(diags, n)
        for i in 1:nreplicates
            println(i, " / ", nreplicates)
            sample!(rng, s, disorder(rng, g), 0, nsamples, ng_only=true)
        end
    end
    s
end

max_stars_repo_name: moshelooks/eugrid

max_issues_repo_name: moshelooks/eugrid

max_forks_repo_name: moshelooks/eugrid


text:
function SlideJoint(a::Body, b::Body, anchorA::Vect, anchorB::Vect, min::Real, max::Real)
    SlideJoint(ccall(dlsym(libchipmunk, :cpSlideJointNew), Ptr{Void}, (Ptr{Void}, Ptr{Void}, Vect, Vect, Cdouble, Cdouble,), a.ptr, b.ptr, anchorA, anchorB, min, max))
end

function get_anchorA(joint::SlideJoint)
    ccall(dlsym(libchipmunk, :cpSlideJointGetAnchorA), Vect, (Ptr{Void},), joint.ptr)
end

function set_anchorA(joint::SlideJoint, anchor::Vect)
    ccall(dlsym(libchipmunk, :cpSlideJointSetAnchorA), Void, (Ptr{Void}, Vect,), joint.ptr, anchor)
end

function get_anchorB(joint::SlideJoint)
    ccall(dlsym(libchipmunk, :cpSlideJointGetAnchorB), Vect, (Ptr{Void},), joint.ptr)
end

function set_anchorB(joint::SlideJoint, anchor::Vect)
    ccall(dlsym(libchipmunk, :cpSlideJointSetAnchorB), Void, (Ptr{Void}, Vect,), joint.ptr, anchor)
end

function get_min(joint::SlideJoint)
    ccall(dlsym(libchipmunk, :cpSlideJointGetMin), Cdouble, (Ptr{Void},), joint.ptr)
end

function set_min(joint::SlideJoint, min::Real)
    ccall(dlsym(libchipmunk, :cpSlideJointSetMin), Void, (Ptr{Void}, Cdouble,), joint.ptr, min)
end

function get_max(joint::SlideJoint)
    ccall(dlsym(libchipmunk, :cpSlideJointGetMin), Cdouble, (Ptr{Void},), joint.ptr)
end

function set_max(joint::SlideJoint, max::Real)
    ccall(dlsym(libchipmunk, :cpSlideJointSetMin), Void, (Ptr{Void}, Cdouble,), joint.ptr, max)
end

export SlideJoint, get_anchorA, set_anchorA, get_anchorB, set_anchorB, get_min, set_min,
get_max, set_max


text:
using JSON

function part1(input)
    data = JSON.parse(readchomp(input))
    return sum_numbers(data)
end

function part2(input)
    data = JSON.parse(readchomp(input))
    return sum_numbers2(data)
end

sum_numbers(data::Dict) = sum(sum_numbers.(values(data)))
sum_numbers(data::Vector) = sum(sum_numbers.(data))
sum_numbers(data::Number) = data
sum_numbers(data::AbstractString) = 0

function sum_numbers2(data::Dict)
    if any(==("red"), values(data))
        return 0
    end
    return sum(sum_numbers2.(values(data)))
end
sum_numbers2(data::Vector) = sum(sum_numbers2.(data))
sum_numbers2(data::Number) = data
sum_numbers2(data::AbstractString) = 0


text:
Base.@kwdef struct Options
    indent_size::Int = 4
    max_margin::Int = 92
    always_for_in::Bool = false
    whitespace_typedefs::Bool = false
    whitespace_ops_in_indices::Bool = false
    remove_extra_newlines::Bool = false
    import_to_using::Bool = false
    pipe_to_function_call::Bool = false
    short_to_long_function_def::Bool = false
    always_use_return::Bool = false
    whitespace_in_kwargs::Bool = true
    annotate_untyped_fields_with_any::Bool = true
    format_docstrings::Bool = false
end


text:
export iauPn00
"""
Precession-nutation, IAU 2000 model:  a multi-purpose function,
supporting classical (equinox-based) use directly and CIO-based
use indirectly.

This function is part of the International Astronomical Union's
SOFA (Standards Of Fundamental Astronomy) software collection.

Status:  support function.

Given:
   date1,date2  double          TT as a 2-part Julian Date (Note 1)
   dpsi,deps    double          nutation (Note 2)

Returned:
   epsa         double          mean obliquity (Note 3)
   rb           double[3][3]    frame bias matrix (Note 4)
   rp           double[3][3]    precession matrix (Note 5)
   rbp          double[3][3]    bias-precession matrix (Note 6)
   rn           double[3][3]    nutation matrix (Note 7)
   rbpn         double[3][3]    GCRS-to-true matrix (Note 8)

Notes:

   1. The TT date date1+date2 is a Julian Date, apportioned in any
   convenient way between the two arguments.  For example,
   JD(TT)=2450123.7 could be expressed in any of these ways,
   among others:

         date1          date2

      2450123.7           0.0       (JD method)
      2451545.0       -1421.3       (J2000 method)
      2400000.5       50123.2       (MJD method)
      2450123.5           0.2       (date & time method)

   The JD method is the most natural and convenient to use in
   cases where the loss of several decimal digits of resolution
   is acceptable.  The J2000 method is best matched to the way
   the argument is handled internally and will deliver the
   optimum resolution.  The MJD method and the date & time methods
   are both good compromises between resolution and convenience.

   2. The caller is responsible for providing the nutation components;
   they are in longitude and obliquity, in radians and are with
   respect to the equinox and ecliptic of date.  For high-accuracy
   applications, free core nutation should be included as well as
   any other relevant corrections to the position of the CIP.

   3. The returned mean obliquity is consistent with the IAU 2000
   precession-nutation models.

   4. The matrix rb transforms vectors from GCRS to J2000.0 mean
   equator and equinox by applying frame bias.

   5. The matrix rp transforms vectors from J2000.0 mean equator and
   equinox to mean equator and equinox of date by applying
   precession.

   6. The matrix rbp transforms vectors from GCRS to mean equator and
   equinox of date by applying frame bias then precession.  It is
   the product rp x rb.

   7. The matrix rn transforms vectors from mean equator and equinox of
   date to true equator and equinox of date by applying the nutation
   (luni-solar + planetary).

   8. The matrix rbpn transforms vectors from GCRS to true equator and
   equinox of date.  It is the product rn x rbp, applying frame
   bias, precession and nutation in that order.

   9. It is permissible to re-use the same array in the returned
   arguments.  The arrays are filled in the order given.

Called:
   iauPr00      IAU 2000 precession adjustments
   iauObl80     mean obliquity, IAU 1980
   iauBp00      frame bias and precession matrices, IAU 2000
   iauCr        copy r-matrix
   iauNumat     form nutation matrix
   iauRxr       product of two r-matrices

Reference:

   Capitaine, N., Chapront, J., Lambert, S. and Wallace, P.,
   "Expressions for the Celestial Intermediate Pole and Celestial
   Ephemeris Origin consistent with the IAU 2000A precession-
   nutation model", Astron.Astrophys. 400, 1145-1154 (2003)

   n.b. The celestial ephemeris origin (CEO) was renamed "celestial
         intermediate origin" (CIO) by IAU 2006 Resolution 2.

This revision:  2013 June 18

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# void iauPn00(double date1, double date2, double dpsi, double deps,
#              double *epsa,
#              double rb[3][3], double rp[3][3], double rbp[3][3],
#              double rn[3][3], double rbpn[3][3])

function iauPn00(date1::Real, date2::Real, dpsi::Real, deps::Real)
   # Allocate return values
   ref_epsa = Ref{Float64}(0.0)
   rb       = zeros(Float64, 3, 3)
   rp       = zeros(Float64, 3, 3)
   rbp      = zeros(Float64, 3, 3)
   rn       = zeros(Float64, 3, 3)
   rbpn     = zeros(Float64, 3, 3)

   ccall((:iauPn00, libsofa_c), Cvoid,
        (Cdouble, Cdouble, Cdouble, Cdouble,
         Ref{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble},
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}),
         convert(Float64, date1),
         convert(Float64, date2),
         convert(Float64, dpsi),
         convert(Float64, deps),
         ref_epsa, rb, rp, rbp, rn, rbpn)

   return ref_epsa[], SMatrix{3,3}(rb'), SMatrix{3,3}(rp'),
      SMatrix{3,3}(rbp'), SMatrix{3,3}(rn'), SMatrix{3,3}(rbpn')
end


text:
using TimeZones: TimeZone, localzone, parse_tz_format

# Parse the TZ environment variable format
# Should mirror the behaviour of running:
# `date -u; env TZ="..." date`

@test_throws ArgumentError parse_tz_format("AB")
@test_throws ArgumentError parse_tz_format("+")
@test_throws ArgumentError parse_tz_format("12")

@test parse_tz_format("") == FixedTimeZone("UTC", 0, 0)
@test parse_tz_format("ABC") == FixedTimeZone("ABC", 0, 0)
@test parse_tz_format("ABC+") == FixedTimeZone("ABC", 0, 0)
@test parse_tz_format("ABC-") == FixedTimeZone("ABC", 0, 0)

@test parse_tz_format("ABC1") == FixedTimeZone("ABC", -1 * 3600, 0)
@test parse_tz_format("ABC+1") == FixedTimeZone("ABC", -1 * 3600, 0)
@test parse_tz_format("ABC-1") == FixedTimeZone("ABC", 1 * 3600, 0)

@test parse_tz_format("ABC-24")  == FixedTimeZone("ABC", 24 * 3600, 0)
@test parse_tz_format("ABC-25")  == FixedTimeZone("ABC", 24 * 3600, 0)
@test parse_tz_format("ABC-100") == FixedTimeZone("ABC", 24 * 3600, 0)

@test parse_tz_format("ABC-00:59")  == FixedTimeZone("ABC", 59 * 60, 0)
@test parse_tz_format("ABC-00:99")  == FixedTimeZone("ABC", 59 * 60, 0)
@test parse_tz_format("ABC-00:100") == FixedTimeZone("ABC", 59 * 60, 0)

@test parse_tz_format("ABC-00:00:59")  == FixedTimeZone("ABC", 59, 0)
@test parse_tz_format("ABC-00:00:99")  == FixedTimeZone("ABC", 59, 0)
@test parse_tz_format("ABC-00:00:100") == FixedTimeZone("ABC", 59, 0)

@test_throws ArgumentError parse_tz_format("ABC+00:")  # Terminal test is valid and equal to "ABC"
@test_throws ArgumentError parse_tz_format("ABC+12:")

# Ensure that the current system's local time zone is supported. If this test fails make
# sure to report it as an issue.
@test isa(localzone(), TimeZone)

if Sys.islinux()
    # Bad TZ environment variable formats
    withenv("TZ" => "+12:00") do
        @test_throws ArgumentError localzone()
    end
    withenv("TZ" => "Europe/Warsaw") do
        @test_throws ArgumentError localzone()
    end

    # Currently unsupported TZ environment variable formats
    withenv("TZ" => "NZST-12:00:00NZDT-13:00:00,M10.1.0,M3.3.0") do
        @test_throws ArgumentError localzone()
    end

    # Valid TZ formats
    withenv("TZ" => "") do
        @test localzone() == utc
    end
    withenv("TZ" => "UTC") do
        @test localzone() == utc
    end

    # Absolute filespec
    warsaw_path = joinpath(TZFILE_DIR, "Europe", "Warsaw")
    warsaw_from_file = open(warsaw_path) do f
        TimeZones.read_tzfile(f, "local")
    end
    withenv("TZ" => ":" * abspath(warsaw_path)) do
        @test localzone() == warsaw_from_file
    end

    # Relative filespec
    warsaw = TimeZone("Europe/Warsaw")
    withenv("TZ" => ":Europe/Warsaw") do
        @test localzone() == warsaw
    end

    # Set TZDIR and use time zone unrecognized by TimeZone
    @test_throws ArgumentError TimeZone("Etc/UTC")
    utc = open(joinpath(TZFILE_DIR, "Etc", "UTC")) do f
        TimeZones.read_tzfile(f, "Etc/UTC")
    end
    withenv("TZ" => ":Etc/UTC", "TZDIR" => TZFILE_DIR) do
        @test localzone() == utc
    end

    # Attempt to use system installed time zone files. On some minimal systems no time zone
    # information will be available.
    @test_throws ArgumentError TimeZone("Etc/GMT-9")
    gmt_minus_9 = FixedTimeZone("Etc/GMT-9", 9 * 3600)
    withenv("TZ" => ":Etc/GMT-9") do
        try
            tz = localzone()
            @test tz == gmt_minus_9
        catch err
            msg = sprint(showerror, err)
            if isa(err, ErrorException) && msg == "unable to locate tzfile: Etc/GMT-9"
                warn("Skipping test: missing system time zone information")
            else
                rethrow()
            end
        end
    end

    # Unable to locate time zone on system
    withenv("TZ" => ":") do
        @test_throws ErrorException localzone()
    end
    withenv("TZ" => ":Etc/Foo") do
        @test_throws ErrorException localzone()
    end
end


text:
const REGISTER_SIZE = 16
const REGISTER_COUNT = 16

const FMA_FAST = false
const FMA = false
const AVX2 = false
const AVX512F = false
const AVX512ER = false
const AVX512PF = false
const AVX512VL = false
const AVX512BW = false
const AVX512DQ = false
const AVX512CD = false
const SIMD_INTEGER_REGISTER_SIZE = REGISTER_SIZE


text:
abstract type BayesFactorModel <: ProbabilisticModel end

"""
NormalEffectSize <: EffectSizeModel

A standard effect size model has two hypotheses: ``H_0``(null) an ``H_1``(alternative):

1. ``H_0``: ``\\mu = m_0``
2. ``H_1``: ``\\mu ≠ m_0``

with the population mean ``\\mu`` and pre-specified standard deviation ``\\sigma``. We want to test whether
``\\mu`` is equal to ``\\mu_0`` or not.

The prior of the standard effect size is

``
\\delta | H_1 \\sim \\text{Normal}(0, \\sigma_0^2)
``

where ``\\delta`` is the standard effect size. When ``\\sigma_0 = 1``the prior is called
*unit-information* prior.

The standard effect size ``\\delta`` is defined as

``
\\delta = \\frac{\\mu - \\mu_0}{\\sigma}.
``

In practice, the standard deviations are unknown but in large sample scenario we assume
they are known and use their estimates.

## Fileds

- `μ0`: mean of null hypothesis.
- `σ0`: The prior standard deviation of the effect size.
- `p0`: prior belief of ``H_0``. This is used to calculate the prior odds. Default ``0.5``.

## Methods

```julia
bayesfactor(model, stats)    # calculate Bayes factor from one group statistics
bayesfactor(model, twostats) # calculate Bayes factor from two group's statistics
```

## References

1. [Chapter 5 hypothesis Testing with Normal Populations](https://statswithr.github.io/book/hypothesis-testing-with-normal-populations.html)
   in *An Introduction to Bayesian Thinking*.
2. Deng, Alex, Jiannan Lu, and Shouyuan Chen. "Continuous monitoring of A/B tests without pain:
   Optional stopping in Bayesian testing." 2016 IEEE international conference on data science
   and advanced analytics (DSAA). IEEE, 2016.
"""
struct NormalEffectSize <: BayesFactorModel
    μ0::Float64
    σ0::Float64
    p0::Float64
end

NormalEffectSize(μ0, σ0) = NormalEffectSize(μ0, σ0, 0.5)
NormalEffectSize(;μ0=0.0, σ0=1.0, p0=0.5) = NormalEffectSize(μ0, σ0, p0)

function bayesfactor(model::NormalEffectSize, stats::NormalStatistics)
    n = stats.n
    σ0 = model.σ0
    δ = effectsize(stats, μ0=model.μ0)
    bf10 = pdf(Normal(0, sqrt(σ0^2+1/n)), δ)/pdf(Normal(0, sqrt(1/n)), δ)
    priorodds = _priorodds(model)
    return priorodds*bf10
end

function bayesfactor(model::NormalEffectSize, twostats::TwoNormalStatistics)
    stats = merge(twostats)
    return bayesfactor(model, stats)
end

"""
    StudentTEffectSize <: BayesFactorModel

A model with Bayes factor from the Student's t distributions.
We have a standard effect size model has two hypotheses: ``H_0``(null) an ``H_1``(alternative):

1. ``H_0``: ``\\mu = m_0``
2. ``H_1``: ``\\mu ≠ m_0``

The model uses the Jeffreys-Zellener-Siow (JZS) prior.
More specifically, we use a Cauchy prior on ``\\mu`` for ``H_1``

``
\\mu | \\sigma^2 \\sim \\text{Cauchy}(0, r^2 \\sigma^2)
``

and a Jeffrey's prior on ``\\sigma``:

``
p(\\sigma^2) \\propto \\frac{1}{\\sigma2}
``

for both ``H_0`` and ``H_1``.

## Fields

- ``r``: Prior standard deviation of the effect size.
- ``rtol``: Numerical tolerence fo the ``quadgk`` function used in the denominator calculation.
- ``p0``: prior belief of ``H_0``. This is used to calculate the prior odds. Default ``0.5``.

## References

- Rouder, J. N., Speckman, P. L., Sun, D., Morey, R. D., & Iverson, G. (2009). Bayesian t tests
  for accepting and rejecting the null hypothesis. Psychonomic bulletin & review, 16(2), 225-237.

"""
struct StudentTEffectSize <: BayesFactorModel
    r::Real
    rtol::Real
    p0::Real
end

StudentTEffectSize(;r=0.707, rtol=1e-8, p0=0.5) = StudentTEffectSize(r, rtol, p0)

function bayesfactor(model::StudentTEffectSize, stats::StudentTStatistics)
    t = stats.t
    v = stats.dof
    n = stats.n

    r = model.r
    rtol = model.rtol

    numerator = (1 + t^2/v)^(-(v+1)/2)
    denominator, _ = quadgk(
        g -> (1+n*g*r^2)^(-0.5) * (1+t^2/((1+n*g*r^2)*v))^(-(v+1)/2)
            * (2π)^(-0.5) * g^(-1.5) * exp(-1/(2*g)),
        0, Inf,
        rtol=rtol)

    # we take invserse to get b_1_0, instead of b_0_1
    priorodds = _priorodds(model)
    return priorodds * denominator/numerator
end

function _priorodds(model::BayesFactorModel)
    p0 = model.p0
    return (1-p0)/p0
end


text:
module TaskDAGAnalyzersTests

include("test_examples.jl")

end  # module TaskDAGAnalyzersTests


text:
using AbstractMCMC
using DynamicPPL
using Random
using Turing

export StaticConditional

"""
    StaticConditional(m, sym)

A "pseudo-sampler" to use automatically extracted conditionals within `Gibbs`.  This works by
calculating the conditionals once on the model `m` for variable `sym`, and reusing them in every
sampling step; hence, the trace of the model must keep the same at every step.

`StaticConditional(m, :x)` will sample the Gibbs conditional of discrete variables `x` with a
defined support, or all such variables subsumed by `x` (i.e., `x[1]`, `x[2]`, ...).

# Examples

```julia

@model function test(x)
    w ~ Dirichlet(2, 1.0)
    p ~ DiscreteNonParametric([0.3, 0.7], w)
    x ~ Bernoulli(p)
end

m = test(false)
cond_p = StaticConditional(m, :p)
sample(m, Gibbs(cond_p, MH(:w)), 10)
```
"""
struct StaticConditional{S, TCond}
    conditionals::TCond

    function StaticConditional(model::Turing.Model, sym::Symbol)
        vn = VarName(sym)
        graph = trackdependencies(model)
        conds = conditionals(graph, vn)
        return new{sym, typeof(conds)}(conds)
    end
end

DynamicPPL.getspace(::StaticConditional{S}) where {S} = (S,)
DynamicPPL.alg_str(::StaticConditional) = "StaticConditional"
Turing.isgibbscomponent(::StaticConditional) = true

function Turing.Sampler(
    alg::StaticConditional,
    model::Turing.Model,
    s::DynamicPPL.Selector=DynamicPPL.Selector()
)
    state = Turing.Inference.SamplerState(Turing.VarInfo(model))
    return Turing.Sampler(alg, Dict{Symbol, Any}(), s, state)
end

function AbstractMCMC.step!(
    rng::AbstractRNG,
    model::Turing.Model,
    spl::Turing.Sampler{<:StaticConditional},
    N::Integer,
    transition;
    kwargs...
)
    vi = spl.state.vi
    values = sampled_values(model, vi)
    for (vn, cond) in spl.alg.conditionals
        updated = rand(rng, cond(values))
        vi[vn] = [updated;]
    end

    model(spl.state.vi)

    return transition
end

function sampled_values(model::Turing.Model, vi::DynamicPPL.UntypedVarInfo)
    assumed_vars = Dict(vn => vi[vn] for vn in keys(vi))
    observed_vars = Dict(VarName(arg) => value for (arg, value) in pairs(model.args))
    return merge(assumed_vars, observed_vars)
end

function sampled_values(model::Turing.Model, vi::DynamicPPL.TypedVarInfo)
    assumed_vars = Dict(vn => vi[vn]
                        for name_meta in values(vi.metadata)
                        for vn in name_meta.vns)
    observed_vars = Dict(VarName(arg) => value for (arg, value) in pairs(model.args))
    return merge(assumed_vars, observed_vars)
end


text:
module DotProductTest

using Test
using ForneyLab
using ForneyLab: outboundType, isApplicable
using ForneyLab: SPDotProductOutNPG, SPDotProductOutNGP, SPDotProductIn1GNP, SPDotProductIn2GPN

#-------------
# Update rules
#-------------

@testset "SPDotProductOutNPG" begin
    @test SPDotProductOutNPG <: SumProductRule{DotProduct}
    @test outboundType(SPDotProductOutNPG) == Message{GaussianMeanVariance}
    @test isApplicable(SPDotProductOutNPG, [Nothing, Message{PointMass}, Message{Gaussian}])

    @test ruleSPDotProductOutNPG(nothing, Message(Multivariate, PointMass, m=[4.0]), Message(Multivariate, GaussianMeanVariance, m=[3.0], v=mat(2.0))) == Message(Univariate, GaussianMeanVariance, m=12.0, v=32.0)
end

@testset "SPDotProductOutNGP" begin
    @test SPDotProductOutNGP <: SumProductRule{DotProduct}
    @test outboundType(SPDotProductOutNGP) == Message{GaussianMeanVariance}
    @test isApplicable(SPDotProductOutNGP, [Nothing, Message{Gaussian}, Message{PointMass}])

    @test ruleSPDotProductOutNGP(nothing, Message(Multivariate, GaussianMeanVariance, m=[3.0], v=mat(2.0)), Message(Multivariate, PointMass, m=[4.0])) == Message(Univariate, GaussianMeanVariance, m=12.0, v=32.0)
end

@testset "SPDotProductIn1GNP" begin
    @test SPDotProductIn1GNP <: SumProductRule{DotProduct}
    @test outboundType(SPDotProductIn1GNP) == Message{GaussianWeightedMeanPrecision}
    @test isApplicable(SPDotProductIn1GNP, [Message{Gaussian}, Nothing, Message{PointMass}])

    @test ruleSPDotProductIn1GNP(Message(Univariate, GaussianWeightedMeanPrecision, xi=3.0, w=2.0), nothing, Message(Multivariate, PointMass, m=[4.0])) == Message(Multivariate, GaussianWeightedMeanPrecision, xi=[12.0], w=mat(32.000000000001))
end

@testset "SPDotProductIn2GPN" begin
    @test SPDotProductIn2GPN <: SumProductRule{DotProduct}
    @test outboundType(SPDotProductIn2GPN) == Message{GaussianWeightedMeanPrecision}
    @test isApplicable(SPDotProductIn2GPN, [Message{Gaussian}, Message{PointMass}, Nothing])

    @test ruleSPDotProductIn2GPN(Message(Univariate, GaussianWeightedMeanPrecision, xi=3.0, w=2.0), Message(Multivariate, PointMass, m=[4.0]), nothing) == Message(Multivariate, GaussianWeightedMeanPrecision, xi=[12.0], w=mat(32.000000000001))
end

end # module


text:
using ClusterManagers
using Distributed

OnCluster = true #set to false to run locally
addWorkers = true #set to false to run serially
println("OnCluster = $(OnCluster)")

# Current number of workers
#--------------------------
currentWorkers = nworkers()
println("Initial number of workers = $(currentWorkers)")

# Increase the number of workers available
#-----------------------------------------
maxNumberWorkers = 10
if addWorkers == true
	if OnCluster == true
	  addprocs(SlurmManager(maxNumberWorkers))
	else
	  addprocs(maxNumberWorkers)
	end
end

# Sanity checks
#-------------
hosts = []
pids = []
for i in workers()
	host, pid = fetch(@spawnat i (gethostname(), getpid()))
	println("Hello I am worker $(i), my host is $(host)")
	push!(hosts, host)
	push!(pids, pid)
end

currentWorkers = nworkers()
println("Number of workers = $(currentWorkers)")

using Plots
using LaTeXStrings
using ParallelDataTransfer

@everywhere using MSM
@everywhere using DataStructures
@everywhere using OrderedCollections
@everywhere using Distributions
@everywhere using Random
@everywhere using DataStructures
@everywhere using Statistics
@everywhere using LinearAlgebra

# Generate simulated data
Random.seed!(1234)  #for replicability reasons
T = 100000          #number of periods
P = 2               #number of dependent variables
beta0 = rand(P)     #choose true coefficients by drawing from a uniform distribution on [0,1]
alpha0 = rand(1)[]  #intercept
theta0 = 0.0        #coefficient to create serial correlation in the error terms
println("True intercept = $(alpha0)")
println("True coefficient beta0 = $(beta0)")
println("Serial correlation coefficient theta0 = $(theta0)")

# Generation of error terms
# row = individual dimension
# column = time dimension
U = zeros(T)
d = Normal()
U[1] = rand(d, 1)[] #first error term
# loop over time periods
for t = 2:T
    U[t] = rand(d, 1)[] + theta0*U[t-1]
end

# Let's simulate the dependent variables x_t
x = zeros(T, P)

d = Uniform(0, 5)
for p = 1:P
    x[:,p] = rand(d, T)
end

# Let's calculate the resulting y_t
y = zeros(T)

for t=1:T
    y[t] = alpha0 + x[t,1]*beta0[1] + x[t,2]*beta0[2] + U[t]
end

# Send simulated data to workers
sendto(workers(), y=y)
sendto(workers(), x=x)

# Visualize data
p1 = scatter(x[1:100,1], y[1:100], xlabel = "x1", ylabel = "y", legend=:none, smooth=true)
p2 = scatter(x[1:100,2], y[1:100], xlabel = "x2", ylabel = "y", legend=:none, smooth=true)
plot(p1, p2)

# Define locally
optionsSMM = MSMOptions(maxFuncEvals=2000, globalOptimizer = :dxnes, localOptimizer = :NelderMead)
myProblem = MSMProblem(options = optionsSMM);

# Send to workers
sendto(workers(), optionsSMM=optionsSMM)
sendto(workers(), myProblem=myProblem)

# Priors
dictPriors = OrderedDict{String,Array{Float64,1}}()
dictPriors["alpha"] = [0.5, 0.001, 1.0]
dictPriors["beta1"] = [0.5, 0.001, 1.0]
dictPriors["beta2"] = [0.5, 0.001, 1.0]
sendto(workers(), dictPriors=dictPriors)

# Empirical moments
dictEmpiricalMoments = OrderedDict{String,Array{Float64,1}}()
dictEmpiricalMoments["mean"] = [mean(y)] #informative on the intercept
dictEmpiricalMoments["mean^2"] = [mean(y.^2)] #informative on the intercept
dictEmpiricalMoments["mean_x1y"] = [mean(x[:,1] .* y)] #informative on betas
dictEmpiricalMoments["mean_x2y"] = [mean(x[:,2] .* y)] #informative on betas
dictEmpiricalMoments["mean_x1y^2"] = [mean((x[:,1] .* y).^2)] #informative on betas
dictEmpiricalMoments["mean_x2y^2"] = [mean((x[:,2] .* y).^2)] #informative on betas

W = Matrix(1.0 .* I(length(dictEmpiricalMoments)))#initialization
#Special case: diagonal matrix
#(you may choose something else)
for (indexMoment, k) in enumerate(keys(dictEmpiricalMoments))
    W[indexMoment,indexMoment] = 1.0/(dictEmpiricalMoments[k][1])^2
end

# Send to workers
sendto(workers(), dictPriors=dictPriors)
sendto(workers(), dictEmpiricalMoments=dictEmpiricalMoments)
sendto(workers(), W=W)

@everywhere set_priors!(myProblem, dictPriors)
@everywhere set_empirical_moments!(myProblem, dictEmpiricalMoments)
@everywhere set_weight_matrix!(myProblem, W)

# x[1] corresponds to the intercept, x[2] corresponds to beta1, x[3] corresponds to beta2
@everywhere function functionLinearModel(x; uniform_draws::Array{Float64,1}, simX::Array{Float64,2}, nbDraws::Int64 = length(uniform_draws), burnInPerc::Int64 = 0)
    T = nbDraws
    P = 2       #number of dependent variables

    alpha = x[1]
    beta = x[2:end]
    theta = 0.0     #coefficient to create serial correlation in the error terms

    # Creation of error terms
    # row = individual dimension
    # column = time dimension
    U = zeros(T)
    d = Normal()
    # Inverse cdf (i.e. quantile)
    gaussian_draws = quantile.(d, uniform_draws)
    U[1] = gaussian_draws[1] #first error term

    # loop over time periods
    for t = 2:T
        U[t] = gaussian_draws[t] + theta*U[t-1]
    end

    # Let's calculate the resulting y_t
    y = zeros(T)

    for t=1:T
        y[t] = alpha + simX[t,1]*beta[1] + simX[t,2]*beta[2] + U[t]
    end

	# Get rid of the burn-in phase:
    #------------------------------
    startT = max(1, Int(nbDraws * (burnInPerc / 100)))

    # Moments:
    #---------
    output = OrderedDict{String,Float64}()
    output["mean"] = mean(y[startT:nbDraws])
    output["mean^2"] = mean(y[startT:nbDraws].^2)
    output["mean_x1y"] = mean(simX[startT:nbDraws,1] .* y[startT:nbDraws])
    output["mean_x2y"] = mean(simX[startT:nbDraws,2] .* y[startT:nbDraws])
    output["mean_x1y^2"] = mean((simX[startT:nbDraws,1] .* y[startT:nbDraws]).^2)
    output["mean_x2y^2"] = mean((simX[startT:nbDraws,2] .* y[startT:nbDraws]).^2)

    return output
end

# Let's freeze the randomness during the minimization
d_Uni = Uniform(0,1)
nbDraws = 2*T #number of draws in the simulated data
uniform_draws = rand(d_Uni, nbDraws)
simX = zeros(length(uniform_draws), 2)
d = Uniform(0, 5)
for p = 1:2
  simX[:,p] = rand(d, length(uniform_draws))
end

# Send to workers
sendto(workers(), simX=simX)
sendto(workers(), uniform_draws=uniform_draws)

# Construct the objective function everywhere
@everywhere set_simulate_empirical_moments!(myProblem, x -> functionLinearModel(x, uniform_draws = uniform_draws, simX = simX))
@everywhere construct_objective_function!(myProblem)

# Safety check: value on the master node == values on slave nodes?
using Test
val_local = myProblem.objective_function(ones(3)); #local execution
val_workers = [];
for w in workers() #Execution on workers
    push!(val_workers, @fetchfrom w myProblem.objective_function(ones(3)))
end
for (wIndex, w) in enumerate(workers())
    @test abs(val_local - val_workers[wIndex]) < 10e-10
end

println("Global Algorithm")
# Choose a global optimizer that supports parallel evaluations (e.g. xnes or dxnes)
# (see the documentation: https://github.com/robertfeldt/BlackBoxOptim.jl)
msm_optimize!(myProblem, verbose = false)

# Access the results using best_candidate
minimizer = msm_minimizer(myProblem)
minimum_val = msm_minimum(myProblem)

println("Minimum objective function = $(minimum_val)")
println("Estimated value for alpha = $(minimizer[1]). True value for beta1 = $(alpha0[1]) \n")
println("Estimated value for beta1 = $(minimizer[2]). True value for beta1 = $(beta0[1]) \n")
println("Estimated value for beta2 = $(minimizer[3]). True value for beta2 = $(beta0[2]) \n")

println("Multistart Algorithm")
# Start several local optimization algorithms in parallel
# Choose algorithms from the package Optim.jl (https://github.com/JuliaNLSolvers/Optim.jl)
# The "global" mininimum is the minimum of the local minima.
msm_multistart!(myProblem, nums = nworkers(), verbose = false)

minimizer_multistart = msm_multistart_minimizer(myProblem)
minimum_multistart = msm_multistart_minimum(myProblem)

println("Minimum objective function = $(minimum_multistart)")
println("Estimated value for alpha = $(minimizer_multistart[1]). True value for beta1 = $(alpha0[1]) \n")
println("Estimated value for beta1 = $(minimizer_multistart[2]). True value for beta1 = $(beta0[1]) \n")
println("Estimated value for beta2 = $(minimizer_multistart[3]). True value for beta2 = $(beta0[2]) \n")

# Empirical Series
#-----------------
X = zeros(T, length(dictEmpiricalMoments))
X[:,1] = y
X[:,2] = y.^2
X[:,3] = (x[:,1] .* y)
X[:,4] = (x[:,2] .* y)
X[:,5] = (x[:,1] .* y).^2
X[:,6] = (x[:,2] .* y).^2

# "Distance Matrix" (see Duffie and Singleton, 1993)
Sigma0 = cov(X)

# Heatmap to visualize correlation
xs = [string("x", i) for i = 1:length(dictEmpiricalMoments)]
ys = [string("x", i) for i = 1:length(dictEmpiricalMoments)]
z = cor(X)
hh = heatmap(xs, ys, z, aspect_ratio = 1)

set_Sigma0!(myProblem, Sigma0)
# nbDraws = number of draws in the simulated data
# To decrease standard errors, increase nbDraws
calculate_Avar!(myProblem, minimizer_multistart, tau = T/nbDraws)

df = summary_table(myProblem, minimizer_multistart, T, 0.05)
println(df)

# Check the rank condition
# Local identification requires D to be full column rank (in a neighborhood of the solution)
D = calculate_D(myProblem, minimizer_multistart)
println("number of parameters: $(size(D,2))")
println("rank of D is: $(rank(D))")

# Slices
vXGrid, vYGrid = msm_slices(myProblem, minimizer_multistart, nbPoints = 7);

p1 = plot(vXGrid[:, 1],vYGrid[:, 1],title = L"\alpha", label = "",linewidth = 3, xrotation = 45)
plot!(p1, [minimizer_multistart[1]], seriestype = :vline, label = "",linewidth = 1)
p2 = plot(vXGrid[:, 2],vYGrid[:, 2],title = L"\beta_1", label = "",linewidth = 3, xrotation = 45)
plot!(p2, [minimizer_multistart[2]], seriestype = :vline, label = "",linewidth = 1)
p3 = plot(vXGrid[:, 3],vYGrid[:, 3],title = L"\beta_2", label = "",linewidth = 3, xrotation = 45)
plot!(p3, [minimizer_multistart[3]], seriestype = :vline, label = "",linewidth = 1)
plot_combined = plot(p1, p2, p3)
display(plot_combined)

# Compare results with GLM
using DataFrames, GLM
data = DataFrame(x1=x[:,1], x2=x[:,2], y= y[:]);
ols = lm(@formula(y ~ x1 + x2), data)
println(ols)


text:
"""
	grdinterpolate(cmd0::String="", arg1=nothing, arg2=nothing; kwargs...)

Interpolate a 3-D cube, 2-D grids or 1-D series from a 3-D data cube or stack of 2-D grids.

Full option list at [`grdinterpolate`]($(GMTdoc)/grdinterpolate.html)

Parameters
----------

- **D** | **meta** | **metadata** :: [Type => Str | NamedTuple]

    Give one or more combinations for values xname, yname, zname (3rd dimension in cube), and dname
    (data value name) and give the names of those variables and in square bracket their units
    ($(GMTdoc)grdinterpolate.html#d)
- **E** | **crossection** :: [Type => Str | GMTdtaset | NamedTuple]

    Specify a crossectinonal profile via a file or from specified line coordinates and modifiers. If a file,
    it must be contain a single segment with either lon lat or lon lat dist records. These must be equidistant.
    ($(GMTdoc)grdinterpolate.html#e)
- **F** | **interp_type** :: [Type => Str]   ``Arg = l|a|c|n[+1|+2]``

    Choose from l (Linear), a (Akima spline), c (natural cubic spline), and n (no interpolation:
    nearest point) [Default is Akima].
- **G** | **outfile** | **outgrid** :: [Type => Str]

    Output file name. If `range` only selects a single layer then the data cube collapses to a regular 2-D grid file
    ($(GMTdoc)grdinterpolate.html#g)
- $(GMT.opt_R)
- **S** | **stack** :: [Type => Str | Tuple | Dataset]	`Arg = x/y|pointfile[+hheader]`

    Rather than compute gridded output, create tile/spatial series through the stacked grids at the given point (x/y)
    or the list of points in pointfile.
    ($(GMTdoc)grdinterpolate.html#s)
- **T** | **range** :: [Type => Str]			`Arg = [min/max/]inc[+i|n] |-Tfile|list`

    Make evenly spaced time-steps from min to max by inc [Default uses input times].
    ($(GMTdoc)grdinterpolate.html#t)
- **Z** | **levels** :: [Type => range]			`Arg = [levels]`

    The `levels` may be specified the same way as in `range`. If not given then we default to an integer
    levels array starting at 0.
    ($(GMTdoc)grdinterpolate.html#z)
- $(GMT.opt_V)
- $(GMT.opt_bi)
- $(GMT.opt_bo)
- $(GMT.opt_di)
- $(GMT.opt_e)
- $(GMT.opt_f)
- $(GMT.opt_g)
- $(GMT.opt_h)
- $(GMT.opt_i)
- $(GMT.opt_n)
- $(GMT.opt_o)
- $(GMT.opt_q)
- $(GMT.opt_s)
- $(GMT.opt_swap_xy)

When using two numeric inputs and no G option, the order of the x,y and grid is not important.
That is, both of this will work: D = grdinterpolate([0 0], G);  or  D = grdinterpolate(G, [0 0]);
"""
function grdinterpolate(cmd0::String="", arg1=nothing, arg2=nothing, arg3=nothing; kwargs...)

	length(kwargs) == 0 && arg1 === nothing && return monolitic("grdinterpolate", cmd0, arg1, arg2, arg3)

	d = init_module(false, kwargs...)[1]		# Also checks if the user wants ONLY the HELP mode

	cmd, = parse_common_opts(d, "", [:R :V_params :bi :bo :di :e :f :g :h :i :n :o :q :s :yx])
	cmd  = parse_these_opts(cmd, d, [[:G :outfile :outgrid], [:Z :levels]])

	cmd  = add_opt(d, cmd, 'D', [:D :meta :metadata],
           (xname="+x", yname="+y", zname="+z", dname="+d", scale="+s", offset="+o", nodata="+n", title="+t", remark="+r", varname="+v"))
	cmd  = add_opt(d, cmd, 'F', [:F :interp_type],
           (linear="_l", akima="_a", cubic="_c", nearest="_n", first_derivative="+1", second_derivative="+2"))

	cmd, _, arg1 = find_data(d, cmd0, cmd, arg1)

	cmd, args, n1, = add_opt(d, cmd, 'E', [:E :crossection], :line, Vector{Any}([arg1, arg2]), (azim="+a", great_circ="_+g", parallel="_+p", inc="+i", length="+l", npoints="+n", middpoint="+o", radius="+r", loxodrome="_+x"))
	if (n1 > 0)
		arg1, arg2 = args[:]
		_Vt = Vector{Any}([arg2, arg3])
	else
		_Vt = Vector{Any}([arg1, arg2])
	end

	cmd, args, n2, = add_opt(d, cmd, 'S', [:S :track], :line, _Vt, (header="+h",))
	if     (n1 == 0 && n2 > 0)  arg1, arg2 = args[:]
	elseif (n2 > 0)             arg2, arg3 = args[:]
	end

	cmd = parse_opt_range(d, cmd, "T")[1]

	#!occursin("-G", cmd) && (cmd *= " -G")
	if (isa(arg1, Tuple))
		for k = 1:length(arg1)  cmd *= " ?"  end		# Need as many ? as numel(arg1)
		common_grd(d, "grdinterpolate " * cmd, arg1..., arg2, arg3)
	else
		common_grd(d, "grdinterpolate " * cmd, arg1, arg2, arg3)
	end

end

# ---------------------------------------------------------------------------------------------------
grdinterpolate(arg1, arg2=nothing, arg3=nothing; kw...) = grdinterpolate("", arg1, arg2, arg3; kw...)


text:
# ensure in right directory and environment
cd(dirname(@__FILE__))
using Pkg
if Pkg.project().name == "MuseInference"
    using TestEnv # you'll need TestEnv in your global environment
    TestEnv.activate()
end

using Documenter, MuseInference, Turing

# change to true to skip running doc code
Documenter.Utilities.Selectors.disable(::Type{Documenter.Expanders.ExampleBlocks}) = false

makedocs(
    sitename = "MuseInference",
    format = Documenter.HTML(
        assets = ["assets/muse.css"],
        disable_git = true,
    ),
    pages = [
        "index.md",
        "userapi.md",
        "devapi.md",
    ],
    checkdocs = :none,
    doctest = false,
    strict = false,
    modules = [MuseInference]
)

deploydocs(
    repo = "https://github.com/marius311/MuseInference.jl",
    devbranch = "main",
    devurl = "latest",
    push_preview = true,
    forcepush = true
)


text:
# wrapper to convert LPQP solver into Conic solver

# To enable Conic support from an LPQP solver, define, e.g.,
# ConicModel(s::GurobiSolver) = LPQPtoConicBridge(LinearQuadraticModel(s))
# Must also implement supportedcones(). List SOC as a supported cone if it can
# be passed as a quadratic constraint.

mutable struct LPQPtoConicBridge <: AbstractConicModel
    lpqpmodel::AbstractLinearQuadraticModel
    qcp::Bool
    c
    A
    b
    Asoc
    Arsoc
    constr_cones
    var_cones
    cbvec::Vector{Float64}
    cbvec2::Vector{Float64}
end

LPQPtoConicBridge(m::AbstractLinearQuadraticModel) = LPQPtoConicBridge(m, false, nothing, nothing, nothing, nothing, nothing, nothing, nothing, Float64[], Float64[])

export LPQPtoConicBridge

numvar(m::LPQPtoConicBridge) = size(m.A,2)
numconstr(m::LPQPtoConicBridge) = size(m.A,1)

# To transform Conic problems into LinearQuadratic problems
function loadproblem!(m::LPQPtoConicBridge, c, A, b, constr_cones, var_cones)
    m.c = c
    m.A = A
    m.b = b
    m.constr_cones = constr_cones
    m.var_cones = var_cones

    # Conic form        LP form
    # min  c'x          min      c'x
    #  st b-Ax ∈ K_1     st lb <= Ax <= b
    #        x ∈ K_2         l <=  x <= u

    # If a cone is anything other than [:Free,:Zero,:NonNeg,:NonPos,:SOC,:SOCRotated], give up.
    bad_cones = [:SDP, :ExpPrimal, :ExpDual]
    # for each SOC and SOCRotated affine constraint, we need to add auxiliary variables
    num_orig = length(c)
    num_aux = 0
    linconstr_coneidx = Int[]
    linconstr_idx = Int[]
    socconstr_idx = Int[]
    rsocconstr_idx = Int[]
    rsoc_start_idx = Int[]
    for i in 1:length(constr_cones)
        cone, idxs = constr_cones[i]
        cone in bad_cones && error("Cone type $(cone) not supported")
        if cone == :SOC
            num_aux += length(idxs)
            append!(socconstr_idx, idxs)
        elseif cone == :SOCRotated
            num_aux += length(idxs)
            push!(rsoc_start_idx, length(rsocconstr_idx)+1)
            push!(rsoc_start_idx, length(rsocconstr_idx)+2)
            append!(rsocconstr_idx, idxs)
        else
            push!(linconstr_coneidx, i)
            if isa(idxs,Number)
                push!(linconstr_idx, idxs)
            else
                append!(linconstr_idx, idxs)
            end
        end
    end
    @assert num_aux == length(socconstr_idx) + length(rsocconstr_idx)
    for (cone,idxs) in var_cones
        cone in bad_cones && error("Cone type $(cone) not supported")
    end

    c = vcat(c,zeros(num_aux))
    # Variable bounds
    l = fill(-Inf, length(c))
    u = fill(Inf, length(c))
    for (cone,idxs) in var_cones
        if cone == :SOC
            l[idxs[1]] = 0
        elseif cone == :SOCRotated
            l[idxs[1]] = 0
            l[idxs[2]] = 0
        else
            cone_l = (cone == :Free || cone == :NonPos) ? -Inf : 0.0
            cone_u = (cone == :Free || cone == :NonNeg) ?  Inf : 0.0
            for idx in idxs
                l[idx] = cone_l
                u[idx] = cone_u
            end
        end
    end

    # set bounds for auxiliary variables
    ksoc = 1
    krsoc = 1
    for (cone,idx) in constr_cones
        if cone == :SOC
            l[num_orig + ksoc] = 0
            ksoc += length(idx)
        elseif cone == :SOCRotated
            l[num_orig + length(socconstr_idx) + krsoc] = 0
            l[num_orig + length(socconstr_idx) + krsoc + 1] = 0
            krsoc += length(idx)
        end
    end

    # matrix for linear constraints
    Alin = A[linconstr_idx,:]

    # Linear constraint bounds
    lb = Array{Float64}(undef, length(linconstr_idx))
    ub = Array{Float64}(undef, length(linconstr_idx))
    k = 1
    for (cone,idxs) in constr_cones
        if cone != :SOC && cone != :SOCRotated
            # :Zero         b - Ax = s == 0 -> Ax == b
            # :NonPos       b - Ax = s <= 0 -> Ax >= b
            # :NonNeg       b - Ax = s >= 0 -> Ax <= b
            # :Free         b - Ax = s free ->  free
            for idx in idxs
                lb[k] = (cone == :Zero || cone == :NonPos) ? b[idx] : -Inf
                ub[k] = (cone == :Zero || cone == :NonNeg) ? b[idx] :  Inf
                k += 1
            end
        end
    end

    if length(socconstr_idx) > 0
        m.Asoc = A[socconstr_idx,:]
        # linear constraints for aux variables
        # for each ||b - Ax|| <= c - d^Tx,
        # introduce y = b - Ax, z = c-d^Tx, and say
        # y^Ty <= z^2.
        # Ax + y = b, so we just need to append some identity columns
        Alin = [ Alin spzeros(length(linconstr_idx),length(socconstr_idx))
        m.Asoc  sparse(SparseArrays.I, length(socconstr_idx), length(socconstr_idx))]
        lbaux = b[socconstr_idx]
        ubaux = lbaux
        lb = [lb; lbaux]
        ub = [ub; ubaux]
    else
        m.Asoc = nothing
    end

    if length(rsocconstr_idx) > 0
        m.Arsoc = A[rsocconstr_idx,:]
        # same thing for rotated SOC
        # note we adjust for factor of 2:
        # rsoc has x'x <= 2pq
        # quadratic form is x'x <= pq
        diagvec = ones(length(rsocconstr_idx))
        diagvec[rsoc_start_idx] .= 1/sqrt(2)
        Alin = [ Alin spzeros(size(Alin,1),length(rsocconstr_idx))
        [m.Arsoc spzeros(size(m.Arsoc,1),length(socconstr_idx))] SparseArrays.spdiagm(0 => diagvec) ]
        lbaux = b[rsocconstr_idx]
        ubaux = lbaux
        lb = [lb; lbaux]
        ub = [ub; ubaux]
    else
        m.Arsoc = nothing
    end

    loadproblem!(m.lpqpmodel, Alin, l, u, c, lb, ub, :Min)
    m.cbvec = zeros(length(c))
    m.cbvec2 = fill(NaN,length(c)-num_aux)

    # Add conic constraints

    for (cone, idx) in var_cones
        if cone == :SOC
            m.qcp = true
            addquadconstr!(m.lpqpmodel, Int[], Float64[], vcat(idx), vcat(idx), [-1.0; ones(length(idx)-1)], '<', 0.0)
        elseif cone == :SOCRotated
            m.qcp = true
            idx1 = idx[1]
            idx2 = idx[2]
            idxrest = idx[3:end]
            addquadconstr!(m.lpqpmodel, Int[], Float64[], [idx1;idxrest], [idx2;idxrest], [-2.0; ones(length(idxrest))], '<', 0.0)
        end
    end

    ksoc = 1
    krsoc = 1
    for (cone,idx) in constr_cones
        if cone == :SOC
            m.qcp = true
            idx1 = num_orig + ksoc
            idxrest = (num_orig+ksoc+1):(num_orig+ksoc+length(idx)-1)
            addquadconstr!(m.lpqpmodel, Int[], Float64[], [idx1; idxrest], [idx1; idxrest], [-1.0;ones(length(idxrest))], '<', 0.0)
            ksoc += length(idx)
        elseif cone == :SOCRotated
            m.qcp = true
            idx1 = num_orig + length(socconstr_idx) + krsoc
            idx2 = num_orig + length(socconstr_idx) + krsoc + 1
            idxrest = (idx2+1):(idx1+length(idx)-1)
            addquadconstr!(m.lpqpmodel, Int[], Float64[], [idx1; idxrest], [idx2; idxrest], [-1.0;ones(length(idxrest))], '<', 0.0)
            krsoc += length(idx)
        end
    end
end

# extend a solution in the conic space to the LPQP space with extra variables
# could be made more efficient by avoiding recomputing vectors
function extend_solution(model::LPQPtoConicBridge,x)
    socconstr_idx = Int[]
    rsocconstr_idx = Int[]
    rsoc_start_idx = Int[]
    num_aux = 0
    for i in 1:length(model.constr_cones)
        cone, idxs = model.constr_cones[i]
        if cone == :SOC
            num_aux += length(idxs)
            append!(socconstr_idx, idxs)
        elseif cone == :SOCRotated
            num_aux += length(idxs)
            push!(rsoc_start_idx, length(rsocconstr_idx)+1)
            push!(rsoc_start_idx, length(rsocconstr_idx)+2)
            append!(rsocconstr_idx, idxs)
        end
    end
    @assert num_aux == length(socconstr_idx) + length(rsocconstr_idx)

    if length(socconstr_idx) > 0
        soc_aux = model.b[socconstr_idx] - model.Asoc*x
    else
        soc_aux = eltype(x)[]
    end

    if length(rsocconstr_idx) > 0
        diagvec = ones(length(rsocconstr_idx))
        diagvec[rsoc_start_idx] = sqrt(2)
        rsoc_aux = diagvec .* (model.b[rsocconstr_idx] - model.Arsoc*x)
    else
        rsoc_aux = eltype(x)[]
    end
    return [x;soc_aux;rsoc_aux]
end

setwarmstart!(model::LPQPtoConicBridge,v) = setwarmstart!(model.lpqpmodel,extend_solution(model,v))

for f in [:optimize!, :status, :getobjval, :getvartype]
    @eval $f(model::LPQPtoConicBridge) = $f(model.lpqpmodel)
end

getsolution(model::LPQPtoConicBridge) = getsolution(model.lpqpmodel)[1:length(model.c)]

function getdual(model::LPQPtoConicBridge)
    if model.qcp
        error("getdual not supported for SOC and SOCRotated cones")
    elseif status(model.lpqpmodel) == :Infeasible
        # Needed to pass LIN3 of coniclineartest
        # "-infeasibility ray" is not guaranteed to be dual feasible but
        # "-getconstrduals - λ getinfeasibilityray" is dual feasible for any λ
        -getconstrduals(model.lpqpmodel) - getinfeasibilityray(model.lpqpmodel)
    else
        -getconstrduals(model.lpqpmodel)
    end
end

function getvardual(model::LPQPtoConicBridge)
    if model.qcp
        error("getvardual not supported for SOC and SOCRotated cones")
    else
        getreducedcosts(model.lpqpmodel)
    end
end

setvartype!(model::LPQPtoConicBridge, vtype) = setvartype!(model.lpqpmodel, vtype)

for f in methods_by_tag[:rewrap]
    @eval $f(model::LPQPtoConicBridge) = $f(model.lpqpmodel)
end

mutable struct LPQPWrapperCallbackData <: MathProgCallbackData
    lpqpcb::MathProgCallbackData
    model::LPQPtoConicBridge
    solvec::Vector{Float64}
    heurvec::Vector{Float64}
end

wrapcb(f,m::LPQPtoConicBridge) = cb -> f(LPQPWrapperCallbackData(cb,m,m.cbvec,m.cbvec2))

function setlazycallback!(m::LPQPtoConicBridge, f)
    if applicable(setlazycallback!, m.lpqpmodel, f)
        setlazycallback!(m.lpqpmodel, wrapcb(f,m))
    else
        error("Solver does not support lazy callbacks")
    end
end

function setcutcallback!(m::LPQPtoConicBridge, f)
    if applicable(setcutcallback!, m.lpqpmodel, f)
        setcutcallback!(m.lpqpmodel, wrapcb(f,m))
    else
        error("Solver does not support cut callbacks")
    end
end

function setheuristiccallback!(m::LPQPtoConicBridge, f)
    if applicable(setheuristiccallback!, m.lpqpmodel, f)
        setheuristiccallback!(m.lpqpmodel, wrapcb(f,m))
    else
        error("Solver does not support heuristic callbacks")
    end
end

function setinfocallback!(m::LPQPtoConicBridge, f)
    if applicable(setinfocallback!, m.lpqpmodel, f)
        setinfocallback!(m.lpqpmodel, wrapcb(f,m))
    else
        error("Solver does not support info callbacks")
    end
end

for f in [:cbgetmipsolution,:cbgetlpsolution,:cbgetobj,
          :cbgetbestbound,:cbgetexplorednodes,
          :cbgetstate]
    @eval ($f)(cb::LPQPWrapperCallbackData) = ($f)(cb.lpqpcb)
end

for f in [:cbgetmipsolution,:cbgetlpsolution]
    @eval function ($f)(cb::LPQPWrapperCallbackData,output)
        ($f)(cb.lpqpcb,cb.solvec)
        copyto!(output,1,cb.solvec,1,length(output))
    end
end

for f in [:cbaddcut!,:cbaddcutlocal!,:cbaddlazy!,:cbaddlazylocal!]
    @eval function ($f)(cb::LPQPWrapperCallbackData,varidx,varcoef,sense,rhs)
        return ($f)(cb.lpqpcb,varidx,varcoef,sense,rhs)
    end
end

function cbsetsolutionvalue!(cb::LPQPWrapperCallbackData, varidx, value)
    cb.heurvec[varidx] = value
end

function cbaddsolution!(cb::LPQPWrapperCallbackData)
    newsol = extend_solution(cb.model,cb.heurvec)
    for i in 1:length(newsol)
        if isfinite(newsol[i])
            cbsetsolutionvalue!(cb.lpqpcb,i,newsol[i])
        end
    end
    cbaddsolution!(cb.lpqpcb)
    fill!(newsol,NaN)
end


text:
argtail(_, args...) = args

split_tuple_type(T) = fieldtype(T, 1), Tuple{argtail(T.parameters...)...}

eltypes(nt::NamedTuple{names}) where {names} = NamedTuple{names, eltypes(values(nt))}
eltypes(t::Tuple) = Tuple{map(eltype, t)...}

alwaysfalse(t) = false

"""
    StructArrays.map_params(f, T)

Apply `f` to each field type of `Tuple` or `NamedTuple` type `T`, returning a
new `Tuple` or `NamedTuple` object.

```julia-repl
julia> StructArrays.map_params(T -> Complex{T}, Tuple{Int32,Float64})
(Complex{Int32}, Complex{Float64})
```
"""
map_params(f::F, ::Type{NamedTuple{names, types}}) where {F, names, types} =
    NamedTuple{names}(map_params(f, types))

function map_params(f::F, ::Type{T}) where {F, T<:Tuple}
    if @generated
        types = fieldtypes(T)
        args = map(t -> :(f($t)), types)
        Expr(:tuple, args...)
    else
        map_params_fallback(f, T)
    end
end

map_params_fallback(f, ::Type{T}) where {T<:Tuple} = map(f, fieldtypes(T))

buildfromschema(initializer::F, ::Type{T}) where {F, T} = buildfromschema(initializer, T, staticschema(T))

"""
    StructArrays.buildfromschema(initializer, T[, S])

Construct a [`StructArray{T}`](@ref) with a function `initializer`, using a schema `S`.

`initializer(T)` is a function applied to each field type of `S`, and should return an `AbstractArray{S}`

`S` is a `Tuple` or `NamedTuple` type. The default value is [`staticschema(T)`](@ref).
"""
function buildfromschema(initializer::F, ::Type{T}, ::Type{NT}) where {F, T, NT<:Tup}
    nt = map_params(initializer, NT)
    StructArray{T}(nt)
end

array_names_types(::Type{StructArray{T, N, C, I}}) where {T, N, C, I} = array_names_types(C)
array_names_types(::Type{NamedTuple{names, types}}) where {names, types} = zip(names, types.parameters)
array_names_types(::Type{T}) where {T<:Tuple} = enumerate(T.parameters)

function apply_f_to_vars_fields(names_types, vars)
    exprs = Expr[]
    for (name, type) in names_types
        sym = QuoteNode(name)
        args = [Expr(:call, :component, var, sym) for var in vars]
        expr = if type <: StructArray
            apply_f_to_vars_fields(array_names_types(type), args)
        else
            Expr(:call, :f, args...)
        end
        push!(exprs, expr)
    end
    return Expr(:block, exprs...)
end

function _foreachfield(names_types, L)
    vars = ntuple(i -> gensym(), L)
    exprs = Expr[]
    for (i, v) in enumerate(vars)
        push!(exprs, Expr(:(=), v, Expr(:call, :getfield, :xs, i)))
    end
    push!(exprs, apply_f_to_vars_fields(names_types, vars))
    push!(exprs, :(return nothing))
    return Expr(:block, exprs...)
end

@generated foreachfield_gen(::S, f, xs::Vararg{Any, L}) where {S<:StructArray, L} =
    _foreachfield(array_names_types(S), L)

foreachfield(f::F, x::StructArray, xs::Vararg{Any, N}) where {F, N} = foreachfield_gen(x, f, x, xs...)

"""
    StructArrays.iscompatible(::Type{S}, ::Type{V}) where {S, V<:AbstractArray}

Check whether element type `S` can be pushed to a container of type `V`.
"""
iscompatible(::Type{S}, ::Type{<:AbstractArray{T}}) where {S, T} = S<:T
iscompatible(::Type{S}, ::Type{<:StructArray{<:Any, <:Any, C}}) where {S, C} = iscompatible(astuple(staticschema(S)), astuple(C))

iscompatible(::Type{Tuple{}}, ::Type{T}) where {T<:Tuple} = false
iscompatible(::Type{T}, ::Type{Tuple{}}) where {T<:Tuple} = false
iscompatible(::Type{Tuple{}}, ::Type{Tuple{}}) = true

function iscompatible(::Type{T}, ::Type{T′}) where {T<:Tuple, T′<:Tuple}
    (f, ls), (f′, ls′) = split_tuple_type(T), split_tuple_type(T′)
    iscompatible(f, f′) && iscompatible(ls, ls′)
end

iscompatible(::S, ::V) where {S, V<:AbstractArray} = iscompatible(S, V)

function _promote_typejoin(::Type{T}, ::Type{T′}) where {T<:NTuple{N, Any}, T′<:NTuple{N, Any}} where N
    (f, ls), (f′, ls′) = split_tuple_type(T), split_tuple_type(T′)
    return Tuple{_promote_typejoin(f, f′), _promote_typejoin(ls, ls′).parameters...}
end

_promote_typejoin(::Type{Tuple{}}, ::Type{Tuple{}}) = Tuple{}
function _promote_typejoin(::Type{NamedTuple{names, types}}, ::Type{NamedTuple{names, types′}}) where {names, types, types′}
    T = _promote_typejoin(types, types′)
    return NamedTuple{names, T}
end

_promote_typejoin(::Type{T}, ::Type{T′}) where {T, T′} = Base.promote_typejoin(T, T′)

function _promote_typejoin(::Type{Pair{A, B}}, ::Type{Pair{A′, B′}}) where {A, A′, B, B′}
    C = _promote_typejoin(A, A′)
    D = _promote_typejoin(B, B′)
    return Pair{C, D}
end

function replace_storage(f, v::AbstractArray{T, N})::AbstractArray{T, N} where {T, N}
    f(v)
end

"""
    StructArrays.replace_storage(f, s::StructArray)

Change storage type for components: each array `v` is replaced by `f(v)`. `f(v)` is expected to have the same
`eltype` and `size` as `v`.

## Examples

If PooledArrays is loaded, we can pool all columns of non `isbitstype`:

```jldoctest
julia> using StructArrays, PooledArrays

julia> s = StructArray(a=1:3, b = fill("string", 3));

julia> s_pooled = StructArrays.replace_storage(s) do v
           isbitstype(eltype(v)) ? v : convert(PooledArray, v)
       end
$(if VERSION < v"1.6-"
    "3-element StructArray(::UnitRange{Int64}, ::PooledArray{String,UInt32,1,Array{UInt32,1}}) with eltype NamedTuple{(:a, :b),Tuple{Int64,String}}:"
else
        "3-element StructArray(::UnitRange{Int64}, ::PooledVector{String, UInt32, Vector{UInt32}}) with eltype NamedTuple{(:a, :b), Tuple{Int64, String}}:"
end)
 (a = 1, b = "string")
 (a = 2, b = "string")
 (a = 3, b = "string")
```
"""
function replace_storage(f, s::StructArray{T}) where T
    cols = components(s)
    newcols = map(t -> replace_storage(f, t), cols)
    StructArray{T}(newcols)
end

astuple(::Type{NamedTuple{names, types}}) where {names, types} = types
astuple(::Type{T}) where {T<:Tuple} = T

strip_params(::Type{<:Tuple}) = Tuple
strip_params(::Type{<:NamedTuple{names}}) where {names} = NamedTuple{names}

hasfields(::Type{<:Tup}) = false
hasfields(::Type{<:NTuple{N, Any}}) where {N} = true
hasfields(::Type{<:NamedTuple{names}}) where {names} = true
hasfields(::Type{T}) where {T} = !isabstracttype(T)
hasfields(::Union) = false

"""
    StructArrays.bypass_constructor(T, args)

Create an instance of type `T` from a tuple of field values `args`, bypassing
possible internal constructors. `T` should be a concrete type.
"""
@generated function bypass_constructor(::Type{T}, args) where {T}
    vars = ntuple(_ -> gensym(), fieldcount(T))
    assign = [:($var::$(fieldtype(T, i)) = getfield(args, $i)) for (i, var) in enumerate(vars)]
    construct = Expr(:new, :T, vars...)
    Expr(:block, assign..., construct)
end


text:
"""
イエーツの連続性の補正による p 値

2×2 分割表の場合，PowerDivergenceTest() で得られる検定統計量はイエーツの連続性の補正をしないものである。

イエーツの連続性の補正をする場合には，以下の関数を定義しておくとよい。
"""

using Distributions, HypothesisTests
function yates(x)
    result = PowerDivergenceTest(x, lambda=1.0)
    n = sum(x)
    a, c, b, d = x
    stat = n*max(0, abs(a*d - b*c) - 0.5n)^2 / ((a+b)*(c+d)*(a+c)*(b+d))
    p = ccdf(Chisq(1), stat)
    println("corrected χ-sq. = $stat,  df = 1,  p value = $p")
end;

"""
使用法

julia> x = [12 35
            43 56]
2×2 Matrix{Int64}:
 12  35
 43  56

julia> yates(x)
corrected χ-sq. = 3.621119907386832,  df = 1,  p value = 0.057050457416995715

"""

max_stars_repo_name: r-de-r/stats

max_issues_repo_name: r-de-r/stats

max_forks_repo_name: r-de-r/stats