text:
##################################################################################
# This file is part of ModelBaseEcon.jl
# BSD 3-Clause License
# Copyright (c) 2020, Bank of Canada
# All rights reserved.
##################################################################################

export Model

const defaultoptions = Options(
    shift=10,
    substitutions=false,
    tol=1e-10,
    maxiter=20,
    verbose=false,
    warn=Options(no_t=true)
)

mutable struct ModelFlags
    linear::Bool
    ssZeroSlope::Bool
    ModelFlags() = new(false, false)
end

Base.show(io::IO, ::MIME"text/plain", flags::ModelFlags) = show(io, flags)
function Base.show(io::IO, flags::ModelFlags)
    names, values = [], []
    for f in fieldnames(ModelFlags)
        push!(names, string(f))
        push!(values, getfield(flags, f))
    end
    align = maximum(length, names) + 3
    println(io, "ModelFlags")
    for (n, v) in zip(names, values)
        println(io, lpad(n, align), " = ", v)
    end
end

"""
    Model <: AbstractModel

Data structure that represents a macroeconomic state space model.

"""
mutable struct Model <: AbstractModel
    "Options are various hyper-parameters for tuning the algorithms"
    options::Options
    "Flags contain meta information about the type of model"
    flags::ModelFlags
    sstate::SteadyStateData
    #### Inputs from user
    # transition variables
    variables::Vector{ModelVariable}
    # shock variables
    shocks::Vector{ModelVariable}
    # transition equations
    equations::Vector{Equation}
    # parameters
    parameters::Parameters
    # auto-exogenize mapping of variables and shocks
    autoexogenize::Dict{Symbol,Symbol}
    #### Things we compute
    maxlag::Int64
    maxlead::Int64
    # auxiliary variables
    auxvars::Vector{ModelVariable}
    # auxiliary equations
    auxeqns::Vector{Equation}
    # ssdata::SteadyStateData
    evaldata::AbstractModelEvaluationData
    #
    # constructor of an empty model
    Model(opts::Options) = new(merge(defaultoptions, opts),
        ModelFlags(), SteadyStateData(), [], [], [], Parameters(), Dict(), 0, 0, [], [], NoMED)
    Model() = new(deepcopy(defaultoptions),
        ModelFlags(), SteadyStateData(), [], [], [], Parameters(), Dict(), 0, 0, [], [], NoMED)
end

@inline auxvars(m::Model) = getfield(m, :auxvars)
@inline nauxvars(m::Model) = length(auxvars(m))

# We have to specialize allvars() nallvars() because we have auxvars here
@inline allvars(m::Model) = vcat(variables(m), shocks(m), auxvars(m))
@inline nallvars(m::Model) = length(variables(m)) + length(shocks(m)) + length(auxvars(m))

@inline alleqns(m::Model) = vcat(equations(m), getfield(m, :auxeqns))
@inline nalleqns(m::Model) = length(equations(m)) + length(getfield(m, :auxeqns))

################################################################
# Specialize Options methods to the Model type

OptionsMod.getoption(model::Model; kwargs...) = getoption(model.options; kwargs...)
OptionsMod.getoption(model::Model, name::Symbol, default) = getoption(model.options, name, default)
OptionsMod.getoption(model::Model, name::AS, default) where AS <: AbstractString = getoption(model.options, name, default)

OptionsMod.getoption!(model::Model; kwargs...) = getoption!(model.options; kwargs...)
OptionsMod.getoption!(model::Model, name::Symbol, default) = getoption!(model.options, name, default)
OptionsMod.getoption!(model::Model, name::AS, default) where AS <: AbstractString = getoption!(model.options, name, default)

OptionsMod.setoption!(model::Model; kwargs...) = setoption!(model.options; kwargs...)
OptionsMod.setoption!(model::Model, name::Symbol, value) = setoption!(model.options, name, value)
OptionsMod.setoption!(model::Model, name::AS, value) where AS <: AbstractString = setoption!(model.options, name, value)
OptionsMod.setoption!(f::Function, model::Model) = (f(model.options); model.options)

################################################################
# Implement access to options and flags and a few other computed properties

function Base.getproperty(model::Model, name::Symbol)
    if name ∈ fieldnames(Model)
        return getfield(model, name)
    end
    if name == :nvars
        return length(getfield(model, :variables))
    elseif name == :nshks
        return length(getfield(model, :shocks))
    elseif name == :nauxs
        return length(getfield(model, :auxvars))
    elseif name == :allvars
        return vcat(getfield(model, :variables), getfield(model, :shocks), getfield(model, :auxvars))
    elseif name == :varshks
        return vcat(getfield(model, :variables), getfield(model, :shocks))
    elseif name == :alleqns
        return vcat(getfield(model, :equations), getfield(model, :auxeqns))
    elseif haskey(getfield(model, :parameters), name)
        return getproperty(getfield(model, :parameters), name)
    elseif name ∈ getfield(model, :options)
        return getoption(model, name, nothing)
    elseif name ∈ fieldnames(ModelFlags)
        return getfield(getfield(model, :flags), name)
    else
        ind = indexin([name], getfield(model, :variables))[1]
        if ind !== nothing
            return getindex(getfield(model, :variables), ind)
        end
        ind = indexin([name], getfield(model, :shocks))[1]
        if ind !== nothing
            return getindex(getfield(model, :shocks), ind)
        end
        ind = indexin([name], getfield(model, :auxvars))[1]
        if ind !== nothing
            return getindex(getfield(model, :auxvars), ind)
        end
        error("This Model doesn't have property $name")
    end
end

function Base.propertynames(model::Model, private::Bool=false)
    return (fieldnames(Model)..., :nvars, :nshks, :nauxs, :allvars, :varshks, :alleqns,
    keys(getfield(model, :options))..., fieldnames(ModelFlags)...,
    Symbol[getfield(model, :variables)...]...,
    Symbol[getfield(model, :shocks)...]...,
    keys(getfield(model, :parameters))...,)
end

function Base.setproperty!(model::Model, name::Symbol, val::Any)
    if name ∈ fieldnames(Model)
        return setfield!(model, name, val)
    elseif haskey(getfield(model, :parameters), name)
        return setproperty!(getfield(model, :parameters), name, val)
    elseif name ∈ getfield(model, :options)
        return setoption!(model, name, val)
    elseif name ∈ fieldnames(ModelFlags)
        return setfield!(getfield(model, :flags), name, val)
    else
        ind = indexin([name], getfield(model, :variables))[1]
        if ind !== nothing
            if getindex(getfield(model, :variables), ind) != val
                throw(ArgumentError("Cannot replace variable with a different name. Use `m.var = update(m.var, ...)` to update variable."))
            end
            return setindex!(getfield(model, :variables), val, ind)
        end
        ind = indexin([name], getfield(model, :shocks))[1]
        if ind !== nothing
            if getindex(getfield(model, :shocks), ind) != val
                throw(ArgumentError("Cannot replace shock with a different name. Use `m.shk = update(m.shk, ...)` to update shock."))
            end
            return setindex!(getfield(model, :shocks), val, ind)
        end
        ind = indexin([name], getfield(model, :auxvars))[1]
        if ind !== nothing
            if getindex(getfield(model, :auxvars), ind) != val
                throw(ArgumentError("Cannot replace aux variable with a different name. Use `m.aux = update(m.aux, ...)` to update aux variable."))
            end
            return setindex!(getfield(model, :auxvars), val, ind)
        end
        setfield!(model, name, val)  # will throw an error since Model doesn't have field `$name`
    end
end

################################################################
# Pretty printing the model and summary (TODO)

export fullprint
fullprint(model::Model) = fullprint(Base.stdout, model)
function fullprint(io::IO, model::Model)
    io = IOContext(io, :compact => true, :limit => false)
    nvar = length(model.variables)
    nshk = length(model.shocks)
    nprm = length(model.parameters)
    neqn = length(model.equations)
    nvarshk = nvar + nshk
    function print_things(io, things...; len=0, maxlen=40, last=false)
        s = sprint(print, things...; context=io, sizehint=0)
        print(io, s)
        len += length(s) + 2
        last && (println(io), return 0)
        (len > maxlen) ? (print(io, "\n    "); return 4) : (print(io, ", "); return len)
    end
    let len = 15
        print(io, length(model.variables), " variable(s): ")
        if nvar == 0
            println(io)
        else
            for v in model.variables[1:end - 1]
                len = print_things(io, v; len=len)
            end
            print_things(io, model.variables[end]; last=true)
        end
    end
    let len = 15
        print(io, length(model.shocks), " shock(s): ")
        if nshk == 0
            println(io)
        else
            for v in model.shocks[1:end - 1]
                len = print_things(io, v; len=len)
            end
            print_things(io, model.shocks[end]; last=true)
        end
    end
    let len = 15
        print(io, length(model.parameters), " parameter(s): ")
        if nprm == 0
            println(io)
        else
            params = collect(model.parameters)
            for (k, v) in params[1:end - 1]
                len = print_things(io, k, " = ", v; len=len)
            end
            k, v = params[end]
            len = print_things(io, k, " = ", v; len=len, last=true)
        end
    end
    print(io, length(model.equations), " equations(s) with ", length(model.auxeqns), " auxiliary equations: \n")
    function print_aux_eq(bi)
        v = model.auxeqns[bi]
        for (_, ai) in filter(tv -> tv[2] > nvarshk, v.vinds)
            ci = ai - nvarshk
            ci < bi && print_aux_eq(ci)
        end
        println(io, "   |->A$bi:   ", v)
    end
    for (i, v) in enumerate(model.equations)
        println(io, "   E$i:   ", v)
        for (_, ai) in filter(tv -> tv[2] > nvarshk, v.vinds)
            print_aux_eq(ai - nvarshk)
        end
    end
end

function Base.show(io::IO, model::Model)
    nvar = length(model.variables)
    nshk = length(model.shocks)
    nprm = length(model.parameters)
    neqn = length(model.equations)
    nvarshk = nvar + nshk
    if nvar == nshk == nprm == neqn == 0
        print(io, "Empty model")
    elseif get(io, :compact, false) || nvar + nshk > 20 || neqn > 20
        # compact print
        print(io, nvar, " variable(s), ")
        print(io, nshk, " shock(s), ")
        print(io, nprm, " parameter(s), ")
        print(io, neqn, " equations(s) with ", length(model.auxeqns), " auxiliary equations.")
    else
        # full print
        fullprint(io, model)
        println(io, "Maximum lag: ", model.maxlag)
        println(io, "Maximum lead: ", model.maxlead)
    end
    return nothing
end

################################################################
# The macros used in the model definition.

# Note: These macros simply store the information into the corresponding
# arrays within the model instance. The actual processing is done in @initialize

export @variables, @logvariables, @neglogvariables, @steadyvariables, @exogenous, @shocks
export @parameters, @equations, @autoshocks, @autoexogenize

"""
    @variables model names...
    @variables model begin
        names...
    end

Define the names of transition variables in the model.

### Example
```jldoctest
@variables model a b c

# If the list is long, use a begin-end block separating names with newline or semicolon
@variables model begin
    a; b
    c
end
````
"""
macro variables(model, block::Expr)
    vars = filter(a -> !isa(a, LineNumberNode), block.args)
    return esc(:(unique!(append!($(model).variables, $vars)); nothing ))
end
macro variables(model, vars::Symbol...)
    return esc(:( unique!(append!($(model).variables, $vars)); nothing ))
end

macro logvariables(model, block::Expr)
    vars = filter(a -> !isa(a, LineNumberNode), block.args)
    return esc(:(unique!(append!($(model).variables, to_log.($vars))); nothing ))
end
macro logvariables(model, vars::Symbol...)
    return esc(:( unique!(append!($(model).variables, to_log.($vars))); nothing ))
end

macro neglogvariables(model, block::Expr)
    vars = filter(a -> !isa(a, LineNumberNode), block.args)
    return esc(:(unique!(append!($(model).variables, to_neglog.($vars))); nothing ))
end
macro neglogvariables(model, vars::Symbol...)
    return esc(:( unique!(append!($(model).variables, to_neglog.($vars))); nothing ))
end

macro steadyvariables(model, block::Expr)
    vars = filter(a -> !isa(a, LineNumberNode), block.args)
    return esc(:(unique!(append!($(model).variables, to_steady.($vars))); nothing ))
end
macro steadyvariables(model, vars::Symbol...)
    return esc(:( unique!(append!($(model).variables, to_steady.($vars))); nothing ))
end

macro exogenous(model, block::Expr)
    vars = filter(a -> !isa(a, LineNumberNode), block.args)
    return esc(:(unique!(append!($(model).variables, to_exog.($vars))); nothing ))
