text:
import sbp
using sbp.Source: source_smoothness, source_discretize_1d, source_discretize_2d
using Test

tol=1e-12;
n = 8;
x = collect(1:n);
N = zeros(n)
for i=1:n
        N[i] = (-1) ^ i
end

moments = (coef, stencil, p, xs) -> abs(sum(coef.*x[stencil].^p) - xs^p);
smoothness = (coef, stencil, s, xs) -> abs(sum(coef.*N[stencil].*x[stencil].^s));

"""
 Source stencil spanning 3 cells (nearest stencil point to the left)

             |     *     |
 o---o---o---o---o---o---o---o
 1   2   3   4   5   6   7   8
 * : source location
"""

@testset "Source stencil spanning 3 cells, left" begin
p = 3
s = 0
xs = 5.5
stencil, coef, alpha = source_smoothness(x, xs, p, s)
for i=0:p
        @test (moments(coef, stencil, i, xs) < tol)
end
@test (alpha == -0.5)
@test (length(stencil) == 4)
@test (stencil[1] == 4)
end

"""
 Source stencil spanning 3 cells (nearest stencil point to the right)
             |     *     |
 o---o---o---o---o---o---o---o
 1   2   3   4   5   6   7   8
 * : source location
"""

@testset "Source stencil spanning 3 cells, right" begin
p = 3
s = 0
xs = 5.6
stencil, coef, alpha = source_smoothness(x, xs, p, s)
for i=0:p
        @test (moments(coef, stencil, i, xs) < tol)
end
@test (abs(alpha - 0.4) < tol)
@test (length(stencil) == 4)
@test (stencil[1] == 4)
end

"""
 Source stencil spanning 3 cells (nearest stencil point close to the right
 boundary)

                 |          *|
 o---o---o---o---o---o---o---o
 1   2   3   4   5   6   7   8
 * : source location
"""

@testset "Source stencil spanning 3 cells, right boundary" begin
p = 3
s = 0
xs = 7.7
stencil, coef, alpha = source_smoothness(x, xs, p, s)
for i=0:p
        @test (moments(coef, stencil, i, xs) < tol)
end
@test (abs(alpha - 0.3) < tol)
@test (length(stencil) == 4)
@test (stencil[1] == 5)
end

"""
 Test of smoothness conditions
 Source stencil spanning 4 cells (nearest stencil point to the right)
     |         *         |
 o---o---o---o---o---o---o---o
 1   2   3   4   5   6   7   8
 * : source location
"""

@testset "Source stencil spanning 3 cells, smoothness conditions" begin
p = 3
s = 2;
xs = 4.5;
stencil, coef, alpha = source_smoothness(x, xs, p, s);
for i=0:p
        @test (moments(coef, stencil, i, xs) < tol)
end
for i=0:(s-1)
        @test (smoothness(coef, stencil, i, xs) < tol)
end
end

@testset "1D source discretization" begin
nx = 10
xp, h = sbp.Grid.grid_xp(nx)
xs = 0.5
p = 3
s = 2
d = source_discretize_1d(xs, p, s, xp, h)
@test isapprox(sum(d) * h, 1)

end

@testset "2D source discretization" begin
nx = 10
ny = 10
xp, h1 = sbp.Grid.grid_xp(nx)
yp, h2 = sbp.Grid.grid_xp(ny)
xm, h1 = sbp.Grid.grid_xm(nx + 1)
ym, h2 = sbp.Grid.grid_xm(ny + 1)
xs = 0.5
ys = 0.5
p = 3
s = 2
d = source_discretize_2d(xs, ys, p, s, xp, yp, h1, h2)
@test isapprox(sum(d) * h1 * h2, 1)

end

max_stars_repo_name: ooreilly/sbpjl

max_issues_repo_name: ooreilly/sbpjl

max_forks_repo_name: ooreilly/sbpjl


text:
@testset "Glyph Collections" begin
    using Makie.FreeTypeAbstraction

    # Test whether Makie's padded signed distance field text matches
    # FreeTypeAbstraction characters in terms of boundingbox
    str = "^_lg"
    chars = collect(str)
    font = Makie.defaultfont()

    scene = Scene()
    campixel!(scene)
    p = text!(scene, str, position = Point2f(30, 37), align = (:left, :baseline))
    glyph_collection = p.plots[1][1][]

    # This doesn't work well because FreeTypeAbstraction doesn't quite scale
    # linearly
    # fta_glyphs = map(char -> renderface(font, char, 64), chars)
    # unit_extents = map(fta_glyphs) do (img, extent)
    #     FontExtent(
    #         extent.vertical_bearing * 20f0 / 64f0,
    #         extent.horizontal_bearing * 20f0 / 64f0,
    #         extent.advance * 20f0 / 64f0,
    #         extent.scale * 20f0 / 64f0
    #     )
    # end
    # origins = let
    #     glyph_scale = p.textsize[] / 64
    #     cumsum(vcat(
    #         - glyph_scale * fta_glyphs[1][2].horizontal_bearing[1],
    #         [glyph_scale * fta_glyphs[i][2].advance[1] for i in 1:3]
    #     ))
    # end

    # This is just repeating code from Makie
    unit_extents = [FreeTypeAbstraction.get_extent(font, char) for char in chars]
    origins = cumsum(20f0 * Float32[
        0,
        unit_extents[1].advance[1],
        unit_extents[2].advance[1],
        unit_extents[3].advance[1]
    ])

    @test glyph_collection isa Makie.GlyphCollection
    @test glyph_collection.glyphs == chars
    @test glyph_collection.fonts == [font for _ in 1:4]
    @test glyph_collection.origins == [Point3f(x, 0, 0) for x in origins]
    @test glyph_collection.extents == unit_extents
    @test glyph_collection.scales.sv == [Vec2f(p.textsize[]) for _ in 1:4]
    @test glyph_collection.rotations.sv == [Quaternionf(0,0,0,1) for _ in 1:4]
    @test glyph_collection.colors.sv == [RGBAf(0,0,0,1) for _ in 1:4]
    @test glyph_collection.strokecolors.sv == [RGBAf(0,0,0,0) for _ in 1:4]
    @test glyph_collection.strokewidths.sv == Float32[0, 0, 0, 0]

    # Test quad data
    positions, char_offsets, quad_offsets, uvs, scales = Makie.text_quads(
        to_ndim(Point3f, p.position[], 0), glyph_collection,
        Vec2f(0), Makie.transform_func_obs(scene)[]
    )

    # Also doesn't work
    # fta_offsets = map(fta_glyphs) do (img, extent)
    #     (extent.horizontal_bearing .- Makie.GLYPH_PADDING[]) * p.textsize[] /
    #         Makie.PIXELSIZE_IN_ATLAS[]
    # end
    # fta_scales = map(fta_glyphs) do (img, extent)
    #     (extent.scale .+ 2 * Makie.GLYPH_PADDING[]) * p.textsize[] /
    #         Makie.PIXELSIZE_IN_ATLAS[]
    # end

    fta_quad_offsets = map(chars) do c
        mini = FreeTypeAbstraction.metrics_bb(c, font, 20.0)[1] |> minimum
        Vec2f(mini .- Makie.GLYPH_PADDING[] * 20.0 / Makie.PIXELSIZE_IN_ATLAS[])
    end

    fta_scales = map(chars) do c
        mini = FreeTypeAbstraction.metrics_bb(c, font, 20.0)[1] |> widths
        Vec2f(mini .+ 2 * Makie.GLYPH_PADDING[] * 20.0 / Makie.PIXELSIZE_IN_ATLAS[])
    end

    @test all(isequal(to_ndim(Point3f, p.position[], 0f0)), positions)
    @test char_offsets == glyph_collection.origins
    @test quad_offsets == fta_quad_offsets
    @test scales  == fta_scales
end

ext: jl


text:
# # Even reduction

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Even reduction.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Even reduction.ipynb)

using Test #src
using DynamicPolynomials
@polyvar x

# We would like to find the minimum value of the polynomial

poly = x^4 - 2x^2

using SumOfSquares

# Symmetry reduction is still a work in progress in SumOfSquares, so we include the following files that will be incorporated into SumOfSquares.jl once SymbolicWedderburn.jl is released:
using SumOfSquares
include(joinpath(dirname(dirname(pathof(SumOfSquares))), "examples", "symmetry.jl"))
include(joinpath(dirname(dirname(pathof(SumOfSquares))), "examples", "scaled_perm.jl"))

# We define the custom action as follows:

using PermutationGroups
function action(mono::MP.AbstractMonomial, p::Permutation)
    if p.perm == perm"(1)(2)" || iseven(MP.degree(mono))
        return 1 * mono
    else
        @assert p.perm == perm"(1,2)"
        return -1 * mono
    end
end
G = PermGroup([perm"(1,2)"])

# We can exploit the symmetry as follows:

import CSDP
solver = CSDP.Optimizer
model = Model(solver)
@variable(model, t)
@objective(model, Max, t)
certificate = SymmetricIdeal(Certificate.MaxDegree(SOSCone(), MonomialBasis, maxdegree(poly)), G, action)
con_ref = @constraint(model, poly - t in SOSCone(), ideal_certificate = certificate)
optimize!(model)
@test value(t) ≈ -1 #src
value(t)

# We indeed find `-1`, let's verify that symmetry was exploited:

@test length(gram_matrix(con_ref).sub_gram_matrices) == 2 #src
@test gram_matrix(con_ref).sub_gram_matrices[1].basis.polynomials == [x^2, 1] #src
@test gram_matrix(con_ref).sub_gram_matrices[2].basis.polynomials == [x] #src
gram_matrix(con_ref)

max_stars_repo_path: examples/Even reduction.jl

max_issues_repo_path: examples/Even reduction.jl

max_forks_repo_path: examples/Even reduction.jl


text:
module BaselineAlgos

using TSML.TSMLTypes
using TSML.Utils

import TSML.TSMLTypes.fit!
import TSML.TSMLTypes.transform!

export Baseline,Identity
       fit!,
       transform!

using StatsBase: mode

mutable struct Baseline <: TSLearner
    model
    args

    function Baseline(args=Dict())
        default_args = Dict(
            :output => :class,
            :strat => mode
        )
        new(nothing,mergedict(default_args,args))
    end
end

function fit!(bsl::Baseline,x::Matrix,y::Vector)
    bsl.model = bsl.args[:strat](y)
end

function transform!(bsl::Baseline,x::Matrix)
    fill(bsl.model,size(x,1))
end

mutable struct Identity <: Transformer
    model
    args

    function Identity(args=Dict())
        default_args = Dict{Symbol,Any}()
        new(nothing,mergedict(default_args,args))
    end
end

function fit!(idy::Identity,x::Matrix,y::Vector)
    nothing
end

function transform!(idy::Identity,x::Matrix)
    return x
end

end


text:
abstract type AbstractModel{T} end

function Base.show(io::IO, m::AbstractModel)
    @printf io "Model \n"
    @printf io "no. parameters:             %i\n" n_parameters(m)
    @printf io "description:\n %s\n"          description(m)
end

@inline function Base.getindex(m::AbstractModel, i::Integer)
    if i <= (j = length(get_parameters(m)))
        return get_parameters(m)[i]
    else
        return m.steady_state[i-j]
    end
end

# Need to define like this so as to disable bounds checking
@inline function Base.getindex(m::AbstractModel, k::Symbol)
    i = m.keys[k]
    @inbounds if i <= (j = length(get_parameters(m)))
        return get_parameters(m)[i]
    else
        return m.steady_state[i-j]
    end
end

@inline function Base.setindex!(m::AbstractModel, value::Number, i::Integer)
    if i <= (j = length(get_parameters(m)))
        param = get_parameters(m)[i]
        param.value = value
        if isa(param, ScaledParameter)
            param.scaledvalue = param.scaling(value)
        end
        return param
    else
        steady_state_param = m.steady_state[i-j]
        steady_state_param.value = value
        return steady_state_param
    end
end

Base.setindex!(m::AbstractModel, value::Array, k::Symbol) = Base.setindex!(m, value, m.keys[k])

@inline function Base.setindex!(m::AbstractModel, value::Array, i::Integer)
    if i <= (j = length(get_parameters(m)))
        param = get_parameters(m)[i]
        param.value = value
        if isa(param, ScaledParameter)
            param.scaledvalue = param.scaling(value)
        end
        return param
    else
        steady_state_param = m.steady_state[i-j]
        steady_state_param.value = value
        return steady_state_param
    end
end

"""
```
setindex!(m::AbstractModel, param::AbstractParameter, i::Integer)
```

If `i`<length(get_parameters(m)), overwrites get_parameters(m)[i] with
param. Otherwise, overwrites m.steady_state[i-length(get_parameters(m)).
"""
@inline function Base.setindex!(m::AbstractModel, param::AbstractParameter, i::Integer)
    if i <= (j = length(get_parameters(m)))
        get_parameters(m)[i] = param
    else
        m.steady_state[i-j] = param
    end
    return param
end

Base.setindex!(m::AbstractModel, value, k::Symbol) = Base.setindex!(m, value, m.keys[k])

"""
```
(<=)(m::AbstractModel{T}, p::AbstractParameter{T}) where T
```

Syntax for adding a parameter to a model: m <= parameter.
NOTE: If `p` is added to `m` and length(m.steady_state) > 0, `keys(m)` will not generate the
index of `p` in `get_parameters(m)`.
"""
function (<=)(m::AbstractModel{T}, p::AbstractParameter{T}) where T

    if !in(p.key, keys(m.keys))

        new_param_index = length(m.keys) + 1

        # grow parameters and add the parameter
        push!(get_parameters(m), p)

        # add parameter location to dict
        setindex!(m.keys, new_param_index, p.key)
    else
        # overwrite the previous parameter with the new one
        setindex!(m, p, p.key)
    end
