text:
"""
`Model`

A struct representing a model.

**Fields**
- `variable_names`: Variable names (e.g. `["X", "Y", "Z"]`).
- `parameter_names`: Parameter names (e.g. `["v", "d1", "d2"]`).
- `problem`: `ODEProblem` or `SDEProblem` or `JumpProblem`.
- `solver_algorithm`: Solver used to solve the `problem`` (e.g. `SOSRI()`).
- `solver_parameters`: Parameters passed to the solver (e,g, `(saveat=0.1,)`)
- `input`: Input to the model representedas a tuple with two elements. The first
    element is a matrix representing the input events and the second element is
    a string that specifies, which parameter is manipulated by the events.
- `output`: Function that transforms `DifferentialEqualtions` solution
    into a matrix. Rows are time steps and columns are variables.

**Functions**
- `set_initial_conditions`
- `set_timespan`
- `create_model`
- `set_solver`
- `get_parameter_index`
- `get_parameter_value`
- `set_parameter_value`
- `set_input`
- `simulate_population`
"""
struct Model
    variable_names::Vector{String}
    parameter_names::Vector{String}
    problem::Union{ODEProblem, SDEProblem, JumpProblem}
    solver_algorithm::Any
    solver_parameters::NamedTuple
    input::Tuple{Matrix{Float64}, String}
    output::Function
end

"""
`_set_model_property(model::Model; kwargs...)`

Create a new copy of `model` with replaced fields as specified by the keyword
    arguments.
"""
function _set_model_property(model::Model; kwargs...)

    # Iterate all possible fields of Model struct
    fields = []
    for field in fieldnames(Model)

        # If the field is in kwargs, use it, otherwise use the original value
        if field in keys(kwargs)
            push!(fields, deepcopy(kwargs[field]))
        else
            push!(fields, deepcopy(getproperty(model, field)))
        end

    end

    # Create a model with new fields
    return Model(fields...)

end

function _get_problem_property(problem, property_name)

    # Property name must be a symbol
    property_name = Symbol(property_name)

    # For JumpProblem access problem.prob, otherwise problem directly
    if problem isa JumpProblem
        property = getproperty(problem.prob, property_name)
    else
        property = getproperty(problem, property_name)
    end

    # Return the found property
    return property

end

function _get_problem_property(model::Model, property_name)

    return _get_problem_property(model.problem, property_name)

end

function _set_problem_property(model::Model; kwargs...)

    # Remake problem with the new properties
    new_problem = remake(model.problem; kwargs...)

    # Replace the problem in the model
    new_model = _set_model_property(model, problem=new_problem)

    return new_model
end

"""
`set_initial_conditions(model::Model, u0)`

Set initial conditions for a model.
"""
function set_initial_conditions(model::Model, u0)
    return _set_problem_property(model; u0=u0)
end

"""
`set_timespan(model::Model, tspan)`

Set timespan for a model.
"""
function set_timespan(model::Model, tspan)
    return _set_problem_property(model; tspan=tspan)
end

function set_output(model::Model, output)
    return _set_model_property(model; output=output)
end