end
macro exogenous(model, vars::Symbol...)
    return esc(:( unique!(append!($(model).variables, to_exog.($vars))); nothing ))
end

"""
    @shocks model names...
    @shocks model begin
        names...
    end

Define the names of transition shocks in the model.

### Example
```jldoctest
@shocks model a_shk b_shk c_shk

# If the list is long, use a begin-end block separating names with newline or semicolon
@shocks model begin
    a_shk; b_shk
    c_shk
end
````
"""
macro shocks(model, block::Expr)
    shks = filter(a -> !isa(a, LineNumberNode), block.args)
    return esc(:( unique!(append!($(model).shocks, to_shock.($shks))); nothing ))
end
macro shocks(model, shks::Symbol...)
    return esc(:( unique!(append!($(model).shocks, to_shock.($shks))); nothing ))
end

"""
    @autoshocks model

Create a list of shocks that matches the list of variables.  Each shock name is
created from a variable name by appending "_shk".
"""
macro autoshocks(model, suf="_shk")
    esc( quote
        $(model).shocks = ModelVariable[
            to_shock(Symbol(v.name, $(QuoteNode(suf)))) for v in $(model).variables if !isexog(v) && !isshock(v)
        ]
        push!($(model).autoexogenize, (
            v.name => Symbol(v.name, $(QuoteNode(suf))) for v in $(model).variables if !isexog(v)
        )...)
        nothing
    end)
end

"""
    @parameters model begin
        name = value
        ...
    end

Declare and define the model parameters.

The parameters must have values. Provide the information in a series of assignment
statements wrapped inside a begin-end block. The names can be used in equations
as if they were regular variables.
"""
macro parameters(model, args::Expr...)
    if length(args) == 1 && args[1].head == :block
        args = args[1].args
    end
    ret = Expr(:block, :($(model).parameters.mod[] = $__module__))
    for a in args
        if a isa LineNumberNode
            continue
        end
        if Meta.isexpr(a, :(=), 2)
            key, value = a.args
            key = QuoteNode(key)
            # value = Meta.quot(value)
            push!(ret.args, :(push!($(model).parameters, $(key) => $(value))))
            continue
        end
        throw(ArgumentError("Parameter definitions must be assignments, not\n  $a"))
    end
    return esc(ret)
end

"""
    @autoexogenize model begin
        varname = shkname
        ...
    end

Define a mapping between variables and shocks that can be used to
conveniently  swap exogenous and endogenous variables.
"""
macro autoexogenize(model, args::Expr...)
    if length(args) == 1 && args[1].head == :block
        args = args[1].args
    end
    args = filter(a -> a isa Expr && a.head == :(=), [args...])
    autoexos = Dict{Symbol,Any}([ex.args for ex in args])
    esc(:( merge!($(model).autoexogenize, $(autoexos)); nothing ))
end

"""
Usage example:
```
@equations model begin
    y[t] = a * y[t-1] + b * y[t+1] + y_shk[t]
```
"""
macro equations(model, block::Expr)
    if block.head != :block
        error("A list of equations mush be within a begin-end block")
    end
    ret = Expr(:block)
    eqn = Expr(:block)
    for expr in block.args
        if isa(expr, LineNumberNode)
            push!(eqn.args, expr)
        else
            push!(eqn.args, expr)
            push!(ret.args, :(push!($model.equations, $(Meta.quot(eqn)))))
            eqn = Expr(:block)
        end
    end
    return esc(ret)
end

################################################################
# The processing of equations during model initialization.

const log_eqn_type = Symbol("@", :log)

export islog
@inline islog(eq::AbstractEquation) = eq.type === log_eqn_type

error_process(msg, expr) = begin
    throw(ArgumentError("$msg\n  During processing of\n  $(expr)"))
end

warn_process(msg, expr) = begin
    @warn "$msg\n  During processing of\n  $(expr)"
end

"""
    process_equation(model::Model, expr; <keyword arguments>)

Process the given expression in the context of the given model and create
an Equation() instance for it.

Internal function. There should be no need to call directly.

"""
function process_equation end
# export process_equation
process_equation(model::Model, expr::String; kwargs...) = process_equation(model, Meta.parse(expr); kwargs...)
# process_equation(model::Model, val::Number; kwargs...) = process_equation(model, Expr(:block, val); kwargs...)
# process_equation(model::Model, val::Symbol; kwargs...) = process_equation(model, Expr(:block, val); kwargs...)
function process_equation(model::Model, expr::Expr;
    modelmodule::Module=moduleof(model),
    line=LineNumberNode(0),
    type=default_eqn_type,
    doc="")

    # a list of all known time series
    allvars = [model.variables; model.shocks; model.auxvars]

    # keep track of model parameters used in expression
    parameters = Set{Symbol}()
    # keep track of references to known time series in the expression
    references = Dict{Tuple{Int64,Int64},Symbol}()
    # keep track of the source code location where the equation was defined
    #  (helps with tracking the locations of errors)
    source = []

    add_reference(sym, tind) = add_reference(sym, tind, indexin([sym], allvars)[1])
    add_reference(sym::Symbol, tind::Int, vind::Int) = begin
        if islog(allvars[vind])
            vsym = Symbol("#log", sym, "#", tind, "#")
        elseif isneglog(allvars[vind])
            vsym = Symbol("#logm", sym, "#", tind, "#")
        else
            vsym = Symbol("#", sym, "#", tind, "#")
        end
        push!(references, (tind, vind) => vsym) # keep track of indexes and dummy symbol
    end

    ###################
    #    process(expr)
    #
    # Process the expression, performing various tasks.
    #  * keep track of mentions of parameters and variables (including shocks)
    #  * remove line numbers from expression, but keep track so we can insert it into the residual functions
    #  * for each time-referenece of variable, create a dummy symbol that will be used in constructing the residual functions
    #
    # leave numbers alone
    process(num::Number) = num
    # store line number and discard it from the expression
    function process(line::LineNumberNode)
        push!(source, line)
        return nothing
    end
    # Symbols are left alone.
    # Mentions of parameters are tracked and left in place
    # Mentions of time series throw errors (they must always have a t-reference)
    function process(sym::Symbol)
        if sym ∈ model.variables
            if model.warn.no_t
                warn_process("Variable `$(sym)` without `t` reference. Assuming `$(sym)[t]`", expr)
            end
            add_reference(sym, 0)
            return Expr(:ref, sym, :t)
        elseif sym ∈ model.shocks
            if model.warn.no_t
                warn_process("Shock `$(sym)` without `t` reference. Assuming `$(sym)[t]`", expr)
            end
            add_reference(sym, 0)
            return Expr(:ref, sym, :t)
        elseif sym ∈ model.auxvars
            error_process("Auxiliary `$(sym)` without `t` reference.", expr)
        elseif haskey(model.parameters, sym)
            push!(parameters, sym)
        else
            # is this symbol valid in the model module?
            try
                modelmodule.eval(sym)
            catch
                error_process("Unknown symbol `$(sym)`.", expr)
            end
        end
        return sym
    end
    # Main version of process() - it's recursive
    function process(ex::Expr)
        if ex.head == :macrocall && ex.args[1] == doc_macro
            push!(source, ex.args[2])
            doc *= ex.args[3]
            return process(ex.args[4])
        end
        if ex.head == :macrocall
            push!(source, ex.args[2])
            if length(ex.args) == 3 && MacroTools.isexpr(ex.args[3], :(=))
                type = ex.args[1]
                return process(ex.args[3])
            end
            mfunc = Symbol("at_", string(ex.args[1])[2:end]) # replace leading @ with at_
            if isdefined(modelmodule, mfunc)
                mfunc = :( $modelmodule.$mfunc )
            elseif isdefined(ModelBaseEcon, mfunc)
                mfunc = :( ModelBaseEcon.$mfunc)
            else
                error_process("Unknown meta function $(ex.args[1]).", ex)
            end
            margs = map(filter(x -> !isa(x, LineNumberNode), ex.args[3:end])) do arg
                arg = process(arg)
                arg isa Expr ? Meta.quot(arg) :
                arg isa Symbol ? QuoteNode(arg) :
                arg
            end
            ex = modelmodule.eval(Expr(:call, mfunc, margs...))
        end
        if ex.head == :ref
            # expression is an indexing expression
            name, index = ex.args
            if haskey(model.parameters, name)
                # indexing in a parameter - leave it alone, but keep track
                push!(parameters, name)
                return Expr(:ref, name, modelmodule.eval(index))
            end
            vind = indexin([name], allvars)[1]  # the index of the variable
            if vind !== nothing
                # indexing in a time series
                tind = modelmodule.eval(:(let t = 0; $index end))  # the lag or lead value
                add_reference(name, tind, vind)
                return normal_ref(name, tind)
            end
            error_process("Undefined reference $(ex).", expr)
        end
        if ex.head == :(=)
            # expression is an equation
            # recursively process the two sides of the equation
            lhs, rhs = ex.args
            lhs = process(lhs)
            rhs = process(rhs)
            return Expr(:(=), lhs, rhs)
        end
        # if we're still here, recursively process the arguments
        args = map(process, ex.args)
        # remove `nothing`
        filter!(args) do a
            a !== nothing
        end
        if ex.head == :if
            if length(args) == 3
                return Expr(:call, :ifelse, args...)
            else
                error_process("Unable to process an `if` statement with a single branch. Use function `ifelse` instead.", expr)
            end
        end
        if ex.head == :call
            return Expr(:call, args...)
        end
        if ex.head == :block && length(args) == 1
            return args[1]
        end
        if ex.head == :incomplete
            # for incomplete expression, args[1] contains the error message
            error_process(ex.args[1], expr)
        end
        error_process("Can't process $(ex).", expr)
    end

    ##################
    #    make_residual_expression(expr)
    #
    # Convert a processed equation into an expression that evaluates the residual.
    #
    #  * each mention of a time-reference is replaced with its symbol
    make_residual_expression(any) = any
    function make_residual_expression(ex::Expr)
        if ex.head == :ref
            name, index = ex.args
            vind = indexin([name], allvars)[1]
            if vind !== nothing
                # The index expression is either t, or t+n or t-n. We made sure of that in process() above.
                if isa(index, Symbol) && index == :t
                    tind = 0
                elseif isa(index, Expr) && index.head == :call && index.args[1] == :- && index.args[2] == :t
                    tind = -index.args[3]
                elseif isa(index, Expr) && index.head == :call && index.args[1] == :+ && index.args[2] == :t
                    tind = +index.args[3]
                else
                    error_process("Unrecognized t-reference expression $index.", expr)
                end
                if need_transform(allvars[vind])
                    func = inverse_transformation(allvars[vind])
                    arg = references[(tind, vind)]
                    return :($func($arg))
                else
                    return references[(tind, vind)]
                end
            end
        elseif ex.head == :(=)
            if type == log_eqn_type
                return Expr(:call, :log, Expr(:call, :/, map(make_residual_expression, ex.args)...))
            else
                return Expr(:call, :-, map(make_residual_expression, ex.args)...)
            end
        end
        return Expr(ex.head, map(make_residual_expression, ex.args)...)
    end

    # call process() to gather information
    new_expr = process(expr)
    MacroTools.isexpr(new_expr, :(=)) || error_process("Expected equation.", expr)
    type ∈ (log_eqn_type, default_eqn_type) || error_process("Unknown equation type $(type).", expr)
    # if source information missing, set from argument
    push!(source, line)
    # collect the indices and dummy symbols of the mentioned variables
    # NOTE: Julia documentation assures us that keys() and values() iterate elements of the Dict in the same order!
    vinds = collect(keys(references))
    vsyms = collect(values(references))
    # make a residual expressoin for the eval function
    residual = make_residual_expression(new_expr)
    # add the source information to residual expression
    residual = Expr(:block, source[1], residual)
    resid, RJ = let mparams = model.parameters
        # create a list of expressions that assign the values of model parameters to
        # variables of the same name
        param_assigments = Expr(:block)
        for p in parameters
            push!(param_assigments.args, :( local $(p) = $(mparams).$(p) ))
        end
        funcs_expr = makefuncs(residual, vsyms, param_assigments; mod=modelmodule)
        modelmodule.eval(funcs_expr)
    end
    return Equation(doc, type, expr, residual, vinds, vsyms, resid, RJ)
end

# we must export this because we call it in the module where the model is being defined
export add_equation!

"""
    add_equation!(model::Model, expr::Expr; modelmodule::Module)

Process the given expression in the context of the given module, create
the Equation() instance for it and add it to the model instance.
"""