end

"""
```
(<=)(m::AbstractModel{T}, ssp::Union{SteadyStateParameter,SteadyStateParameterArray}) where {T}
```

Add a new steady-state value to the model by appending `ssp` to the `m.steady_state` and
adding `ssp.key` to `m.keys`.
"""
function (<=)(m::AbstractModel{T}, ssp::Union{SteadyStateParameter, SteadyStateParameterArray}) where {T}

    if !in(ssp.key, keys(m.keys))
        new_param_index = length(m.keys) + 1

        # append ssp to steady_state vector
        push!(m.steady_state, ssp)

        # add parameter location to dict
        setindex!(m.keys, new_param_index, ssp.key)
    else
        # overwrite the previous parameter with the new one
        setindex!(m, ssp, ssp.key)
    end
end

"""
```
(<=)(m::AbstractModel{T}, ssp::SteadyStateParameterGrid) where {T}

```

Add a new steady-state value to the model by appending `ssp` to the `m.steady_state` and
adding `ssp.key` to `m.keys`.
"""
function (<=)(m::AbstractModel{T}, ssp::SteadyStateParameterGrid) where {T}

    if !in(ssp.key, keys(m.keys))
        new_param_index = length(m.keys) + 1

        # append ssp to steady_state vector
        push!(m.steady_state, ssp)

        # add parameter location to dict
        setindex!(m.keys, new_param_index, ssp.key)
    else
        # overwrite the previous parameter with the new one
        setindex!(m, ssp, ssp.key)
    end
end

Distributions.logpdf(m::AbstractModel) = logpdf(get_parameters(m))
Distributions.pdf(m::AbstractModel) = exp(logpdf(m))

# Convenience functions
n_states(m::AbstractModel)                  = length(m.endogenous_states)
n_states_augmented(m::AbstractModel)        = n_states(m) + length(m.endogenous_states_augmented)
n_shocks_exogenous(m::AbstractModel)        = length(m.exogenous_shocks)
n_shocks_expectational(m::AbstractModel)    = length(m.expected_shocks)
n_observables(m::AbstractModel)             = length(m.observables)
n_pseudo_observables(m::AbstractModel)      = length(m.pseudo_observables)
n_equilibrium_conditions(m::AbstractModel)  = length(m.equilibrium_conditions)
n_parameters(m::AbstractModel)              = length(get_parameters(m))
n_parameters_steady_state(m::AbstractModel) = length(m.steady_state)
n_parameters_free(m::AbstractModel)         = length(get_free_para_inds(get_parameters(m); regime_switching = true))

function n_parameters_regime_switching(m::AbstractModel)
    return n_parameters_regime_switching(get_parameters(m))
end

get_parameters(m::AbstractModel) = m.parameters

function get_fixed_para_inds(m::AbstractModel; regime_switching::Bool = false,
                             toggle::Bool = true)
    return get_fixed_para_inds(m.parameters; regime_switching = regime_switching,
                               toggle = toggle)
end

function get_free_para_inds(m::AbstractModel; regime_switching::Bool = false,
                             toggle::Bool = true)
    return get_free_para_inds(m.parameters; regime_switching = regime_switching,
                               toggle = toggle)
end

"""
```
get_dict(m, class, index)
```
"""
function get_dict(m::AbstractModel, class::Symbol)
    if class == :states
        m.endogenous_states
    elseif class == :obs
        m.observables
    elseif class == :pseudo
        m.pseudo_observables
    elseif class in [:shocks, :stdshocks]
        m.exogenous_shocks
    else
        throw(ArgumentError("Invalid class: $class. Must be :states, :obs, :pseudo, :shocks, or :stdshocks"))
    end
end

"""
```
get_key(m, class, index)
```

Returns the name of the state (`class = :states`), observable (`:obs`),
pseudo-observable (`:pseudo`), or shock (`:shocks` or `:stdshocks`)
corresponding to the given `index`.
"""
function get_key(m::AbstractModel, class::Symbol, index::Int)
    dict = get_dict(m, class)
    out = Base.filter(key -> dict[key] == index, collect(keys(dict)))
    if length(out) == 0
        error("Key corresponding to index $index not found for class: $class")
    elseif length(out) > 1
        error("Multiple keys corresponding to index $index found for class: $class")
    else
        return out[1]
    end
end

# Interface for I/O settings
spec(m::AbstractModel)         = m.spec
subspec(m::AbstractModel)      = m.subspec
saveroot(m::AbstractModel)     = get_setting(m, :saveroot)
dataroot(m::AbstractModel)     = get_setting(m, :dataroot)

# Interface for data
data_vintage(m::AbstractModel) = get_setting(m, :data_vintage)
data_id(m::AbstractModel)      = get_setting(m, :data_id)

#=
Build paths to where input/output/results data are stored.

Description:
Creates the proper directory structure for input and output files, treating the DSGE/save
    directory as the root of a savepath directory subtree. Specifically, the following
    structure is implemented:

    dataroot/

    savepathroot/
                 output_data/<spec>/<subspec>/log/
                 output_data/<spec>/<subspec>/<out_type>/raw/
                 output_data/<spec>/<subspec>/<out_type>/work/
                 output_data/<spec>/<subspec>/<out_type>/tables/
                 output_data/<spec>/<subspec>/<out_type>/figures/

Note: we refer to the savepathroot/output_data/<spec>/<subspec>/ directory as saveroot.
=#
# """
# ```
# logpath(model)
# ```
# Returns path to log file. Path built as
# ```
# <output root>/output_data/<spec>/<subspec>/log/log_<filestring>.log
# ```
# """
# function logpath(m::AbstractModel)
#     return savepath(m, "log", "log.log")
# end

strs = [:work, :raw, :tables, :figures, :log]
fns = [Symbol(x, "path") for x in strs]
for (str, fn) in zip(strs, fns)
    @eval begin
        # First eval function
        function $fn(m::AbstractModel,
                     out_type::String,
                     file_name::String = "",
                     filestring_addl::Vector{String}=Vector{String}())
            return savepath(m, out_type, $(string(str)), file_name, filestring_addl)
        end

        # Then, add docstring to it
        @doc $(
        """
        ```
        $fn(m::AbstractModel, out_type::String, file_name::String="",
            filestring_addl::Vector{String}=Vector{String}())
        ```

        Returns path to specific $str output file, creating containing directory as needed. If
        `file_name` not specified, creates and returns path to containing directory only.
        Path built as
        ```
        <output root>/output_data/<spec>/<subspec>/<out_type>/$str/<file_name>_<filestring>.<ext>
        ```
        """
        ) $fn
    end
end

# Not exposed to user. Actually create path and insert model string to file name.
function savepath(m::AbstractModel,
                  out_type::String,
                  sub_type::String,
                  file_name::String = "",
                  filestring_addl::Vector{String} = Vector{String}())
    # Containing directory
    dir = String(joinpath(saveroot(m), "output_data", spec(m), subspec(m), out_type, sub_type))

    if !isempty(file_name)
        base = filestring_base(m)
        return savepath(dir, file_name, base, filestring_addl)
    else
        return dir
    end
end

function savepath(dir::String,
                  file_name::String = "",
                  filestring_base::Vector{String} = Vector{String}(),
                  filestring_addl::Vector{String} = Vector{String}())
    if !isdir(dir)
        mkpath(dir)
    end

    if !isempty(file_name)
        (base, ext) = splitext(file_name)
        myfilestring = filestring(filestring_base, filestring_addl)
        file_name_detail = base * myfilestring * ext

        return joinpath(dir, file_name_detail)
    else
        return dir
    end
end

# Input data handled slightly differently, because it is not model-specific.
"""
```
inpath(m::AbstractModel, in_type::T, file_name::T="") where T<:String
```

Returns path to specific input data file, creating containing directory as needed. If
`file_name` not specified, creates and returns path to containing directory only. Valid
`in_type` includes:

* `\"raw\"`: raw input series
* `\"data\"`: transformed data in model units
* `\"cond\"`: conditional data - nowcasts for the current forecast quarter, or related
* `\"user\"`: user-supplied data for starting parameter vector, hessian, or related
* `\"scenarios\"`: alternative scenarios

Path built as
```
<data root>/<in_type>/<file_name>
```
"""
function inpath(m::AbstractModel, in_type::String, file_name::String="")
    path = dataroot(m)
    # Normal cases.
    if in_type in ["raw", "data", "cond", "scenarios"]
        path = joinpath(path, in_type)
    # User-provided inputs. May treat this differently in the future.
    elseif in_type == "user"
        path = joinpath(path, "user")
    else
        error("Invalid in_type: ", in_type)
    end

    # Containing dir
    if !isdir(path)
        mkpath(path)
    end

    # If file_name provided, return full path
    if !isempty(file_name)
        path = joinpath(path, file_name)
    end

    return path
end

function filestring_base(m::AbstractModel)
    if !m.testing
        base = Vector{String}()
        for (skey, sval) in m.settings
            if sval.print
                push!(base, to_filestring(sval))
            end
        end
        return base
    else
        return ["test"]
    end
end

filestring(m::AbstractModel) = filestring(m, Vector{String}(undef, 0))
filestring(m::AbstractModel, d::String) = filestring(m, [String(d)])
function filestring(m::AbstractModel, d::Vector{String})
    base = filestring_base(m)
    return filestring(base, d)
end

function filestring(base::Vector{String}, d::Vector{String})
    filestrings = vcat(base, d)
    sort!(filestrings)
    return "_" * join(filestrings, "_")
end

function filestring(d::Vector{String})
    sort!(d)
    return "_" * join(d, "_")
end

"""
```
rand(d::Union{DegenerateMvNormal,MvNormal}, m::AbstractModel; cc::AbstractFloat = 1.0)
```

Generate a draw from `d` with variance optionally scaled by `cc^2`.
"""
function rand(d::Union{DegenerateMvNormal,MvNormal}, m::AbstractModel; cc::AbstractFloat = 1.0)
    return d.μ + cc*d.σ*randn(m.rng, length(d))
end

"""
```
rand(d::Union{DegenerateMvNormal,MvNormal}, rng::MersenneTwister; cc::AbstractFloat = 1.0)
```

Generate a draw from `d` with variance optionally scaled by `cc^2`.
"""
function rand(d::Union{DegenerateMvNormal,MvNormal}, rng::MersenneTwister; cc::AbstractFloat = 1.0)
    return d.μ + cc*d.σ*randn(rng, length(d))
end

"""
`rand_prior(m::AbstractModel; ndraws::Int = 100_000)`

Draw a random sample from the model's prior distribution.
"""
function rand_prior(m::AbstractModel; ndraws::Int = 100_000)
    T = typeof(get_parameters(m)[1].value)
    npara = length(get_parameters(m))
    priorsim = Array{T}(undef, ndraws, npara)

    for i in 1:ndraws
        priodraw = Array{T}(undef, npara)

        # Parameter draws per particle
        for j in 1:length(get_parameters(m))

            priodraw[j] = if !get_parameters(m)[j].fixed
                prio = rand(get_parameters(m)[j].prior.value)

                # Resample until all prior draws are within the value bounds
                while !(get_parameters(m)[j].valuebounds[1] < prio < get_parameters(m)[j].valuebounds[2])
                    prio = rand(get_parameters(m)[j].prior.value)
                end

                prio
            else
                get_parameters(m)[j].value
            end
        end
        priorsim[i,:] = priodraw'
    end

    priorsim
end

@inline unpack(m::AbstractModel, ::Val{k}) where {k} = get_untransformed_values(m[k])

function parameters2namedtuple(m::AbstractModel)
    return parameters2namedtuple(get_parameters(m))
end

max_stars_repo_name: FRBNY-DSGE/ModelConstructors

max_issues_repo_name: FRBNY-DSGE/ModelConstructors

max_forks_repo_name: FRBNY-DSGE/ModelConstructors


text:
# ALL TESTS THAT CONCERN UNITCELLS

# include
using LatticePhysics
using Base.Test

# begin the Unitcell testblock
lattice_testset = @testset "Lattices" begin

################################################################################
#
#   TYPE DEFINITIONS OF OBJECT CLASSES IN JULIA
#
################################################################################

# begin the testset
@testset "General Lattice functions" begin

    # get a test unitcell
    unitcell = getUnitcellHoneycomb()
    # get a test lattice
    lattice = getLattice(unitcell, -4)

    @testset "File access with JLD" begin
        # test saving
        @test typeof(saveLattice(lattice))==String
        # test loading
        @test typeof(loadLattice(lattice))==String
        @test typeof(loadLattice(lattice.filename))==Lattice
    end;

    @testset "Test and Information functions" begin
        # test printing information
        @test printInfo(lattice)==nothing
        @test printInfo(lattice, detailed=true)==nothing
        # test testing method
        @test typeof(testLattice(lattice, 2, 2))==Bool
    end;

    @testset "Connectivity information functions" begin
        # testing connectivity information
        @test typeof(getConnectivityList(lattice)) == Array{Array,1}
        # testing connection information
        @test typeof(getConnectionList(lattice)) == Array{Array,1}
        # testing connection strength information
        @test typeof(getConnectionStrengthList(lattice)) == Array{Any,1}
    end;

# end the testset here
end;

################################################################################
#
#   IMPLEMENTATIONS OF DIFFERENT LATTICE CONSTRUCTION RELATED STUFF
#   AS WELL AS DIFFERNT LATTICE MODIFICATION STUFF
#
################################################################################