"""
`create_model`

Create a model based on the `DifferentialEqualtions` problem.

**Arguments**
- `variable_names`: Variable names (e.g. `["X", "Y", "Z"]`).
- `parameter_names`: Parameter names (e.g. `["v", "d1", "d2"]`).
- `problem`: `ODEProblem` or `SDEProblem` or `JumpProblem`.

**Returns**
- `model`: Model represented as a `Model` struct.
"""
function create_model(variable_names, parameter_names, problem)

    # Check that variable names match the length of the initial state
    u0 = _get_problem_property(problem, "u0")
    if length(variable_names) != length(u0)
        error("Incorrect number of variable names!")
    end

    # Check that parameter names match the length of the parameter vector
    p = _get_problem_property(problem, "p")
    if length(parameter_names) != length(p)
        error("Incorrect number of parameter names!")
    end

    # Set up a default algorithm based on the problem type
    if problem isa ODEProblem
        solver_algorithm = Tsit5()
    elseif problem isa SDEProblem
        solver_algorithm = SOSRI()
    elseif problem isa JumpProblem
        solver_algorithm = SSAStepper()
    else
        error("Unknown problem type!")
    end

    # Set up the other default parameters the do not depend on the problem type
    solver_parameters = (saveat=0.1,)
    input = (Matrix{Float64}(undef, 0, 0), "")
    output = sol -> Matrix(sol[:, :]')

    # Create the Model struct
    model = Model(
        variable_names,
        parameter_names,
        problem,
        solver_algorithm,
        solver_parameters,
        input,
        output
    )

    return model

end

"""
`set_solver(model::Model, algorithm=nothing; merge_kwargs=true, kwargs...)`

Set solver algorithm and its properties.

**Arguments**
- `model`: Model.

**Optional Arguments**
- `algorithm`: Solver algorithm (e.g. `Tsit5()` or `SOSRI()`).

**Keyword Arguments**
- `merge_kwargs`: If `false`, the original arguments are deleted.
- `kwargs...`: Solver parameters.

**Returns**
- `new_model`: Model with updated solver algorithm and parameters.
"""
function set_solver(model::Model, algorithm=nothing; merge_kwargs=true, kwargs...)

    if :callback in keys(kwargs)
        error("Callback is reserved for the input function!")
    end

    if isnothing(algorithm)
        # If the algorithm was not specified, use to original one
        solver_algorithm = model.solver_algorithm
    else
        # If the algorithm was specified, use the new one
        solver_algorithm = algorithm
    end

    if length(kwargs) > 0
        # If any kwargs were passed, add them to the solver_parameters

        if merge_kwargs
            # Merge kwargs with the existing parameters
            solver_parameters = merge(model.solver_parameters, NamedTuple(kwargs))
        else
            # Replace the old parameters with the new ones
            if :callback in keys(model.solver_parameters)
                # Preserve callback
                callback = model.solver_parameters[:callback]
                solver_parameters = merge(NamedTuple(kwargs), (callback=callback,))
            else
                solver_parameters = NamedTuple(kwargs)
            end
        end

    else
        # If no kwargs were passed, just copy the original ones
        solver_parameters = model.solver_parameters

    end

    # Build a new model
    new_model = _set_model_property(model; solver_algorithm=solver_algorithm,
        solver_parameters=solver_parameters)

    return new_model

end

"""
Set callback to the solver parameters.
"""
function _set_callback(model::Model, callback)

    # Merge callback with existing solver parameters
    solver_parameters = merge(model.solver_parameters, (callback=callback,))

    # Build a new model
    new_model = _set_model_property(model; solver_parameters=solver_parameters)

    return new_model

end

"""
`get_parameter_index(model::Model, parameter_name::String)`

Find index of the specified parameter.

**Arguments**
- `model`: Model.
- `parameter_name`: Name of the parameter.

**Returns**
- `index`: Index indicating the position of the parameter.
"""
function get_parameter_index(model::Model, parameter_name::String)

    parameter_names = model.parameter_names
    index = findfirst(parameter_names .== parameter_name)
    if isnothing(index)
        error("Parameter '$(parameter_name)' is not in the model!")
    end
    return index

end

"""
`get_parameter_value(model::Model, parameter_name::String)`

Get a parameter value.

**Arguments**
- `model`: Model.
- `parameter_name`: Name of the parameter.

**Returns**
- `parameter_value`: Value of the parameter.
"""
function get_parameter_value(model::Model, parameter_name::String)

    index = get_parameter_index(model, parameter_name)
    parameter_values = _get_problem_property(model, "p")
    return parameter_values[index]

end

"""
`set_parameter_value(model::Model, name::String, value::Number)`

Set a parameter value.

**Arguments**
- `model`: Model.
- `name`: Name of the parameter. Can also be in the form of "d1=d2=d3", which
    sets all parameters d1, d2, and d3 to the same value.
- `value`: New value for the parameter.

**Returns**
- `new_model`: Copy of `model` with updated parameter value.
"""
function set_parameter_value(model::Model, name::String, value::Number)

    # Get array of parameter values
    parameter_values = _get_problem_property(model, "p")

    if occursin("=", name)
        # If "=" occurs in `name`, split the string into individual parameters
        name_array = String.(split(name, "="))
        for name_split in name_array
            # Set each parameter to `value`
            index = get_parameter_index(model, name_split)
            parameter_values[index] = value
        end
    else
        # Set the desired parameter to `value`
        index = get_parameter_index(model, name)
        parameter_values[index] = value
    end

    # Rebuild the model with the new parameter value
    new_model = _set_problem_property(model, p=parameter_values)
    return new_model

end

"""
`model::Model, name_array::Array, value_array::Array`

Set parameter values.

**Arguments**
- `model`: Model.
- `name_array`: Parameter names. Can also be in the form of "d1=d2=d3", which
    sets all parameters d1, d2, and d3 to the same value.
- `value_array`: New values for the parameters.

**Returns**
- `new_model`: Copy of `model` with updated parameter values.
"""
function set_parameter_value(model::Model, name_array::Array, value_array::Array)

    # The input arrays must have the same length
    if length(name_array) != length(value_array)
        error("The array of parameter names and values must have the same length!")
    end

    # Set the parameters
    new_model = deepcopy(model)
    for (name, value) in zip(name_array, value_array)
        new_model = set_parameter_value(new_model, name, value)
    end

    return new_model

end

"""
`_create_callback(events::Matrix, i)`

Generate `PresetTimeCallback` from `events` modyfing `i`-th parameter.
"""
function _create_callback(events::Matrix, i)

    p = NaN  # holds the original value of the parameter

    # Initialization at the beginning of the integration
    initialize = function (c, u, t, integrator)

        # Protect the original parameters from overwriting
        integrator.p = copy(integrator.p)

        # Save the default parameter
        p = integrator.p[i]

        # Set the parameter to 0
        integrator.p[i] = 0.0

    end

    # Set the parameter to its default value at the beginning of the event
    tstops_on = events[:, 1]
    affect_on! = (integrator) -> integrator.p[i] = p
    callback_on = PresetTimeCallback(tstops_on, affect_on!,
        save_positions=(false, false), initialize=initialize)

    # Set the parameter to zero at the end of the event
    tstops_off = events[:, 2]
    affect_off! = (integrator) -> integrator.p[i] = 0.0
    callback_off = PresetTimeCallback(tstops_off, affect_off!,
        save_positions=(false, false))

    return CallbackSet(callback_on, callback_off)

end

"""
`_events_to_function(events::Matrix)`

Convert `events` to a function.
"""
function _events_to_function(events::Matrix)
    N = size(events, 1)
    function fun(t)
        for i = 1:N
            if events[i, 1] <= t < events[i, 2]
                return 1.0
            end
        end
        return 0.0
    end
    return fun
end

"""
`_create_discrete_callback(events::Matrix, i)`

Generate `DiscreteCallback` from `events` modyfing `i`-th parameter.
"""
function _create_discrete_callback(events::Matrix, i)

    p = NaN  # holds the original value of the parameter

    # Initialization at the beginning of the integration
    initialize = function (c, u, t, integrator)

        # Protect the original parameters from overwriting
        integrator.p = copy(integrator.p)

        # Save the default parameter
        p = integrator.p[i]

        # Set the parameter to 0
        integrator.p[i] = 0.0

    end

    # Convert the events to a function
    fun = _events_to_function(events)

    # Call the function at every step of the integration
    affect! = (integrator) -> integrator.p[i] = fun(integrator.t)*p
    condition = (u, t, integrator) -> true
    callback = DiscreteCallback(condition, affect!; initialize=initialize,
        save_positions=(false, false))

    return callback

end

"""
`set_input(model::Model, events::Matrix, parameter_name="I")`

Set input to the model.

**Arguments**
- `model`: Model.
- `events`: Matrix representing the input square function.
- `parameter_name`: Parameter that is being modified by the input.
"""
function set_input(model::Model, events::Matrix, parameter_name="I")

    parameter_index = get_parameter_index(model, parameter_name)

    # Create a callback
    if model.problem isa Union{ODEProblem, SDEProblem}
        # Use PresetTimeCallback for ODE and SDE models
        callback = _create_callback(events, parameter_index)
    elseif model.problem isa JumpProblem
        # Use DiscreteCallback for Jump models
        callback = _create_discrete_callback(events, parameter_index)
    end

    # Build a new model with the specified input
    new_model = _set_callback(model, callback)
    new_model = _set_model_property(new_model, input=(events, parameter_name))

    return new_model

end

"""
`simulate_population(model::Model, trajectories=1; save_trajectories=true,
    show_progress=false)`

Simulate a population.

**Arguments**
- `model`: Model.

**Optional Arguments**
- `trajectories`: Number of trajectories to simulate.

**Keyword Arguments**
- `save_trajectories`: If `false`, only mean is saved.
- `show_progress`: If `true`, show a progress bar in the terminal.
"""
function simulate_population(model::Model, trajectories=1;
    save_trajectories=true, show_progress=false)

    # Make sure that the original model is not modified
    model = deepcopy(model)

    # Prepare variables for PopulationSolution fields
    t = Vector{Float64}(undef, 0)  # time
    m = Matrix{Float64}(undef, 0, 0)  # mean
    U = Array{Float64, 3}(undef, 0, 0, 0)  # trajectories
    events = model.input[1]  # events

    # Initialize the porgress meter
    if show_progress
        progress_meter = ProgressMeter.Progress(trajectories; barlen=20)
    end

    # Simulate trajectories
    success_arr = fill(false, trajectories)  # success for individual trajectories
    lk = ReentrantLock()
    Threads.@threads for i = 1:trajectories

        model2 = deepcopy(model)

        # Solve the problem
        sol2 = solve(model2.problem, model2.solver_algorithm; model2.solver_parameters...)

        # Check if the simulation was successful
        if model2.problem isa Union{ODEProblem, SDEProblem}
            success_arr[i] = sol2.retcode == :Success
        elseif model2.problem isa JumpProblem
            success_arr[i] = sol2.retcode == :Default
        else
            error("Unknown model type!")
        end

        # Save the solution if the current simulation was successful
        if success_arr[i]

            lock(lk) do

                # Save time
                if isempty(t)
                    t = sol2.t
                end

                # Map solution to a matrix
                x = model2.output(sol2)

                # Add the trajectory contribution to the mean
                if isempty(m)
                    # Initialize the mean vector
                    m = (x ./ trajectories)
                else
                    # Add solution to the mean vector
                    m .+= (x ./ trajectories)
                end

                # Save the individual trajectory
                if save_trajectories
                    if isempty(U)
                        # Initialize the matrix for individual trajectories
                        nsamples = size(x, 1)
                        nvariables = size(x, 2)
                        U = fill(NaN, nsamples, nvariables, trajectories)
                    end
                    # Add the current trajectory to the trajectory list
                    U[:, :, i] = x
                end

            end

        else
            # Stop the simulation if the current simulation was unsuccessful
            break
        end

        if show_progress
            lock(lk) do
                ProgressMeter.next!(progress_meter)
            end
        end

    end

    # Build the output struct
    success = all(success_arr)
    solution = PopulationSolution(t, m, U, events, success)

    return solution

end

"""
`scan`

Scan parameters of a model.

**Arguments**
- `model`: Model.
- `parameter_names`: Parameter names to scan.
- `parameter_values`: Parameter values for each parameter to scan over.
- `summary_function`: A function that takes in a model, evaluates it, and
    returns some parameters. If the function is called without any arguments,
    it returns the names of the parameters as a vector of strings. For example
    ```
    julia> summary_function(model)
    [1.1, 2.2, 3.3]
    julia> summary_function()
    ["Period", "Amplitude", "Phase"]
    ```

**Keyword Arguments**
- `show_progress`: If `true`, show a progress bar in the terminal.

**Returns**
- `df`: Dataframe with results.
"""
function scan(model::Model, parameter_names, parameter_values,
    summary_function; show_progress=false)

    # Protect the original model from overwriting
    model = deepcopy(model)

    # Find all possible parameter combinations
    parameter_value_combinations = find_all_combinations(parameter_values)
    n = size(parameter_value_combinations, 1)

    # Initialize summary matrix
    summary_names = summary_function()
    scan_summary = fill(NaN, n, length(summary_names))

    # Initialize progress bar
    if show_progress
        progressmeter = ProgressMeter.Progress(n; barlen=20)
    end

    # Iterate all parameter value combinations
    lk = ReentrantLock()
    Threads.@threads for i = 1:n

        # Set model parameters
        pmodel = set_parameter_value(model, parameter_names, parameter_value_combinations[i, :])

        # Evaluate the sumary function
        pmodel_summary = summary_function(pmodel)

        # Save the model summary to the overall matrix
        scan_summary[i, :] = pmodel_summary

        # Update progress bar
        if show_progress
            lock(lk) do
                ProgressMeter.next!(progressmeter)
            end
        end

    end

    # Build the output dataframe
    matrix = hcat(parameter_value_combinations, scan_summary)
    names = vcat(parameter_names, summary_names)
    df = DataFrame(matrix, names)

    return df

end

"""
`estimate_binary_arnold_tongue`

Use binary search to estimate arnold tongue for a model.

**Arguments**
- `model`: Model.
- `periods`: A two-element array like `[minimal_period, maximal_period]`.
- `amplitudes`: Input amplitudes to scan.

**Keyword Arguments**
- `trajectories`: Number of trajectories in a population.
- `i_solution`: Number of the state to use to estimate the period.
- `n_lags`: Number of lags used to estimate the autocorrelation functions.
- `input_parameter_name`: Name of the parameter that control the input amplitude.
- `min_time`: Minimal time to consider (to remove the effect of initial conditions).
- `max_time`: Maximal time of the simulation.
- `show_progress`: If true, show a progress bar in the terminal.
- `period_error`: Maximal allowed output period deviation from the input period.
- `boundary_error`: Maximal allowed error to estimate the boundary between the
    entrained and unentrained regions.
- `maximal_error_attempts`: How many attempts should be given to an errored simulation.
"""
function estimate_binary_arnold_tongue(model::Model, periods, amplitudes;
    trajectories=1, i_solution=1, n_lags=1000, input_parameter_name="I",
    min_time=nothing, max_time=nothing, show_progress=false, period_error=0.01,
    boundary_error=0.01, maximal_error_attempts=0)

    # Protect the original model from overwriting
    model = deepcopy(model)

    # Set the maximal time of the simulation
    if !isnothing(max_time)
        model = set_timespan(model, (0.0, max_time))
    end

    # Get minimal and maximal amplitude and the number of periods to scan
    minimal_amplitude = amplitudes[1]
    maximal_amplitude = amplitudes[end]
    n_periods = length(periods)

    # Each element is the minimal amplitude of entrainment for the given period
    entrainment_amplitudes = fill(NaN, n_periods)

    # Initialize a progress bar
    if show_progress
        progressmeter = ProgressMeter.Progress(n_periods; barlen=20)
    end

    # Iterate all period-amplitude combinations
    lk = ReentrantLock()
    Threads.@threads for i = 1:n_periods  # Threads.@threads

        # Set the forcing (input) period
        forcing_period = periods[i]

        # Create events for the current forcing period
        if model.problem isa JumpProblem
            end_time = model.problem.prob.tspan[2]
        else
            end_time = model.problem.tspan[2]
        end
        nevents = end_time / forcing_period + 1
        events = create_events([(:DD, forcing_period*0.5),
            (:LD, nevents, forcing_period*0.5, forcing_period*0.5)])

        # Set offset (solution before this time is discarted)
        if isnothing(min_time)
            offset = end_time*0.5
        else
            offset = min_time
        end

        # pmodel = deepcopy(model)
        pmodel = set_input(model, events, input_parameter_name)

        """
        This function will be optimized using the binary boundary search. It
        simulates the model for the given forcing amplitude and returns `true`
        if the model is entrained and `false` if the model is unentrained.
        """
        function fun(forcing_amplitude)

            # Set the forcing amplitude and events to the model
            pmodel = set_parameter_value(pmodel, input_parameter_name, forcing_amplitude)

            # Simulate the population
            solution = PopulationSolution()
            success = false
            counter = 0
            while !success
                try
                    # lock(lk) do
                    solution = simulate_population(pmodel, trajectories;
                        save_trajectories=false)
                    # end
                    success = true
                catch err
                    counter += 1
                    if counter >= maximal_error_attempts
                        @warn "[LAST] Simulation threw an error: $err"
                        solution = PopulationSolution()
                        success = true
                    else
                        @warn "[$counter] Simulation threw an error: $err"
                    end
                end
            end

            # If the simulation was successful, estimate the period
            if solution.success

                # Discard the simulation start
                solution = select_time(solution, mintime=offset)
                t = solution.time
                m = solution.mean[:, i_solution]

                R_period, _ = estimate_period(t, m; n_lags=n_lags)

                return abs(R_period - forcing_period) <= period_error

            else

                return false

            end

        end

        # Perform binary search
        entrainment_amplitudes[i] = binary_boundary_search(fun,
            [minimal_amplitude, maximal_amplitude], boundary_error)

        # Update the progress bar
        if show_progress
            ProgressMeter.next!(progressmeter)
        end

    end

    # Create the output dataframe
    arnold_tongue = DataFrame(periods=periods, amplitudes=entrainment_amplitudes)

    return arnold_tongue

end

max_stars_repo_name: vkumpost/stoosc

max_issues_repo_name: vkumpost/stoosc

max_forks_repo_name: vkumpost/stoosc


text:
export takens_best_estimate

"""
    takens_best_estimate(X, εmax, metric = Chebyshev(),εmin = 0) → D_C, D_C_95u, D_C_95l
Use the so-called "Takens' best estimate" [^Takens1985][^Theiler1988]
method for estimating the correlation dimension
`D_C` and the upper (`D_C_95u`) and lower (`D_C_95l`) confidence limit for the given dataset `X`.

The original formula is
```math
D_C \\approx \\frac{C(\\epsilon_\\text{max})}{\\int_0^{\\epsilon_\\text{max}}(C(\\epsilon) / \\epsilon) \\, d\\epsilon}
```
where ``C`` is the [`correlationsum`](@ref) and ``\\epsilon_\\text{max}`` is an upper cutoff.
Here we use the later expression
```math
D_C \\approx - \\frac{1}{\\eta},\\quad \\eta = \\frac{1}{(N-1)^*}\\sum_{[i, j]^*}\\log(||X_i - X_j|| / \\epsilon_\\text{max})
```
where the sum happens for all ``i, j`` so that ``i < j`` and ``||X_i - X_j|| < \\epsilon_\\text{max}``.
In the above expression, the bias in the original paper has already been corrected, as suggested in [^Borovkova1999].

The confidence limits are estimated from the log-likelihood function by finding
the values of `D_C` where the function has fallen by 2 from its maximum, see e.g.
[^Barlow] chapter 5.3
Because the CLT does not apply (no independent measurements), the limits are not
neccesarily symmetric.

According to [^Borovkova1999], introducing a lower cutoff `εmin` can make the
algorithm more stable (no divergence), this option is given but defaults to zero.

If `X` comes from a delay coordinates embedding of a timseries `x`, a recommended value
for ``\\epsilon_\\text{max}`` is `std(x)/4`.

[^Takens1985]: Takens, On the numerical determination of the dimension of an attractor, in: B.H.W. Braaksma, B.L.J.F. Takens (Eds.), Dynamical Systems and Bifurcations, in: Lecture Notes in Mathematics, Springer, Berlin, 1985, pp. 99–106.
[^Theiler1988]: Theiler, [Lacunarity in a best estimator of fractal dimension. Physics Letters A, 133(4–5)](https://doi.org/10.1016/0375-9601(88)91016-X)
[^Borovkova1999]: Borovkova et al., [Consistency of the Takens estimator for the correlation dimension. The Annals of Applied Probability, 9, 05 1999.](https://doi.org/10.1214/aoap/1029962747)
[^Barlow]: Barlow, R., Statistics - A Guide to the Use of Statistical Methods in the Physical Sciences. Vol 29. John Wiley & Sons, 1993
"""
function takens_best_estimate(X, εmax, metric = Chebyshev(); εmin=0)
    n, η, N = 0, zero(eltype(X)), length(X)
    @inbounds for i in 1:N
        for j in i+1:N
            d = evaluate(metric, X[i], X[j])
            if εmin < d < εmax
                n += 1
                η += log(d/εmax)
            end
        end
    end
    # bias-corrected version (log-likelihood function shifted on x-axis)
    α = -(n-1)/η
    # biased version (maximum of original log-likelihood function)
    α_b = -n/η
    # value of maximum of original log-likelihood function
    mxl = n*log(α_b) + α_b * η
    # at the 95%-confidence interval, the log-l function has dropped by 2
    # -> log_l(x) - mxl + 2 = 0
    # this is a result of the invariance of the MLE, a really nice property
    # these limits are not going to be perfectly symmetric (CLT does not apply)
    mn, mx = fzeros(x-> n * log(x) + η * x - mxl +2 , 0,2*α)

    # Since the bias-correction is just a shift of the log-l function on the
    # x-axis, we can easily shift the confidence limits by the bias α-α_b
    α95u = α - α_b + mn
    α95l = α - α_b + mx

    return α, α95u, α95l
end

ext: jl


text:
function enqueue_transmissions!(state::SimState, ::Type{Val{ConstantKernelContact}}, source_id::Integer, params::SimParams)
  progression = progressionof(params, source_id)

  start_time = progression.incubation_time
  end_time =  if      !ismissing(progression.mild_symptoms_time);   progression.mild_symptoms_time
              elseif  !ismissing(progression.severe_symptoms_time); progression.severe_symptoms_time
              elseif  !ismissing(progression.recovery_time);        progression.recovery_time
              else    error("no recovery nor symptoms time defined")
              end

  total_infection_rate = (end_time - start_time) * params.constant_kernel_param

  total_infection_rate *= spreading(params, source_id)

  num_infections = rand(state.rng, Poisson(total_infection_rate))

  if num_infections == 0
    return
  end
  @assert start_time != end_time "pathologicaly short time for infections there shouldn't be any infections but are $num_infections, progression=$progression"

  time_dist = Uniform(state.time, end_time - start_time + state.time) # in global time reference frame

  num_individuals = size(params.progressions, 1)

  for _ in 1:num_infections
    subject_id = sample(state.rng, 1:(num_individuals-1)) # exclude the source itself
    if subject_id >= source_id # restore the indexing
      subject_id +=1
    end

    if Healthy == health(state, subject_id)
      infection_time::TimePoint = rand(state.rng, time_dist) |> TimePoint
      @assert state.time <= infection_time <= (end_time-start_time + state.time)
      push!(state.queue, Event(Val(TransmissionEvent), infection_time, subject_id, source_id, ConstantKernelContact))
    end
  end
  nothing
end

function enqueue_transmissions!(state::SimState, ::Type{Val{HouseholdContact}}, source_id::Integer, params::SimParams)
  household = householdof(params,source_id)

  if 1==length(household)
    return
  end

  progression = progressionof(params, source_id)

  start_time = progression.incubation_time
  end_time =  if      !ismissing(progression.severe_symptoms_time); progression.severe_symptoms_time
              elseif  !ismissing(progression.recovery_time);        progression.recovery_time
              else    error("no recovery nor severe symptoms time defined")
              end

  max_time = time(state) - start_time + end_time

  mean_infection_time = (length(household)-1) / params.household_kernel_param
  time_dist = Exponential(mean_infection_time)

  for subject_id in household
    if subject_id == source_id || Healthy != health(state, subject_id)
      continue
    end

    infection_time = time(state) + rand(state.rng, time_dist)
    if infection_time > max_time
      continue
    end

    @assert time(state) <= infection_time <= (end_time - start_time + time(state))
    push!(state.queue, Event(Val(TransmissionEvent),
      infection_time,
      subject_id,
      source_id,
      HouseholdContact)
    )
  end
  nothing
end

function enqueue_transmissions!(state::SimState, ::Type{Val{HospitalContact}}, source_id::Integer, params::SimParams)
  if nothing === params.hospital_kernel_params
    return
  end
  progression = progressionof(params, source_id)

  @assert progression.severity in SA[Severe, Critical]
  @assert !ismissing(progression.severe_symptoms_time)
  @assert !ismissing(progression.recovery_time) || !ismissing(progression.death_time)

  start_time = progression.severe_symptoms_time
  end_time = ismissing(progression.recovery_time) ? progression.death_time : progression.recovery_time

  total_infection_rate = (end_time - start_time) * params.hospital_kernel_params.kernel_constant

  num_infections = rand(state.rng, Poisson(total_infection_rate))

  if num_infections == 0
    return
  end
  @assert start_time != end_time "pathologicaly short time for infections there shouldn't be any infections but are $num_infections, progression=$progression"

  time_dist = Uniform(state.time, end_time - start_time + state.time) # in global time reference frame

  for _ in 1:num_infections
    subject_id = sample(state.rng, params.hospital_kernel_params.hospital_staff_ids)
    if subject_id == source_id # self infection not possible
      continue
    end

    if Healthy == health(state, subject_id)
      infection_time::TimePoint = rand(state.rng, time_dist) |> TimePoint
      @assert state.time <= infection_time <= (end_time-start_time + state.time)
      push!(state.queue, Event(Val(TransmissionEvent), infection_time, subject_id, source_id, HospitalContact))
    end
  end
end

function enqueue_transmissions!(state::SimState, ::Type{Val{FriendshipContact}}, source_id::Integer, params::SimParams)
  if nothing == params.friendship_kernel_params
    return
  end
  progression = progressionof(params, source_id)

  start_time = progression.incubation_time

  end_time =  if      !ismissing(progression.mild_symptoms_time);   progression.mild_symptoms_time
              elseif  !ismissing(progression.severe_symptoms_time); progression.severe_symptoms_time
              elseif  !ismissing(progression.recovery_time);        progression.recovery_time
              else    error("no recovery nor symptoms time defined")
              end

  total_infection_rate = (end_time - start_time) * params.friendship_kernel_params.kernel_constant * socialcompetence(params, source_id)

  total_infection_rate *= spreading(params, source_id)

  num_infections = rand(state.rng, Poisson(total_infection_rate))

  if num_infections == 0
    return
  end
  @assert start_time != end_time "pathologicaly short time for infections there shouldn't be any infections but are $num_infections, progression=$progression"

  time_dist = Uniform(state.time, end_time - start_time + state.time) # in global time reference frame

  num_individuals = size(params.progressions, 1)

  for _ in 1:num_infections
    subject_id = friend_sample(params.friendship_kernel_params.sampler, age(params, source_id), gender(params, source_id), state.rng)

    if Healthy == health(state, subject_id)
      infection_time::TimePoint = rand(state.rng, time_dist) |> TimePoint
      @assert state.time <= infection_time <= (end_time-start_time + state.time)
      push!(state.queue, Event(Val(TransmissionEvent), infection_time, subject_id, source_id, ConstantKernelContact))
    end
  end
  nothing
end


text:
# ---
# title: 74. Search a 2D Matrix
# id: problem74
# author: Qling
# date: 2021-03-30
# difficulty: Medium
# categories: Array, Binary Search
# link: <https://leetcode.com/problems/search-a-2d-matrix/description/>
# hidden: true
# ---
#
# Write an efficient algorithm that searches for a value in an `m x n` matrix.
# This matrix has the following properties:
#
#   * Integers in each row are sorted from left to right.
#   * The first integer of each row is greater than the last integer of the previous row.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/10/05/mat.jpg)
#
#
#
#     Input: matrix = [[1,3,5,7],[10,11,16,20],[23,30,34,50]], target = 3
#     Output: true
#
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2020/10/05/mat2.jpg)
#
#
#
#     Input: matrix = [[1,3,5,7],[10,11,16,20],[23,30,34,50]], target = 13
#     Output: false
#
#
# **Example 3:**
#
#
#
#     Input: matrix = [], target = 0
#     Output: false
#
#
#
#
# **Constraints:**
#
#   * `m == matrix.length`
#   * `n == matrix[i].length`
#   * `0 <= m, n <= 100`
#   * `-104 <= matrix[i][j], target <= 104`
#
#
## @lc code=start
using LeetCode