function add_equation!(model::Model, expr::Expr; modelmodule::Module=moduleof(model))
    source = LineNumberNode[]
    auxeqns = Expr[]
    type = default_eqn_type
    doc = ""

    ##################################
    # We process() the expression looking for substitutions.
    # If we find one, we create an auxiliary variable and equation.
    # We also keep track of line number, so we can label the aux equation as
    # defined on the same line.
    #
    # We make sure to make a copy of the expression and not to overwrite it.
    #
    process(any) = any
    function process(line::LineNumberNode)
        push!(source, line)
        return line
    end
    function process(expr::Expr)
        if expr.head == :block && expr.args[1] isa LineNumberNode && length(expr.args) == 2
            push!(source, expr.args[1])
            return process(expr.args[2])
        end
        if expr.head == :macrocall
            mname, mline = expr.args[1:2]
            margs = expr.args[3:end]
            push!(source, mline)
            if mname == doc_macro
                doc = margs[1]
                return process(margs[2])
            end
            if mname == log_eqn_type && length(margs) == 1
                type = mname
                return process(margs[1])
            end
            return Expr(:macrocall, mname, nothing, process.(margs)...)
        end
        # recursively process all arguments
        ret = Expr(expr.head)
        for i in eachindex(expr.args)
            push!(ret.args, process(expr.args[i]))
        end
        if getoption!(model; substitutions=true)
            local arg
            matched = @capture(ret, log(arg_))
            # is it log(arg)
            if matched && isa(arg, Expr)
                local var1, var2, ind1, ind2
                # is it log(x[t]) ?
                matched = @capture(arg, var1_[ind1_])
                if matched
                    mv = model.:($var1)
                    if islog(mv)
                        # log variable is always positive, no need for substitution
                        @goto skip_substitution
                    elseif isshock(mv) || isexog(mv)
                        if model.verbose
                            @info "Found log($var1), which is a shock or exogenous variable. Make sure $var1 data is positive."
                        end
                        @goto skip_substitution
                    elseif islin(mv) && model.verbose
                        @info "Found log($var1). Consider making $var1 a log variable."
                    end
                else
                    # is it log(x[t]/x[t-1]) ?
                    matched2 = @capture(arg, op_(var1_[ind1_], var2_[ind2_]))
                    if matched2 && op ∈ (:/, :+, :*) && has_t(ind1) && has_t(ind2) && islog(model.:($var1)) && islog(model.:($var2))
                        @goto skip_substitution
                    end
                end
                aux_expr = process_equation(model, Expr(:(=), arg, 0); modelmodule=modelmodule)
                if length(aux_expr.vinds) == 0
                    # arg doesn't contain any variables, no need for substitution
                    @goto skip_substitution
                end
                # substitute log(something) with auxN and add equation exp(auxN) = something
                push!(model.auxvars, :new)
                model.auxvars[end] = auxs = Symbol("aux", model.nauxs)
                push!(auxeqns, Expr(:(=), Expr(:call, :exp, Expr(:ref, auxs, :t)), arg))
                return Expr(:ref, auxs, :t)
                @label skip_substitution
                nothing
            end
        end
        return ret
    end

    new_expr = process(expr)
    if isempty(source)
        push!(source, LineNumberNode(0))
    end
    eqn = process_equation(model, new_expr; modelmodule=modelmodule, line=source[1], type=type, doc=doc)
    push!(model.equations, eqn)
    model.maxlag = max(model.maxlag, eqn.maxlag)
    model.maxlead = max(model.maxlead, eqn.maxlead)
    for i ∈ eachindex(auxeqns)
        eqn = process_equation(model, auxeqns[i]; modelmodule=modelmodule, line=source[1])
        push!(model.auxeqns, eqn)
        model.maxlag = max(model.maxlag, eqn.maxlag)
        model.maxlead = max(model.maxlead, eqn.maxlead)
    end
    model.evaldata = NoMED
    return model
end
@assert precompile(add_equation!, (Model, Expr))

############################
### Initialization routines

export @initialize

function initialize!(model::Model, modelmodule::Module)
    # Note: we cannot use moduleof here, because the equations are not initialized yet.
    if model.evaldata !== NoMED
        error("Model already initialized.")
    end
    initfuncs(modelmodule)
    model.parameters.mod[] = modelmodule
    varshks = model.varshks
    model.variables = varshks[.!isshock.(varshks)]
    model.shocks = varshks[isshock.(varshks)]
    empty!(model.auxvars)
    eqns = [e.expr for e in model.equations]
    empty!(model.equations)
    empty!(model.auxeqns)
    for e in eqns
        add_equation!(model, e; modelmodule=modelmodule)
    end
    for (i, v) in enumerate(model.allvars)
        model.:($(v.name)) = update(v, index=i)
    end
    model.evaldata = ModelEvaluationData(model)
    initssdata!(model)
    update_links!(model.parameters)
    return nothing
end

"""
    @initialize model

Prepare a model instance for analysis. Call this macro after all
variable names, shock names and equations have been defined.
"""
macro initialize(model::Symbol)
    # thismodule = @__MODULE__
    # modelmodule = __module__
    return quote
        $(@__MODULE__).initialize!($(model), $(__module__))
    end |> esc
end

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

eval_RJ(x::AbstractMatrix{Float64}, m::Model) = eval_RJ(x, m.evaldata)
eval_R!(r::AbstractVector{Float64}, x::AbstractMatrix{Float64}, m::Model) = eval_R!(r, x, m.evaldata)
# @inline printsstate(io::IO, m::Model) = printsstate(io, m.sstate)
# @inline printsstate(m::Model) = printsstate(m.sstate)
@inline issssolved(m::Model) = issssolved(m.sstate)

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

# export update_auxvars
"""
    update_auxvars(point, model; tol=model.tol, default=0.0)

Calculate the values of auxiliary variables from the given values of regular variables and shocks.

Auxiliary variables were introduced as substitutions, e.g. log(expression) was replaced by aux1 and
equation was added exp(aux1) = expression, where expression contains regular variables and shocks.

This function uses the auxiliary equation to compute the value of the auxiliary variable for the
given values of other variables. Note that the given values of other variables might be inadmissible,
in the sense that expression is negative. If that happens, the auxiliary variable is set to the given
`default` value.

If the `point` array does not contain space for the auxiliary variables, it is extended appropriately.

If there are no auxiliary variables/equations in the model, return *a copy* of `point`.

!!! note
    The current implementation is specialized only to log substitutions.
    TODO: implement a general approach that would work for any substitution.
"""
function update_auxvars(data::AbstractArray{Float64,2}, model::Model;
                tol::Float64=model.options.tol, default::Float64=0.0
)
    nauxs = length(model.auxvars)
    if nauxs == 0
        return copy(data)
    end
    (nt, nv) = size(data)
    nvarshk = length(model.variables) + length(model.shocks)
    if nv ∉ (nvarshk, nvarshk + nauxs)
        error("Incorrect number of columns $nv. Expected $nvarshk or $(nvarshk + nauxs).")
    end
    mintimes = 1 + model.maxlag + model.maxlead
    if nt < mintimes
        error("Insufficient time periods $nt. Expected $mintimes or more.")
    end
    result = [data[:,1:nvarshk] zeros(nt, nauxs)]
    for (i, eqn) in enumerate(model.auxeqns)
        for t in (eqn.maxlag + 1):(nt - eqn.maxlead)
            idx = [CartesianIndex((t + ti, vi)) for (ti, vi) in eqn.vinds]
            res = eqn.eval_resid(result[idx])
            if res < 1.0
                result[t, nvarshk + i] = log(1.0 - res)
            else
                result[t, nvarshk + i] = default
            end
        end
    end
    return result
end


text:
using QXContexts
using MPI
using Logging
using Random

import QXContexts.Logger: @perf

MPI.Init()
comm = MPI.COMM_WORLD
rank = MPI.Comm_rank(comm)

#io = nothing
if MPI.Comm_size(comm) > 1
    io = MPI.File.open(comm, "mpi_io.dat", read=true,  write=true, create=true)
else
    io = stdout
end
global_logger(QXContexts.Logger.QXLogger(io, Logging.Error, Dict{Any, Int64}(), false))

@info "Hello world, I am $(rank) of $(MPI.Comm_size(comm)) world-size"
@warn "Rank $(rank) doesn't like what you are doing"
@error "Rank $(rank) is very unhappy!"

if rank == 2
    @info "Hello again from rank $(rank)"
    a = @perf sum(rand(100,100) * rand(100,100))
    println(a)
else
    @warn "I AM RANK $(rank)"
end

ext: jl


text:
input = read(stdin, String)

constraints, myticket, nearby = split(input, "\n\n")
myticket = [parse(Int, s) for s in split(split(myticket, '\n')[2], ',', keepempty=false)]

# Constraints : (min1, max1, min2, max2)
constr = []
for c in split(constraints, '\n', keepempty=false)
    colon = findfirst(':', c)
    part1, _, part2 = split(c[colon + 1: end], keepempty=false)
    part1 = strip(part1)
    part2 = strip(part2)
    (min1, max1), (min2, max2) = split(part1, '-', keepempty=false), split(part2, '-', keepempty=false)

    push!(constr, (parse(Int, min1), parse(Int, max1), parse(Int, min2), parse(Int, max2)))
end

errorrate = 0

# poss[index in tickets][index in constraint] = Possible that ticket matches constraint index
poss = ones(Bool, 20, 20)
for ticket in split(nearby, '\n', keepempty=false)[2:end]
    global errorrate, part2prod

    nb = [parse(Int, s) for s in split(ticket, ',', keepempty=false)]

    for (ni, n) in enumerate(nb)
        valid = false

        for (ci, (min1, max1, min2, max2)) in enumerate(constr)
            if min1 <= n <= max1 || min2 <= n <= max2
                valid = true

                break
            end
        end

        if !valid
            errorrate += n
        else
            # Part 2
            for (ci, (min1, max1, min2, max2)) in enumerate(constr)
                if !(min1 <= n <= max1 || min2 <= n <= max2)
                    global poss[ni, ci] = false
                end
            end
        end
    end
end

# Part 2
class2idx = Dict()

# Sums (all values in range 1 to 20) to start from min sum to max sum
sums = [(i, sum(poss[i, begin : end])) for i in 1:20]
sort!(sums, lt=((_, a), (_, b)) -> a < b)

for i in 1:20
    ci = sums[i][1]

    # There is only one 1
    class = findfirst(poss[ci, begin : end])

    # Remove all 1s of this class (already taken)
    for j in 1:20
        poss[j, class] = false
    end

    # Set class
    class2idx[class] = ci
end

part2prod = reduce(*, [myticket[class2idx[i]] for i in 1:6])

println("Part 1 : $errorrate")
println("Part 2 : $part2prod")


text:
struct PrefixMatch <: MethodType
end

computeDocFreq(method::PrefixMatch, docFreq) = docFreq.doc_freq

function similarity(method::PrefixMatch, query::Vector{UInt32}, candidate::Vector{UInt32}, doc_freq::Vector{Float64})
    q_len = length(query)
    c_len = length(candidate)
    prefix_len = 0.0

    for i=1:min(q_len, c_len)
        @inbounds query[i] != candidate[i] && break
        prefix_len += 1.0
    end

    return prefix_len / max(q_len, c_len)
end

max_stars_repo_name: irving-muller/FaST

max_issues_repo_name: irving-muller/FaST

max_forks_repo_name: irving-muller/FaST


text:
"""
    associated_legendre(x; l_max, m)

Returns the associated Legendre polynomials `Pₗᵐ(x)` and their
first derivative `(d/dx)Pₗᵐ(x)`, for `l=1,...,l_max`. The polynomials
are unnormalized and include the Condon–Shortley phase.
"""
function associated_legendre(x; l_max, m)
    abs(x) < 1 || @error "Argument 'x' must be between -1 and 1."
    plm = Legendre.Plm(0:l_max, m, x)
    # compute first derivate using recurrence
    # (x^2 - 1)(d/dx)p_nm(x) = n * x * p_nm(x) - (n + m) * p_{n-1, m}(x)
    plm_d = similar(plm)
    for n in 1:l_max
        n_index = n + 1
        plm_d[n_index] = (n*x*plm[n_index] - (n+m)*plm[n_index-1]) / (x^2 - 1)
    end
    # return values that correspond to l=1,...,l_max
    result_plm = plm[2:end]
    result_plm_d = plm_d[2:end]
    @assert length(result_plm) == length(result_plm_d) == l_max
    return result_plm, result_plm_d
end

max_stars_repo_name: Riarrieta/MieSeries

max_issues_repo_name: Riarrieta/MieSeries

max_forks_repo_name: Riarrieta/MieSeries