# begin the testset
@testset "Lattice construction functions" begin

    # get a test unitcell
    unitcell_2d = getUnitcellHoneycomb()
    unitcell_3d = getUnitcellDiamond()

    # BASED ON UNITCELLS
    @testset "Construction based on Unitcells" begin

        # PERIODIC
        @testset "Periodic" begin
            # SPECIFIC FUNCTION
            # - just getting lattice
            @test testLattice(getLatticePeriodic(unitcell_2d, [-4,-4]), 2, 2)
            @test testLattice(getLatticePeriodic(unitcell_3d, [-4,-4,-4]), 3, 3)
            # - getting lattice and saving
            @test testLattice(getLatticePeriodic(unitcell_2d, [-4,-4], save=true), 2, 2)
            @test testLattice(getLatticePeriodic(unitcell_3d, [-4,-4,-4], save=true), 3, 3)
            # - getting lattice and saving and loading
            @test testLattice(getLatticePeriodic(unitcell_2d, [-4,-4], save=true, load=true), 2, 2)
            @test testLattice(getLatticePeriodic(unitcell_3d, [-4,-4,-4], save=true, load=true), 3, 3)
            # GENERAL FUNCTION
            # - just getting lattice
            @test testLattice(getLattice(unitcell_2d, [-4,-4]), 2, 2)
            @test testLattice(getLattice(unitcell_3d, [-4,-4,-4]), 3, 3)
            # - getting lattice and saving
            @test testLattice(getLattice(unitcell_2d, [-4,-4], save=true), 2, 2)
            @test testLattice(getLattice(unitcell_3d, [-4,-4,-4], save=true), 3, 3)
            # - getting lattice and saving and loading
            @test testLattice(getLattice(unitcell_2d, [-4,-4], save=true, load=true), 2, 2)
            @test testLattice(getLattice(unitcell_3d, [-4,-4,-4], save=true, load=true), 3, 3)
        end;

        # SEMI PERIODIC
        @testset "Semi-Periodic" begin
            # SPECIFIC FUNCTION
            # - just getting lattice
            @test testLattice(getLatticeSemiperiodic(unitcell_2d, [4,-4]), 1, 2)
            @test testLattice(getLatticeSemiperiodic(unitcell_3d, [4,-4,-4]), 2, 3)
            # - getting lattice and saving
            @test testLattice(getLatticeSemiperiodic(unitcell_2d, [4,-4], save=true), 1, 2)
            @test testLattice(getLatticeSemiperiodic(unitcell_3d, [4,-4,-4], save=true), 2, 3)
            # - getting lattice and saving and loading
            @test testLattice(getLatticeSemiperiodic(unitcell_2d, [4,-4], save=true, load=true), 1, 2)
            @test testLattice(getLatticeSemiperiodic(unitcell_3d, [4,-4,-4], save=true, load=true), 2, 3)
            # GENERAL FUNCTION
            # - just getting lattice
            @test testLattice(getLattice(unitcell_2d, [4,-4]), 1, 2)
            @test testLattice(getLattice(unitcell_3d, [4,-4,-4]), 2, 3)
            # - getting lattice and saving
            @test testLattice(getLattice(unitcell_2d, [4,-4], save=true), 1, 2)
            @test testLattice(getLattice(unitcell_3d, [4,-4,-4], save=true), 2, 3)
            # - getting lattice and saving and loading
            @test testLattice(getLattice(unitcell_2d, [4,-4], save=true, load=true), 1, 2)
            @test testLattice(getLattice(unitcell_3d, [4,-4,-4], save=true, load=true), 2, 3)
        end;

        # OPEN
        @testset "Open" begin
            # SPECIFIC FUNCTION
            # - just getting lattice
            @test testLattice(getLatticeOpen(unitcell_2d, [4,4]), 0, 2)
            @test testLattice(getLatticeOpen(unitcell_3d, [4,4,4]), 0, 3)
            # - getting lattice and saving
            @test testLattice(getLatticeOpen(unitcell_2d, [4,4], save=true), 0, 2)
            @test testLattice(getLatticeOpen(unitcell_3d, [4,4,4], save=true), 0, 3)
            # - getting lattice and saving and loading
            @test testLattice(getLatticeOpen(unitcell_2d, [4,4], save=true, load=true), 0, 2)
            @test testLattice(getLatticeOpen(unitcell_3d, [4,4,4], save=true, load=true), 0, 3)
            # GENERAL FUNCTION
            # - just getting lattice
            @test testLattice(getLattice(unitcell_2d, [4,4]), 0, 2)
            @test testLattice(getLattice(unitcell_3d, [4,4,4]), 0, 3)
            # - getting lattice and saving
            @test testLattice(getLattice(unitcell_2d, [4,4], save=true), 0, 2)
            @test testLattice(getLattice(unitcell_3d, [4,4,4], save=true), 0, 3)
            # - getting lattice and saving and loading
            @test testLattice(getLattice(unitcell_2d, [4,4], save=true, load=true), 0, 2)
            @test testLattice(getLattice(unitcell_3d, [4,4,4], save=true, load=true), 0, 3)
        end;

    # end of unitcell based construction
    end;

    # BASED ON BOND DISTANCE
    @testset "Construction based on Bond Distance" begin

        # - just getting lattice
        @test testLattice(getLatticeByBondDistance(unitcell_2d, 4), 0, 2)
        @test testLattice(getLatticeByBondDistance(unitcell_3d, 4), 0, 3)
        @test testLattice(getLatticeByBondDistance(unitcell_2d, 4, origin=2), 0, 2)
        @test testLattice(getLatticeByBondDistance(unitcell_3d, 4, origin=2), 0, 3)
        # - getting lattice and saving
        @test testLattice(getLatticeByBondDistance(unitcell_2d, 4, save=true), 0, 2)
        @test testLattice(getLatticeByBondDistance(unitcell_3d, 4, save=true), 0, 3)
        # - getting lattice and saving and loading
        @test testLattice(getLatticeByBondDistance(unitcell_2d, 4, save=true, load=true), 0, 2)
        @test testLattice(getLatticeByBondDistance(unitcell_3d, 4, save=true, load=true), 0, 3)

    # end of bond distance based construction
    end;

    # BASED ON SHAPE
    @testset "Construction based on Shape" begin

        # SPHERE / CIRCLE
        @testset "Sphere" begin
            # - just getting lattice
            @test testLattice(getLatticeInSphere(unitcell_2d, 2.0), 0, 2)
            @test testLattice(getLatticeInSphere(unitcell_3d, 2.0), 0, 3)
            @test testLattice(getLatticeInSphere(unitcell_2d, 2.0, origin=2), 0, 2)
            @test testLattice(getLatticeInSphere(unitcell_3d, 2.0, origin=2), 0, 3)
            # - getting lattice and saving
            @test testLattice(getLatticeInSphere(unitcell_2d, 2.0, save=true), 0, 2)
            @test testLattice(getLatticeInSphere(unitcell_3d, 2.0, save=true), 0, 3)
            # - getting lattice and saving and loading
            @test testLattice(getLatticeInSphere(unitcell_2d, 2.0, save=true, load=true), 0, 2)
            @test testLattice(getLatticeInSphere(unitcell_3d, 2.0, save=true, load=true), 0, 3)
        # end testset
        end;

        # BOX
        @testset "Box" begin
            # defining a box that workds for 2d and 3d
            box = [4.0, 2.0, 3.0]
            # - just getting lattice
            @test testLattice(getLatticeInBox(unitcell_2d, box), 0, 2)
            @test testLattice(getLatticeInBox(unitcell_3d, box), 0, 3)
            @test testLattice(getLatticeInBox(unitcell_2d, box, origin=2), 0, 2)
            @test testLattice(getLatticeInBox(unitcell_3d, box, origin=2), 0, 3)
            # - getting lattice and saving
            @test testLattice(getLatticeInBox(unitcell_2d, box, save=true), 0, 2)
            @test testLattice(getLatticeInBox(unitcell_3d, box, save=true), 0, 3)
            # - getting lattice and saving and loading
            @test testLattice(getLatticeInBox(unitcell_2d, box, save=true, load=true), 0, 2)
            @test testLattice(getLatticeInBox(unitcell_3d, box, save=true, load=true), 0, 3)
        # end testset
        end;

        # GENERAL
        @testset "General shape" begin
            # defining a box that workds for 2d and 3d
            shape_function(point) = sum(point.*point) < 2.0
            # - just getting lattice
            @test testLattice(getLatticeInShape(unitcell_2d, shape_function), 0, 2)
            @test testLattice(getLatticeInShape(unitcell_3d, shape_function), 0, 3)
            @test testLattice(getLatticeInShape(unitcell_2d, shape_function, "myshape"), 0, 2)
            @test testLattice(getLatticeInShape(unitcell_3d, shape_function, "myshape"), 0, 3)
            @test testLattice(getLatticeInShape(unitcell_2d, shape_function, "myshape", origin=2), 0, 2)
            @test testLattice(getLatticeInShape(unitcell_3d, shape_function, "myshape", origin=2), 0, 3)
            # - getting lattice and saving
            @test testLattice(getLatticeInShape(unitcell_2d, shape_function, "myshape", save=true), 0, 2)
            @test testLattice(getLatticeInShape(unitcell_3d, shape_function, "myshape", save=true), 0, 3)
            # - getting lattice and saving and loading
            @test testLattice(getLatticeInShape(unitcell_2d, shape_function, "myshape", save=true, load=true), 0, 2)
            @test testLattice(getLatticeInShape(unitcell_3d, shape_function, "myshape", save=true, load=true), 0, 3)
        # end testset
        end;

    # end of shape based construction
    end;

# end the testset here
end;

# end the lattice test block here
end;


text:
# Test all the various ways of specifying a function

@testset "iface" begin

# Using functions modifying in-place

function f!(F, x)
    F[1] = (x[1]+3)*(x[2]^3-7)+18
    F[2] = sin(x[2]*exp(x[1])-1)
end

function j!(J, x)
    J[1, 1] = x[2]^3-7
    J[1, 2] = 3*x[2]^2*(x[1]+3)
    u = exp(x[1])*cos(x[2]*exp(x[1])-1)
    J[2, 1] = x[2]*u
    J[2, 2] = u
end

function fj!(F, J, x)
    F[1] = (x[1]+3)*(x[2]^3-7)+18
    F[2] = sin(x[2]*exp(x[1])-1)
    J[1, 1] = x[2]^3-7
    J[1, 2] = 3*x[2]^2*(x[1]+3)
    u = exp(x[1])*cos(x[2]*exp(x[1])-1)
    J[2, 1] = x[2]*u
    J[2, 2] = u
end

r = nlsolve(OnceDifferentiable(f!, [ -0.5; 1.4], [ -0.5; 1.4]), [ -0.5; 1.4])
@test converged(r)

r = nlsolve(OnceDifferentiable(f!, j!, [ -0.5; 1.4], [ -0.5; 1.4]), [ -0.5; 1.4])
@test converged(r)

r = nlsolve(OnceDifferentiable(f!, j!, fj!, [ -0.5; 1.4], [ -0.5; 1.4]),  [ -0.5; 1.4])
@test converged(r)

r = nlsolve(f!, [ -0.5; 1.4])
@test converged(r)

r = nlsolve(f!, j!, [ -0.5; 1.4])
@test converged(r)

#r = nlsolve(only_f!_and_fj!(f!, fj!), [ -0.5; 1.4])
#@test converged(r)

#r = nlsolve(only_fj!(fj!), [ -0.5; 1.4])
#@test converged(r)

# Use not-in-place forms
function f(x)
    F = Array{eltype(x)}(undef, 2)
    F[1] = (x[1]+3)*(x[2]^3-7)+18
    F[2] = sin(x[2]*exp(x[1])-1)
    return F
end

function g(x)
    J = Array{eltype(x)}(undef, 2, 2)
    J[1, 1] = x[2]^3-7
    J[1, 2] = 3*x[2]^2*(x[1]+3)
    u = exp(x[1])*cos(x[2]*exp(x[1])-1)
    J[2, 1] = x[2]*u
    J[2, 2] = u
    return J
end

function fg(x)
    F = Array{eltype(x)}(undef, 2)
    J = Array{eltype(x)}(undef, 2, 2)
    F[1] = (x[1]+3)*(x[2]^3-7)+18
    F[2] = sin(x[2]*exp(x[1])-1)
    J[1, 1] = x[2]^3-7
    J[1, 2] = 3*x[2]^2*(x[1]+3)
    u = exp(x[1])*cos(x[2]*exp(x[1])-1)
    J[2, 1] = x[2]*u
    J[2, 2] = u
    return F, J
end

r = nlsolve(f, [ -0.5; 1.4]; inplace = false)
@test converged(r)

r = nlsolve(f, [ -0.5; 1.4]; inplace = false, autodiff = :forward)
@test converged(r)

r = nlsolve(f, g, [ -0.5; 1.4]; inplace = false)
@test converged(r)

r = nlsolve(f, g, fg, [ -0.5; 1.4]; inplace = false)
@test converged(r)

# Using functions taking scalar as inputs

function f(x, y)
    fx = (x+3)*(y^3-7)+18
    fy = sin(y*exp(x)-1)
    return(fx,fy)
end

r = nlsolve(n_ary(f), [ -0.5; 1.4])
@test converged(r)
r = nlsolve(n_ary(f), [ -0.5; 1.4], autodiff = :forward)
@test converged(r)

@testset "anderson trace issue #160" begin

    function f_2by2!(F, x)
        F[1] = (x[1]+3)*(x[2]^3-7)+18
        F[2] = sin(x[2]*exp(x[1])-1)
    end

    function g_2by2!(J, x)
        J[1, 1] = x[2]^3-7
        J[1, 2] = 3*x[2]^2*(x[1]+3)
        u = exp(x[1])*cos(x[2]*exp(x[1])-1)
        J[2, 1] = x[2]*u
        J[2, 2] = u
    end

    df = OnceDifferentiable(f_2by2!, g_2by2!, [ -0.5; 1.4], [ -0.5; 1.4])

    r = nlsolve(df, [ 0.01; .99], method = :anderson, m = 10, beta=.01, show_trace=true)