function search_matrix(matrix::Vector{Vector{Int}}, target::Int)::Bool
    m, n = length(matrix), length(matrix[1])

    (m == 0 || n == 0) && return false
    left, right = 1, m * n

    while left <= right
        mid = left + (right - left) ÷ 2
        num = matrix[mid ÷ n][mod1(mid, n)]
        if num == target
            return true
        elseif num < target
            left = mid + 1
        else
            right = mid - 1
        end
    end

    return false
end
## @lc code=end


text:
using MultiWriters
using Test

const io = IOBuffer()
const rd, wr = redirect_stderr()
@testset "MultiWriters.jl" begin
    mw = MultiWriter(io, stderr)
    mwprint(mw, "Hello", "World")
    helloworld = "HelloWorld"
    @test String(take!(io)) == helloworld
    @test String(readavailable(rd)) == helloworld

    mwwrite(mw, "Hello", "World")
    helloworld = "HelloWorld"
    @test String(take!(io)) == helloworld
    @test String(readavailable(rd)) == helloworld

    mwprintln(mw, "Hello", "World")
    helloworld = "HelloWorld\n"
    @test String(take!(io)) == helloworld
    @test String(readavailable(rd)) == helloworld

    mwshow(mw, "Hello")
    hello = "Hello"
    @test String(take!(io)) == sprint(show, hello)
    @test String(readavailable(rd)) == sprint(show, hello)
end


text:
struct GetIndex{T,Args} <: UnaryNodeOp{T}
    args::Args
end

stateless_operator(::GetIndex) = true
time_agnostic(::GetIndex) = true
always_ticks(::GetIndex) = true

operator!(op::GetIndex, x) = getindex(x, op.args...)

"""
    getindex(x::Node, args...)

A node whose values are generated by calling `getindex(value, args...)` on each `value`
obtained from the node `x`.
"""
function Base.getindex(x::Node, args...)
    T = output_type(getindex, value_type(x), map(typeof, args)...)
    return obtain_node((x,), GetIndex{T,typeof(args)}(args))
end

# Special-case optimisations for a single colon.
Base.getindex(x::Node{<:AbstractArray}, ::Colon) = vec(x)
Base.getindex(x::Node{<:Tuple}, ::Colon) = x

"""
    vec(x::Node{<:AbstractArray}) -> Node{<:AbstractVector}

Return a vector whose values are those of `x`, but flattened into a single vector.

If `x` has values which are already an `AbstractVector`, this will be a no-op.
"""
Base.vec(x::Node{<:AbstractArray}) = apply(vec, x)
Base.vec(x::Node{<:AbstractVector}) = x


text:
using PQEcolaPoint
using DataFrames

# setup DTE -
T = 149.0
DTE = range(T, stop=0.0,step=-1.0) |> collect

# setup underlying range -
Sₒ = 69.41
IV = 0.4693
σ = Sₒ*(IV)*(√(T/365))
L = round(Sₒ - 1.96*σ)  # 95%
U = round(Sₒ + 1.96*σ)  # 95%
S = range(max.(L, 0.0),stop=U, step=0.01) |> collect

# setup IV array -
IV_array = range(0.1, stop = 1.2, step = 0.005) |> collect

# setup strike coordinates -
K = range(45.0, stop = 90.0, step = 5.0) |> collect

# build JIT lattice model -
JIT_lattice = build(CRRJITContractPremiumLatticeModel, PutContractModel, S, K, IV_array, DTE; number_of_levels = 80)

# origin version -
origin_point = encode(JIT_lattice; S = 70.12, DTE = 53.0, IV = 0.5129, K = 65.0)
start_point = encode(JIT_lattice; S = 68.41, DTE = 56.0, IV = 0.5957, K = 65.0)
end_point = encode(JIT_lattice; S = 68.41 , DTE = 84.0, IV = 0.5737, K = 55.0)
C = roll(JIT_lattice, origin_point, start_point=>end_point)

# build my points -
test = roll(JIT_lattice, start_point=>end_point)