text:
"""
    readLFT(fname::AbstractString)

Import a leadfield matrix in the BESA LFT format

#### Arguments

* `fname`   = location and name of file to import
* `verbose` = should the function output user feedback

#### Returns

* Array representing leadfield values

#### Example

L = readLFT(verbose=true)
"""
function readLFT(fname::AbstractString = Pkg.dir("Private", "test", "data", "FEM_Default50MRI-Aniso3_CR80.lft");
        verbose::Bool = false)

    debug("Leadfield file: $(fname)")

    # Open file
    fid = open(fname, "r")

    # Read header information from file
    version         = read(fid, Int32, 1)[1]
    numSensors      = read(fid, Int32, 1)[1] # n_s
    numNodes        = read(fid, Int32, 1)[1] # n_n
    numOrientations = read(fid, Int32, 1)[1] # n_o

    # Read in maximum leadfield values
    numMaxLfValues = numNodes * numOrientations
    Lmax_jk = read(fid, Float32, numMaxLfValues)

    # Read in leadfield fractions
    numLfFracs = numSensors * numNodes[1] * numOrientations
    Lc_ijk = read(fid, Int16, numLfFracs)

    # Check that everything has been read
    if ~eof(fid)
        println("!! Something went wrong. File is longer than expected")
    end

    close(fid)

    ##  Convert to true leadfield values from fraction

    # Reshape Lmax
    # reshape requires particular inputs. Hopefully this is relaxed to save conversion
    Lmax_jk = reshape(transpose(Lmax_jk), (
        convert(Int64,numNodes),
        convert(Int64,numOrientations)))

    # Reshape Lc
    Lc_ijk = reshape(transpose(Lc_ijk), (
        convert(Int64, numNodes),
        convert(Int64, numOrientations),
        convert(Int64, numSensors)))

    # Scale leadfield
    # Lc_ijk = L_ijk * (SHRT_MAX / Lmax_jk) -> L_ijk = Lc_ijk / (SHRT_MAX / Lmax_jk)
    #                                          L_ijk = Lc_ijk * (Lmax_jk / SHRT_MAX)
    SHRT_MAX = 2^15-1

    L_ijk = Array(Float32, (
        convert(Int64, numNodes),
        convert(Int64, numOrientations),
        convert(Int64, numSensors)
        ))

    i = 1
    while i <= numNodes

        j = 1
        while j <= numOrientations

            k = 1
            while k <= numSensors

                L_ijk[i,j,k] =  Lc_ijk[i,j,k] * (Lmax_jk[i,j] / SHRT_MAX)

                k = k + 1
            end
            j = j + 1
        end
        i = i + 1
    end

    if verbose
        println("LFT version:            $(version)")
        println("Number of sensors:      $(numSensors)")
        println("Number of nodes:        $(numNodes)")
        println("Number of orientations: $(numOrientations)")
        println("Number of maximums:     $numMaxLfValues")
        println("Number of fractions:    $numLfFracs")
        println("Lc_ijk  = [$(size(Lc_ijk)[1]), $(size(Lc_ijk)[2]), $(size(Lc_ijk)[3])]")
        println("Lmax_jk = [$(size(Lmax_jk)[1]), $(size(Lmax_jk)[2])]")
        println("L_ijk   = [$(size(L_ijk)[1]), $(size(L_ijk)[2]), $(size(L_ijk)[3])]")
    end

    return L_ijk
end


text:
using BenchmarkTools
using CUDA
using Oceananigans
using Benchmarks

# Benchmark function

function benchmark_incompressible_model(Arch, FT, N)
    grid = RegularRectilinearGrid(FT, size=(N, N, N), extent=(1, 1, 1))
    model = IncompressibleModel(architecture=Arch(), float_type=FT, grid=grid)

    time_step!(model, 1) # warmup

    trial = @benchmark begin
        @sync_gpu time_step!($model, 1)
    end samples=10

    return trial
end

# Benchmark parameters

Architectures = has_cuda() ? [CPU, GPU] : [CPU]
Float_types = [Float32, Float64]
Ns = [32, 64, 128, 256]

# Run and summarize benchmarks

print_system_info()
suite = run_benchmarks(benchmark_incompressible_model; Architectures, Float_types, Ns)

df = benchmarks_dataframe(suite)
sort!(df, [:Architectures, :Float_types, :Ns], by=(string, string, identity))
benchmarks_pretty_table(df, title="Incompressible model benchmarks")

if GPU in Architectures
    df_Δ = gpu_speedups_suite(suite) |> speedups_dataframe
    sort!(df_Δ, [:Float_types, :Ns], by=(string, identity))
    benchmarks_pretty_table(df_Δ, title="Incompressible model CPU -> GPU speedup")
end


text:
module TheEqualiser
#
# This module automatically adjusts taxes (it and ni, optionally)
# so the net cost of benefit or other changes
# is close to zero.
#
# TODO needs a lot of work:
#
# - more options - so basic rate only etc;
# - use passed-in functions to equalise (like op_tax! below);
# - check results are in sensible bounds (e.g. >= 0 < 100 tax rates).
#
using Roots
using UUIDs
using Observables

using ScottishTaxBenefitModel
using .Definitions
using .Monitor
using .Results
using .Runner
using .RunSettings
using .STBOutput
using .STBParameters
using .Utils

@enum EqTargets eq_it eq_ni eq_it_ni
export EqTargets,eq_it,eq_ni,eq_it_ni
export equalise
#
# Roots only allows 1 parameter, I think, so:
#
mutable struct RunParameters{T<:AbstractFloat}
    params :: TaxBenefitSystem{T}
    settings :: Settings
	base_cost :: T
    target :: EqTargets
    obs    :: Observable
end

# TODO another possible approach is to pass in editing
# functions such as:
function op_tax!( sys :: TaxBenefitSystem{T}, r :: T ) where T <: AbstractFloat
    sys.it.non_savings_rates .+= r
end

function run( x :: T, rparams :: RunParameters{T} ) where T <: AbstractFloat
    nsr = deepcopy( rparams.params.it )
    nsi = deepcopy( rparams.params.ni )
    if rparams.target in [eq_it, eq_it_ni]
        rparams.params.it.non_savings_rates .+= x
    end
    # TODO check sensible it rates
    if rparams.target in [eq_ni, eq_it_ni]
        rparams.params.ni.primary_class_1_rates .+= x
        rparams.params.ni.class_4_rates .+= x
    end
    # TODO check sensible ni rates
    results = do_one_run(rparams.settings, [rparams.params], rparams.obs )
    # restore
    rparams.params.it = nsr
    rparams.params.ni = nsi
    summary = summarise_frames(results, rparams.settings)
    nc = summary.income_summary[1][1,:net_cost]
    return round( nc - rparams.base_cost, digits=0 )
end

"""
Adjust the thing in `target` so that the net cost of the changes in `sys`
is close to `base_cost`
"""
function equalise(
    target :: EqTargets,
    sys :: TaxBenefitSystem{T},
    settings :: Settings,
    base_cost :: T,
    observer :: Observable ) :: T where T<:AbstractFloat
    zerorun = ZeroProblem( run, 0.0 ) # fixme guess at 0.0 ?
    rparams = RunParameters( sys, settings, base_cost, target, observer )
    incch = solve( zerorun, rparams )
    #
    # TODO test incch is sensible
    return incch
end

end # module

max_stars_repo_name: grahamstark/ScottishTaxBenefitModel

max_issues_repo_name: grahamstark/ScottishTaxBenefitModel

max_forks_repo_name: grahamstark/ScottishTaxBenefitModel


text:
module Doppelkeks

using Logging

export Literal, Clause, Formula, solve

include("structs.jl")

abstract type SAT_State end
struct SAT <: SAT_State end
struct UNSAT <: SAT_State end

function solve(formula::Formula)
    trail = Trail(Tuple{Literal, Bool}[])
    while true
        if doBacktrack(trail, formula)
            if isempty(decisions(trail))
                return UNSAT, nothing
            else
                trail = applyBacktrack(trail)
            end
        else
            if vars(trail) == vars(formula)
                return SAT, trail
            else
                applyDecide!(trail, formula)
            end
        end
    end

    return UNSAT, nothing
end

###### Utiliy

function doBacktrack(trail, formula)
    isempty(trail) && return false

    for clause in formula
        all(!l in trail for l in clause) && return true
    end

    return false
end

function vars(f::Formula)
    res = Set{Int}()

    for c in f.clauses
        for l in c.lits
            push!(res, l.name)
        end
    end

    return res
end

function vars(trail)
    res = Set{Int}()

    for (l,_) in trail
        push!(res, l.name)
    end

    return res
end

###### Trail functions

function decisions(trail)
    filter(x -> x[2], trail)
end

function lastDecision(trail)
    for i in length(trail):-1:1
        name, decision = trail[i]
        decision && return name
    end
    return nothing
end

function decisionsTo(trail, lit)
    decisions(trail[1:findfirst(x -> x[1] == lit, trail)])
end

function currentLevel(trail)
    length(decisions(trail))
end

function level(lit, trail)
    length(decisionsTo(trail, lit))
end

function prefixToLevel(trail, level)
    filter(x -> x[2] <= level, trail)
end

function prefixBeforeLastDecision(trail)
    for i in length(trail):-1:1
        _, decision = trail[i]
        decision && return Trail(trail[1:i-1])
    end
    return nothing
end

function lastAssertedLiteral(clause, trail)
    lastLit = first(clause)
    lastLevel = level(lastLit, trail)

    for i in 2:length(clause)
        tmp = clause[i]
        tmpLvl = level(tmp, trail)
        if tmpLvl > lastLevel
            lastLevel = tmpLvl
            lastLit = tmp
        end
    end

    return lastLit
end

maxLevel(clause, trail) = level(lastAssertedLiteral(clause, trail), trail)

function assertLiteral!(trail, literal, decision)
    push!(trail, (literal, decision))
end

##### Decisions

function applyDecide!(trail, formula)
    l = selectLiteral(trail, formula)
    assertLiteral!(trail, l, true)
    @info string(rpad("applyDecide", 16), l, " -> ", trail)
end

function applyBacktrack(trail)
    l = lastDecision(trail)
    M = prefixBeforeLastDecision(trail)
    assertLiteral!(M, !l, false)
    @info string(rpad("applyBacktrack", 16), !l, " -> ", M)
    M
end

function selectLiteral(trail, formula)
    t_vars = vars(trail)
    f_vars = vars(formula)
    @assert t_vars != f_vars
    pos_vars = setdiff(f_vars, t_vars)
    Literal(rand(pos_vars), rand(Bool))
end

end # module


text:
using Pkg
Pkg.activate(joinpath(@__DIR__, ".."))
Pkg.instantiate()

# ## Setup
using Dojo

# ## Mechanism
mechanism = get_mechanism(:pendulum,
    timestep=0.01,
    gravity=-9.81,
    damper=5.0,
    spring=0.0)

# ## Controller
function controller!(mechanism, t)
    ## Target state
    x_goal = [1.0 * π; 0.0]

    ## Current state
    x = get_minimal_state(mechanism)

    ## Gains
    K = [5.0 0.5] * 0.1

    # Control inputs
    u = -K * (x - x_goal)
    set_input!(mechanism, u)
end

# ##Simulate
initialize!(mechanism, :pendulum,
    angle=0.0 * π,
    angular_velocity=0.0);

storage = simulate!(mechanism, 2.0, controller!,
    record=true,
    verbose=true);

# ## Visualize
vis = Visualizer()
render(vis)
visualize(mechanism, storage,
    vis=vis);