end
end


text:
import Base.copy

using JSON
using DataStructures

type Map
    name
    nodes
    edges
end

copy(m::Map) = Map(m.name, copy(m.nodes), copy(m.edges))

Items = Dict("stool" => 1, "chair" => 2, "easel" => 3,
    "hatrack" => 4, "lamp" => 5, "sofa" => 6, "" => 7)

Walls = Dict("butterfly" => 1, "fish" => 2, "tower" => 3)

Floors = Dict("blue" => 1, "brick" => 2, "concrete" => 3, "flower" => 4,
    "grass" => 5, "gravel" => 6, "wood" => 7, "yellow" => 8)

MapColors = Dict("black" => 1, "blue" => 2, "brown" => 3,
    "green" => 4, "grey" => 5, "pink" => 6, "red" => 7, "yellow" => 8)

ColorMapping = Dict(1 => ["blue"], 2 => ["red", "brown"], 3 => ["grey"],
    4 => ["pink"], 5 => ["green"], 6 => ["black", "grey"],
    7 => ["brown"], 8 => ["yellow", "green"])

function getmap(fname)
    j = JSON.parsefile(fname; dicttype=DataStructures.OrderedDict, use_mmap=true)
    gridmap = j["map"]
    #println(keys(gridmap))
    #println(typeof(gridmap["nodes"]["node"]))

    name = gridmap["_name"]
    nodes = Dict()
    edges = Dict()
    for n in gridmap["nodes"]["node"]
        x = parse(Int, n["_x"])
        y = parse(Int, n["_y"])
        item = Items[n["_item"]]
        get!(nodes, (x,y), item)
    end

    for e in gridmap["edges"]["edge"]
        arr = map(a -> parse(Int, a), split(e["_node1"], ","))
        n1 = (arr[1], arr[2])
        arr = map(a -> parse(Int, a), split(e["_node2"], ","))
        n2 = (arr[1], arr[2])
        wall = Walls[e["_wall"]]
        flr = Floors[e["_floor"]]
        d = get!(edges, n1, Dict(n2 => (wall, flr)))
        get!(d, n2, (wall, flr))
        d = get!(edges, n2, Dict(n1 => (wall, flr)))
        get!(d, n1, (wall, flr))
    end
    return Map(name, nodes, edges)
end

function readmapsjson(fname)
    maps = JSON.parsefile(fname)

    converted = Dict()

    for k in keys(maps)
        m = maps[k]
        nodes = Dict()
        edges = Dict()

        for kn in keys(m["nodes"])
            nodes[eval(parse(kn))] = m["nodes"][kn]
        end

        for ke in keys(m["edges"])
            n1 = eval(parse(ke))
            for kke in keys(m["edges"][ke])
                n2 = eval(parse(kke))
                wall, flr = m["edges"][ke][kke]
                d = get!(edges, n1, Dict(n2 => (wall, flr)))
                get!(d, n2, (wall, flr))
            end
        end

        converted[k] = Map(k, nodes, edges)
    end
    return converted
end

function getlocation(map, current, action)
    if action == 1#move
        if current[3] == 0
            next = (current[1], current[2] - 1, current[3])
        elseif current[3] == 90
            next = (current[1] + 1, current[2], current[3])
        elseif current[3] == 180
            next = (current[1], current[2] + 1, current[3])
        elseif current[3] == 270
            next = (current[1] - 1, current[2] , current[3])
        end
    elseif action == 2#right
        next = (current[1], current[2], (current[3] + 90) % 360)
    elseif action == 3#left
        next = (current[1], current[2], (current[3] - 90 + 360) % 360)
    else#stop
        next = current
    end
    return next
end

max_stars_repo_name: ozanarkancan/navi-app

max_issues_repo_name: ozanarkancan/navi-app

max_forks_repo_name: ozanarkancan/navi-app


text:
function my_retry(f::Function, num_retries::Integer = 1)
    result = retry(f, delays=ExponentialBackOff(n=num_retries))()
    return result
end


text:
function revers(side :: HorizonSide)
    side = HorizonSide((Int(side)+2)%4)
    return side
end

function count_move!(r, side :: HorizonSide, count)
    for i = 1:count
        move!(r,side)
    end
end

function movements_to_side!(r, side :: HorizonSide)
    count = 0
    while isborder(r,side) ==  false
        move!(r,side)
        count = count + 1
    end
    return count
end

function home_path!(r)

    home = Int[]

    while (isborder(r, Nord) == false) || (isborder(r, West) == false)
        s = movements_to_side!(r, Nord)
        w = movements_to_side!(r, West)
        push!(home, s)
        push!(home, w)
    end
    return home
end

function go_home!(r, home)

    home_path!(r)

    reverse!(home)

    for i in 1: length(home)
        if i % 2 == 0
            count_move!(r, Sud, home[i])
        else
            count_move!(r, Ost, home[i])
        end
    end
end

function find_move!(r, side :: HorizonSide, border_side :: HorizonSide)
    while (isborder(r,side) ==  false) && (isborder(r,border_side) ==  false)
        move!(r,side)
    end
    return isborder(r,border_side)
end

function find(r)

    side = Sud
    border_side = Ost
    flag = false

    while flag != true
        flag = find_move!(r, side, border_side);
        if flag
            break
        end
        move!(r,Ost)
        side = revers(side)
    end

end

function perimetr(r)

    for side in [Nord, Ost, Sud, West]
        putmarker!(r)
        move!(r, side)
        while isborder(r, HorizonSide((Int(side)+3)%4)) == true
            putmarker!(r)
            move!(r, side)
        end
    end

    while ismarker(r) ==  false
        putmarker!(r)
        move!(r,Nord)
    end

end

function all_perimetr(r)
    home = Int[]

    home = home_path!(r)
    find(r)
    perimetr(r)
    go_home!(r, home)

end

all_perimetr(r)
show(r)

max_stars_repo_name: AntonDudDm/mirea-progs

max_issues_repo_name: AntonDudDm/mirea-progs

max_forks_repo_name: AntonDudDm/mirea-progs