text:
function triangulate(polygon::Vector{Vector{Point{2, Float64}}})
    lengths = map(x-> UInt32(length(x)), polygon)
    len = UInt32(length(lengths))
    array = ccall(
        (:u32_triangulate_f64, libearcut),
        Tuple{Ptr{GLTriangleFace}, Cint},
        (Ptr{Ptr{Float64}}, Ptr{UInt32}, UInt32),
        polygon, lengths, len
    )
    return unsafe_wrap(Vector{GLTriangleFace}, array[1], array[2])
end

function triangulate(polygon::Vector{Vector{Point{2, Float32}}})
    lengths = map(x-> UInt32(length(x)), polygon)
    len = UInt32(length(lengths))
    array = ccall(
        (:u32_triangulate_f32, libearcut),
        Tuple{Ptr{GLTriangleFace}, Cint},
        (Ptr{Ptr{Float32}}, Ptr{UInt32}, UInt32),
        polygon, lengths, len
    )
    return unsafe_wrap(Vector{GLTriangleFace}, array[1], array[2])
end

function triangulate(polygon::Vector{Vector{Point{2, Int64}}})
    lengths = map(x-> UInt32(length(x)), polygon)
    len = UInt32(length(lengths))
    array = ccall(
        (:u32_triangulate_i64, libearcut),
        Tuple{Ptr{GLTriangleFace}, Cint},
        (Ptr{Ptr{Int64}}, Ptr{UInt32}, UInt32),
        polygon, lengths, len
    )
    return unsafe_wrap(Vector{GLTriangleFace}, array[1], array[2])
end

function triangulate(polygon::Vector{Vector{Point{2, Int32}}})
    lengths = map(x-> UInt32(length(x)), polygon)
    len = UInt32(length(lengths))
    array = ccall(
        (:u32_triangulate_i32, libearcut),
        Tuple{Ptr{GLTriangleFace}, Cint},
        (Ptr{Ptr{Int32}}, Ptr{UInt32}, UInt32),
        polygon, lengths, len
    )
    return unsafe_wrap(Vector{GLTriangleFace}, array[1], array[2])
end


text:
module AdvancedHMC

const DEBUG = convert(Bool, parse(Int, get(ENV, "DEBUG_AHMC", "0")))

using Statistics: mean, var, middle
using LinearAlgebra: Symmetric, UpperTriangular, mul!, ldiv!, dot, I, diag, cholesky, UniformScaling
using StatsFuns: logaddexp, logsumexp
import Random
using Random: GLOBAL_RNG, AbstractRNG
using ProgressMeter: ProgressMeter
using UnPack: @unpack

using Setfield
import Setfield: ConstructionBase

using ArgCheck: @argcheck

using DocStringExtensions

import AbstractMCMC

import StatsBase: sample

include("utilities.jl")

# Notations
# ℓπ: log density of the target distribution
# θ: position variables / model parameters
# ∂ℓπ∂θ: gradient of the log density of the target distribution w.r.t θ
# r: momentum variables
# z: phase point / a pair of θ and r

include("metric.jl")
export UnitEuclideanMetric, DiagEuclideanMetric, DenseEuclideanMetric

include("hamiltonian.jl")
export Hamiltonian

include("integrator.jl")
export Leapfrog, JitteredLeapfrog, TemperedLeapfrog

include("trajectory.jl")
export Trajectory, HMCKernel,
       FixedNSteps, FixedIntegrationTime,
       ClassicNoUTurn, GeneralisedNoUTurn, StrictGeneralisedNoUTurn,
       EndPointTS, SliceTS, MultinomialTS,
       find_good_stepsize

# Useful defaults

struct NUTS{TS, TC} end

"""
$(SIGNATURES)

Convenient constructor for the no-U-turn sampler (NUTS).
This falls back to `HMCKernel(Trajectory{TS}(int, TC(args...; kwargs...)))` where

- `TS<:Union{MultinomialTS, SliceTS}` is the type for trajectory sampler
- `TC<:Union{ClassicNoUTurn, GeneralisedNoUTurn, StrictGeneralisedNoUTurn}` is the type for termination criterion.

See [`ClassicNoUTurn`](@ref), [`GeneralisedNoUTurn`](@ref) and [`StrictGeneralisedNoUTurn`](@ref) for details in parameters.
"""
NUTS{TS, TC}(int::AbstractIntegrator, args...; kwargs...) where {TS, TC} =
    HMCKernel(Trajectory{TS}(int, TC(args...; kwargs...)))
NUTS(int::AbstractIntegrator, args...; kwargs...) =
    HMCKernel(Trajectory{MultinomialTS}(int, GeneralisedNoUTurn(args...; kwargs...)))
NUTS(ϵ::AbstractScalarOrVec{<:Real}) =
    HMCKernel(Trajectory{MultinomialTS}(Leapfrog(ϵ), GeneralisedNoUTurn()))

export NUTS

# Deprecations for trajectory.jl

abstract type AbstractTrajectory end

struct StaticTrajectory{TS} end
@deprecate StaticTrajectory{TS}(int::AbstractIntegrator, L) where {TS} HMCKernel(Trajectory{TS}(int, FixedNSteps(L)))
@deprecate StaticTrajectory(int::AbstractIntegrator, L) HMCKernel(Trajectory{EndPointTS}(int, FixedNSteps(L)))
@deprecate StaticTrajectory(ϵ::AbstractScalarOrVec{<:Real}, L) HMCKernel(Trajectory{EndPointTS}(Leapfrog(ϵ), FixedNSteps(L)))

struct HMCDA{TS} end
@deprecate HMCDA{TS}(int::AbstractIntegrator, λ) where {TS} HMCKernel(Trajectory{TS}(int, FixedIntegrationTime(λ)))
@deprecate HMCDA(int::AbstractIntegrator, λ) HMCKernel(Trajectory{EndPointTS}(int, FixedIntegrationTime(λ)))
@deprecate HMCDA(ϵ::AbstractScalarOrVec{<:Real}, λ) HMCKernel(Trajectory{EndPointTS}(Leapfrog(ϵ), FixedIntegrationTime(λ)))

@deprecate find_good_eps find_good_stepsize

export StaticTrajectory, HMCDA, find_good_eps

include("adaptation/Adaptation.jl")
using .Adaptation
import .Adaptation: StepSizeAdaptor, MassMatrixAdaptor, StanHMCAdaptor, NesterovDualAveraging

# Helpers for initializing adaptors via AHMC structs

StepSizeAdaptor(δ::AbstractFloat, stepsize::AbstractScalarOrVec{<:AbstractFloat}) =
    NesterovDualAveraging(δ, stepsize)
StepSizeAdaptor(δ::AbstractFloat, i::AbstractIntegrator) = StepSizeAdaptor(δ, nom_step_size(i))

MassMatrixAdaptor(m::UnitEuclideanMetric{T}) where {T} =
    UnitMassMatrix{T}()
MassMatrixAdaptor(m::DiagEuclideanMetric{T}) where {T} =
    WelfordVar{T}(size(m); var=copy(m.M⁻¹))
MassMatrixAdaptor(m::DenseEuclideanMetric{T}) where {T} =
    WelfordCov{T}(size(m); cov=copy(m.M⁻¹))

MassMatrixAdaptor(
    m::Type{TM},
    sz::Tuple{Vararg{Int}}=(2,)
) where {TM<:AbstractMetric} = MassMatrixAdaptor(Float64, m, sz)

MassMatrixAdaptor(
    ::Type{T},
    ::Type{TM},
    sz::Tuple{Vararg{Int}}=(2,)
) where {T, TM<:AbstractMetric} = MassMatrixAdaptor(TM(T, sz))

# Deprecations

@deprecate StanHMCAdaptor(n_adapts, pc, ssa) initialize!(StanHMCAdaptor(pc, ssa), n_adapts)
@deprecate NesterovDualAveraging(δ::AbstractFloat, i::AbstractIntegrator) StepSizeAdaptor(δ, i)
@deprecate Preconditioner(args...) MassMatrixAdaptor(args...)

export StepSizeAdaptor, NesterovDualAveraging,
       MassMatrixAdaptor, UnitMassMatrix, WelfordVar, WelfordCov,
       NaiveHMCAdaptor, StanHMCAdaptor

include("diagnosis.jl")

include("sampler.jl")
export sample

include("abstractmcmc.jl")
export DifferentiableDensityModel

include("contrib/ad.jl")

### Init

using Requires

function __init__()
    @require OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" begin
        export DiffEqIntegrator
        include("contrib/diffeq.jl")
    end

    @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" begin
        include("contrib/forwarddiff.jl")
    end

    @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin
        include("contrib/zygote.jl")
    end

    @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" begin
        include("contrib/cuda.jl")
    end
end

end # module


text:
module IMDBGraphs

using ..ProblemBenchmarks: PROBLEM_DATA_DIR
using Base.Iterators
using Statistics

mutable struct IMDBNode
    name::String # actor//film name
    neighbors::Set{IMDBNode} # adjacent nodes
    IMDBNode(name::AbstractString) = new(name, Set{IMDBNode}())
end

const IMDBGraph = Dict{String, IMDBNode}

function fetch_node!(G::IMDBGraph, name::AbstractString)
    if haskey(G, name)
        return G[name]
    else # if node doesn't exist, create it
        node = IMDBNode(name)
        G[name] = node
        return node
    end
end

function centrality_mean(G::IMDBGraph, start_name::AbstractString)
    distances = Dict{String,UInt}()
    next_names = Set{String}([String(start_name)])
    current_distance = 0
    while !(isempty(next_names))
        neighbor_names = Set{String}()
        for name in next_names
            if !haskey(distances, name)
                distances[name] = current_distance
                for node in fetch_node!(G, name).neighbors
                    push!(neighbor_names, node.name)
                end
            end
        end
        current_distance += 1
        next_names = neighbor_names
    end
    return mean(values(distances))
end

function read_graph(path)
    G, actor_names = IMDBGraph(), Set{String}()
    open(path, "r") do file
        while !(eof(file))
            k = split(strip(readline(file)), "\t")
            actor_name, movie_name = k[1], join(k[2:3], "_")
            actor_node = fetch_node!(G, actor_name)
            movie_node = fetch_node!(G, movie_name)
            push!(actor_names, actor_node.name)
            push!(actor_node.neighbors, movie_node)
            push!(movie_node.neighbors, actor_node)
        end
    end
    return G, sort!(collect(actor_names))
end

function perf_imdb_centrality(n_actors)
    path = joinpath(PROBLEM_DATA_DIR, "imdb.tsv")
    G, actor_names = read_graph(path)
    results = Dict{String, Float64}()
    for name in take(actor_names, n_actors)
        results[name] = centrality_mean(G, name)
    end
    return sort!([Pair{Float64,String}(v, k) for (k, v) in results])
end

end # module


text:
using AxisAlgorithms
using Test

@testset "matmul" begin
    n = 5
    m = 3
    src = rand(n,n,n)
    for M in (rand(m,n), sprand(m,n,0.2))
        for dim = 1:3
            dest1 = mapslices(b->M*b, src, dims=dim)
            sz = fill(n,3)
            sz[dim] = m
            dest2 = rand(sz...)
            AxisAlgorithms.A_mul_B_md!(dest2, M, src, dim)
            @test dest1 ≈ dest2
            if !issparse(M)
                rand!(dest2)
                AxisAlgorithms.A_mul_B_perm!(dest2, M, src, dim)
                @test dest1 ≈ dest2
            end
        end
    end

    # Test size-checking
    for dim = 1:3
        M = sprand(m,n,0.2)
        dest1 = mapslices(b->M*b, src, dims=dim)
        sz = fill(n,3)
        sz[dim] = m+1
        dest2 = rand(sz...)
        @test_throws DimensionMismatch AxisAlgorithms.A_mul_B_md!(dest2, M, src, dim)
        @test_throws DimensionMismatch AxisAlgorithms.A_mul_B_perm!(dest2, M, src, dim)
        sz[dim] = m
        sz[mod1(dim+1,3)] = 1
        dest2 = rand(sz...)
        @test_throws DimensionMismatch AxisAlgorithms.A_mul_B_md!(dest2, M, src, dim)
        @test_throws DimensionMismatch AxisAlgorithms.A_mul_B_perm!(dest2, M, src, dim)
    end

    # Test 1x1 and 2x2 cases
    n = 5
    for dim = 1:3
        sz = fill(n,3)
        sz[dim] = 1
        src2 = rand(sz...)
        M = fill(3.2,1,1)
        dest1 = mapslices(b->M*b, src2, dims=dim)
        dest2 = AxisAlgorithms.A_mul_B_md(M, src2, dim)
        @test dest1 ≈ dest2
        sz[dim] = 2
        src2 = rand(sz...)
        M = rand(2,2)
        dest1 = mapslices(b->M*b, src2, dims=dim)
        dest2 = AxisAlgorithms.A_mul_B_md(M, src2, dim)
        @test dest1 ≈ dest2
    end