text:
@testset "test distributed dataset creation" begin
    @testset "5-bus case file, number of samples" begin
		# Create worker processes
		#nproc = Sys.CPU_THREADS # The desired number of CPUs to run with
		# Clear any existing worker processes if present
		Distributed.nprocs() > 1 && Distributed.rmprocs(Distributed.workers())
		# Create worker processes
		Distributed.addprocs(nproc - 1; exeflags="--project")

		# Import functions used on all worker processes
		Distributed.@everywhere using OPFLearn

		net_file = "pglib_opf_case5_pjm.m"
		K = 100
		results = dist_create_samples(net_file, K, net_path="data")
        @test size(results["inputs"]["pd"],1) == K
		@test size(results["inputs"]["pd"],2) == 3
		@test size(results["outputs"]["p_gen"],1) == K
		@test size(results["outputs"]["p_gen"],2) == 5
		@test size(results["duals"]["v_max"],1) == K
		@test size(results["duals"]["v_max"],2) == 5
    end

    @testset "5-bus case file, specified variables" begin
		# Create worker processes
		#nproc = Sys.CPU_THREADS # The desired number of CPUs to run with
		# Clear any existing worker processes if present
		Distributed.nprocs() > 1 && Distributed.rmprocs(Distributed.workers())
		# Create worker processes
		Distributed.addprocs(nproc - 1; exeflags="--project")

		# Import functions used on all worker processes
		Distributed.@everywhere using OPFLearn

		input_vars = ["pd", "qd"]
		output_vars = ["p_gen", "vm_gen"]
		dual_vars = ["v_min", "qg_min", "pto_max"]

		net_file = "pglib_opf_case5_pjm.m"
		K = 100
		results = dist_create_samples(net_file, K, net_path="data", input_vars=input_vars,
									  output_vars=output_vars, dual_vars=dual_vars)
        @test length(results["inputs"]) == length(input_vars)
		@test length(results["outputs"]) == length(output_vars)
		@test length(results["duals"]) == length(dual_vars) + 3  # Duals includes k, uas, and variances
		@test all([x in keys(results["inputs"]) for x in input_vars])
		@test all([x in keys(results["outputs"]) for x in output_vars])
		@test all([x in keys(results["duals"]) for x in dual_vars])
    end

	@testset "all save options distributed" begin
		# Create worker processes
		#nproc = Sys.CPU_THREADS # The desired number of CPUs to run with
		# Clear any existing worker processes if present
		Distributed.nprocs() > 1 && Distributed.rmprocs(Distributed.workers())
		# Create worker processes
		Distributed.addprocs(nproc - 1; exeflags="--project")

		# Import functions used on all worker processes
		Distributed.@everywhere using OPFLearn

		net_file = "pglib_opf_case5_pjm.m"
		K = 100
		results = dist_create_samples(net_file, K, net_path="data",
									  save_certs=true, stat_track=1, save_infeasible=true)
		@test haskey(results, "inputs")
		@test haskey(results, "outputs")
		@test haskey(results, "duals")
		@test haskey(results, "infeasible_inputs")
		@test haskey(results, "polytope")
		@test haskey(results, "stats")
    end

end


text:
module fastmaxent

import Base: show

import Graphs: degree
import LinearAlgebra: diagind
import StatsBase: countmap
#=
import Base:
    +

import Graphs:
    degree, strength
=#
using Graphs
using SimpleWeightedGraphs
using NLsolve
using IndirectArrays

# MODELS
export UBCM, UBCMCompact, solve!

# METRICS
export degree

include("models.jl")
include("metrics.jl")
end # module


text:
@static if VERSION < v"1.4"

update_project() = @msg updateProject(false)
handle(()->false, "allProjects")

else  # if VERSION < v"1.4"

update_project() = @msg updateProject(project_info())

project_info(path = Pkg.project().path) = (name = Pkg.REPLMode.projname(path), path = path)

# adapted from https://github.com/JuliaLang/Pkg.jl/blob/eb3726d8f9c68bb91707a5c0e9809c95f1c1eee7/src/API.jl#L331-L726
# but here we only look at "user-depot" and collect usages of each Manifest.toml
using Pkg.TOML, Dates

function allprojects()
  manifest_usage = Dict{String,DateTime}()
  usagefile_path = joinpath(Base.DEPOT_PATH[1], "logs", "manifest_usage.toml")
  isfile(usagefile_path) || return (projects = [], active = "")
  for (filename, infos) in TOML.parse(read(usagefile_path, String))
    for info in infos
      manifest_usage[filename] = max(get(manifest_usage, filename, DateTime(0)), DateTime(info["time"]))
    end
  end

  # IDEA: maybe filtering dated manifests ?
  all_project_files = map(collect(keys(manifest_usage))) do f
    replace(f, "Manifest.toml" => "Project.toml")
  end
  filter!(isfile, all_project_files)
  projects = project_info.(all_project_files)
  active = project_info().name
  return (projects = projects, active = active)
end

handle(allprojects, "allProjects")

end  # if VERSION < v"1.4"


text:
# Bogumił Kamiński, 2019-03-25

# run as
# julia installation.jl

# Tested under Julia 1.1.0 (but any Julia 1.x version should work).
# In order to install Julia go to https://julialang.org/downloads/
# and follow the instructions given there.

# I give package versions under which the scripts were tested
# but installing their latest versions should be OK.

# In some environments installation of Cbc or PyPlot might fail
# due to clashes with external dependencies on OS level.
# Please contact me if you encounter such problems.

# "4. Logparser" example assumes that git is installed on the host machine
# and its executable is in search path.

using Pkg

Pkg.add("Distributions") # 0.17.0
Pkg.add("DataFrames") # 0.17.1
Pkg.add("ForwardDiff") # v0.10.3
Pkg.add("CSV") # 0.4.3
Pkg.add("GLM") # 1.1.1
Pkg.add("Flux") # 0.8.1
Pkg.add("HTTP") # 0.8.0
Pkg.add("JuMP") # 0.19.0
Pkg.add("Cbc") # 0.6.0
Pkg.add("PyPlot") # 2.8.0


text:
module Example

using Gadfly, Colors, DataFrames
import Gadfly.plot

import Base.rand
export AbstractPoint, AbstractLine,AbstractPlane, Point, Line, Plane
export area, point, distance, convexHull, boundingbox
export layer, layerPolygon

include("types.jl")
include("functions.jl")
include("plotting.jl")

end # module


text:
using ExplicitQPs
using Base.Test
using JuMP
using Gurobi

@testset "trivial model" begin
    m = Model(solver=GurobiSolver(OutputFlag=0))
    @variable m x >= 0.5
    @variable m y
    @constraint m y >= x
    JuMP.fix(x, 0.5)
    @objective m Min x^2 + y^2 + 0.1 * x * y
    solve(m)
    @test getvalue(y) ≈ 0.5
    @test getvalue(x) ≈ 0.5
    ex = ExplicitQPs.explicit_solution(m, [x])
    s = ExplicitQPs.solution(ex, y)
    @test getvalue(s) ≈ 0.5
    @test ExplicitQPs.parameter(ex, x) == getvalue(x)
    @test ExplicitQPs.gradient(ex, y) ≈ [1.0]
    @test ExplicitQPs.jacobian(ex, [y]) ≈ [1.0]
end

@testset "models with affine cost" begin
    @testset "model 1" begin
        m = Model(solver=GurobiSolver(OutputFlag=0))
        @variable m x
        @variable m y
        @constraint m y <= x
        JuMP.fix(x, 0.5)
        @objective m Min (y - 0.6)^2
        solve(m)
        @test getvalue(y) ≈ 0.5
        @test getvalue(x) ≈ 0.5
        ex = ExplicitQPs.explicit_solution(m, [x])
        s = ExplicitQPs.solution(ex, y)
        @test getvalue(s) ≈ 0.5
        @test ExplicitQPs.parameter(ex, x) == getvalue(x)
        @test ExplicitQPs.gradient(ex, y) ≈ [1.0]
        @test ExplicitQPs.jacobian(ex, [y]) ≈ [1.0]
    end

    @testset "model 2" begin
        m = Model(solver=GurobiSolver(OutputFlag=0))
        @variable m x
        @variable m y
        @constraint m y <= 10
        JuMP.fix(x, 0.5)
        @objective m Min (y - 0.6)^2
        solve(m)
        @test getvalue(y) ≈ 0.6
        @test getvalue(x) ≈ 0.5
        ex = ExplicitQPs.explicit_solution(m, [x])
        s = ExplicitQPs.solution(ex, y)
        @test getvalue(s) ≈ 0.6
        @test ExplicitQPs.parameter(ex, x) == getvalue(x)
        @test ExplicitQPs.gradient(ex, y) ≈ [0.0]
        @test ExplicitQPs.jacobian(ex, [y]) ≈ [0.0]
    end
end

@testset "simple mpc" begin
    env = Gurobi.Env()
    function run_mpc(x0)
        m = Model(solver=GurobiSolver(env, OutputFlag=0))
        @variable(m, x[1:2])
        JuMP.fix.(x, x0)
        U = []
        X = [x]
        N = 3
        Δt = 0.1
        for i in 1:N
            u = @variable(m, lowerbound=-1, upperbound=1)
            xi = @variable(m, [1:2])
            @constraint(m, xi[2] == X[end][2] + Δt * u - 0.1)
            @constraint(m, xi[1] == X[end][1] + Δt * xi[2])
            push!(U, u)
            push!(X, xi)
        end

        @objective m Min sum([x[1]^2 + 0.01 * x[2]^2 for x in X]) + 0.01 * sum([(u - 0.5)^2 for u in U])
        solve(m)
        m, X, U
    end

    srand(1)
    correct = 0
    incorrect = 0
    for i in 1:100
        x0 = 0.1 .* randn(2)
        m, X, U = run_mpc(x0);
        ex = ExplicitQPs.explicit_solution(m, X[1]);
        s = ExplicitQPs.solution(ex, U[1])
        J = ExplicitQPs.jacobian(ex, [U[1]])

        if isapprox(getvalue(s), getvalue(U[1]), atol=1e-3, rtol=1e-3)
            correct += 1
        else
            incorrect += 1
        end

        for j in 1:2
            Δx = zeros(2)
            eps = 1e-9
            Δx[j] = eps
            m2, X2, U2 = run_mpc(x0 .+ Δx)
            if isapprox((getvalue(U2[1]) - getvalue(U[1])) / eps, J[1, j], rtol=1e-4)
                correct += 1
            else
                incorrect += 1
            end
        end
    end
    @test (correct / (correct + incorrect)) >= 0.85
end

max_stars_repo_name: rdeits/ExplicitQPs

max_issues_repo_name: rdeits/ExplicitQPs

max_forks_repo_name: rdeits/ExplicitQPs


text:
module Plots

using DataFrames, Plots

default(dpi = 300, titlefont = (16, "times"), guidefont = (12, "times"), tickfont = (12, "times"), legendfont = (6, "times"))

function runchart(raw::DataFrame, data::Symbol, batch::Symbol;
                center::Function = mean,
                spec::Real = 100,
                USL::Real = 110,
                LSL::Real = 90,
                sigma::Real = 3,
                title::String = "Run chart",
                xlabel::String = "Batch",
                ylabel::String = "Assay (%)")

    gdf = groupby(raw, batch)
    point = length(unique(raw[!, batch]))
    run = size(raw, 1) / point
    plot([1, point], repeat([USL], 2),color = :green, linestyle = :dash, label = "USL/LSL", legend = :topright)
    plot!([1, point] ,repeat([LSL], 2),color = :green, linestyle = :dash, label = false)
    plot!([1, point], repeat([median(combine(gdf, data => center)[!,2])], 2), color = :black, label = "Median", linestyle = :dash )
    if sigma > 0
        plot!([1, point], repeat([spec + sigma * std(raw[!, data])], 2), color = :chartreuse, linestyle = :dash, label = "±$(sigma)σ")
        plot!([1, point], repeat([spec - sigma * std(raw[!, data])], 2), color = :chartreuse, linestyle = :dash, label = false)
    end
    scatter!(raw[!, batch], raw[!, data], label = false, color = :deepskyblue, markeralpha = .3)
    plot!(1:point, combine(gdf, data => center)[!,2], color = :black, width = 3, label = "Batch $center")
    up = max(USL, raw[!, data]...)
    down = min(LSL, raw[!, data]...)
    Δ = up - down
    μ = (up + down) / 2
    ylims!(μ - 1.5 * Δ, μ + 2 * Δ)
    xlabel!(xlabel)
    ylabel!(ylabel)
    title!(title)
end