text:
# Map operations
I(A) = one(A)
F(∇u) = ∇u + I(∇u)
J(∇u) = det(F(∇u))
Finv(∇u) = inv(F(∇u))
FinvT(∇u) = (Finv(∇u)')
E(∇u) = 0.5 * ((F(∇u)')⋅F(∇u) - I(∇u))

# Map operations derivatives
dF(∇du) = ∇du
dJ(∇u,∇du) = J(∇u)*tr(Finv(∇u)⋅dF(∇du))
dFinv(∇u,∇du) = -Finv(∇u) ⋅ dF(∇du) ⋅ Finv(∇u)
dFinvT(∇u,∇du) = (dFinv(∇u,∇du)')
dE(∇u,∇du) = 0.5 * ((dF(∇du)')⋅F(∇u) + (F(∇u)')⋅dF(∇du))

# Fluid constitutive laws
# Cauchy stress
σᵥ_Ωf(μ,u) = 2.0*μ*ε(u)
σᵥ_Ωf(μ) = (∇v,Finv) -> μ*(∇v⋅Finv + (Finv')⋅(∇v'))
# First Piola-Kirchhoff stress
Pᵥ_Ωf(μ,u,v) = (J∘∇(u)) * (σᵥ_Ωf(μ)∘(∇(v),Finv∘∇(u)))' ⋅ (FinvT∘∇(u))
function Pₚ_Ωf(u,p)
  typeof(- (J∘∇(u)) * p * tr((FinvT∘∇(u))))
  - (J∘∇(u)) * p * tr((FinvT∘∇(u)))
end
# First Piola-Kirchhoff stress Jacobian
dPᵥ_Ωf_du(μ,u,du,v) = (dJ∘(∇(u),∇(du))) * (σᵥ_Ωf(μ)∘(∇(v),(Finv∘∇(u))))' ⋅ (FinvT∘∇(u)) +
                      (J∘∇(u)) * (σᵥ_Ωf(μ)∘(∇(v),(dFinv∘(∇(u),∇(du)))))' ⋅ (FinvT∘∇(u)) +
                      (J∘∇(u)) * (σᵥ_Ωf(μ)∘(∇(v),(Finv∘∇(u))))' ⋅ (dFinvT∘(∇(u),∇(du)))
dPₚ_Ωf_du(u,du,p) = - p * ( (dJ∘(∇(u),∇(du))) * tr((FinvT∘∇(u))) +
                           (J∘∇(u)) * tr((dFinvT∘(∇(u),∇(du)))) )
dPᵥ_Ωf_dv(μ,u,dv) = Pᵥ_Ωf(μ,u,dv)
dPₚ_Ωf_dp(u,dp) = Pₚ_Ωf(u,dp)
# Convective term
conv(c,∇v) = (∇v') ⋅ c
dconv(dc,∇dv,c,∇v) = conv(c,∇dv) + conv(dc,∇v)

# Solid constitutive laws
# Second Piola-Kirchhoff stress (Saint-Venant)
Sₛᵥ_Ωs(λ,μ,u) = 2*μ*(E∘∇(u)) + λ*tr((E∘∇(u)))*(I∘∇(u))
dSₛᵥ_Ωs_du(λ,μ,u,du) = 2*μ*(dE∘(∇(u),∇(du))) + λ*tr((dE∘(∇(u),∇(du))))*(I∘∇(u))
# First Piola-Kirchhoff stress (Saint-Venant)
Pₛᵥ_Ωs(λ,μ,u) = (F∘∇(u)) ⋅ Sₛᵥ_Ωs(λ,μ,u)
dPₛᵥ_Ωs_du(λ,μ,u,du) = (dF∘∇(du)) ⋅ Sₛᵥ_Ωs(λ,μ,u) + (F∘∇(u)) ⋅ dSₛᵥ_Ωs_du(λ,μ,u,du)
# C(F) = (F')⋅F
# function S_NH(∇u,λ,μ)
#   Cinv = invLaw(C(F(∇u)))
#   μ*(I(∇u)-Cinv) + λ*logLaw(J(∇u))*Cinv
# end
# function dS_NH(∇u,∇du,λ,μ)
#   Cinv = invLaw(C(F(∇u)))
#   _dE = dE(∇du,∇u)
#   λ*(Cinv⊙_dE)*Cinv + 2*(μ-λ*logLaw(J(∇u)))*Cinv⋅_dE⋅(Cinv')
# end

# Mesh constitutive laws
αₘ(J) = 1.0e-5 / J
σₘ(λ,μ) = ε -> λ*tr(ε)*I(ε) + 2.0*μ*ε
dαₘ(J,dJ) = - 1.0e-5 * 1.0 / (J*J) * dJ
dσₘ(λ,dλ,μ,dμ,ε,dε) = σ_m(λ,μ,dε) + σ_m(dλ,dμ,ε)


text:
# Text provided under a Creative Commons Attribution license, CC-BY.  All code is made available under the FSF-approved BSD-3 license.  (c) Lorena A. Barba, Gilbert F. Forsyth 2017. Thanks to NSF for support via CAREER award #1149784.
# [@LorenaABarba](https://twitter.com/LorenaABarba)
# Translation to Julia by [@miguelraz_](https://twitter.com/miguelraz_)

# 12 steps to Navier–Stokes
# =====
# ***

# Up to now, all of our work has been in one spatial dimension (Steps [1](./01_Step_1.ipynb) to [4](./05_Step_4.ipynb)). We can learn a lot in just 1D, but let's grow up to flatland: two dimensions.

# In the following exercises, you will extend the first four steps to 2D. To extend the 1D finite-difference formulas to partial derivatives in 2D or 3D, just apply the definition: a partial derivative with respect to $x$ is the variation in the $x$ direction *at constant* $y$.

# In 2D space, a rectangular (uniform) grid is defined by the points with coordinates:

# $$x_i = x_0 +i \Delta x$$

# $$y_i = y_0 +i \Delta y$$

# Now, define $u_{i,j} = u(x_i,y_j)$ and apply the finite-difference formulas on either variable $x,y$ *acting separately* on the $i$ and $j$ indices. All derivatives are based on the 2D Taylor expansion of a mesh point value around $u_{i,j}$.

# Hence, for a first-order partial derivative in the $x$-direction, a finite-difference formula is:

# $$ \frac{\partial u}{\partial x}\biggr\rvert_{i,j} = \frac{u_{i+1,j}-u_{i,j}}{\Delta x}+\mathcal{O}(\Delta x)$$

# and similarly in the $y$ direction. Thus, we can write backward-difference, forward-difference or central-difference formulas for Steps 5 to 12. Let's get started!

# Step 5: 2-D Linear Convection
# ----
# ***

# The PDE governing 2-D Linear Convection is written as

# $$\frac{\partial u}{\partial t}+c\frac{\partial u}{\partial x} + c\frac{\partial u}{\partial y} = 0$$

# This is the exact same form as with 1-D Linear Convection, except that we now have two spatial dimensions to account for as we step forward in time.

# Again, the timestep will be discretized as a forward difference and both spatial steps will be discretized as backward differences.

# With 1-D implementations, we used $i$ subscripts to denote movement in space (e.g. $u_{i}^n-u_{i-1}^n$).  Now that we have two dimensions to account for, we need to add a second subscript, $j$, to account for all the information in the regime.

# Here, we'll again use $i$ as the index for our $x$ values, and we'll add the $j$ subscript to track our $y$ values.

# With that in mind, our discretization of the PDE should be relatively straightforward.

# $$\frac{u_{i,j}^{n+1}-u_{i,j}^n}{\Delta t} + c\frac{u_{i, j}^n-u_{i-1,j}^n}{\Delta x} + c\frac{u_{i,j}^n-u_{i,j-1}^n}{\Delta y}=0$$

# As before, solve for the only unknown:

# $$u_{i,j}^{n+1} = u_{i,j}^n-c \frac{\Delta t}{\Delta x}(u_{i,j}^n-u_{i-1,j}^n)-c \frac{\Delta t}{\Delta y}(u_{i,j}^n-u_{i,j-1}^n)$$

# We will solve this equation with the following initial conditions:

# $$u(x,y) = \begin{cases}
# \begin{matrix}
# 2\ \text{for} & 0.5 \leq x, y \leq 1 \cr
# 1\ \text{for} & \text{everywhere else}\end{matrix}\end{cases}$$

# and boundary conditions:

# $$u = 1\ \text{for } \begin{cases}
# \begin{matrix}
# x =  0,\ 2 \cr
# y =  0,\ 2 \end{matrix}\end{cases}$$

using Plots

### Variable declarations

nₓ = 81
ny = 81 # note we can't call this n\_y because of the Unicode standard, not because of Julia
nₜ = 100
c = 1
Δx = 2 / (nx - 1)
Δy = 2 / (ny - 1)
σ = 0.2
Δt = σ * Δx
rangex = range(start = 0, stop = 2, length = nₓ)
rangey = range(start = 0, stop = 2, length = ny)

### TODO
u = ones(ny, nₓ) ##create a 1xn vector of 1's
uₙ = ones(ny, nₓ)

###Assign initial conditions

##set hat function I.C. : u(.5<=x<=1 && .5<=y<=1 ) is 2
# We're going to use an array comprehension for this:
u = [ .5 <= x <= 1 && .5 <= y <= 1 ? 2.0 : 1.0 for x in rangex, y in rangey ]

# Array comprehensions are super useful when initializing data, but might not be as efficient
# They work by putting a `for` expression inside square brackets

firstten = [i for i in 1:10]
squares = [i^2 for i in 1:10]

# You can also filter the elements with a predicate
evensquares = [i^2 for i in 1:10 if iseven(i)]

# and have multiple loops
squaremat = [i^2 for i in 1:4, j in 1:4]

###Plot Initial Condition
### TODO
##the figsize parameter can be used to produce different sized images
fig = pyplot.figure(figsize=(11, 7), dpi=100)
ax = fig.gca(projection='3d')
X, Y = numpy.meshgrid(x, y)
surf = ax.plot_surface(X, Y, u[:], cmap=cm.viridis)

### 3D Plotting Notes

# To plot a projected 3D result, make sure that you have added the Axes3D library.

#     from mpl_toolkits.mplot3d import Axes3D

# The actual plotting commands are a little more involved than with simple 2d plots.

# ```python
# fig = pyplot.figure(figsize=(11, 7), dpi=100)
# ax = fig.gca(projection='3d')
# surf2 = ax.plot_surface(X, Y, u[:])
# ```

# The first line here is initializing a figure window.  The **figsize** and **dpi** commands are optional and simply specify the size and resolution of the figure being produced.  You may omit them, but you will still require the

#     fig = pyplot.figure()

# The next line assigns the plot window the axes label 'ax' and also specifies that it will be a 3d projection plot.  The final line uses the command

#     plot_surface()

# which is equivalent to the regular plot command, but it takes a grid of X and Y values for the data point positions.

# ##### Note

# The `X` and `Y` values that you pass to `plot_surface` are not the 1-D vectors `x` and `y`.  In order to use matplotlibs 3D plotting functions, you need to generate a grid of `x, y` values which correspond to each coordinate in the plotting frame.  This coordinate grid is generated using the numpy function `meshgrid`.

#     X, Y = numpy.meshgrid(x, y)

# ### Iterating in two dimensions

# To evaluate the wave in two dimensions requires the use of several nested for-loops to cover all of the `i`'s and `j`'s.
# ~~~Since Python is not a compiled language there can be noticeable slowdowns in the execution of code with multiple for-loops.~~~
# Since Julia is great at for loops, we'll just write simple for loops
# First try evaluating the 2D convection code and see what results it produces.

# TODO
u = ones((ny, nx))
# u[int(.5 / dy):int(1 / dy + 1), int(.5 / dx):int(1 / dx + 1)] = 2

for n in 1:nₜ ##loop across number of time steps
    uₙ = copy(u)
    row, col = size(u)
    # TODO C vs FORTRAN order
    for j in 2:row
        for i in 2:col
            u[j, i] = (uₙ[j, i] - (c * Δt / Δx * (uₙ[j, i] - uₙ[j, i - 1])) -
                                  (c * Δt / Δy * (uₙ[j, i] - uₙ[j - 1, i])))
            # Can you see why using begin/end can help to see the rectangular boundary conditions?
            u[begin, :] .= 1.0
            u[end, :]   .= 1.0
            u[:, begin] .= 1.0
            u[:, end]   .= 1.0
        end
    end
end

# TODO
fig = pyplot.figure(figsize=(11, 7), dpi=100)
ax = fig.gca(projection='3d')
surf2 = ax.plot_surface(X, Y, u[:], cmap=cm.viridis)

# Array Operations
# ----------------

# Here the same 2D convection code is implemented, but instead of using nested for-loops, the same calculations are evaluated using array views.

# TODO timing
u = ones((ny, nx))
u[CartesianIndices((rangex, rangey))] .= 2

for n in 1:nₜ ##loop across number of time steps
    uₙ = copy(u)
    u[2:end, 2:end] = @views (uₙ[2:end, 2:end] - (c * Δt / Δx * (uₙ[2:end, 2:end] - uₙ[2:end, begin:(end-1)])) -
                              (c * Δt / Δy * (uₙ[2:end, 2:end] - uₙ[begin:(end-1), 2:end])))
    u[begin, :] .= 1.0
    u[end,   :] .= 1.0
    u[:, begin] .= 1.0
    u[:,   end] .= 1.0
end

### Index agnosticism
# The expression `(begin+1):end` represents the same as `2:end`, but only when all the indices are linear. If someone were to shift the indices of `u`
# because it's a special type of matrix (as in OffsetArrays.jl), our code would still work without a hitch or performance hit.
# This is precisely the sort of composability that lets you reuse your code and drag-and-drop algorithms to solve interesting problems
# without having to change your code internals.
# #TODO Unreasonable effectiveness of multiple dispatch.

# TODO
fig = pyplot.figure(figsize=(11, 7), dpi=100)
ax = fig.gca(projection='3d')
surf2 = ax.plot_surface(X, Y, u[:], cmap=cm.viridis)

## Learn More

The video lesson that walks you through the details for Step 5 (and onwards to Step 8) is **Video Lesson 6** on You Tube:

# TODO
# YouTubeVideo('tUg_dE3NXoY')
# src="https://www.youtube.com/embed/tUg_dE3NXoY"

max_stars_repo_name: mkitti/CFDJulia

max_issues_repo_name: mkitti/CFDJulia

max_forks_repo_name: mkitti/CFDJulia


text:
# this contains stuff that is common for VAE, WAE and AAE models
"""
	AEModel

A Union of VAE and AAE types.
"""
const AEModel = Union{VAE, AAE}

# these functions need to be overloaded for convolutional models
function ConditionalDists.condition(p::ConditionalMvNormal, z::AbstractArray{T,4}) where T
	(μ,σ) = ConditionalDists.mean_var(p.mapping(z))
	if size(σ,1) == 1
		σ = dropdims(σ, dims=1)
	end
	ConditionalDists.BatchMvNormal(μ,σ)
end
Distributions.logpdf(d::ConditionalDists.BMN, x::AbstractArray{T,4}) where T<:Real =
	Distributions.logpdf(d, reshape(x,:,size(x,4)))
Distributions.mean(p::ConditionalMvNormal, z::AbstractArray{T,4}) where T<:Real = mean(condition(p,z))
Distributions.var(p::ConditionalMvNormal, z::AbstractArray{T,4}) where T<:Real = var(condition(p,z))
Distributions.rand(p::ConditionalMvNormal, z::AbstractArray{T,4}) where T<:Real = rand(condition(p,z))
Distributions.logpdf(p::ConditionalMvNormal, x::AbstractArray{T,4}, z::AbstractArray{T,2}) where T<:Real =
	logpdf(condition(p,z), x)

"""
	reconstruct(model::AEModel, x)

Data reconstruction.
"""
reconstruct(model::AEModel, x) = mean(model.decoder, rand(model.encoder, x))
reconstruct(model::AEModel, x::AbstractArray{T,4}) where T =
	reshape(mean(model.decoder, rand(model.encoder, x)), size(x)...)

"""
	generate(model::AEModel, N::Int[, outdim])

Data generation. Support output dimension if the output needs to be reshaped, e.g. in convnets.
"""
generate(model::AEModel, N::Int) = mean(model.decoder, rand(model.prior, N))
generate(model::AEModel, N::Int, outdim) = reshape(generate(model, N), outdim..., :)

"""
	encode_mean(model::AEModel, x)

Produce data encodings.
"""
encode_mean(model::AEModel, x) = mean(model.encoder, x)
"""
	encode_mean_gpu(model::AEModel, x[, batchsize])

Produce data encodings. Works only on 4D tensors.
"""
encode_mean_gpu(model, x::AbstractArray{T,4}) where T = encode_mean(model, gpu(Array(x)))
function encode_mean_gpu(model, x::AbstractArray{T,4}, batchsize::Int) where T
	# this has to be done in a loop since doing cat(map()...) fails if there are too many arguments
	dl = Flux.Data.DataLoader(x, batchsize=batchsize)
	z = encode_mean_gpu(model, iterate(dl,1)[1])
	N = size(x, ndims(x))
	encodings = gpu(zeros(eltype(z), size(z,1), N))
	for (i,batch) in enumerate(dl)
		encodings[:,1+(i-1)*batchsize:min(i*batchsize,N)] .= encode_mean_gpu(model, batch)
	end
	encodings
end

"""
	reconstruction_score(model::AEModel, x[, L=1])

Anomaly score based on the reconstruction probability of the data.
"""
function reconstruction_score(model::AEModel, x)
	p = condition(model.decoder, rand(model.encoder, x))
	-logpdf(p, x)
end
reconstruction_score(model::AEModel, x, L::Int) =
	mean([reconstruction_score(model, x) for _ in 1:L])

"""
	reconstruction_score_mean(model::AEModel, x)

Anomaly score based on the reconstruction probability of the data. Uses mean of encoding.
"""
function reconstruction_score_mean(model::AEModel, x)
	p = condition(model.decoder, mean(model.encoder, x))
	-logpdf(p, x)
end
"""
	latent_score(model::AEModel, x[, L=1])

Anomaly score based on the similarity of the encoded data and the prior.
"""
function latent_score(model::AEModel, x)
	z = rand(model.encoder, x)
	-logpdf(model.prior, z)
end
latent_score(model::AEModel, x, L::Int) =
	mean([latent_score(model, x) for _ in 1:L])

"""
	latent_score_mean(model::AEModel, x)

Anomaly score based on the similarity of the encoded data and the prior. Uses mean of encoding.
"""
function latent_score_mean(model::AEModel, x)
	z = mean(model.encoder, x)
	-logpdf(model.prior, z)
end

"""
	aae_score(model::AAE, x, alpha::Real)

A combination of reconstruction and discriminator score.
"""
aae_score(model::AAE, x, alpha::Real) =
	alpha*GenerativeAD.Models.reconstruction_score_mean(model, x) .+
	(1-alpha)*vec(model.discriminator(mean(model.encoder, x)))

# JacoDeco score
# see https://arxiv.org/abs/1905.11890
"""
	jacobian(f, x)

Jacobian of f given due to x.
"""
function jacobian(f, x)
	y = f(x)
	n = length(y)
	m = length(x)
	T = eltype(y)
	j = Array{T, 2}(undef, n, m)
	for i in 1:n
		j[i, :] .= gradient(x -> f(x)[i], x)[1]
	end
	return j
end

"""
	lJacoD(m,x)

Jacobian decomposition JJ(m,x).
"""
function lJacoD(m,x)
	JJ = zeros(eltype(x),size(x,ndims(x)));
	zg = mean(m.encoder,x);
	for i=1:size(x,ndims(x))
		jj,J = jacobian(y->mean(m.decoder,reshape(y,:,1))[:],zg[:,i]);
		(U,S,V) = svd(J);
		JJ[i]= sum(2*log.(S));
	end
	JJ
end

"""
	lpx(m,x)

p(x|g(x))
"""
lpx(m,x) = logpdf(m.decoder,x,mean(m.encoder,x))

"""
	lpz(m,x)

p(z|e(x))
"""
lpz(m,x) = logpdf(m.prior,mean(m.encoder,x)) #

"""
	lp_orthD(m,x)

JacoDeco score: p(x|g(x)) + p(z|e(x)) - JJ(m,x)
"""
jacodeco(m,x) = (lpx(m,x) .+ lpz(m,x) .- lJacoD(m,x));


text:
"""
pn_methods.jl
"""

function solve!(solver::ProjectedNewtonSolver)
    # initialize
    if solver.opts.verbose_pn
        println("\nProjection:")
    end
    ϵ_feas = solver.opts.pn_vtol
    # compute value and jacobian of all constraints
    evaluate_copy_constraints!(solver, solver.X, solver.U)
    viol = norm(solver.d[solver.a], Inf)
    # run solve
    step_count = 0
    while (step_count < solver.opts.n_steps) && (viol > ϵ_feas)
        viol = _projection_solve!(solver)
        if solver.opts.multiplier_projection
            res = multiplier_projection!(solver)
        else
            res = Inf
        end
        step_count += 1
        # compute cost
        J = 0.
        for k = 1:solver.N
            J += cost(solver.obj, solver.X, solver.U, k)
        end
        J_prev = solver.stats.cost[solver.stats.iterations]
        dJ = J_prev - J
        # log
        record_iteration!(solver.stats, cost=J, c_max=viol, is_pn=true,
                          dJ=dJ, gradient=res, penalty_max=NaN)
    end
    # terminate
    terminate!(solver)
    return solver
end

function _projection_solve!(solver::ProjectedNewtonSolver)
    # initialize
    max_refinements = 10
    convergence_rate_threshold = solver.opts.r_threshold
    ρ_chol = solver.opts.ρ_chol
    ρ_primal = solver.opts.ρ_primal
    # update constraints and derivatives c, D, d, λ, a
    evaluate_copy_constraints!(solver, solver.X, solver.U)
    D = solver.D[solver.a, :]
    d = solver.d[solver.a]
    # update costs and derivatives H, g
    evaluate_copy_costs!(solver)
    H = Diagonal(solver.H)
    # check violation
    viol0 = norm(d, Inf)
    if solver.opts.verbose_pn
        println("feas0: $viol0")
    end
    # regularize Hessian
    if ρ_primal > 0.0
        for i = 1:solver.N
            H[i, i] += ρ_primal
        end
    end
    # compute projection operators
    HinvD = H \ D'
    S = Symmetric(D * HinvD)
    Sreg = cholesky(S + ρ_chol * I)
    # line search until convergence
    viol_prev = viol0
    refine_count = 0
    while refine_count < max_refinements
        # line search
        viol = _projection_linesearch!(solver, (S, Sreg), HinvD)
        # log conv rate
        convergence_rate = log10(viol) / log10(viol_prev)
        viol_prev = viol
        if solver.opts.verbose_pn
            println("conv rate: $convergence_rate")
        end
        # exit if converged
        if ((convergence_rate < convergence_rate_threshold) ||
            (viol < solver.opts.pn_vtol))
            break
        end
        refine_count += 1
    end
    return viol_prev
end

function _projection_linesearch!(solver::ProjectedNewtonSolver, S, HinvD)
    # initialize
    viol = Inf
    solve_tol = 1e-8
    refinement_iters = 25
    α = 1.
    ϕ = 0.5
    viol_decreased = false
    max_iter_count = 10
    # grab variables
    N = solver.N
    convals = solver.convals
    a = solver.a
    d = solver.d[a]
    viol0 = norm(d, Inf)
    X = solver.X
    X_tmp = solver.X_tmp
    x_ginds = solver.x_ginds
    U = solver.U
    U_tmp = solver.U_tmp
    u_ginds = solver.u_ginds
    # run linesearch until violation decreases or
    # maximum number of searches
    count = 1
    while true
        δλ = reg_solve(S[1], d, S[2], solve_tol, refinement_iters)
        δZ = -HinvD * δλ
        δZ .*= α
        # build temporary trajectory from update
        for k = 1:N - 1
            solver.X_tmp[k] .= solver.X[k]
            solver.X_tmp[k] .+= δZ[x_ginds[k]]
            solver.U_tmp[k] .= solver.U[k]
            solver.U_tmp[k] .+= δZ[u_ginds[k]]
        end
        solver.X_tmp[N] .= solver.X[N]
        solver.X_tmp[N] .+= δZ[x_ginds[N]]
        # check constraint violation
        evaluate_copy_constraints!(solver, X_tmp, U_tmp)
        max_violation, _ = max_violation_penalty(convals)
        d = solver.d[a]
        viol = norm(d, Inf)
        # log
        if solver.opts.verbose_pn
            println("feas: $(viol) (α = $(α))")
        end
        # evaluate convergence
        if viol < viol0
            viol_decreased = true
            break
        elseif count > max_iter_count
            break
        end
        count += 1
        α *= ϕ
    end
    # accept updated trajectory if the violation decreased
    if viol_decreased
        for k = 1:solver.N - 1
            solver.X[k] .= solver.X_tmp[k]
            solver.U[k] .= solver.U_tmp[k]
        end
        solver.X[N] .= solver.X_tmp[N]
    end
    return viol
end

function multiplier_projection!(solver::ProjectedNewtonSolver)
    D = solver.D[solver.a, :]
    d = solver.d[solver.a]
    λ = solver.λ[solver.a]
    g = solver.g
    res0 = g + D'λ
    A = D * D'
    Areg = A + I * solver.opts.ρ_primal
    b = D * res0
    δλ = -reg_solve(A, b, Areg)
    λ += δλ
    res = g + D'λ  # primal residual
    return norm(res)
end

function reg_solve(A, b, B, tol=1e-10, max_iters=10)
    x = B \ b
    count = 0
    while count < max_iters
        r = b - A * x
        if norm(r) < tol
            break
        else
            x += B \ r
            count += 1
        end
    end
    return x
end

"""
compute the value and jacobian of all constraints along the specified
trajectory, storing them in the proper linearized constraint
"""
function evaluate_copy_constraints!(solver::ProjectedNewtonSolver,
                                    X::Vector{<:AbstractVector},
                                    U::Vector{<:AbstractVector}; eval=true, jac=true,
                                    override_const_jac=false)
    for k = 1:solver.N
        for conval in solver.convals[k]
            # evaluate, active, copy
            if eval
                evaluate!(conval.c, conval.con, X, U, k)
                update_active!(conval)
                solver.d[conval.c_ginds] .= conval.c
                solver.λ[conval.c_ginds] .= conval.λ
                solver.a[conval.c_ginds] .= conval.a
            end
            # jacobian, copy
            if jac && !conval.con.const_jac
                jacobian_copy!(solver.D, conval.con, X, U, k, conval.c_ginds, solver.x_ginds,
                               solver.u_ginds)
            end
        end
    end
    return nothing
end

function evaluate_copy_costs!(solver::ProjectedNewtonSolver)
    H = solver.H
    g = solver.g
    E = solver.E
    x_ginds = solver.x_ginds
    u_ginds = solver.u_ginds
    N = solver.N
    for k = 1:N-1
        cost_derivatives!(E, solver.obj, solver.X, solver.U, k)
        H[x_ginds[k], x_ginds[k]] .= E.Q
        H[u_ginds[k], u_ginds[k]] .= E.R
        H[u_ginds[k], x_ginds[k]] .= E.H
        g[x_ginds[k]] .= E.q
        g[u_ginds[k]] .= E.r
    end
    cost_derivatives!(E, solver.obj, solver.X, solver.U, N)
    H[x_ginds[N], x_ginds[N]] .= E.Q
    g[x_ginds[N]] .= E.q
end


text:
# module Lupin

if VERSION < v"1.7.0-DEV.611"
Base.:∘(f::Function, g::Function) = (x...; kwargs...) -> f(g(x...; kwargs...))
end

# module Lupin


text:
@inline nrows(v::V) where {T, V<:AbstractVector{T}} = length(v)
@inline ncols(v::V) where {T, V<:AbstractVector{T}} = 1

@inline nrows(m::M) where {T, M<:AbstractMatrix{T}} = size(m)[1]
@inline ncols(m::M) where {T, M<:AbstractMatrix{T}} = size(m)[2]

@inline nrows(a::A) where {T, N, A<:AbstractArray{T,N}} = size(m)[1]
@inline ncols(a::A) where {T, N, A<:AbstractArray{T,N}} = N > 1 ? size(m)[2] : throw(BoundsError("$(size(a))[2]"))
@inline nstrata(a::A) where {T, N, A<:AbstractArray{T,N}} = N > 2 ? size(m)[3] : throw(BoundsError("$(size(a))[3]"))

#=
     these benchmark favorably relative to these canoncial formulations

@inline linear_to_cartesian(m::Array{T,N}, linearidx::Int) where {N, T} =
    Tuple(CartesianIndices(m)[linearidx])

@inline cartesian_to_linear(m::Array{T,N}, cartesianidx::C) where {T,N, C<:Tuple} =
     (LinearIndices(m))[cartesianidx...,]

@inline cartesian_to_linear(m::Array{T,N}, cartesianidx::C) where {N, T, C<:Vector{T}} =
     (LinearIndices(m))[cartesianidx...,]
=#

@inline cartesian_to_linear(nrows::Int, row::Int, col::Int) = row + (col-1)*nrows

@inline cartesian_to_linear(m::M, row::Int, col::Int) where {T,M<:AbstractMatrix{T}} =
    row + (col-1)*size(m)[1]

@inline linear_to_cartesian(m::M, linearidx::Int) where {T,M<:AbstractMatrix{T}} =
    linear_to_cartesian(size(m)[1], linearidx)

@inline function linear_to_cartesian(nrows::Int, linearidx::Int)
    row = rem(linearidx, nrows)
    row = ifelse(row === 0, nrows, row)
    col = div(linearidx-row, nrows) + 1
    return row,col
end


text:
module Workers

const WORK_QUEUE = Channel{Task}(0)

macro async(thunk)
    esc(quote
        tsk = @task $thunk
        tsk.storage = current_task().storage
        put!(Workers.WORK_QUEUE, tsk)
        tsk
    end)
end

function init()
    tids = Threads.nthreads() == 1 ? (1:1) : 2:Threads.nthreads()
    Threads.@threads for tid in 1:Threads.nthreads()
        if tid in tids
            Base.@async begin
                for task in WORK_QUEUE
                    schedule(task)
                    wait(task)
                end
            end
        end
    end
    return
end

end # module


text:
using Oceananigans.Fields: KernelComputedField
using KernelAbstractions: @kernel, @index

grid = RegularRectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1),
                            topology=(Periodic, Periodic, Bounded))

@kernel function double!(doubled_field, grid, single_field)
    i, j, k = @index(Global, NTuple)
    @inbounds doubled_field[i, j, k] = 2 * single_field[i, j, k]
end

@kernel function multiply!(multiplied_field, grid, field, multiple)
    i, j, k = @index(Global, NTuple)
    @inbounds multiplied_field[i, j, k] = multiple * field[i, j, k]
end

for arch in archs
    @testset "KernelComputedField [$(typeof(arch))]" begin
        @info "  Testing KernelComputedField..."

        single_field = Field(Center, Center, Center, arch, grid)

        doubled_field = KernelComputedField(Center, Center, Center,
                                            double!, arch, grid,
                                            computed_dependencies = single_field)

        # Test that the constructor worked
        @test doubled_field isa KernelComputedField

        multiple = 3
        multiplied_field = KernelComputedField(Center, Center, Center,
                                               multiply!, arch, grid,
                                               computed_dependencies = doubled_field,
                                               parameters = multiple)

        # Test that the constructor worked
        @test multiplied_field isa KernelComputedField

        set!(single_field, π)
        @test single_field[1, 1, 1] == convert(eltype(single_field), π)

        compute!(doubled_field)
        @test doubled_field[1, 1, 1] == 2π

        # Test boundary conditions / halo filling
        @test doubled_field[0, 1, 1] == doubled_field[2, 1, 1] # periodic
        @test doubled_field[1, 0, 1] == doubled_field[1, 2, 1] # periodic
        @test doubled_field[1, 1, 0] == doubled_field[1, 1, 1] # no flux
        @test doubled_field[1, 1, 1] == doubled_field[1, 1, 2] # no flux

        set!(doubled_field, 0)
        compute!(multiplied_field)

        @test doubled_field[1, 1, 1] == 2π
        @test multiplied_field[1, 1, 1] == multiple * 2 * π

        doubled_face_field = KernelComputedField(Center, Center, Face,
                                                 double!, arch, grid,
                                                 computed_dependencies = single_field)

        # Test that nothing happens for fields on faces in bounded directions
        compute!(doubled_face_field)
        @test doubled_face_field[1, 1, 0] == 0
        @test doubled_face_field[1, 1, 3] == 0
    end
end


text:
module NamedVectors

using Printf: @printf
using ..Traits: efficientoperations

export NamedVector, HomoNamedVector

"""
    NamedVector

Abstract type for all named vectors.
"""
abstract type NamedVector end
@inline Base.:(==)(nv₁::NamedVector, nv₂::NamedVector) = keys(nv₁)==keys(nv₂) && values(nv₁)==values(nv₂)
@inline Base.isequal(nv₁::NamedVector, nv₂::NamedVector) = isequal(keys(nv₁), keys(nv₂)) && isequal(values(nv₁), values(nv₂))
@inline Base.getindex(nv::NamedVector, index::Int) = getfield(nv, index)
@inline Base.setindex!(nv::NamedVector, value, index::Int) = setfield!(nv, index, value)
@inline Base.:<(nv₁::NamedVector, nv₂::NamedVector) = <(efficientoperations, nv₁, nv₂)
@inline Base.isless(nv₁::NamedVector, nv₂::NamedVector) = isless(efficientoperations, nv₁, nv₂)
@inline Base.hash(nv::NamedVector, h::UInt) = hash(values(nv), h)
@inline Base.length(::Type{NV}) where {NV<:NamedVector} = fieldcount(NV)
@inline Base.length(nv::NamedVector) = length(typeof(nv))
@inline Base.iterate(nv::NamedVector, state=1) = (state > length(nv)) ? nothing : (getfield(nv, state), state+1)
@inline Base.iterate(rv::Iterators.Reverse{<:NamedVector}, state=length(rv.itr)) = (state < 1) ? nothing : (getfield(rv.itr, state), state-1)
@inline Base.keys(nv::NamedVector) = fieldnames(typeof(nv))
@inline Base.values(nv::NamedVector) = convert(Tuple, nv)
@inline Base.pairs(nv::NamedVector) = Base.Generator(=>, keys(nv), values(nv))
Base.show(io::IO, nv::NamedVector) = @printf io "%s(%s)" nameof(typeof(nv)) join(repr.(values(nv)), ", ")

"""
    convert(::Type{Tuple}, nv::NamedVector) -> Tuple
    convert(::Type{NV}, nv::Tuple) where {NV<:NamedVector} -> NV

Convert a named vector to tuple and vice versa.
"""
@generated Base.convert(::Type{Tuple}, nv::NamedVector) = Expr(:tuple, (:(getfield(nv, $i)) for i = 1:fieldcount(nv))...)
function Base.convert(::Type{NV}, nv::Tuple) where NV<:NamedVector
    @assert fieldcount(NV) == length(nv) "convert error: mismatched length between $NV($(fieldcount(NV))) and input tuple($(length(nv)))."
    return NV(nv...)
end

"""
    zero(::Type{NV}) where NV<:NamedVector -> NV
    zero(nv::NamedVector) -> typeof(nv)

Get a concrete `NamedVector` with all values being zero.
"""
@inline @generated function Base.zero(::Type{NV}) where NV<:NamedVector
    zeros = (zero(fieldtype(NV, i)) for i = 1:fieldcount(NV))
    return :(NV($(zeros...)))
end
@inline Base.zero(nv::NamedVector) = zero(typeof(nv))

"""
    replace(nv::NamedVector; kwargs...) -> typeof(nv)

Return a copy of a concrete `NamedVector` with some of the field values replaced by the keyword arguments.
"""
@inline Base.replace(nv::NamedVector; kwargs...) = replace(efficientoperations, nv; kwargs...)

"""
    map(f, nvs::NV...) where NV<:NamedVector -> NV

Apply function `f` element-wisely on the input named vectors.
"""
@generated function Base.map(f, nvs::NV...) where NV<:NamedVector
    exprs = Vector{Expr}(undef, fieldcount(NV))
    for i = 1:length(exprs)
        tmp = [:(nvs[$j][$i]) for j = 1:length(nvs)]
        exprs[i] = :(f($(tmp...)))
    end
    return :(($NV)($(exprs...)))
end

"""
    HomoNamedVector{T}

Abstract type for all homogeneous named vectors.
"""
abstract type HomoNamedVector{T} <: NamedVector end
@inline Base.eltype(::Type{<:HomoNamedVector{T}}) where T = T
@inline Base.eltype(nv::HomoNamedVector) = eltype(typeof(nv))

end #module


text:
#=
# Shared memory

**Shared memory** is the memory in a SM(symmetric multiprocessor) which is accessable to all threads running on the SM. It is much faster than global memory being much closer in proximity. The amount of shared memory available depends on the compute capability of the GPU. Increasing the amount of shared memory may reduce occupancy.

`sync_threads()` is a function which adds a [barrier](https://en.wikipedia.org/wiki/Barrier_(computer_science)) for all threads in a thread block. A **barrier** ensures that all threads which belong to it stall once they reach it until all other threads reach the barrier. This is commonly used as a synchronization mechanism to eliminate race conditions.

### Array Reversal

Our job is to reverse an array, i.e $[1, 2, 3] \rightarrow [3, 2, 1]$.
=#

using CUDA, BenchmarkTools

function reverse(input, output = similar(input))
    len = length(input)
    for i = 1:cld(len,2)
        output[i], output[len - i + 1] = input[len - i + 1], input[i]
    end
    output
end

#-

reverse([1, 2, 3, 4, 5])

#-

function gpu_reverse(input, output)
    tid = threadIdx().x
    len = length(input)
    if tid <= cld(len, 2)
        output[tid], output[len - tid + 1] = input[len - tid + 1], input[tid]
    end
    return
end

#-

A = CuArray(collect(1:5))
B = similar(A)
@cuda blocks=1 threads=length(A) gpu_reverse(A, B)
B

#=
There are two ways to declare shared memory: Statically and Dynamically. We declaring it statically when the amount we need is the same amount for all kernel launches and known when writing the kernel. In the other case we declare it dynamically and specify it while launching with the `@cuda` macro using the `shmem` argument.
=#

@doc @cuStaticSharedMem

#-

function gpu_stshmemreverse(input, output)
    ## Maximum size of array is 64
    shmem = @cuStaticSharedMem(eltype(output), 64)
    tid = threadIdx().x
    len = length(input)
    shmem[tid] = input[len - tid + 1]
    output[tid] = shmem[tid]
    return
end

#-

A = CuArray(collect(1:32))
B = similar(A)
@cuda blocks=1 threads=length(A) gpu_stshmemreverse(A, B)
print(B)

#=
When the amount of shared memory required isn't known while writing the kernel overallocating is not a good idea because that may potentially reduce our occupancy. As an SM's resource usage increases it's occupancy goes down. Hence, it's best to use dynamically allocated memory when memory usage can only be known at launch time.
=#

@doc @cuDynamicSharedMem

#-

function gpu_dyshmemreverse(input, output)
    shmem = @cuDynamicSharedMem(eltype(output), (length(output),))
    tid = threadIdx().x
    len = length(input)
    shmem[tid] = input[len - tid + 1]
    output[tid] = shmem[tid]
    return
end

#-

C = CuArray(collect(1:32))
D = similar(C)
@cuda blocks=1 threads=length(C) shmem=length(C) gpu_dyshmemreverse(C, D)
print(D)

#=
## Matrix Transpose
Matrix transpose is an operation which flips a matrix along its main diagonal.

![transpose](../assets/transpose.png)
=#

A = reshape(1:9, (3, 3))

#-

A'

#-

A'' ## (A')'

#=
Ignore the types that Julia returns. `LinearAlgebra.Adjoint` is a wrapper that uses [lazy evaluation](https://en.wikipedia.org/wiki/Lazy_evaluation) to compute the result as required.
=#

## CPU implementation
function cpu_transpose(input, output = similar(input, (size(input, 2), size(input, 1))))
    ## the dimensions of the resultant matrix are reversed
    for index = CartesianIndices(input)
            output[index[2], index[1]] = input[index]
    end
    output
end

#-

A = reshape(1:20, (4, 5))

#-

cpu_transpose(A)

#=
Before we begin working on the GPU consider the following code.
=#

A = CuArray(reshape(1.:9, (3, 3)))

println("A => ", pointer(A))
CUDA.@allowscalar begin
    for i in eachindex(A)
        println(i, " ", A[i], " ", pointer(A, i))
    end
end

#=
Notice how consecutive elements are in the same column rather than the same row. This is because Julia stores its multidimensional arrays in a column major order like Fortran. In contrast to C/C++ which are row-major languages. The reason for making Julia's arrays column major is because a lot of linear algebra libraries are column major to begin with (https://discourse.julialang.org/t/why-column-major/24374/3).
=#

## To index our 2-D array we will split the input into tiles of 32x32 elements.
## Each thread block will launch with 32x8 = 256 threads
## Each thread will work on 4 elements.
const TILE_DIM = 32

function gpu_transpose_kernel(input, output)
    tile_index = ((blockIdx().y, blockIdx().x) .- 1) .* TILE_DIM

    ## each thread manages 4 rows (8x4 = 32)
    for i in 1:4
        thread_index = (threadIdx().y + (i - 1)*8, threadIdx().x)
        index = CartesianIndex(tile_index .+ thread_index)
        (index[1] > size(input, 1) || index[2] > size(input, 2)) && continue
        @inbounds output[index] = input[index[2], index[1]]
    end

    return
end

#-

function gpu_transpose(input, output = similar(input, (size(input, 2), size(input, 1))))
    threads = (32, 8)
    blocks = cld.(size(input), (32, 32))
    @cuda blocks=blocks threads=threads gpu_transpose_kernel(input, output)
    output
end

#-

A = CuArray(reshape(1f0:1089, 33, 33))

#-

gpu_transpose(A)

#-

A = CUDA.rand(10000, 10000)
B = similar(A)
@benchmark CUDA.@sync gpu_transpose($A, $B)

#-

@benchmark CUDA.@sync $B .= $A

#=
## Coalescing Memory Access

Compared to a simple elementwise copy we are roughly at 60% performance. Both kernels have a single load and store for each value. If all loads and stores were independent of each other then this should not have happened.

Consider a thread accessing(load or store) a single value in global memory. Instead of transferring just the one value the GPU will instead transfer a larger chunk of memory as a single transaction. For example on NVIDIA's K20 GPU this size was 128 bytes.
When threads in a warp access consecutive memory addresses the GPU can service multiple threads in the same transaction. This is known as **coalesced memory access**. Access time is effectively reduced by minimizing the number of transactions. However when threads access non-sequentially or sparse data then transactions are serialised.

We want consecutive threads of a warp to access consecutive elements in memory. When the thread block is one-dimensional it is straightforward to determine a thread's`warpId` i.e. `warpId = threadId().x % warpsize()`.
According to NVIDIA's documentation on [thread hierarchy](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#thread-hierarchy).
> The index of a thread and its thread ID relate to each other in a straightforward way: For a one-dimensional block, they are the same; for a two-dimensional block of size (Dx, Dy),the thread ID of a thread of index (x, y) is (x + y Dx); for a three-dimensional block of size (Dx, Dy, Dz), the thread ID of a thread of index (x, y, z) is (x + y Dx + z Dx Dy).

In our kernel there are four loads and stores per thread.
- `tile_index = ((blockIdx().y, blockIdx().x) .- 1) .* TILE_DIM`
- `thread_index = (threadIdx().y + (i - 1)*8, threadIdx().x)`
- `index = CartesianIndex(tile_index .+ thread_index)`
- `Load: input[index[2], index[1]]`
- `Store: output[index[1], index[2]]`

The loads are coalesced because the column is indexed by `index[2]` which has `threadIdx().x` and the stores are non-coalesced because they are indexed by `index[1]` which has `threadIdx().y`.

To ensure coalescing during both loads and stores we will use shared memory. We will load from global memory a column and store it in shared memory as a row, effectively transposing it. Once all threads have written to shared memory we can write back to global memory column wise.

![Coalesced transpose](../assets/coalesced_transpose.png)
=#

function gpu_transpose_kernel2(input, output)
    ## Declare shared memory
    shared = @cuStaticSharedMem(eltype(input), (TILE_DIM, TILE_DIM))

    ## Modify thread index so threadIdx().x dominates the column
    block_index = ((blockIdx().y, blockIdx().x) .- 1) .* TILE_DIM

    for i in 1:4
        thread_index = (threadIdx().x, threadIdx().y + (i - 1)*8)
        index = CartesianIndex(block_index .+ thread_index)

        (index[1] > size(input, 1) || index[2] > size(input, 2)) && continue
        @inbounds shared[thread_index[2], thread_index[1]] = input[index]
    end

    ## Barrier to ensure all threads have completed writing to shared memory
    sync_threads()

    ## swap tile index
    block_index = ((blockIdx().x, blockIdx().y) .- 1) .* TILE_DIM

    for i in 1:4
        thread_index = (threadIdx().x, threadIdx().y + (i - 1)*8)
        index = CartesianIndex(block_index .+ thread_index)

        (index[1] > size(output, 1) || index[2] > size(output, 2)) && continue
        @inbounds output[index] = shared[thread_index...]
    end
    return
end

function gpu_transpose_shmem(input, output = similar(input, (size(input, 2), size(input, 1))))
    threads = (32, 8)
    blocks = cld.(size(input), (32, 32))
    @cuda blocks=blocks threads=threads gpu_transpose_kernel2(input, output)
    output
end

#-

X = CuArray(reshape(1f0:1089, (33, 33)))
Y = similar(X)
gpu_transpose_shmem(X, Y)

#-

@benchmark CUDA.@sync gpu_transpose_shmem($A, $B)

#-

@benchmark CUDA.@sync B .= A

#=
### Shared Memory Bank conflicts

Inside a SM, shared memory is divided into banks. Modern NVIDIA GPUs have 32 banks which have a 4-byte boundary. This means addresses 1-4 of shared memory are serviced by bank 1, addresses 5-8 are serviced by bank two and so on. When multiple threads access memory from the same bank then their requests are serialised.

Nsight compute gives statistics about shared memory usage. Running the profiler on `gpu_transpose_shmem` for an input of 33x33 of `Float32` we get:

![shared-memory-conflicts](../assets/shmem-conflicts.png)

It reports zero conflicts during shared loads because of we load columnwise. The 1023 store conflicts can be explained as follows. When an entire column is read it is stored to a row. Consecutive elements in a row differ in address by `column_length*sizeof(datatype)`. In 33 tile columns we write directly to a complete row where 32 elements are written hence there are 31 write conflicts (33*31 = 1023). `CUDA.jl` docs have a brief Nsight compute usage guide [here](https://juliagpu.github.io/CUDA.jl/dev/development/profiling/#NVIDIA-Nsight-Compute).

The fix is quite simple, pad the column length in shared memory by 1. Now consecutive elements in a row will differ by 33 % 32 = 1 hence no more bank conflicts.

i.e. `shared = @cuStaticSharedMem(eltype(input), (TILE_DIM + 1, TILE_DIM))`
=#

function gpu_transpose_kernel3(input, output)
    ## Declare shared memory
    shared = @cuStaticSharedMem(eltype(input), (TILE_DIM + 1, TILE_DIM))

    ## Modify thread index so threadIdx().x dominates the column
    block_index = ((blockIdx().y, blockIdx().x) .- 1) .* TILE_DIM

    for i in 1:4
        thread_index = (threadIdx().x, threadIdx().y + (i - 1)*8)
        index = CartesianIndex(block_index .+ thread_index)

        (index[1] > size(input, 1) || index[2] > size(input, 2)) && continue
        @inbounds shared[thread_index[2], thread_index[1]] = input[index]
    end

    ## Barrier to ensure all threads have completed writing to shared memory
    sync_threads()

    ## swap tile index
    block_index = ((blockIdx().x, blockIdx().y) .- 1) .* TILE_DIM

    for i in 1:4
        thread_index = (threadIdx().x, threadIdx().y + (i - 1)*8)
        index = CartesianIndex(block_index .+ thread_index)

        (index[1] > size(output, 1) || index[2] > size(output, 2)) && continue
        @inbounds output[index] = shared[thread_index...]
    end
    return
end

function gpu_transpose_noconf(input, output = similar(input, (size(input, 2), size(input, 1))))
    threads = (32, 8)
    blocks = cld.(size(input), (32, 32))
    @cuda blocks=blocks threads=threads gpu_transpose_kernel3(input, output)
    output
end

#-

@benchmark CUDA.@sync gpu_transpose_noconf($A, $B)

#-

@benchmark CUDA.@sync gpu_transpose_shmem($A, $B)

#=
An obvious improvement, we can also confirm with Nsight compute if there are no bank conflicts.

![shmem-noconf](../assets/shmem-noconf.png)
=#


text:
# Transit structures
abstract type TransitOutput{T} <: AbstractOutput{T} end

"""
    TransitTiming{T<:AbstractFloat} <: TransitOutput{T}

Transit times and derivatives.

# (User-facing) Fields
- `tt::Matrix{T}` : The transit times of each body.
- `dtdq0::Array{T,4}` : Derivatives of the transit times with respect to the initial Cartesian coordinates and masses.
- `dtdelements::Array{T,4}` : Derivatives of the transit times with respect to the initial orbital elements and masses.
"""
struct TransitTiming{T<:AbstractFloat} <: TransitOutput{T}
    tt::Matrix{T}
    dtdq0::Array{T,4}
    dtdelements::Array{T,4}

    # Internal
    count::Vector{Int64}
    ntt::Int64
    ti::Int64
    occs::Vector{Int64}
    dtdq::Array{T,3}
    gsave::Vector{T}
    s_prior::State{T}
    s_transit::State{T}
end

"""
    TransitTiming(tmax, ic; ti)

Constructor for [`TransitTiming`](@ref) type.

# Arguments
- `tmax::T` : Expected total elapsed integration time. (Allocates arrays accordingly)
- `ic::ElementsIC{T}` : Initial conditions for the system

## Optional
- `ti::Int64=1` : Index of the body with respect to which transits are measured. (Default is the central body)
"""
function TransitTiming(tmax::T,ic::ElementsIC{T},ti::Int64=1) where T<:AbstractFloat
    n = ic.nbody
    ind = isfinite.(tmax./ic.elements[:,2])
    ntt = maximum(ceil.(Int64,abs.(tmax./ic.elements[ind,2])).+3)
    tt = zeros(T,n,ntt)
    dtdq0 = zeros(T,n,ntt,7,n)
    dtdelements = zeros(T,n,ntt,7,n)
    count = zeros(Int64,n)
    occs = setdiff(collect(1:n),ti)
    dtdq = zeros(T,1,7,n)
    gsave = zeros(T,n)
    s_prior = State(ic)
    s_transit = State(ic)
    return TransitTiming(tt,dtdq0,dtdelements,count,ntt,ti,occs,dtdq,gsave,s_prior,s_transit)
end

"""
    TransitParameters{T<:AbstractFloat} <: TransitOutput{T}

Transit times, impact parameters, sky-velocities, and derivatives.

# (User-facing) Fields
- `ttbv::Matrix{T}` : The transit times, impact parameter, and sky-velocity of each body.
- `dtbvdq0::Array{T,5}` : Derivatives of the transit times, impact parameters, and sky-velocities with respect to the initial Cartesian coordinates and masses.
- `dtbvdelements::Array{T,5}` : Derivatives of the transit times, impact parameters, and sky-velocities with respect to the initial orbital elements and masses.
"""
struct TransitParameters{T<:AbstractFloat} <: TransitOutput{T}
    ttbv::Array{T,3}
    dtbvdq0::Array{T,5}
    dtbvdelements::Array{T,5}

    # Internal
    count::Vector{Int64}
    ntt::Int64
    ti::Int64
    occs::Vector{Int64}
    dtbvdq::Array{T,3}
    gsave::Vector{T}
    s_prior::State{T}
    s_transit::State{T}
end

"""
    TransitParameters(tmax, ic; ti)

Constructor for [`TransitParameters`](@ref) type.

# Arguments
- `tmax::T` : Expected total elapsed integration time. (Allocates arrays accordingly)
- `ic::ElementsIC{T}` : Initial conditions for the system

## Optional
- `ti::Int64=1` : Index of the body with respect to which transits are measured. (Default is the central body)
"""
function TransitParameters(tmax::T,ic::ElementsIC{T},ti::Int64=1) where T<:AbstractFloat
    n = ic.nbody
    ind = isfinite.(tmax./ic.elements[:,2])
    ntt = maximum(ceil.(Int64,abs.(tmax./ic.elements[ind,2])).+3)
    ttbv = zeros(T,3,n,ntt)
    dtbvdq0 = zeros(T,3,n,ntt,7,n)
    dtbvdelements = zeros(T,3,n,ntt,7,n)
    count = zeros(Int64,n)
    occs = setdiff(collect(1:n),ti)
    dtbvdq = zeros(T,3,7,n)
    gsave = zeros(T,n)
    s_prior = State(ic)
    s_transit = State(ic)
    return TransitParameters(ttbv,dtbvdq0,dtbvdelements,count,ntt,ti,occs,dtbvdq,gsave,s_prior,s_transit)
end

struct TransitSnapshot{T<:AbstractFloat} <: TransitOutput{T}
    nt::Int64
    times::Vector{T}
    bsky2::Matrix{T}
    vsky::Matrix{T}
    dbvdq0::Array{T,5}
    dbvdelements::Array{T,5}
end

function TransitSnapshot(times::Vector{T},ic::ElementsIC{T}) where T<:AbstractFloat
    n = ic.nbody
    nt = length(times)
    return TransitSeries(nt,times,zeros(T,n,nt),zeros(T,n,nt),zeros(T,2,n,nt,7,n),zeros(T,2,n,nt,7,n))
end

function zero_out!(tt::TransitOutput{T}) where T
    for i in 1:length(fieldnames(typeof(tt)))
        if typeof(getfield(tt,i)) <: Array{T}
            getfield(tt,i) .= zero(T)
        end
    end
    tt.count .= 0
end

"""

Main integrator method for Transit calculations.
"""
function (intr::Integrator)(s::State{T}, tt::TransitOutput{T}, d::Derivatives{T}; grad::Bool=true) where T<:AbstractFloat

    t0 = s.t[1] # Initial time
    nsteps = abs(round(Int64,intr.tmax/intr.h))
    h = intr.h * check_step(t0, intr.tmax+t0) # get direction of integration

    for i in tt.occs
        # Compute the relative sky velocity dotted with position:
        tt.gsave[i] = g!(i,tt.ti,s.x,s.v)
    end

    istep = 0
    for _ in 1:nsteps

        # Take an integration step
        if grad
            intr.scheme(s,d,h)
        else
            intr.scheme(s,h)
        end
        istep += 1
        s.t[1] = t0 + (istep * h)

        # Check if a transit occured; record time.
        detect_transits!(s,d,tt,intr,grad=grad)
    end
    # Calculate derivatives
    if grad
        calc_dtdelements!(s,tt)
    end
end

"""Wrapper so the user doesn't need to create a `Derivatives` type."""
function (intr::Integrator)(s::State{T}, tt::TransitOutput{T}; grad::Bool=true, return_arrays::Bool=false) where T<:AbstractFloat
    # Preallocate arrays
    d = Derivatives(T, s.n)

    intr(s, tt, d, grad=grad)
    if return_arrays; return d; end # Return preallocated arrays
    return
end

"""

Integrator method for outputting `TransitSnapshot`.
"""
function (intr::Integrator)(s::State{T},ts::TransitSnapshot{T};grad::Bool=true) where T<:AbstractFloat
    # Integrate to, and output b and v_sky for each body, for a list of times.
    # Should be sorted times.
    # NOTE: Need to fix so that if initial time is a 0, s.dqdt isn't 0s.
    if grad; dbvdq = zeros(T,2,7,s.n); end

    # Integrate to each time, using intr.h, and output b and vsky (only want primary transits for now)
    t0 = s.t[1]
    for (j,ti) in enumerate(ts.times)
        intr(s,ti;grad=grad)
        for i in 2:s.n
            if grad
                ts.vsky[i,j],ts.bsky2[i,j] = calc_dbvdq!(s,dbvdq,1,i)
                for ibv=1:2, k=1:7, p=1:s.n
                    ts.dbvdq0[ibv,i,j,k,p] = dbvdq[ibv,k,p]
                end
            else
                ts.vsky[i,j],ts.bsky2[i,j] = calc_bv(s,1,i)
            end
        end
        # Step to the next full time step, as measured from t0.
        #steps = round(Int64 ,(s.t[1] - t0) / intr.h, RoundUp)
        #intr(s,t0 + (steps * intr.h); grad=grad)
    end
    calc_dbvdelements!(s,ts)
    return
end

# Includes for source
files = ["timing.jl","snapshot.jl"]
include.(files)


text:
struct CanalConfig
    H::Real
    hd::Real
    Ld::Real
    x0::Real
    L::Real
    d::Real
end

CanalConfig(H::Real, hd::Real, Ld::Real, x0::Real, L::Real) = CanalConfig(H, hd, Ld, x0, L, 2*H)
CanalConfig(H::Real) = CanalConfig(H, 0.55*H, 2*H, 4*H, 16*H)

function δ(x::Real, c::CanalConfig)
    diff = x - c.x0
    l = c.Ld * 0.5
    if (-l <= diff) && (diff <= l)
        tr = cos(2pi * (x - c.x0) / c.Ld)
        c.hd * ( (1 + tr) * 0.5 )^2
    else
        0.
    end
end

δ(x::AbstractVector, c::CanalConfig) = δ.(x, Ref(c))

function inbump(c::CanalConfig, x, y)
    diff = abs(x - c.x0)
    l = c.Ld * 0.5
    if diff <= l
        tr = cos(2pi * (x - c.x0) / c.Ld)
        yc = c.hd * ( (1 + tr) * 0.5 )^2
        if c.H - abs(y) <= yc
            true
        else
            false
        end
    else
        false
    end
end


text:
using Emporos
using Test

@testset "Emporos.jl" begin
    # Write your tests here.
end


text:
st1 = vcat([10i .+ (0:9) for i in 0:9]...)
# a epoch of spike train
y,n,w,i = epochspiketrain(st1,10,20)
@test y==10:19
@test n==10
@test w==(10,20)
@test i==11:20
# two epochs of spike train
ys,ns,ws,is = epochspiketrain(st1,[10,20],[20,30])
@test ys[1]==y
@test ns[1]==n
@test ws[1]==w
@test is[1]==i
@test ys[2]==20:29
@test ns[2]==10
@test ws[2]==(20,30)
@test is[2]==21:30
ys1,ns1,ws1,is1 = epochspiketrain(st1,[10 20;20 30])
@test ys1==ys
@test ns1==ns
@test ws1==ws
@test is1==is
ys2,ns2,ws2,is2 = epochspiketrain(st1,[10,20,30])
@test ys2==ys
@test ns2==ns
@test ws2==ws
@test is2==is
# two epochs of two spike trains
st1st1 = [st1,st1]
yss,nss,wss,iss = epochspiketrains(st1st1,[10,20,30])
@test yss[1]==yss[2]
@test nss[1]==nss[2]
@test iss[1]==iss[2]
@test yss[1]==ys
@test nss[1]==ns
@test wss == ws
@test iss[1]==is
# two epoch responses of spike train
r = epochspiketrainresponse(st1,[10,20],[20,30],israte=false)
@test r[1]==r[2]==10
r1 = epochspiketrainresponse(st1,[10 20;20 30],israte=false)
@test r1[1]==r1[2]==10
# epochspiketrainresponse and epochspiketrainresponse_ono
spike = sort(rand(1:100000,1000))
on = collect(0:150:95000)
off = collect(100:150:95100)
n1 = epochspiketrainresponse(spike,on,off,israte=false)
n2 = epochspiketrainresponse_ono(spike,on,off,israte=false)
@test n1==n2
@btime epochspiketrainresponse($(spike),$(on),$(off),israte=false)
@btime epochspiketrainresponse_ono($(spike),$(on),$(off),israte=false)
# mean psth of two spike trains
m,se,x = psthspiketrains(st1st1,[10,20,30],israte=false)
@test m==[10,10]
@test se==[0,0]
@test x==[15,25]
m1,se1,x1 = psthspiketrains(st1st1,10:10:30,israte=false)
@test m1==m
@test se1==se
@test x1==x

## Poisson Model of Spike Generation, [David Heeger(2000), Poisson Model of Spike Generation](http://www.cns.nyu.edu/~david/handouts/poisson.pdf)
rate = 50
duration = 1000
refractoryperiod = 2
ntrain = 100

## Homogeneous Poisson Process, where instantaneous ﬁring rate is constant

# Inter-Spike-Interval Exponential Distribution Method
sts = [poissonspiketrain(rate,duration,rp=refractoryperiod) for _ in 1:ntrain]
plotspiketrain(sts,timeline=[0,duration])

# Instantaneous Firing Rate Method
sts = [poissonspiketrain(t->rate,duration,rp=refractoryperiod) for _ in 1:ntrain]
plotspiketrain(sts,timeline=[0,duration])

## Inhomogeneous Poisson Process, where instantaneous ﬁring rate is changing
sts = [poissonspiketrain(t->rate*(sin(0.05t)+1),duration,rp=refractoryperiod) for _ in 1:ntrain]
plotspiketrain(sts,timeline=[0,duration])