end

max_stars_repo_name: timholy/FilterMD

max_issues_repo_name: timholy/FilterMD

max_forks_repo_name: timholy/FilterMD


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

# Checkers #
# ========================== #
check_generators(stochasticprogram::StochasticProgram{N}) where N = _check_generators(stochasticprogram, Val(N))
function _check_generators(stochasticprogram::StochasticProgram, ::Val{N}) where N
    _check_stage_generator(stochasticprogram, N)
    _check_generators(stochasticprogram, Val(N-1))
    return nothing
end
function _check_generators(stochasticprogram::StochasticProgram, ::Val{1})
    has_generator(stochasticprogram, :stage_1) || error("First-stage problem not defined in stochastic program. Consider @stage 1.")
    return nothing
end
function _check_stage_generator(stochasticprogram::StochasticProgram{N}, stage::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    stage_key = Symbol(:stage_, stage)
    has_generator(stochasticprogram, stage_key) || error("Stage problem $stage not defined in stochastic program. Consider @stage $stage.")
    return nothing
end
# Model generation #
# ========================== #
"""
    stage_one_model(stochasticprogram::StochasticProgram; optimizer = nothing)

Return a generated copy of the first stage model in `stochasticprogram`. Optionally, supply a capable `optimizer` to the stage model.
"""
function stage_one_model(stochasticprogram::StochasticProgram; optimizer = nothing)
    # Return possibly cached model
    cache = problemcache(stochasticprogram)
    if haskey(cache, :stage_1)
        m = cache[:stage_1]
        optimizer != nothing && set_optimizer(m, optimizer)
        return m
    end
    # Check that the required generators have been defined
    has_generator(stochasticprogram, :stage_1) || error("First-stage problem not defined in stochastic program. Consider @stage 1.")
    model = optimizer == nothing ? Model() : Model(optimizer)
    # Prepare decisions
    model.ext[:decisions] = Decisions(Val{1}(); is_node = true)
    add_decision_bridges!(model)
    # Generate and cache first-stage model
    generator(stochasticprogram, :stage_1)(model, stage_parameters(stochasticprogram, 1))
    cache[:stage_1] = model
    return model
end
"""
    stage_model(stochasticprogram::StochasticProgram, stage::Integer, scenario::AbstractScenario; optimizer = nothing)

Return a generated stage model corresponding to `scenario`, in `stochasticprogram`. Optionally, supply a capable `optimizer` to the stage model.
"""
function stage_model(stochasticprogram::StochasticProgram{N},
                     stage::Integer,
                     scenario::AbstractScenario;
                     optimizer = nothing) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N")
    stage == 1 && return stage_one_model(stochasticprogram, optimizer)
    # Check generators
    stage_key = Symbol(:stage_, stage)
    decision_key = Symbol(:stage_, stage - 1, :_decisions)
    has_generator(stochasticprogram, stage_key) || error("Stage problem $stage not defined in stochastic program. Consider @stage $stage")
    has_generator(stochasticprogram, decision_key) || error("No decision variables defined in stage problem $(stage-1).")
    # Create stage model
    stage_model = optimizer == nothing ? Model() : Model(optimizer)
    # Prepare decisions
    stage_model.ext[:decisions] = Decisions(stage; is_node = true)
    add_decision_bridges!(stage_model)
    # Generate and return the stage model
    generator(stochasticprogram, decision_key)(stage_model, decision_params)
    generator(stochasticprogram, stage_key)(stage_model, stage_params, scenario)
    return stage_model
end
function generate_stage_one!(stochasticprogram::StochasticProgram)
    haskey(stochasticprogram.problemcache, :stage_1) && return nothing
    has_generator(stochasticprogram, :stage_1) || error("First-stage problem not defined in stochastic program. Consider @stage 1.")
    stochasticprogram.problemcache[:stage_1] = JuMP.Model()
    generator(stochasticprogram, :stage_1)(stochasticprogram.problemcache[:stage_1], stage_parameters(stochasticprogram, 1))
    return nothing
end

function generate_proxy!(stochasticprogram::StochasticProgram{N}) where N
    # First-stage decisions are unique (reuse)
    proxy(stochasticprogram, 1).ext[:decisions] = Decisions((stochasticprogram.decisions[1],); is_node = true)
    # Generate first stage
    has_generator(stochasticprogram, :stage_1) || error("First-stage problem not defined in stochastic program. Consider @stage 1.")
    generator(stochasticprogram, :stage_1)(proxy(stochasticprogram, 1), stage_parameters(stochasticprogram, 1))
    # Generate remaining stages
    for s in 2:N
        # Initialize decisions
        decisions = ntuple(Val{s}()) do i
            if i == s - 1
                # Known decisions from the previous stages are
                # the same everywhere.
                return stochasticprogram.decisions[s]
            end
            return DecisionMap()
        end
        proxy(stochasticprogram, s).ext[:decisions] = Decisions(decisions)
        # Check generator
        stage_key = Symbol(:stage_, s)
        decision_key = Symbol(:stage_, s - 1, :_decisions)
        has_generator(stochasticprogram, stage_key) || error("Stage problem $stage not defined in stochastic program. Consider @stage $stage.")
        has_generator(stochasticprogram, decision_key) || error("No decision variables defined in stage problem $(stage-1).")
        # Generate
        generator(stochasticprogram, decision_key)(proxy(stochasticprogram, s), stage_parameters(stochasticprogram, s-1))
        generator(stochasticprogram, stage_key)(proxy(stochasticprogram, s), stage_parameters(stochasticprogram, s), scenario(stochasticprogram, s, 1))
    end
    return nothing
end
"""
    clear!(stochasticprogram::StochasticProgram)

Clear the `stochasticprogram`, removing all model definitions.
"""
function clear!(stochasticprogram::StochasticProgram{N}) where N
    # Clearing invalidates the cache
    invalidate_cache!(stochasticprogram)
    # Dispatch clearup on stochastic structure
    clear!(structure(stochasticprogram))
    # Clear proxy
    for s in 1:N
        proxy_ = proxy(stochasticprogram, s)
        # Clear decisions
        if haskey(proxy_.ext, :decisions)
            clear!(proxy_.ext[:decisions])
        end
        # Clear model
        empty!(proxy_)
    end
    return nothing
end
"""
    generate!(stochasticprogram::StochasticProgram)

Generate the `stochasticprogram` using the model definitions from @stage and available data.
"""
function generate!(stochasticprogram::StochasticProgram{N}) where N
    # Clear stochasticprogram before re-generation
    clear!(stochasticprogram)
    # Abort early if there are no scenarios
    num_scenarios(stochasticprogram) == 0 && return nothing
    # Check generators
    check_generators(stochasticprogram::StochasticProgram{N})
    # Generate proxy
    generate_proxy!(stochasticprogram)
    # Dispatch generation on stochastic structure
    generate!(stochasticprogram, structure(stochasticprogram))
    return nothing
end
# Outcome model generation #
# ========================== #
function _outcome_model!(outcome_model::JuMP.Model,
                         decision_generator::Function,
                         generator::Function,
                         decision_params::Any,
                         stage_params::Any,
                         stage_one_model::JuMP.Model,
                         decisions::AbstractVector,
                         scenario::AbstractScenario)
    # Prepare decisions
    outcome_model.ext[:decisions] = Decisions(Val{2}(); is_node = true)
    add_decision_bridges!(outcome_model)
    # Generate the outcome model
    decision_generator(outcome_model, decision_params)
    generator(outcome_model, stage_params, scenario)
    # Copy first-stage objective
    copy_decision_objective!(stage_one_model,
                             outcome_model,
                             all_known_decision_variables(outcome_model, 1))
    # Update the known decision values
    update_known_decisions!(outcome_model.ext[:decisions][1], decisions)
    update_known_decisions!(outcome_model)
    return nothing
end
"""
    outcome_model(stochasticprogram::TwoStageStochasticProgram,
                  decision::AbstractVector,
                  scenario::AbstractScenario;
                  optimizer = nothing)

Return the resulting second stage model if `decision` is the first-stage decision in the provided `scenario`, in `stochasticprogram`. The supplied `decision` must match the defined decision variables in `stochasticprogram`. Optionally, supply a capable `optimizer` to the outcome model.
"""
function outcome_model(stochasticprogram::TwoStageStochasticProgram,
                       decision::AbstractVector,
                       scenario::AbstractScenario;
                       optimizer = nothing)
    # Check generators
    has_generator(stochasticprogram,:stage_1_decisions) || error("First-stage not defined in stochastic program. Consider @first_stage or @stage 1.")
    has_generator(stochasticprogram,:stage_2) || error("Second-stage problem not defined in stochastic program. Consider @second_stage.")
    # Sanity checks on given decision vector
    length(decision) == num_decisions(stochasticprogram) || error("Incorrect length of given decision vector, has ", length(decision), " should be ", num_decisions(stochasticprogram))
    all(.!(isnan.(decision))) || error("Given decision vector has NaN elements")
    # Create outcome model
    outcome_model = optimizer == nothing ? Model() : Model(optimizer)
    _outcome_model!(outcome_model,
                    generator(stochasticprogram,:stage_1_decisions),
                    generator(stochasticprogram,:stage_2),
                    stage_parameters(stochasticprogram, 1),
                    stage_parameters(stochasticprogram, 2),
                    stage_one_model(stochasticprogram),
                    decision,
                    scenario)
    return outcome_model
end
"""
    outcome_model(stochasticprogram::StochasticProgram{N},
                  decisions::NTuple{N-1,AbstractVector}
                  scenario_path::NTuple{N-1,AbstractScenario},
                  optimizer = nothing)

Return the resulting `N`:th stage model if `decisions` are the decisions taken in the previous stages and `scenario_path` are the realized scenarios up to stage `N` in `stochasticprogram`. Optionally, supply a capable `solver` to the outcome model.
"""
function outcome_model(stochasticprogram::StochasticProgram{N},
                       decisions::NTuple{M,AbstractVector},
                       scenario_path::NTuple{M,AbstractScenario};
                       optimizer = nothing) where {N,M}
    N == M - 1 || error("Inconsistent number of stages $N and number of decisions and scenarios $M")
    # TODO
end
# ========================== #


text:
#### Code and functions specific to the Yodzis-Innes metabolic model
#### "Slow organisms exhibit sudden population disappearances in a reddened world" by Greyson-Gaito, Gellner, & McCann.

@with_kw mutable struct YodInnScalePar
    aₜ = 54.9
    aᵣ = 34.3
    rat = 10^-1
    δ = 0.55
    fⱼ = 0.99
    aⱼ = 89.2
    k = 3.0
    R₀ = 0.8
    ε = 0.1
    x = ( aₜ / aᵣ ) * rat^0.25
    y = fⱼ * aⱼ / aₜ
end

function yod_inn_scale_II!(du, u, p, t)
    @unpack aₜ, aᵣ, y, δ, k, R₀, ε, rat = p
    fₑ = ( (rat * ε) ^ 0.25 ) / ( rat ^ 0.25 )
    xₛ = ( aₜ / aᵣ ) * (rat * ε)^0.25
    R, C = u
    du[1] = R * (1 - R / k) -  ( xₛ * y / ( (1 - δ) * fₑ ) )  * R * C  / (R₀ + R)
    du[2] = ( xₛ * y * R * C  / (R₀ + R) ) - xₛ * C
    return
end #setting up the Yodzis-Innes model

function yod_inn_scale_II(u, par)
    du = similar(u)
    yod_inn_scale_II!(du, u, par, 0.0)
    return du
end #setting up the Yodzis-Innes model

function eq_yodinn_II(p)
    @unpack x, y, δ, k, R₀ = p
    eq_II_R = R₀ / (y - 1)
    eq_II_C = ( ( 1 - δ ) / x ) * eq_II_R * (1 - eq_II_R / k)
    return vcat(eq_II_R, eq_II_C)
end #calculating the interior equilibrium of the Yodzis-Innes model

function con_iso_yodinn(p)
    @unpack R₀, y = p
    R₀/(y - 1)
end #calculating the consumer isocline of the Yodzis-Innes model

function res_iso_yodinn(R, p)
    @unpack x, y, δ, R₀, k  = p
    (R*(R*δ - R + R₀*δ - R₀) + k*(-R*δ + R - R₀*δ + R₀))/(k*x*y)
end #function to help calculate the resource isocline of the Yodzis-Innes model

function iso_plot_YodInn(resrange, par)
    data = [res_iso_yodinn(R, par) for R in resrange]
    plot(collect(resrange), data)
    return vlines(con_iso_yodinn(par), 0, 0.5, colors="orange")
end #plots the isoclines of the Yodzis-Innes model

function YodInn_pert(ep, R0, freq, r, seed, tsend, tvals)
    Random.seed!(seed)
    par = YodInnScalePar()
    par.ε = ep
    par.R₀ = R0
    noise = noise_creation(r, tsend / freq)
    count = 1
    u0 = [eq_yodinn_II(par)[1], eq_yodinn_II(par)[2] + noise[1]]
    tspan = (0.0, tsend)

    function pert_cb2(integrator)
        count += 1
        if isapprox(integrator.u[2], 0.00000000; atol = 1e-8)
            integrator.u[2] = 0.00000000
        else
            integrator.u[2] = maximum([integrator.u[2] + noise[count], 0.0])
        end
    end

    cb = PeriodicCallback(pert_cb2, freq, initial_affect = false) #as of may 29th - initial_affect does not actually do the affect on the first time point
    prob = ODEProblem(yod_inn_scale_II!, u0, tspan, par)
    sol = DifferentialEquations.solve(prob, callback = cb, reltol = 1e-8)
    return solend = sol(tvals)
end #function to numerically solve the Yodzis-Innes consumer-resource model with added noise to the consumer

function pert_YodInn_phase_plot(ep, R0, freq, r, seed, tsend, tvals)
    sol = YodInn_pert(ep, R0, freq, r, seed, tsend, tvals)
    plot(sol[1, :], sol[2, :], color = "green")
    xlabel("Resource")
    return ylabel("Consumer")
end #plot a phase space time series for the Yodzis-Innes model

### Below are functions to help calculate the feasible R₀ values, the isocline functions, and the Hopf bifurcation values for the Yodzis-Innes model
# #Feasible R₀
# function feasibleR0(p)
#     @unpack y, k = p
#     lower = ( (y-1)/(y+1) ) * k
#     upper = (y-1) * k
#     return vcat(lower, upper)
# end

# feasibleR0(YodInnScalePar()) #0.699849 1.82557

## isoclines
# using SymPy
# @vars R C
# @vars x y R₀ k δ

# f(R, C) = R * (1 - R / k) - ( x * y / (1 - δ) ) * R * C / ( R₀ + R )
# g(R ,C) = x * y * R * C / (R + R₀) - x * C

# SymPy.solve(f(R,C),R)
# SymPy.solve(f(R,C),C)
# SymPy.solve(g(R,C),R)
# SymPy.solve(g(R,C),C)

# #find maximum of resource isocline
# using Calculus
# Calculus.simplify(Calculus.differentiate("(R*(R*δ - R + R₀*δ - R₀) + k*(-R*δ + R - R₀*δ + R₀))/(k*x*y)", :R))
# @vars R C
# @vars x y R₀ k δ a
# r(R) = (((((((R * δ - R) + R₀ * δ) - R₀) + R * (δ - 1)) + k * (1 + -δ)) * (k * x * y)) / (k * x * y) ^ 2)

# SymPy.solve(r(R), R)

# function hopf_YodInn(p)
#     @unpack k, R₀, δ, x, y= p
#     hopf_R = -R₀/2 + k/2
#     hopf_C = (hopf_R*(hopf_R*δ - hopf_R + R₀*δ - R₀) + k*(-hopf_R*δ + hopf_R - R₀*δ + R₀))/(k*x*y)
#     return vcat(hopf_R,hopf_C)
# end

# hopf_YodInn(YodInnScalePar()) #R=1.1, C=6.651085519687272


text:
include("test_ops.jl")
include("test_dispatch.jl")


text:
struct PercivalAlg end
const AugLag = PercivalAlg

@params struct PercivalOptions
    nt::NamedTuple
end
function PercivalOptions(; first_order = true, memory = 5, kwargs...)
    return PercivalOptions(
        (;first_order = first_order,
        memory = memory, inity = ones,
        kwargs...),
    )
end
const AugLagOptions = PercivalOptions

@params mutable struct PercivalWorkspace <: Workspace
    model::Model
    problem::Percival.NLPModels.AbstractNLPModel
    x0::AbstractVector
    options::PercivalOptions
    counter::Base.RefValue{Int}
end
function PercivalWorkspace(
    model::Model, x0::AbstractVector = getinit(model);
    options = PercivalOptions(), kwargs...,
)
    problem, counter = getpercival_problem(model, x0)
    return PercivalWorkspace(model, problem, x0, options, counter)
end
@params struct PercivalResult
    minimizer
    minimum
    problem
    result
    fcalls
end

function optimize!(workspace::PercivalWorkspace)
    @unpack problem, options, x0, counter = workspace
    counter[] = 0
    m = getnconstraints(workspace.model)
    result = _percival(problem; options.nt..., inity = options.nt.inity(m))
    result.solution = result.solution[1:length(x0)]
    return PercivalResult(
        copy(result.solution), result.objective, problem, result, counter[],
    )
end

function _percival(nlp;
    μ::Real = eltype(nlp.meta.x0)(10.0),
    max_iter::Int = 1000, max_time::Real = Inf,
    max_eval::Int = 100000, atol::Real = 1e-6,
    rtol::Real = 1e-6, ctol::Real = 1e-6, first_order = true, memory = 5,
    subsolver_logger::Percival.AbstractLogger = Percival.NullLogger(),
    inity = nothing, max_cgiter = 100, subsolver_max_eval = 200, kwargs...,
)
    modifier = m -> NLPModelsModifiers.LBFGSModel(m, mem = memory)
    _kwargs = (
        max_iter = max_iter, max_time = max_time,
        max_eval = max_eval, atol = atol, rtol = rtol,
        subsolver_logger = subsolver_logger,
        subproblem_modifier = first_order ? modifier : identity,
        subsolver_max_eval = subsolver_max_eval,
        subsolver_kwargs = Dict(:max_cgiter => max_cgiter),
    )
    if Percival.unconstrained(nlp) || Percival.bound_constrained(nlp)
        return Percival.percival(
            Val(:tron), nlp; _kwargs...,
        )
    elseif Percival.equality_constrained(nlp)
        return Percival.percival(
            Val(:equ), nlp; inity = inity, _kwargs...,
        )
    else # has inequalities
        return Percival.percival(
            Val(:ineq), nlp; inity = inity, _kwargs...,
        )
    end
end

function Workspace(model::AbstractModel, optimizer::PercivalAlg, args...; kwargs...,)
    return PercivalWorkspace(model, args...; kwargs...)
end

function getpercival_problem(model::Model, x0::AbstractVector)
    eq = if length(model.eq_constraints.fs) == 0
        nothing
    else
        model.eq_constraints
    end
    ineq = if length(model.ineq_constraints.fs) == 0
        nothing
    else
        model.ineq_constraints
    end
    obj = CountingFunction(getobjective(model))
    return getpercival_problem(
        obj,
        ineq,
        eq,
        x0,
        getmin(model),
        getmax(model),
    ), obj.counter
end
function getpercival_problem(obj, ineq_constr, eq_constr, x0, xlb, xub)
    nvars = length(x0)
    if ineq_constr !== nothing
        ineqval = ineq_constr(x0)
        ineq_nconstr = length(ineqval)
    else
        ineqval = Float64[]
        ineq_nconstr = 0
    end
    if eq_constr !== nothing
        eqval = eq_constr(x0)
        eq_nconstr = length(eqval)
    else
        eqval = Float64[]
        eq_nconstr = 0
    end
    c = x -> begin
        if ineq_constr !== nothing
            v1 = ineq_constr(x)
        else
            v1 = eltype(x)[]
        end
        if eq_constr !== nothing
            v2 = eq_constr(x)
        else
            v2 = eltype(x)[]
        end
        return [v1; v2]
    end
    lcon = [fill(-Inf, ineq_nconstr); zeros(eq_nconstr)]
    ucon = zeros(ineq_nconstr + eq_nconstr)
    nlp = ADNLPModels.ADNLPModel(obj, x0, xlb, xub, c, lcon, ucon, adbackend = ADNLPModels.ZygoteAD())
    return nlp
end


text:
@test SS_FILE_D ∩ SS_RANK_3 == SquareSet(SQ_D3)

@test SS_FILE_D ∪ SS_RANK_3 == SquareSet(
    SQ_D1,
    SQ_D2,
    SQ_D3,
    SQ_D4,
    SQ_D5,
    SQ_D6,
    SQ_D7,
    SQ_D8,
    SQ_A3,
    SQ_B3,
    SQ_C3,
    SQ_D3,
    SQ_E3,
    SQ_F3,
    SQ_G3,
    SQ_H3,
)

@test SS_FILE_D - SS_RANK_3 == SquareSet(SQ_D1, SQ_D2, SQ_D4, SQ_D5, SQ_D6, SQ_D7, SQ_D8)

@test SS_RANK_5 - SQ_E5 == SquareSet(SQ_A5, SQ_B5, SQ_C5, SQ_D5, SQ_F5, SQ_G5, SQ_H5)

@test SS_RANK_5 - SQ_E6 == SS_RANK_5

@test SS_RANK_5 + SQ_E6 ==
      SquareSet(SQ_A5, SQ_B5, SQ_C5, SQ_D5, SQ_E5, SQ_F5, SQ_G5, SQ_H5, SQ_E6)

@test issingleton(SquareSet(SQ_D5))
@test !issingleton(SquareSet(SQ_D5, SQ_C5))
@test !issingleton(SS_EMPTY)

@test issubset(SquareSet(SQ_E5, SQ_E3, SQ_E1), SS_FILE_E)
@test !issubset(SquareSet(SQ_E5, SQ_E3, SQ_D1), SS_FILE_E)
@test issubset(SS_EMPTY, SquareSet(SQ_C2))

for s in ALL_SQUARES
    @test s ∈ filesquares(file(s))
    @test s ∈ ranksquares(rank(s))
end

@test squarecount(SS_FILE_A ∪ SS_RANK_4) == 15
@test squarecount(SS_EMPTY) == 0

@test shift_n(SS_RANK_2) == SS_RANK_3
@test shift_n(SquareSet(SQ_D3, SQ_E4, SQ_F8)) == SquareSet(SQ_D4, SQ_E5)
@test shift_s(SS_RANK_2) == SS_RANK_1
@test shift_s(SquareSet(SQ_C3, SQ_D2, SQ_E1)) == SquareSet(SQ_C2, SQ_D1)
@test shift_e(SS_FILE_F) == SS_FILE_G
@test shift_e(SquareSet(SQ_F5, SQ_G6, SQ_H7)) == SquareSet(SQ_G5, SQ_H6)
@test shift_w(SS_FILE_C) == SS_FILE_B
@test shift_w(SquareSet(SQ_C5, SQ_B6, SQ_A7)) == SquareSet(SQ_B5, SQ_A6)

const BLOCKERS = SquareSet(SQ_B2, SQ_E3, SQ_A5, SQ_E7)

@test bishopattacks(BLOCKERS, SQ_E5) == SquareSet(
    SQ_B2,
    SQ_C3,
    SQ_D4,
    SQ_F6,
    SQ_G7,
    SQ_H8,
    SQ_H2,
    SQ_G3,
    SQ_F4,
    SQ_D6,
    SQ_C7,
    SQ_B8,
)

@test rookattacks(BLOCKERS, SQ_E5) ==
      SquareSet(SQ_E3, SQ_E4, SQ_E6, SQ_E7, SQ_A5, SQ_B5, SQ_C5, SQ_D5, SQ_F5, SQ_G5, SQ_H5)

@test queenattacks(BLOCKERS, SQ_E5) ==
      bishopattacks(BLOCKERS, SQ_E5) ∪ rookattacks(BLOCKERS, SQ_E5)

@test knightattacks(SQ_E5) ==
      SquareSet(SQ_D3, SQ_F3, SQ_C4, SQ_G4, SQ_C6, SQ_G6, SQ_D7, SQ_F7)

@test kingattacks(SQ_E5) ==
      SquareSet(SQ_D4, SQ_E4, SQ_F4, SQ_D5, SQ_F5, SQ_D6, SQ_E6, SQ_F6)

@test pawnattacks(WHITE, SQ_E5) == SquareSet(SQ_D6, SQ_F6)
@test pawnattacks(BLACK, SQ_E5) == SquareSet(SQ_D4, SQ_F4)

@test squaresbetween(SQ_A4, SQ_D4) == SquareSet(SQ_B4, SQ_C4)

@test squaresbetween(SQ_F7, SQ_A2) == SquareSet(SQ_B3, SQ_C4, SQ_D5, SQ_E6)


text:
using FreeTypeAbstraction: height_insensitive_boundingbox

function parent_transform(x)
    p = parent(transformation(x))
    isnothing(p) ? Mat4f(I) : p.model[]
end

function boundingbox(x)
    return parent_transform(x) * data_limits(x)
end

function project_widths(matrix, vec)
    pr = project(matrix, vec)
    zero = project(matrix, zeros(typeof(vec)))
    return pr - zero
end

function rotate_bbox(bb::Rect3f, rot)
    points = decompose(Point3f, bb)
    Rect3f(Ref(rot) .* points)
end

function gl_bboxes(gl::GlyphCollection)
    scales = gl.scales.sv isa Vec2f ? (gl.scales.sv for _ in gl.extents) : gl.scales.sv
    map(gl.extents, gl.fonts, scales) do ext, font, scale
        unscaled_hi_bb = height_insensitive_boundingbox_with_advance(ext, font)
        hi_bb = Rect2f(
            Makie.origin(unscaled_hi_bb) * scale,
            widths(unscaled_hi_bb) * scale
        )
    end
end

function height_insensitive_boundingbox_with_advance(ext, font)
    l = 0f0
    r = FreeTypeAbstraction.hadvance(ext)
    b = FreeTypeAbstraction.descender(font)
    t = FreeTypeAbstraction.ascender(font)
    return Rect2f((l, b), (r - l, t - b))
end

function boundingbox(glyphcollection::GlyphCollection, position::Point3f, rotation::Quaternion)

    if isempty(glyphcollection.glyphs)
        return Rect3f(position, Vec3f(0, 0, 0))
    end

    chars = glyphcollection.glyphs
    glyphorigins = glyphcollection.origins
    glyphbbs = gl_bboxes(glyphcollection)

    bb = Rect3f()
    for (char, charo, glyphbb) in zip(chars, glyphorigins, glyphbbs)
        # ignore line breaks
        # char in ('\r', '\n') && continue

        charbb = rotate_bbox(Rect3f(glyphbb), rotation) + charo + position
        if !isfinite_rect(bb)
            bb = charbb
        else
            bb = union(bb, charbb)
        end
    end
    !isfinite_rect(bb) && error("Invalid text boundingbox")
    bb
end

function boundingbox(layouts::AbstractArray{<:GlyphCollection}, positions, rotations)

    if isempty(layouts)
        Rect3f((0, 0, 0), (0, 0, 0))
    else
        bb = Rect3f()
        broadcast_foreach(layouts, positions, rotations) do layout, pos, rot
            if !isfinite_rect(bb)
                bb = boundingbox(layout, pos, rot)
            else
                bb = union(bb, boundingbox(layout, pos, rot))
            end
        end
        !isfinite_rect(bb) && error("Invalid text boundingbox")
        bb
    end
end

function boundingbox(x::Text{<:Tuple{<:GlyphCollection}})
    boundingbox(
        x[1][],
        to_ndim(Point3f, x.position[], 0),
        to_rotation(x.rotation[])
    )
end

function boundingbox(x::Text{<:Tuple{<:AbstractArray{<:GlyphCollection}}})
    boundingbox(
        x[1][],
        to_ndim.(Point3f, x.position[], 0),
        to_rotation(x.rotation[])
    )
end

function text_bb(str, font, size)
    rot = Quaternionf(0,0,0,1)
    layout = layout_text(
        str, size, font, Vec2f(0), rot, 0.5, 1.0,
        RGBAf(0, 0, 0, 0), RGBAf(0, 0, 0, 0), 0f0)
    return boundingbox(layout, Point3f(0), rot)
end

"""
Calculate an approximation of a tight rectangle around a 2D rectangle rotated by `angle` radians.
This is not perfect but works well enough. Check an A vs X to see the difference.
"""
function rotatedrect(rect::Rect{2}, angle)
    ox, oy = rect.origin
    wx, wy = rect.widths
    points = @SMatrix([
        ox oy;
        ox oy+wy;
        ox+wx oy;
        ox+wx oy+wy;
    ])
    mrot = @SMatrix([
        cos(angle) -sin(angle);
        sin(angle) cos(angle);
    ])
    rotated = mrot * points'

    rmins = minimum(rotated, dims = 2)
    rmaxs = maximum(rotated, dims = 2)

    return Rect2(rmins..., (rmaxs .- rmins)...)
end


text:
"""
    module Countries

Tools for handling countries (according to ISO-3116).

The API consists of:
- `Country`
- `country_numeric`
- `country_alpha2`
- `country_alpha3`
- `country_name`
- `country_assigned`
- `new_country`
- `alias_country`
- `each_country`
"""
module Countries

using DelimitedFiles: readdlm

export Country, country_numeric, country_alpha2, country_alpha3, country_name, country_assigned, new_country, alias_country, each_country

### constructors

"""
    Country(id)

A country with the given `id`.

It is canonically represented by its alpha2 code, such as "GB". Two countries with
the same code are identically equal.

The following are all ways to construct the UK:
```julia
# ISO-3166 codes
Country(826)
Country("GBR")
Country("gbr")
Country("GB")

# by name (or unambiguous partial name)
Country("United Kingdom of Great Britain and Northern Ireland")
Country("United Kingdom")
Country("Britain")

# by alias
alias_country("England", "GB")
Country("england")
```
"""
struct Country
    idx::Int16
    Country(::Val{:new}, idx::Int16) = new(idx)
end

const DATA_PATH = joinpath(dirname(@__DIR__), "data", "IP2LOCATION-COUNTRY-INFORMATION.CSV")
const DATA_COLS = (alpha2=1, name=2, alpha3=3, numeric=4)

const STRICT = Ref(false)

const ALPHA2 = [string(c1, c2) for c1 in 'A':'Z' for c2 in 'A':'Z']
const COUNTRIES = [Country(Val(:new), Int16(i)) for (i,_) in enumerate(ALPHA2)]
const ASSIGNED = [false for _ in ALPHA2]
const ALPHA3 = ["" for _ in ALPHA2]
const NUMERIC = [Int16(0) for _ in ALPHA2]
const NAME = ["" for _ in ALPHA2]

const LOOKUP = Dict{String,Country}(zip(ALPHA2, COUNTRIES))
const LOOKUP_NUM = Dict{Int16,Country}()
const LOOKUP_CACHE = Dict{String,Country}()

const AMBIG = Dict{String,Country}()

_lookup_error(x, extra="") = error("$(repr(x)) is not a recognized country name$extra; you can use new_country to define a new country or alias_country to define an alias for an existing country")

Country(x::AbstractString) = get!(LOOKUP_CACHE, x) do
    u = uppercase(x)
    get(LOOKUP, u) do
        if length(u) > 3
            ks = String[]
            cs = Set{Country}()
            for (k, c) in LOOKUP
                if occursin(u, k)
                    push!(ks, k)
                    push!(cs, c)
                end
            end
            if length(cs) == 1 && !STRICT[]
                c = only(cs)
                AMBIG[x] = c
                return c
            elseif isempty(ks)
                _lookup_error(x)
            else
                sort!(ks, by=Country)
                kk = join([repr(k) for k in ks], ", ", " or ")
                _lookup_error(x, " (perhaps you meant $kk)")
            end
        else
            _lookup_error(x)
        end
    end
end

Country(x::Integer) = get(LOOKUP_NUM, x) do
    error("$x is not a recognized country code")
end

Country(x::Country) = x

### inspect countries

"""
    country_numeric(country)

The numeric code of the given country.

Example:
```julia-repl
julia> country_numeric("GBR")
826
```
"""
country_numeric(c) = @inbounds NUMERIC[Country(c).idx]

"""
    country_alpha2(country)

The alpha2 code of the given country.

Example:
```julia-repl
julia> country_alpha2("GBR")
"GB"
```
"""
country_alpha2(c) = @inbounds ALPHA2[Country(c).idx]

"""
    country_alpha3(country)

The alpha3 code of the given country.

Example:
```julia-repl
julia> country_alpha3("United Kingdom")
"GBR"
```
"""
country_alpha3(c) = @inbounds ALPHA3[Country(c).idx]

"""
    country_name(country)

The name of the given country.

Example:
```julia-repl
julia> country_name("GB")
"United Kingdom of Great Britain and Northern Ireland"
```
"""
country_name(c) = @inbounds NAME[Country(c).idx]

"""
    country_assigned(country)

True if the given country is assigned.

Example:
```julia-repl
julia> country_assigned("GB")
true

julia> country_assigned("ZZ")
false
```
"""
country_assigned(c) = @inbounds ASSIGNED[Country(c).idx]

"""
    each_country()

Iterator over each assigned country.

Example:
```julia-repl
julia> collect(each_country())
250-element Vector{Country}:
 AD: Andorra
 AE: United Arab Emirates
 ⋮
 ZM: Zambia
 ZW: Zimbabwe
```
"""
function each_country()
    return (c for c in COUNTRIES if country_assigned(c))
end

### overload Base

function Base.print(io::IO, c::Country)
    print(io, country_alpha2(c))
end

function Base.show(io::IO, c::Country)
    if get(io, :typeinfo, Any) == Country
        print(io, country_alpha2(c))
    else
        show(io, Country)
        print(io, "(")
        show(io, country_alpha2(c))
        print(io, ")")
    end
end

function Base.show(io::IO, ::MIME"text/plain", c::Country)
    if get(io, :compact, false)
        show(io, c)
    else
        print(io, country_alpha2(c), ": ", country_assigned(c) ? country_name(c) : "(Unassigned Country)")
    end
end

function Base.write(io::IO, c::Country)
    return write(io, c.idx)
end

function Base.read(io::IO, ::Type{Country})
    return Country(Val(:new), read(io, Int16))
end

function Base.:(==)(c1::Country, c2::Country)
    c1.idx == c2.idx
end

function Base.hash(c::Country, h::UInt)
    hash(c.idx, hash(Country, h))
end

function Base.isequal(c1::Country, c2::Country)
    isequal(c1.idx, c2.idx)
end

function Base.isless(c1::Country, c2::Country)
    isless(c1.idx, c2.idx)
end

### new countries

function _check_numeric(x)
    x = convert(Int16, x)
    0 ≤ x ≤ 999 || error("numeric must be between 0 and 999")
    return x
end

function _check_alpha2(x)
    x = uppercase(convert(String, x))
    length(x) == 2 || error("alpha2 must be length 2")
    all('A' ≤ c ≤ 'Z' for c in x) || error("alpha2 must consist only of letters")
    return x
end

function _check_alpha3(x)
    x = uppercase(convert(String, x))
    if !isempty(x)
        length(x) == 3 || error("alpha3 must be empty or length 3")
        all('A' ≤ c ≤ 'Z' for c in x) || error("alpha3 must consist only of letters")
    end
    return x
end

"""
    new_country(; alpha2, alpha3="", numeric=0, name="")

Register a new country with the given data.

Example:
```julia-repl
julia> new_country(alpha2="ZZ", alpha3="ZZZ", numeric=999, name="Zedland")

julia> Country("zzz")
ZZ: Zedland
```
"""
function new_country(; alpha2, alpha3="", numeric=0, name="")
    alpha2 = _check_alpha2(alpha2)
    alpha3 = _check_alpha3(alpha3)
    numeric = _check_numeric(numeric)
    name = convert(String, name)
    uname = uppercase(name)
    country = Country(alpha2)
    idx = country.idx
    # check none of the info clashes
    if ASSIGNED[idx]
        error("country=$(repr(country)) already assigned")
    end
    if haskey(LOOKUP, alpha3)
        error("alpha3=$(repr(alpha3)) already used by country=$(Country(alpha3))")
    end
    if haskey(LOOKUP, uname)
        error("name=$(repr(name)) already used by country=$(Country(name))")
    end
    if haskey(LOOKUP_NUM, numeric)
        error("numeric=$numeric already used by country=$(Country(numeric))")
    end
    # assign
    ASSIGNED[idx] = true
    ALPHA3[idx] = alpha3
    NUMERIC[idx] = numeric
    NAME[idx] = name
    # lookup
    for k in [alpha3, uname]
        if !isempty(k)
            LOOKUP[k] = country
        end
    end
    if numeric > 0
        LOOKUP_NUM[numeric] = country
    end
    empty!(LOOKUP_CACHE)
    return
end

"""
    alias_country(alias, country)

Register an alias for the given country so that `Country(alias)` returns `country`.

Example:
```julia-repl
julia> alias_country("England", "GBR")

julia> Country("england")
GB: United Kingdom of Great Britain and Northern Ireland
```
"""
function alias_country(name, country)
    name = convert(String, name)
    uname = uppercase(name)
    country = Country(country)
    if length(name) < 4
        error("alias=$(repr(name)) too short, must have length at least 4")
    end
    if haskey(LOOKUP, uname)
        if LOOKUP[uname] != country
            error("alias=$(repr(name)) already used by country=$(Country(name))")
        end
        return
    end
    LOOKUP[uname] = country
    empty!(LOOKUP_CACHE)
    return
end

### populate default countries at precompile-time

let
    table = readdlm(DATA_PATH, ',', String)
    nrows, ncols = size(table)
    @assert nrows == 250
    for i in 2:nrows
        numeric = parse(Int16, table[i,DATA_COLS.numeric])
        alpha2 = table[i,DATA_COLS.alpha2]
        alpha3 = table[i,DATA_COLS.alpha3]
        name = table[i,DATA_COLS.name]
        new_country(; numeric, alpha2, alpha3, name)
    end
end

### precompile

for country in Any["GB", "GBR", "United Kingdom", 826, Country("GB"), Country("ZZ")]
    Country(country)
    country_alpha2(country)
    country_alpha3(country)
    country_assigned(country)
    country_name(country)
    country_numeric(country)
end
for country in ["GB", "gb", "Gb", "FOO", "foo", "Foo", "FOOO", "fooo", "Fooo"]
    try
        Country(country)
    catch
    end
end
collect(each_country())
empty!(AMBIG)

end # module


text:
"""
    diffusion(g, p, n)

Run diffusion simulation on `g` for `n` steps with spread
probabilities based on `p`. Return a vector with the set of new
vertices reached at each step of the simulation.

### Optional Arguments
- `initial_infections=sample(vertices(g), 1)`: A list of vertices that
are infected at the start of the simulation.
- `watch=Vector()`: While simulation is always run on the full graph,
specifying `watch` limits reporting to a specific set of vertices reached
during the simulation. If left empty, all vertices will be watched.
- `normalize=false`: if `false`, set the probability of spread from a vertex ``i`` to
each of the out_neighbors of ``i`` to ``p``. If `true`, set the probability of spread
from a vertex ``i`` to each of the `out_neighbors` of ``i`` to
``\\frac{p}{outdegreee(g, i)}``.
"""

function diffusion(g::AbstractGraph,
                   p::Real,
                   n::Integer;
                   watch::AbstractVector=Vector{Int}(),
                   initial_infections::AbstractVector=LightGraphs.sample(vertices(g), 1),
                   normalize::Bool=false
                   )

    # Initialize
    T = eltype(g)
    watch_set = Set{T}(watch)
    infected_vertices = IntSet(initial_infections)
    vertices_per_step::Vector{Vector{T}} = [Vector{T}() for i in 1:n]

    # Record initial infection
    if !isempty(watch_set)
        watched_initial_infections = intersect(watch_set, initial_infections)
        vertices_per_step[1] = T.(watched_initial_infections)
    else
        vertices_per_step[1] = T.(initial_infections)
    end

    # Run simulation
    for step in 2:n
        new_infections = Set{T}()

        for i in infected_vertices
            outn = out_neighbors(g, i)
            outd = length(outn)

            if outd > 0
                if normalize
                    local_p = p / outdegree(g, i)
                else
                    local_p = p
                end

                union!(new_infections, randsubseq(outn, local_p))
            end
        end

        # Record only new infections
        setdiff!(new_infections, infected_vertices)
            if !isempty(watch_set)
                vertices_per_step[step] = T.(collect(intersect(new_infections, watch_set)))
            else
                vertices_per_step[step] = collect(new_infections)
            end

        # Add new to master set of infected
        union!(infected_vertices, new_infections)
    end

    return vertices_per_step
end

"""
    diffusion_rate(results)
    diffusion_rate(g, p, n; ...)
Given the results of a `diffusion` output or the parameters
to the `diffusion` simulation itself, (run and) return the rate of
diffusion as a vector representing the cumulative number of vertices
infected at each simulation step, restricted to vertices included
in `watch`, if specified.
"""
diffusion_rate(x::Vector{Vector{T}}) where T <: Integer = cumsum(length.(x))
diffusion_rate(g::AbstractGraph, p::Real, n::Integer;
    initial_infections::AbstractVector=LightGraphs.sample(vertices(g), 1),
    watch::AbstractVector=Vector{Int}(),
    normalize::Bool=false
    ) = diffusion_rate(
           diffusion(g, p, n,
           initial_infections=initial_infections,
           watch=watch, normalize=normalize
           )
)


text:
#= Introduction to Chain and DataFrameMacros packages

This script uses a subset of data reported in Fühner, Golle, Granacher, & Kliegl (2021). Physical fitness in third grade of primary school: A mixed model analysis of 108,295 children and 515 schools.

All children were between 6.0 and 6.99 years at legal keydate (30 September) of school enrollement, that is they were in their ninth year of life in the third grade. To speed up we work with a reduced data set and less complex models than those in the reference publication. This illustrates also drawing a stratified subsample from a large data set.

=#
using DataFrames
using Chain
using DataFrameMacros

using CSV
using Arrow
using Downloads

using StatsBase
using Dates

using MacroTools: prettify

#= Readme for './data/fggk21.rds'

Number of scores: 525126

 1. Cohort: 9 levels; 2011-2019
 2. School: 515 levels
 3. Child: 108295 levels; all children are between 8.0 and 8.99 years old
 4. Sex: "Girls" (n=55,086), "Boys" (n= 53,209)
 5. age: testdate - middle of month of birthdate
 6. Test: 5 levels
     + Endurance (`Run`):  6 minute endurance run [m]; to nearest 9m in 9x18m field
     + Coordination (`Star_r`): star coordination run [m/s]; 9x9m field, 4 x diagonal = 50.912 m
     + Speed(`S20_r`): 20-meters sprint [m/s]
     + Muscle power low (`SLJ`): standing long jump [cm]
     + Muscle power up (`BPT`): 1-kg medicine ball push test [m]
 7. score - see units
 =#

 # Read the data

df = DataFrame(Arrow.Table("./data/fggk21.arrow"))
describe(df)

#= Extract a stratified subsample

We extract a random sample of 5 children from the Sex (2) x Test (5) cells of the design. Cohort and School are random.

=#

dat =
	@chain df begin
  	  @transform(:Sex2 = :Sex == "Girls" ? "female" : "male")
   	  @groupby(:Test, :Sex)
   	  combine(x -> x[sample(1:nrow(x), 5), :])
    end

# Three macros:  @transform, @groupby, and @chain -- one at a time

## transform and @transform -- also note the ternary operator for ifelse

### long
transform(dat, :Sex =>
    ByRow(bla -> bla == "female" ? "girl" : "boy") => :Sex2)

transform(dat, :age => (col -> col .+ 1) => :ageplus)

### short - as used above
@transform(dat, :Sex2 = :Sex == "female" ? "girl" : "boy")

@transform(dat, :ageplus = :age + 1)

@transform(dat, @c :age .- mean(:age))  # @c = columnwise

### version 1 - traditional Julia style
df1 = DataFrame(Arrow.Table(Downloads.download(url)));
describe(df1)

## groupby and @groupby
@groupby(dat, :Age = round(Int, :age))

## @chain  --

### reference to a file
url = "https://github.com/RePsychLing/SMLP2021/raw/main/notebooks/data/fggk21.arrow"

### version 1 - traditional Julia style
df1 = DataFrame(Arrow.Table(Downloads.download(url)));
describe(df1)

### version 2 - with Chain
df2 = @chain url begin
    Downloads.download
    Arrow.Table
    DataFrame
   	@groupby(:Test, :Sex)
   	combine(x -> x[sample(1:nrow(x), 2), :])
    @aside CSV.write("test.csv", _)
    end;
df2
describe(df2)

### look behind the scene
@macroexpand(@chain url begin
    Downloads.download
    Arrow.Table
    DataFrame
    end) |> prettify


text:
using QuadrupedController
using Test

using PyCall

@testset "QuadrupedController.jl" begin
    @test QuadrupedController != PyNULL()
    @test size(names(QuadrupedController)) == (13,)
    @test names(QuadrupedController)[1] == :Command
    @test names(QuadrupedController)[2] == :Configuration
    @test names(QuadrupedController)[3] == :QuadrupedController
    @test names(QuadrupedController)[4] == :Robot
    @test names(QuadrupedController)[5] == :behavior_state
    @test names(QuadrupedController)[6] == :behavior_state_string
    @test names(QuadrupedController)[7] == :end_turn
    @test names(QuadrupedController)[8] == :run!
    @test names(QuadrupedController)[9] == :toggle_activate
    @test names(QuadrupedController)[10] == :toggle_hop
    @test names(QuadrupedController)[11] == :toggle_trot
    @test names(QuadrupedController)[12] == :turn_left
    @test names(QuadrupedController)[13] == :turn_right
end

@testset "Command" begin
    command = Command()
    @test string(command) == "Command([0.0, 0.0], 0.0, -0.16, 0.0, 0.0, 0, false, false, false)"
end

@testset "Configuration" begin
    @test Configuration != PyNULL()
    @test Configuration.__module__ == "Config"
    @test string(Configuration) == "PyObject <class 'Config.Configuration'>"

    config = Configuration()
    @test config.ABDUCTION_OFFSET == 0.03
end

@testset "Robot" begin
    config = Configuration()
    @test config.z_clearance == 0.07
    config.z_clearance = 0.01
    command = Command([0.4, 0.0], 0.5, -0.08, 0.1)
    @test string(command) == "Command([0.4, 0.0], 0.5, -0.08, 0.1, 0.0, 0, false, false, false)"
    robot = Robot(config, command)
    @test behavior_state(robot.state) == -1 # DEACTIVATED
    @test behavior_state_string(robot.state) == "DEACTIVATED: -1"
    toggle_activate(robot)
    @test behavior_state(robot.state) == 0  # REST
    @test behavior_state_string(robot.state) == "REST: 0"
    toggle_trot(robot)
    @test behavior_state(robot.state) == 1  # TROT
    @test behavior_state_string(robot.state) == "TROT: 1"
end


text:
################################################################################

"Reset hidden Markov model object."
function reset!(self::HMM)
  self.tb = [χ .= [0; repeat(-1:-1, length(self.tb[1]) - 1)] for χ = self.tb]
end

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

max_stars_repo_name: DanielRivasMD/HiddenMarkovModelReaders

max_issues_repo_name: DanielRivasMD/HiddenMarkovModelReaders

max_forks_repo_name: DanielRivasMD/HiddenMarkovModelReaders