function runchart(raw::DataFrame, cpkresult::Vextor{Real}, data::Symbol, batch::Symbol;
                center::Function = mean,
                spec::Real = 100,
                USL::Real = 110,
                LSL::Real = 90,
                sigma::Real = 3,
                title::String = "Run chart",
                xlabel::String = "Batch",
                ylabel::String = "Assay (%)")

    gdf = groupby(raw, batch)
    point = length(unique(raw[!, batch]))
    run = size(raw, 1) / point
    plot([1, point], repeat([USL], 2),color = :green, linestyle = :dash, label = "USL/LSL", legend = :topright)
    plot!([1, point] ,repeat([LSL], 2),color = :green, linestyle = :dash, label = false)
    plot!([1, point], repeat([median(combine(gdf, data => center)[!,2])], 2), color = :black, label = "Median", linestyle = :dash )
    if sigma > 0
    plot!([1, point], repeat([spec + sigma * std(raw[!, data])], 2), color = :chartreuse, linestyle = :dash, label = "±$(sigma)σ")
    plot!([1, point], repeat([spec - sigma * std(raw[!, data])], 2), color = :chartreuse, linestyle = :dash, label = false)
    end
    good = findall(x->x > 1.33, cpkresult)
    medium = findall(x->(x < 1.33) & (x > 1), cpkresult)
    bad = findall(x->x < 1, cpkresult)
    isempty(good) || begin
        sub = filter(batch => x->x in good, raw)
        scatter!(sub[!, batch], sub[!, data], label = "Cpk > 1.33", color = :deepskyblue, markeralpha = .3)
    end
    isempty(medium) || begin
        sub = filter(batch => x->x in medium, raw)
        scatter!(sub[!, batch], sub[!, data], label = "1 < Cpk < 1.33", color = :red, markeralpha = .3)
    end
    isempty(bad) || begin
        sub = filter(batch => x->x in bad, raw)
        scatter!(sub[!, batch], sub[!, data], label = "Cpk < 1", color = :purple, markeralpha = .3)
    end
    plot!(1:point, combine(gdf, data => center)[!,2], color = :black, width = 3, label = "Batch $center")
    up = max(USL, raw[!, data]...)
    down = min(LSL, raw[!, data]...)
    Δ = up - down
    μ = (up + down) / 2
    ylims!(μ - 1.5 * Δ, μ + 2 * Δ)
    xlabel!(xlabel)
    ylabel!(ylabel)
    title!(title)
end

# xbarchart
# ichart: n(subgroup) = 1

# rchart: n(subgroup) = 2-8
# schart: n(subgroup) > 8
# mrchart: n(subgroup) = 1

# capability histogram
# capability plot

end


text:
# this is to read alexei's files
using DelimitedFiles, PyPlot
flist = readdlm("filelist", String)
dvec = Array{Array{Float64}, 1}(undef, length(flist))
for (i,f) in enumerate(flist)
    dvec[i] = readdlm(f, Float64, skipstart=1)
end
d = reduce(vcat, dvec)
long = d[:,1]
lat  = d[:,2]
σ = 1 ./d[:,3]
z = d[:,4]
# but I had to write it and reproject in Geosoft
writedlm("d_GDA94.txt", d)
# convert to Lambert conformal conic and then we have
d, h = readdlm("mohodepth_GDA_LCC.csv",',', header=true)
x, y = d[:,5], d[:,6]
f = figure()
s1 = subplot(121)
scatter(long, lat, c=z)
s1.set_aspect(1.)
s2 = subplot(122)
scatter(x, y, c=z)
s2.set_aspect(1.)

max_stars_repo_name: GeoscienceAustralia/HiQGA

max_issues_repo_name: GeoscienceAustralia/HiQGA


text:
mutable struct WasiConfig
    wasi_config_ptr::Ptr{wasi_config_t}
    program_name::String
end
WasiConfig(name) = WasiConfig(wasi_config_new(name), name)

Base.show(io::IO, config::WasiConfig) = write(io, "WasiConfig(", config.program_name, ")")
Base.unsafe_convert(::Type{Ptr{wasi_config_t}}, config::WasiConfig) = config.wasi_config_ptr

mutable struct WasiEnv
    wasi_env_ptr::Ptr{wasi_env_t}

    WasiEnv(wasi_env_ptr::Ptr{wasi_env_t}) =
        finalizer(new(wasi_env_ptr)) do wasi_env
            wasi_env_delete(wasi_env)
        end
end
function WasiEnv(wasi_config_ptr::Ptr{wasi_config_t})
    wasi_env_ptr =
        @ccall libwasmer.wasi_env_new(wasi_config_ptr::Ptr{wasi_config_t})::Ptr{wasi_env_t}
    @assert wasi_env_ptr != C_NULL "Failed to create env"

    WasiEnv(wasi_env_ptr)
end
WasiEnv(config::WasiConfig) = config.wasi_config_ptr |> WasiEnv

Base.unsafe_convert(::Type{Ptr{wasi_env_t}}, wasi_env::WasiEnv) = wasi_env.wasi_env_ptr

function wasi_env_read_stdout(wasi_env, size = 128)
    vec = Vector{Int8}(undef, size)
    nread = LibWasmer.wasi_env_read_stdout(
        wasi_env,
        pointer(vec),
        size,
    )

    vec[begin:nread]
end
function wasi_env_read_stderr(wasi_env, size = 128)
    vec = Vector{Int8}(undef, size)
    nread = LibWasmer.wasi_env_read_stderr(
        wasi_env,
        pointer(vec),
        size,
    )

    vec[begin:nread]
end

function get_imports(store::WasmStore, wasm_module::WasmModule, wasi_env)
    extern_vec = WasmPtrVec(wasm_extern_t)

    res = wasi_get_imports(
        store,
        wasm_module,
        wasi_env,
        extern_vec,
    )
    @assert res "Failed to get WASI imports"

    extern_vec
end


text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

"""
Contains all complex data used for testing and benchmarking.
"""
module TestData

using OpticSim
using OpticSim: tobeziersegments # try to use only exported functions so this list should stay short
using OpticSim.Geometry
using OpticSim.Emitters
using OpticSim.GlassCat

using StaticArrays
using DataFrames: DataFrame
using Unitful
using Images
using Base: @.
using LinearAlgebra

include("curves.jl")
include("surfaces.jl")
include("lenses.jl")
include("systems.jl")
include("other.jl")

end # module TestData
export TestData


text:
# syntax: proto3
using ProtoBuf
import ProtoBuf.meta

mutable struct TopologyProto <: ProtoType
    mesh_shape::Base.Vector{Int32}
    num_tasks::Int32
    num_tpu_devices_per_task::Int32
    device_coordinates::Base.Vector{Int32}
    TopologyProto(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct TopologyProto
const __pack_TopologyProto = Symbol[:mesh_shape,:device_coordinates]
meta(t::Type{TopologyProto}) = meta(t, ProtoBuf.DEF_REQ, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, __pack_TopologyProto, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

export TopologyProto


text:
using Distributed, DataFrames, CSV
addprocs(33)

@everywhere  begin
    using Jedi
    using LinearAlgebra
    using Distributions
    using SharedArrays
    emat = 2 * (ones(4, 4) - Matrix{Float64}(I, 4, 4))
    N = 1000
    f0 = 20/2N
end
rho = [0, 0.1, 0.5, 1., 2]
E = SharedArray{Float64, 3}(length(rho),33, 1000)
L = SharedArray{Float64, 3}(length(rho),33, 1000)
RHO = SharedArray{Float64, 3}(length(rho),33, 1000)

l_arr = collect(8:40)

@sync @distributed for k in 1:33
    for r in 1:length(rho)
        f = fermi_fitness(f0=f0, l=l_arr[k])
        for j in 1:1000
            pop = driver_trailer(N=1000, l=l_arr[k], L=l_arr[k])
            initiate_rand!(pop, 1)
            for i in 1:500000
                Jedi.bp_substitution!(pop, emat, f)
                if rand() < rho[r]/N
                    driver_mutation!(pop)
                end
            end
            E[r, k, j] = get_energy(pop, emat)[1]
            L[r, k, j] = l_arr[k]
            RHO[r, k, j] = rho[r]
        end
        println("Rho $(rho[r]), length $(l_arr[k]) done.")
    end
end

df = DataFrame(gamma=[(E...)...], l=[(L...)...], rho=[(RHO...)...])
CSV.write("script3_results.csv", df)


text:
__precompile__(true)
"""
Strs package

Copyright 2017-2019 Gandalf Software, Inc., Scott P. Jones
Licensed under MIT License, see LICENSE.md
"""
module Strs

using ModuleInterfaceTools
export @api

using PCRE2

import StrTables: lookupname, matchchar, matches, longestmatches, completions
export lookupname, matchchar, matches, longestmatches, completions

@api extend! StrRegex, StrLiterals, Format
@api use! StrFormat, StrEntities

using InternedStrings: intern
export @i_str, @I_str, intern

macro i_str(str) ; Expr(:call, :intern, interpolated_parse(str, StrBase._str)) ; end
macro I_str(str) ; Expr(:call, :intern, interpolated_parse(str, StrBase._str, true)) ; end

@api reexport StrRegex, StrLiterals

@api freeze

end # module Strs


text:
using YaoBase, CacheServers
export CacheFragment, CachedBlock, update_cache
export cache, pull, update!, update_cache, clearall!, iscached, iscacheable

"""
    CacheFragment{BT, K, MT}

A fragment that will be stored for each cached block (of type `BT`) on a cache server.
"""
struct CacheFragment{BT,K,MT}
    ref::BT
    storage::Dict{K,MT}

    function CacheFragment{BT,K,MT}(x::BT) where {BT,K,MT}
        new{BT,K,MT}(x, Dict{K,MT}())
    end

    function CacheFragment{BT,K}(x::BT) where {BT} where {K}
        new{BT,K,cache_type(BT)}(x, Dict{K,cache_type(BT)}())
    end

    function CacheFragment(x::BT) where {BT}
        CacheFragment{BT,typeof(cache_key(x))}(x)
    end
end

# default update rule
function CacheServers.update!(frag::CacheFragment, val)
    if !iscached(frag)
        frag.storage[cache_key(frag.ref)] = val
    end
    frag
end

CacheServers.iscached(frag::CacheFragment) = cache_key(frag.ref) in keys(frag.storage)
CacheServers.pull(frag::CacheFragment) = frag.storage[cache_key(frag.ref)]
CacheServers.clear!(frag::CacheFragment) = (empty!(frag.storage); frag)

"""
    CachedBlock{ST, BT, D} <: TagBlock{BT, D}

A label type that tags an instance of type `BT`. It forwards
every methods of the block it contains, except [`mat`](@ref)
and [`apply!`](@ref), it will cache the matrix form whenever
the program has.
"""
struct CachedBlock{ST,BT,D} <: TagBlock{BT,D}
    server::ST
    content::BT
    level::Int

    function CachedBlock(server::ST, x::BT, level::Int) where {ST,D,BT<:AbstractBlock{D}}
        alloc!(server, x, CacheFragment(x))
        new{ST,BT,D}(server, x, level)
    end
end

CacheServers.iscached(c::CachedBlock) = iscached(c.server, c.content)
iscacheable(c::CachedBlock) = iscacheable(c.server, c.content)
chsubblocks(cb::CachedBlock, blk::AbstractBlock) = CachedBlock(cb.server, blk, cb.level)
occupied_locs(x::CachedBlock) = occupied_locs(content(x))
PropertyTrait(::CachedBlock) = PreserveAll()

function update_cache(::Type{T}, c::CachedBlock) where {T}
    if !iscached(c.server, c.content)
        m = mat(T, c.content)
        push!(c.server, m, c.content)
    end
    return c
end

CacheServers.clear!(x::AbstractBlock) = x
CacheServers.clear!(c::CachedBlock) = (clear!(c.server, c.content); c)

function mat(::Type{T}, c::CachedBlock) where {T}
    if !iscached(c.server, c.content)
        m = mat(T, c.content)
        push!(c.server, m, c.content)
        return m
    end
    return pull(c)
end

function CacheServers.pull(c::CachedBlock)
    return pull(c.server, c.content)
end

function _apply!(r::AbstractArrayReg{D,T}, c::CachedBlock, signal) where {D,T}
    if signal > c.level
        r.state .= mat(T, c) * r
    else
        _apply!(r, c.content)
    end
    return r
end

_apply!(r::AbstractRegister, c::CachedBlock) = _apply!(r, c.content)
_apply!(r::AbstractArrayReg{D,T}, c::CachedBlock) where {D,T} = (r.state .= mat(T, c) * r.state; r)

Base.similar(c::CachedBlock, level::Int) = CachedBlock(c.server, c.content, level)
Base.copy(c::CachedBlock) = CachedBlock(c.server, copy(c.content), c.level)
Base.length(x::CachedBlock) = length(content(x))
Base.getindex(c::CachedBlock, index...) = getindex(content(c), index...)
Base.setindex!(c::CachedBlock, val, index...) = setindex!(content(c), val, index...)
Base.iterate(c::CachedBlock) = iterate(content(c))
Base.iterate(c::CachedBlock, st) = iterate(content(c), st)
Base.eltype(x::CachedBlock) = eltype(content(x))

const DefaultCacheServer = get_server(AbstractBlock, CacheFragment)

cache(x::Function, level::Int = 1; recursive = false) =
    n -> cache(x(n), level; recursive = recursive)

"""
    cache(x[, level=1; recursive=false])

Create a [`CachedBlock`](@ref) with given block `x`, which will cache the matrix of `x` for the first time
it calls [`mat`](@ref), and use the cached matrix in the following calculations.

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> cache(control(3, 1, 2=>X))
nqudits: 3
[cached] control(1)
   └─ (2,) X

julia> chain(cache(control(3, 1, 2=>X)), repeat(H))
nqudits: 3
chain
└─ [cached] control(1)
      └─ (2,) X

```
"""
function cache(x::AbstractBlock, level::Int = 1; recursive = false)
    return cache(DefaultCacheServer, x, level, recursive = recursive)
end

function clearall!(x::CachedBlock)
    return clear!(x)
end

function clearall!(x::CachedBlock{ST,BT}) where {ST,BT<:CompositeBlock}
    for each in subblocks(x.content)
        clearall!(each)
    end
    clear!(x)
    return x
end

function cache(
    server::AbstractCacheServer,
    x::AbstractBlock,
    level::Int;
    recursive::Bool = false,
)
    return CachedBlock(server, x, level)
end

function cache(
    server::AbstractCacheServer,
    x::ChainBlock,
    level::Int;
    recursive::Bool = false,
)
    if recursive
        chain = similar(x)
        for (i, each) in enumerate(block)
            chain[i] = cache(server, each, level, recursive)
        end
    else
        chain = x
    end

    return CachedBlock(server, chain, level)
end

function cache(
    server::AbstractCacheServer,
    block::KronBlock,
    level::Int;
    recursive::Bool = false,
)
    if recursive
        x = similar(block)
        for (k, v) in block
            x[k] = cache(server, v, level, recursive = recursive)
        end
    else
        x = block
    end

    return CachedBlock(server, x, level)
end


text:
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

compute(n::Int)::Int = Int(ceil((log(10) * (n - 1) + 0.5log(5)) / log(0.5(1 + √5))))

compute(3)

compute(1_000)

@benchmark compute(1_000)

max_stars_repo_name: Dynortice/Project-Euler

max_issues_repo_name: Dynortice/Project-Euler

max_forks_repo_name: Dynortice/Project-Euler


text:
# This file was generated, do not modify it. # hide
reuse!(res)
push!(res, 1, 10)
push!(res, 2, 9)
push!(res, 3, 8)
push!(res, 4, 7)
push!(res, 6, 5)
@show res

max_stars_repo_name: sadit/SimilaritySearchDemos

max_issues_repo_name: sadit/SimilaritySearchDemos

max_forks_repo_name: sadit/SimilaritySearchDemos


text:
using VDPTag
using Plots
using POMDPToolbox
using Reel
using ProgressMeter
using PmapProgressMeter
using ContinuousPOMDPTreeSearchExperiments
using MCTS
using ParticleFilters
using POMCPOW
using POMCP

@everywhere begin
    using POMDPs
    using POMDPToolbox
    using MCTS
    using VDPTag
    using POMCPOW
    using POMCP
    using ContinuousPOMDPTreeSearchExperiments
    using ParticleFilters

    pomdp = VDPTagPOMDP()
    seed=27

    rollout_policy = ToNextML(mdp(pomdp))
    solver = POMCPOWSolver(tree_queries=100_000,
                           criterion=MaxUCB(40.0),
                           final_criterion=MaxTries(),
                           max_depth=20,
                           k_action=8.0,
                           alpha_action=1/20,
                           k_observation=4.0,
                           alpha_observation=1/20,
                           estimate_value=FORollout(rollout_policy),
                           # check_repeat_act=false,
                           # check_repeat_obs=false,
                           rng=MersenneTwister(seed)
                          )

    updater = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(100_000), 0.0, MersenneTwister(seed+100))
end

#=
N = 100
# s_rewards = SharedArray(Float64, N)
prog = Progress(N, desc="Simulating...")
rewards = pmap(prog, 1:N) do i
# @showprogress "Simulating..." for i in 1:N
    policy = solve(solver, pomdp)

    hr = HistoryRecorder(max_steps=100, rng=MersenneTwister(i))
    hist = simulate(hr, pomdp, policy, updater)
    discounted_reward(hist)
end
@show mean(rewards)
=#

#=
policy = solve(solver, pomdp)
action(policy, initial_state_distribution(pomdp))
blink(policy)
=#

#=
policy = solve(solver, pomdp)

hr = HistoryRecorder(max_steps=100, rng=MersenneTwister(1), show_progress=true)
hist = simulate(hr, pomdp, policy, updater)
@show discounted_reward(hist)

gr()
frames = Frames(MIME("image/png"), fps=2)
@showprogress "Creating gif..." for i in 1:length(hist)
    push!(frames, plot(pomdp, view(hist, 1:i)))
end

filename = string(tempname(), "_vdprun.gif")
write(filename, frames)
println(filename)
run(`setsid gifview $filename`)
=#


text:
ENV["GKSwstype"] = "nul"  # For StatsPlots
println("pid: ", getpid())

# Read command line args.
if length(ARGS) > 1
  resultsdir = ARGS[1]
  awsbucket = ARGS[2]
  snum = parse(Int, ARGS[3])
  println("ARGS: ", ARGS)
else
  resultsdir = "results/test/"
  awsbucket = nothing
  snum = 4
end
flush(stdout)

# Create results/images dir if needed
imgdir = "$(resultsdir)/img"
mkpath(imgdir)

import Pkg; Pkg.activate(joinpath(@__DIR__, "../../../"))

using CytofDensityEstimation
const CDE = CytofDensityEstimation
using LaTeXStrings
using Distributions
using StatsPlots
using HypothesisTests
import Random
using RCall
using BSON
import CytofDensityEstimation.Model.default_ygrid

include("scenarios.jl")

plotsize = (400, 400)

# Simulate data.
Random.seed!(5);
simdata = scenarios(snum)
yC, yT = simdata[:yC], simdata[:yT]

# Plot true data density.
function plot_trude_data_density()
  plot(default_ygrid(), pdf.(simdata[:mmC], default_ygrid()), lw=3,
       label=L"y_C", color=:blue)
  plot!(default_ygrid(), pdf.(simdata[:mmT], default_ygrid()), lw=3,
        label=L"y_T", color=:red)
end
plot_trude_data_density()
plot!(size=plotsize)
savefig(joinpath(imgdir, "data-true-density.pdf"))
closeall()

# Plot observed data density.
function plot_observed_data_density()
  density(yC[isfinite.(yC)],  lw=3, label=L"y_C", color=:blue)
  density!(yT[isfinite.(yT)], lw=3, label=L"y_T", color=:red)
end
plot_observed_data_density()
plot!(size=plotsize)
savefig(joinpath(imgdir, "data-kde.pdf"))
closeall()

# Define data, prior, and initial state.
Random.seed!(4) # 7
K = 7 # 3
data = CDE.Model.Data(yC, yT)
prior_mu = let
  yfinite = [data.yC_finite; data.yT_finite]
  Normal(mean(yfinite), std(yfinite))
end
prior = CDE.Model.Prior(K, mu=prior_mu, nu=LogNormal(1.6, 0.4),
                        p=Beta(1, 99), omega=InverseGamma(.1, .1),
                        eta=Dirichlet(K, 1 / K), psi=Normal(-1, 1))
state = CDE.Model.State(data, prior)
tuners = CDE.Model.Tuners(K, 0.1)

println("Priors:")
for fn in fieldnames(CDE.Model.Prior)
  println(fn, " => ", getfield(prior, fn))
end

# Run chain.
state.beta = 0
state.p = 0.01
@time chain, laststate, summarystats = CDE.ppfit(state, data, prior, tuners,
                                                 nsamps=[2000],
                                                 nburn=500, thin=2, fix=[:p],
                                                 rep_aux=10, warmup=4000)

# Save results
BSON.bson("$(resultsdir)/results.bson",
          Dict(:chain=>chain, :laststate=>laststate,
               :summarystats=>summarystats, :simdata=>simdata))

# Load via:
out = BSON.load("$(resultsdir)/results.bson")

# Print KS statistic.
@rimport stats as rstats
ks_fit = rstats.ks_test(yC, yT)
println("KS p-value: ", ks_fit["p.value"][1])
flush(stdout)

# Print summary statistics.
CDE.Model.printsummary(out[:chain], out[:summarystats])
flush(stdout)

# Plot results.
@time CDE.Model.plotpostsummary(out[:chain], out[:summarystats], yC, yT,
                                imgdir, simdata=out[:simdata], bw_postpred=.2,
                                xlims_=(-6, 6))
flush(stdout)

# Bayes Factor
VD = Vector{Dict{Symbol, Any}}
chain1 = filter(s -> s[:beta] == 1, out[:chain][1])
chain0 = filter(s -> s[:beta] == 0, out[:chain][1])
@time if length(chain1) > 0 && length(chain0) > 0
  @time lbf = CDE.Model.log_bayes_factor(data, VD(chain0), VD(chain1))
  pp1 = CDE.Model.posterior_prob1(bf, logpriorodds=logit(state.p))
  println("Log Bayes Factor: $(lbf) | Post. prob. for model 1: $(pp1)")

  # DIC
  dic0, dic1 = CDE.dic(VD(chain0), data), CDE.dic(VD(chain1), data)
  println("(DIC0, DIC1): ($(round(dic0, digits=3)), $(round(dic1, digits=3)))")
else
  if length(chain1) > 0
    println("M1 is preferred.")
  else
    println("M0 is preferred.")
  end
end
flush(stdout)

if awsbucket != nothing
  CDE.Util.s3sync(from=resultsdir, to=awsbucket, tags=`--exclude '*.nfs'`)
end

println("Done!")
flush(stdout)

max_stars_repo_name: luiarthur/CytofDensityEstimation

max_issues_repo_name: luiarthur/CytofDensityEstimation

max_forks_repo_name: luiarthur/CytofDensityEstimation


text:
using JSON, Statistics, GitHub, Base64, SHA, Downloads, Dates

function cpu_key()
    cpus = map(Sys.cpu_info()) do cpu
        replace(cpu.model, " " => "")
    end
    return join(unique(cpus), "")
end
julia_key() = "julia-" * replace(string(VERSION), "."=>"-")
tag_commit(ctx, tag) = GitHub.tag(ctx.repo, tag).object["sha"]
current_commit() =  chomp(read(`git rev-parse HEAD`, String))
latest_commit(ctx, branch) = GitHub.branch(ctx.repo, branch)
function github_context()
    owner = "JuliaPlots"
    return (
        owner = owner,
        repo = GitHub.Repo("$(owner)/Makie.jl"),
        auth = GitHub.authenticate(ENV["GITHUB_TOKEN"]),
        scratch_repo = GitHub.Repo("MakieOrg/scratch")
    )
end

struct BenchInfo
    julia::String
    cpu::String
    branch::String
    commit::String
    commit_message::String
    project::String
end

function Base.show(io::IO, info::BenchInfo)
    branch = isempty(info.branch) ? "" : " $(info.branch),"
    msg = info.commit_message[1:min(length(info.commit_message), 20)]
    print(io, "BenchInfo($(info.julia), $(info.cpu),$(branch) $(msg))")
end

function best_name(info::BenchInfo)
    isempty(info.branch) || return info.branch
    return info.commit[1:5]
end

function unique_id(info::BenchInfo)
    isempty(info.commit) || return info.commit
    return info.branch
end

function BenchInfo(;
        julia::String=julia_key(),
        cpu::String=cpu_key(),
        branch::String="",
        commit::String="",
        commit_message::String="",
        project=""
    )
    return BenchInfo(
        julia,
        cpu,
        branch,
        commit,
        commit_message,
        project
    )
end

function BenchInfo(commit::GitHub.Commit; kw...)
    isnothing(commit.commit) && error("Invalid commit")
    return BenchInfo(
        commit=commit.sha,
        commit_message=commit.commit.message,
    )
end

function BenchInfo(repo, commit_sha::AbstractString; kw...)
    commit = GitHub.commit(repo, commit_sha)
    return BenchInfo(commit; kw...)
end

function get_tag_info(tag::GitHub.Tag)
    str = sprint() do io
        Downloads.download(tag.object["url"], io)
    end
    dict = JSON.parse(str)
    parts = splitpath(tag.url.path) # why is this not straight-forward?
    repo = GitHub.Repo(parts[3] * "/" * parts[4])
    return (
        commit = dict["object"]["sha"],
        name = dict["tag"],
        repo = repo,
        message = dict["message"],
    )
end

function BenchInfo(tag::GitHub.Tag; kw...)
    info = get_tag_info(tag)
    return BenchInfo(info.repo, info.commit; branch=info.name, kw...)
end

function BenchInfo(branch::GitHub.Branch; kw...)
    return BenchInfo(branch.commit; branch=branch.name, kw...)
end

function get_file(ctx, repo_path)
    file = GitHub.file(ctx.scratch_repo, repo_path; auth=ctx.auth, handle_error=false)
    if isnothing(file.sha)
        return (sha=nothing, content=nothing)
    else
        return (sha=file.sha, content = base64decode(file.content))
    end
end

upload_file(ctx, file, repo_path) = upload_data(ctx, read(file), repo_path, file)
function upload_data(ctx, bytes_or_str::Union{AbstractString, AbstractVector{UInt8}}, repo_path, file_name="data")
    b64 = Base64.base64encode(bytes_or_str)
    # see if there is an old file, (handle_error=false won't error, and instead returns file with all nothing)
    old_file = get_file(ctx, repo_path)
    params = Dict("content"=>b64, "message"=>"update benchmark for $(file_name)")
    if !isnothing(old_file.sha)
        # if there was a file already, we need to supply its sha
        params["sha"] = old_file.sha
    end
    GitHub.create_file(ctx.scratch_repo, repo_path; auth=ctx.auth, params=params)
    return "https://github.com/MakieOrg/scratch/blob/main/$(repo_path)?raw=true"
end

function make_or_edit_comment(ctx, pr, comment)
    prev_comments, _ = GitHub.comments(ctx.repo, pr; auth=ctx.auth)
    idx = findfirst(c-> c.user.login == "MakieBot", prev_comments)
    if isnothing(idx)
        GitHub.create_comment(ctx.repo, pr; auth=ctx.auth, params=Dict("body"=>comment))
    else
        GitHub.edit_comment(ctx.repo, prev_comments[idx], :pr; auth=ctx.auth, params=Dict("body"=>comment))
    end
end

function run_bench(info::BenchInfo; n=5)
    commit = info.commit
    if info.cpu != cpu_key()
        error("Not running on requested CPU. Request: $(info.cpu), actual: $(cpu_key())")
    end
    if info.julia != julia_key()
        error("Not running on requested Julia. Requested: $(info.julia), actual: $(julia_key())")
    end

    if isdir(info.project)
        project = info.project
        @info("using $(project)")
    else
        @info("Creating project for $(commit)")
        project = "benchmark-project"
        isdir(project) && rm(project; force=true, recursive=true)
        Pkg.activate(project)
        pkgs = ["MakieCore", "Makie", "CairoMakie"]
        pkgs = [PackageSpec(name=pkg, rev=commit) for pkg in pkgs]
        Pkg.add(pkgs)
    end

    results = Vector{Float64}[]
    for i in 1:n
        result = read(`$(Base.julia_cmd()) --project=$(project) ./benchmark-ttfp.jl`, String)
        tup = eval(Meta.parse(result))
        @show tup
        push!(results, [tup...])
    end
    Pkg.activate(".") # activate the main project
    return results
end

bench_info(ctx, info::BenchInfo) = info
bench_info(ctx, gh_typ::GitHub.GitHubType) = BenchInfo(gh_typ)
bench_info(ctx, commit_or_smth) = BenchInfo(ctx.repo, commit_or_smth)

function get_benchmark_data(ctx, info::BenchInfo; n=10, force=false)
    uuid = unique_id(info)
    repo_dir_path = "benchmarks/$(julia_key())/$(cpu_key())"
    repo_base_path = "$(repo_dir_path)/$(uuid)"
    repo_data_path = "$(repo_base_path).json"
    @info("Getting $(repo_data_path)")
    old_file = get_file(ctx, repo_data_path)
    if !force && !isnothing(old_file.content)
        @info("Benchmark already exists, loading from file")
        return JSON.parse(String(old_file.content))
    else
        @info("Benchmark doesn't exist, run benchmark")
        global results = run_bench(info; n=n)

        result_with_info = Dict(
            "branch" => info.branch,
            "message" => info.commit_message,
            "date" => string(now()),
            "results" => results
        )
        upload_data(ctx, JSON.json(result_with_info), repo_data_path, "for commit $(uuid[1:5])")
        return result_with_info
    end
end

function plot_benchmark!(ax, data, pos; width=0.9, height=0.3, alpha=0.2)
    used = first.(data)
    ttfp = last.(data)
    sum_t = used .+ ttfp
    n = length(data)
    bplot!(d, c) = scatter!(ax, fill(pos, n), d, marker=Rect, markerspace=:data, markersize=Vec2f(width, height), color=(c, alpha))

    bplot!(used, :darkred)
    bplot!(ttfp, :blue)
    bplot!(sum_t, :black)
    hwidth = (width / 2) .+ 0.01
    vlines!(ax, [pos - hwidth, pos + hwidth], linewidth=0.5, color=(:black, 0.5))
    return (used, ttfp, sum_t)
end

function plot_benchmarks(benchmarks, benchmark_infos)
    f = Figure()
    ax = Axis(f[1,1]; ylabel="median time (s)")
    mdata = Vector{Float64}[]
    xticks = Int[]
    xticklabels = String[]
    for (i, data) in enumerate(benchmarks)
        pdata = data["results"]
        info = benchmark_infos[i]
        push!(mdata, plot_benchmark!(ax, pdata, i)...)
        push!(xticks, i)
        push!(xticklabels, best_name(info))
    end
    medians = median.(mdata)
    medians = unique(round.(median.(mdata)))
    ax.yticks = medians
    ax.xticks = (xticks, xticklabels)
    legend_elems = [
        PolyElement(color=:darkred),
        PolyElement(color=:blue),
        PolyElement(color=:black)]
    legend_names = ["using time", "plot time", "total time"]
    Legend(f[1, 1], legend_elems, legend_names;
        tellheight = false,
        tellwidth = false,
        margin = (10, 10, 10, 10), bgcolor = (:white, 0.5),
        framewidth = 0.5,
        halign = :right, valign = :bottom)
    hidexdecorations!(ax, ticklabels=false)
    f
end

function upload_data(ctx, fig, repo_path)
    io = IOBuffer()
    show(io, MIME"image/png"(), fig)
    bytes = take!(io)
    upload_data(ctx, bytes, repo_path, "plot")
end

function run_benchmarks(ctx, to_benchmark;
        n=10, force=false, pr_to_comment=get(ENV, "PR_NUMBER", nothing))

    bench_infos = bench_info.(Ref(ctx), to_benchmark)
    @info("benchmarking:")
    display(bench_infos)

    benchmarks = get_benchmark_data.(Ref(ctx), bench_infos)

    @info("done benchmarking, plotting")
    fig = plot_benchmarks(benchmarks, bench_infos)

    name = join(map(best_name, bench_infos), "-vs-")

    @info("uploading plot $(name) to github")
    image_url = upload_data(ctx, fig, "benchmarks/$(name).png")

    comment = """
    ## Compile Times benchmark

    ![]($(image_url))
    """

    if !isnothing(pr_to_comment)
        @info("Commenting plot on PR $(pr_to_comment)")
        pr = GitHub.pull_request(ctx.repo, pr_to_comment)
        make_or_edit_comment(ctx, pr, comment)
    else
        @info("No comment, no PR found")
    end
    return fig
end


text:
module FIB

using POMDPs
using POMDPModelTools
using POMDPPolicies
using Printf

import POMDPs: Solver, Policy, solve

export
    FIBSolver,
    solve

include("vanilla.jl")

end # module


text:
c1 = ConstrainedValueResampling(TruncateStd(1), TruncateStd(1), TruncateQuantiles(0.1, 0.2))
c2 = ConstrainedValueResampling(100, TruncateStd(1), TruncateStd(1), TruncateQuantiles(0.1, 0.2))

@test c1 isa ConstrainedValueResampling{3}
@test c1.n == 1
@test c2 isa ConstrainedValueResampling{3}
@test c2.n == 100


text:
export CuFunction

type CuFunction
    m::CuModule # avoid CuModule gc-ed
    ptr::CUfunction
end

function CuFunction(m::CuModule, name::String)
    p = CUfunction[0]
    cuModuleGetFunction(p, m, name)
    CuFunction(m, p[1])
end

Base.unsafe_convert(::Type{CUfunction}, f::CuFunction) = f.ptr


text:
################################################################################
#  Copyright 2021, Marta Vanin                                                 #
################################################################################
"solves the reduced state estimation in current and voltage rectangular coordinates (ReducedIVR formulation)"
function run_linear_ivr_red_mc_se(data::Union{Dict{String,<:Any},String}, solver; kwargs...)
    return run_mc_selin(data, PowerModelsDistributionStateEstimation.ReducedIVRUPowerModel, solver; kwargs...)
end

function run_mc_selin(data::Union{Dict{String,<:Any},String}, model_type::Type, solver; kwargs...)
    if haskey(data["se_settings"], "criterion")
        _PMDSE.assign_unique_individual_criterion!(data)
    end
    if !haskey(data["se_settings"], "rescaler")
        data["se_settings"]["rescaler"] = 1
        @warn "Rescaler set to default value, edit data dictionary if you wish to change it."
    end
    if !haskey(data["se_settings"], "number_of_gaussian")
        data["se_settings"]["number_of_gaussian"] = 10
        @warn "Estimation criterion set to default value, edit data dictionary if you wish to change it."
    end
    return PowerModelsDistribution.run_mc_model(data, model_type, solver, build_mc_selin; kwargs...)
end

"specification of an exact linear state estimation problem for the reduced IVR Flow formulation
The difference with the regular state estimation available in PowerModelsDistributionStateEstimation is that
variable_mc_bus_voltage and constraint_mc_gen_setpoint_se are unbounded "
function build_mc_selin(pm::PowerModelsDistributionStateEstimation.ReducedIVRUPowerModel)
    # Variables

     PowerModelsDistribution.variable_mc_bus_voltage(pm, bounded = false)
     PowerModelsDistributionStateEstimation.variable_mc_branch_current(pm, bounded = true)
     PowerModelsDistributionStateEstimation.variable_mc_generator_current_se(pm, bounded = true)
     PowerModelsDistributionStateEstimation.variable_mc_load_current(pm, bounded = true)
     PowerModelsDistributionStateEstimation.variable_mc_residual(pm, bounded = true)
     PowerModelsDistributionStateEstimation.variable_mc_measurement(pm, bounded = true)

    # Constraints
    for (i,bus) in PowerModelsDistribution.ref(pm, :ref_buses)
        @assert bus["bus_type"] == 3
        PowerModelsDistribution.constraint_mc_theta_ref(pm, i)
    end

    # gens should be constrained before KCL, or Pd/Qd undefined
    for id in PowerModelsDistribution.ids(pm, :gen)
        PowerModelsDistributionStateEstimation.constraint_mc_generator_power_se(pm, id; bounded = false)
    end

    for (i,bus) in PowerModelsDistribution.ref(pm, :bus)
        PowerModelsDistributionStateEstimation.constraint_mc_current_balance_se(pm, i)
    end

    for i in PowerModelsDistribution.ids(pm, :branch)
        PowerModelsDistributionStateEstimation.constraint_current_to_from(pm, i)
        PowerModelsDistributionStateEstimation.constraint_mc_bus_voltage_drop(pm, i)
    end

    for (i,meas) in PowerModelsDistribution.ref(pm, :meas)
        PowerModelsDistributionStateEstimation.constraint_mc_residual(pm,i)
    end

    PowerModelsDistributionStateEstimation.objective_mc_se(pm)

end


text:
# A Julia implementation of the Affine ciphertext

# Determines if 2 values are coprime
function coprime(x, y)
    xlist = []
    ylist = []
    for i in 2 : x
        if x % i == 0
            push!(xlist, i)
        end
    end
    for j in 2 : y
        if y % j == 0
            push!(ylist, j)
        end
    end
    for a in xlist
        for b in ylist
            if a == b
                return false
            end
        end
    end
    return true
end

# Changes a char to an int value
function charToVal(c)
    alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    val = findfirst(isequal(c), alphabet)
    return val
end

# Changes a value to a char1
function valToChar(i)
    alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    char = alphabet[i]
    return char
end

# Main function
function affine()

    m = 26 # The number of letters in the alphabet used
    println("Please input the plaintext")
    plaintext = readline()
    plaintext = uppercase(plaintext)
    plaintext = replace(plaintext, " " => "")

    # Take keys and ensure they're valid
    println("Please enter the a value")
    a = parse(Int, readline())
    while !coprime(a, m)
        println("a value invalid, must be coprime with m value " * string(m))
        a = parse(Int, readline())
    end
    println("Please enter the b value")
    b = parse(Int, readline())

    # Produce the ciphertext
    ciphertext = ""
    for c in plaintext
        val = charToVal(c) - 1 # Must account for Julia being 1-indexed
        newVal = mod(((a * val) + b), m)
        newchar = valToChar(newVal + 1)
        ciphertext *= newchar
    end
    println(ciphertext)
end

affine()