text:
"""
    abstract type ProductOfSets{T} end

Represents a cartesian product of sets of given types.
"""
abstract type ProductOfSets{T} end

set_index(::ProductOfSets, ::Type{S}) where {S<:MOI.AbstractSet} = nothing

function _sets_code(esc_name, T, type_def, set_types...)
    code = Expr(:block, type_def)
    esc_types = esc.(set_types)
    for (i, esc_type) in enumerate(esc_types)
        push!(
            code.args,
            :(
                function $MOIU.set_index(
                    ::$esc_name{$(T)},
                    ::Type{$(esc_type)},
                ) where {$T}
                    return $i
                end
            ),
        )
    end
    push!(
        code.args,
        :($MOIU.set_types(::$esc_name{$T}) where {$T} = [$(esc_types...)]),
    )
    return code
end

"""
    abstract type MixOfScalarSets{T} <: ProductOfSets{T} end

Product of scalar sets in the order the constraints are added, mixing the
constraints of different types.

Use [`@mix_of_scalar_sets`](@ref) to generate a new subtype.
"""
abstract type MixOfScalarSets{T} <: ProductOfSets{T} end

"""
    @mix_of_scalar_sets(name, set_types...)

Generate a new [`MixOfScalarSets`](@ref) subtype.

## Example

```julia
@mix_of_scalar_sets(
    MixedIntegerLinearProgramSets,
    MOI.GreaterThan{T},
    MOI.LessThan{T},
    MOI.EqualTo{T},
    MOI.Integer,
)
```
"""
macro mix_of_scalar_sets(name, set_types...)
    esc_name = esc(name)
    T = esc(:T)
    type_def = :(struct $(esc_name){$(T)} <: $(MOIU).MixOfScalarSets{$(T)}
        """
        `set_ids[i]` maps the row `i` to the corresponding set type.
        """
        set_ids::Vector{Int}

        $(esc_name){$(T)}() where {$(T)} = new(Int[])
    end)
    return _sets_code(esc_name, T, type_def, set_types...)
end

MOI.is_empty(sets::MixOfScalarSets) = isempty(sets.set_ids)

MOI.empty!(sets::MixOfScalarSets) = empty!(sets.set_ids)

MOI.dimension(sets::MixOfScalarSets) = length(sets.set_ids)

rows(::MixOfScalarSets, ci::MOI.ConstraintIndex) = ci.value

function add_set(sets::MixOfScalarSets, i)
    push!(sets.set_ids, i)
    return length(sets.set_ids)
end

final_touch(::MixOfScalarSets) = nothing

function MOI.get(
    sets::MixOfScalarSets{T},
    ::MOI.ListOfConstraintTypesPresent,
) where {T}
    present = Set(sets.set_ids)
    return Tuple{Type,Type}[
        (_affine_function_type(T, S), S) for
        S in set_types(sets) if set_index(sets, S) in present
    ]
end

function MOI.get(
    sets::MixOfScalarSets,
    ::MOI.NumberOfConstraints{F,S},
) where {F,S}
    i = set_index(sets, S)
    return count(isequal(i), sets.set_ids)
end

function MOI.get(
    sets::MixOfScalarSets,
    ::MOI.ListOfConstraintIndices{F,S},
) where {F,S}
    i = set_index(sets, S)
    return MOI.ConstraintIndex{F,S}[
        MOI.ConstraintIndex{F,S}(ci) for
        (ci, set_type) in enumerate(sets.set_ids) if set_type == i
    ]
end

function MOI.is_valid(
    sets::MixOfScalarSets,
    ci::MOI.ConstraintIndex{F,S},
) where {F,S}
    i = set_index(sets, S)
    if i === nothing
        return false
    end
    return ci.value in eachindex(sets.set_ids) && sets.set_ids[ci.value] == i
end

"""
    abstract type OrderedProductOfSets{T} <: ProductOfSets{T} end

Product of sets in the order the constraints are added, grouping the
constraints of the same types contiguously.

Use [`@product_of_sets`](@ref) to generate new subtypes.
"""
abstract type OrderedProductOfSets{T} <: ProductOfSets{T} end

"""
    @product_of_sets(name, set_types...)

Generate a new [`OrderedProductOfSets`](@ref) subtype.

## Example

```julia
@product_of_sets(
    LinearOrthants,
    MOI.Zeros,
    MOI.Nonnegatives,
    MOI.Nonpositives,
    MOI.ZeroOne,
)
```
"""
macro product_of_sets(name, set_types...)
    esc_name = esc(name)
    T = esc(:T)
    type_def = :(
        mutable struct $(esc_name){$(T)} <: $(MOIU).OrderedProductOfSets{$(T)}
            """
            During the copy, this counts the number of rows corresponding to
            each set. At the end of copy, `final_touch` is called, which
            converts this list into a cumulative ordering.
            """
            num_rows::Vector{Int}

            """
            A dictionary which maps the `set_index` and `offset` of a set to the
            dimension, i.e., `dimension[(set_index,offset)] → dim`.
            """
            dimension::Dict{Tuple{Int,Int},Int}

            """
            A sanity bit to check that we don't call functions out-of-order.
            """
            final_touch::Bool

            function $(esc_name){$(T)}() where {$(T)}
                return new(
                    zeros(Int, $(length(set_types))),
                    Dict{Tuple{Int,Int},Int}(),
                    false,
                )
            end
        end
    )
    return _sets_code(esc_name, T, type_def, set_types...)
end

MOI.is_empty(sets::OrderedProductOfSets) = all(iszero, sets.num_rows)

function MOI.empty!(sets::OrderedProductOfSets)
    fill!(sets.num_rows, 0)
    empty!(sets.dimension)
    sets.final_touch = false
    return
end

function MOI.dimension(sets::OrderedProductOfSets)
    @assert sets.final_touch
    return sets.num_rows[end]
end

function rows(
    sets::OrderedProductOfSets{T},
    ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{T},S},
) where {T,S}
    @assert sets.final_touch
    i = set_index(sets, S)
    return (i == 1 ? 0 : sets.num_rows[i-1]) + ci.value
end

function rows(
    sets::OrderedProductOfSets{T},
    ci::MOI.ConstraintIndex{MOI.VectorAffineFunction{T},S},
) where {T,S}
    @assert sets.final_touch
    i = set_index(sets, S)
    offset = i == 1 ? 0 : sets.num_rows[i-1]
    return (offset + ci.value - 1) .+ (1:sets.dimension[(i, ci.value)])
end

function add_set(sets::OrderedProductOfSets, i)
    @assert !sets.final_touch
    sets.num_rows[i] += 1
    return sets.num_rows[i]
end

function add_set(sets::OrderedProductOfSets, i, dim)
    @assert !sets.final_touch
    ci = sets.num_rows[i] + 1
    sets.dimension[(i, ci)] = dim
    sets.num_rows[i] += dim
    return ci
end

function final_touch(sets::OrderedProductOfSets)
    @assert !sets.final_touch
    for i in 2:length(sets.num_rows)
        sets.num_rows[i] += sets.num_rows[i-1]
    end
    sets.final_touch = true
    return
end

"""
    num_rows(sets::OrderedProductOfSets, ::Type{S}) where {S}

Return the number of rows corresponding to a set of type `S`. That is, it is
the sum of the dimensions of the sets of type `S`.
"""
function num_rows(sets::OrderedProductOfSets, ::Type{S}) where {S}
    i = set_index(sets, S)
    if !sets.final_touch || i == 1
        return sets.num_rows[i]
    end
    return sets.num_rows[i] - sets.num_rows[i-1]
end

function MOI.get(
    sets::OrderedProductOfSets{T},
    ::MOI.ListOfConstraintTypesPresent,
) where {T}
    return Tuple{Type,Type}[
        (_affine_function_type(T, S), S) for
        S in set_types(sets) if num_rows(sets, S) > 0
    ]
end

struct _UnevenIterator
    i::Int
    start::Int
    stop::Int
    dimension::Dict{Tuple{Int,Int},Int}
end

Base.IteratorSize(::_UnevenIterator) = Base.SizeUnknown()

function Base.iterate(it::_UnevenIterator, cur = it.start)
    if cur > it.stop
        return nothing
    end
    return (cur, cur + it.dimension[(it.i, cur)])
end

function Base.in(x::Int64, it::_UnevenIterator)
    return it.start <= x <= it.stop && haskey(it.dimension, (it.i, x))
end

function _range_iterator(
    ::OrderedProductOfSets{T},
    ::Int,
    start::Int,
    stop::Int,
    ::Type{MOI.ScalarAffineFunction{T}},
) where {T}
    return start:stop
end

function _range_iterator(
    sets::OrderedProductOfSets{T},
    i::Int,
    start::Int,
    stop::Int,
    ::Type{MOI.VectorAffineFunction{T}},
) where {T}
    return _UnevenIterator(i, start, stop, sets.dimension)
end

function _range_iterator(
    sets::OrderedProductOfSets{T},
    ::Type{F},
    ::Type{S},
) where {T,F,S}
    i = set_index(sets, S)
    if i === nothing || F != _affine_function_type(T, S)
        return
    end
    return _range_iterator(sets, i, 1, num_rows(sets, S), F)
end

_length(::Nothing) = 0
_length(r::UnitRange) = length(r)
_length(r::_UnevenIterator) = count(_ -> true, r)

function MOI.get(
    sets::OrderedProductOfSets,
    ::MOI.NumberOfConstraints{F,S},
) where {F,S}
    r = _range_iterator(sets, F, S)
    return _length(r)
end

function MOI.get(
    sets::OrderedProductOfSets,
    ::MOI.ListOfConstraintIndices{F,S},
) where {F,S}
    rows = _range_iterator(sets, F, S)
    if rows === nothing
        return MOI.ConstraintIndex{F,S}[]
    end
    return MOI.ConstraintIndex{F,S}.(rows)
end

function MOI.is_valid(
    sets::OrderedProductOfSets,
    ci::MOI.ConstraintIndex{F,S},
) where {F,S}
    r = _range_iterator(sets, F, S)
    return r !== nothing && ci.value in r
end


text:
"""
    classnames() -> Vector{String}

Return the 10 names for the CIFAR10 classes as a vector of strings.
"""
classnames() = CLASSES

"""
    traintensor([T = N0f8], [indices]; [dir]) -> Array{T}

Return the CIFAR-10 **training** images corresponding to the
given `indices` as a multi-dimensional array of eltype `T`. If
the corresponding labels are required as well, it is recommended
to use [`CIFAR10.traindata`](@ref) instead.

The image(s) is/are returned in the horizontal-major
memory layout as a single numeric array. If `T <: Integer`, then
all values will be within `0` and `255`, otherwise the values are
scaled to be between `0` and `1`.

If the parameter `indices` is omitted or an `AbstractVector`, the
images are returned as a 4D array (i.e. a `Array{T,4}`) in
WHCN format (width, height, #channels, #images).
For integer `indices` instead, a 3D array in WHC format is returned.

```julia-repl
julia> CIFAR10.traintensor() # load all training images
32×32×3×50000 Array{N0f8,4}:
[...]

julia> CIFAR10.traintensor(Float32, 1:3) # first three images as Float32
32×32×3×3 Array{Float32,4}:
[...]
```

If `indices` is an `Integer`, a single image is returned as
`Array{T,3}` array.

```julia-repl
julia> CIFAR10.traintensor(1) # load first training image
32×32×3 Array{N0f8,3}:
[...]
```

You can use the utility function
[`convert2image`](@ref) to convert an CIFAR-10 array into a
horizontal-major Julia image with the appropriate `RGB` eltype.

```julia-repl
julia> CIFAR10.convert2image(CIFAR10.traintensor(1)) # convert to column-major colorant array
32×32 Array{RGB{N0f8},2}:
[...]
```

$(download_docstring("CIFAR10", DEPNAME))
"""
function traintensor(args...; dir = nothing)
    traindata(args...; dir = dir)[1]
end

# needed for type inference for some reason
function traintensor(::Type{T}, args...; dir = nothing) where T
    traindata(T, args...; dir = dir)[1]
end

"""
    testtensor([T = N0f8], [indices]; [dir]) -> Array{T}

Return the CIFAR-10 **test** images corresponding to the given
`indices` as a multi-dimensional array of eltype `T`. If the
corresponding labels are required as well, it is recommended to
use [`CIFAR10.testdata`](@ref) instead.

Images are returned in horizontal-major
memory layout as a single numeric array. If `T <: Integer`, then
all values will be within `0` and `255`, otherwise the values are
scaled to be between `0` and `1`.

If the parameter `indices` is omitted or an `AbstractVector`, the
images are returned as a 4D array (i.e. a `Array{T,4}`) in
WHCN format (width, height, #channels, #images).
For integer `indices` instead, a 3D array in WHC format is returned.

```julia-repl
julia> CIFAR10.testtensor() # load all training images
32×32×3×10000 Array{N0f8,4}:
[...]

julia> CIFAR10.testtensor(Float32, 1:3) # first three images as Float32
32×32×3×3 Array{Float32,4}:
[...]
```

If `indices` is an `Integer`, a single image is returned as
`Array{T,3}`.

```julia-repl
julia> CIFAR10.testtensor(1) # load first training image
32×32×3 Array{N0f8,3}:
[...]
```

You can use the utility function
[`convert2image`](@ref) to convert an CIFAR-10 array into a
horizontal-major HW Julia image with the appropriate `RGB` eltype.

```julia-repl
julia> CIFAR10.convert2image(CIFAR10.testtensor(1)) # convert to column-major colorant array
32×32 Array{RGB{N0f8},2}:
[...]
```

$(download_docstring("CIFAR10", DEPNAME))
"""
function testtensor(args...; dir = nothing)
    testdata(args...; dir = dir)[1]
end

function testtensor(::Type{T}, args...; dir = nothing) where T
    testdata(T, args...; dir = dir)[1]
end

"""
    trainlabels([indices]; [dir])

Returns the CIFAR-10 **trainset** labels corresponding to the
given `indices` as an `Int` or `Vector{Int}`. The values of the
labels denote the zero-based class-index that they represent (see
[`CIFAR10.classnames`](@ref) for the corresponding names). If
`indices` is omitted, all labels are returned.

```julia-repl
julia> CIFAR10.trainlabels() # full training set
50000-element Array{Int64,1}:
 6
 9
 ⋮
 1
 1

julia> CIFAR10.trainlabels(1:3) # first three labels
3-element Array{Int64,1}:
 6
 9
 9

julia> CIFAR10.trainlabels(1) # first label
6

julia> CIFAR10.classnames()[CIFAR10.trainlabels(1) + 1] # corresponding name
"frog"
```

$(download_docstring("CIFAR10", DEPNAME))
"""
function trainlabels(args...; dir = nothing)
    traindata(args...; dir = dir)[2]
end

"""
    testlabels([indices]; [dir])

Returns the CIFAR-10 **testset** labels corresponding to the
given `indices` as an `Int` or `Vector{Int}`. The values of the
labels denote the zero-based class-index that they represent (see
[`CIFAR10.classnames`](@ref) for the corresponding names). If
`indices` is omitted, all labels are returned.

```julia-repl
julia> CIFAR10.testlabels() # full training set
10000-element Array{Int64,1}:
 3
 8
 ⋮
 1
 7

julia> CIFAR10.testlabels(1:3) # first three labels
3-element Array{Int64,1}:
 3
 8
 8

julia> CIFAR10.testlabels(1) # first label
3

julia> CIFAR10.classnames()[CIFAR10.testlabels(1) + 1] # corresponding name
"cat"
```

$(download_docstring("CIFAR10", DEPNAME))
"""
function testlabels(args...; dir = nothing)
    testdata(args...; dir = dir)[2]
end

"""
    traindata([T = N0f8], [indices]; [dir]) -> images, labels

Returns the CIFAR-10 **trainingset** corresponding to the given
`indices` as a two-element tuple. If `indices` is omitted the
full trainingset is returned. The first element of the return
values will be the images as a multi-dimensional array, and the
second element the corresponding labels as integers.

The image(s) is/are returned in horizontal-major
memory layout as a single numeric array of eltype `T`. If `T <:
Integer`, then all values will be within `0` and `255`, otherwise
the values are scaled to be between `0` and `1`. The integer
values of the labels correspond 1-to-1 the digit that they
represent.

```julia
train_x, train_y = CIFAR10.traindata() # full datatset
train_x, train_y = CIFAR10.traindata(2) # only second observation
train_x, train_y = CIFAR10.traindata(dir="./CIFAR10") # custom folder
```

$(download_docstring("CIFAR10", DEPNAME))

Take a look at [`CIFAR10.traintensor`](@ref) and
[`CIFAR10.trainlabels`](@ref) for more information.
"""
function traindata(args...; dir = nothing)
    traindata(N0f8, args...; dir = dir)
end

function traindata(::Type{T}; dir = nothing) where T
    # placeholders for the chunks
    Xs = Vector{Array{UInt8,4}}(undef, NCHUNKS)
    Ys = Vector{Vector{Int}}(undef, NCHUNKS)
    # loop over all 5 trainingset files (i.e. chunks)
    for file_index in 1:NCHUNKS
        file_name = filename_for_chunk(file_index)
        file_path = datafile(DEPNAME, file_name, dir)
        # load all the data from each file and append it to
        # the placeholders X and Y
        X, Y = Reader.readdata(file_path)
        Xs[file_index] = X
        Ys[file_index] = Y
    end
    # cat all the placeholders into one image array
    # and one label array. (good enough)
    images = cat(Xs..., dims=4)::Array{UInt8,4}
    labels = vcat(Ys...)::Vector{Int}
    # optionally transform the image array before returning
    bytes_to_type(T, images), labels
end

function traindata(::Type{T}, index::Integer; dir = nothing) where T
    @assert 1 <= index <= NCHUNKS * Reader.CHUNK_SIZE "parameter \"index\" ($index) not in 1:$(NCHUNKS*Reader.CHUNK_SIZE)"
    # we only need to return a single image, so the task breaks
    # down to computing which file that image is in.
    file_index = ceil(Int, index / Reader.CHUNK_SIZE)
    file_name = filename_for_chunk(file_index)
    file_path = datafile(DEPNAME, file_name, dir)
    # once we know the file we just need to compute the approriate
    # offset of the image realtive to that file.
    sub_index = ((index - 1) % Reader.CHUNK_SIZE) + 1
    image, label = Reader.readdata(file_path, sub_index)
    # optionally transform the image array before returning
    bytes_to_type(T, image), label
end

function traindata(::Type{T}, indices::AbstractVector; dir = nothing) where T
    mi, ma = extrema(indices)
    @assert mi >= 1 && ma <= NCHUNKS * Reader.CHUNK_SIZE "not all elements in parameter \"indices\" are in 1:$(NCHUNKS*Reader.CHUNK_SIZE)"
    # preallocate a buffer we will reuse for reading individual
    # images. "buffer" is written to length(indices) times
    buffer = Array{UInt8,3}(undef, Reader.NROW, Reader.NCOL, Reader.NCHAN)
    # we know the types and dimensions of the return values,
    # so we can preallocate them
    images = Array{UInt8,4}(undef, Reader.NROW, Reader.NCOL, Reader.NCHAN, length(indices))
    labels = Array{Int,1}(undef, length(indices))
    # loop over all 5 trainingset files (i.e. chunks)
    for file_index in 1:NCHUNKS
        file_name = filename_for_chunk(file_index)
        file_path = datafile(DEPNAME, file_name, dir)
        open(file_path, "r") do io
            # for each chunk we loop through indices once
            @inbounds for (i, index) in enumerate(indices)
                # check if "index" is part of the current chunk
                cur_file_index = ceil(Int, index / Reader.CHUNK_SIZE)
                cur_file_index == file_index || continue
                # if it is, then compute the relative offset
                sub_index = ((index - 1) % Reader.CHUNK_SIZE) + 1
                # read the corresponding image into "buffer"
                # and the corresponding label into "y"
                _, y = Reader.readdata!(buffer, io, sub_index)
                # write the image into the appropriate position
                # of our preallocated "images" array.
                copyto!(view(images,:,:,:,i), buffer)
                # same with labels
                labels[i] = y
            end
        end
    end
    # optionally transform the image array before returning
    bytes_to_type(T, images::Array{UInt8,4}), labels::Vector{Int}
end

"""
    testdata([T = N0f8], [indices]; [dir]) -> images, labels

Returns the CIFAR-10 **testset** corresponding to the given
`indices` as a two-element tuple. If `indices` is omitted the
full testset is returned. The first element of the return
values will be the images as a multi-dimensional array, and the
second element the corresponding labels as integers.

The image(s) is/are returned in the horizontal-major
memory layout as a single numeric array of eltype `T`. If `T <:
Integer`, then all values will be within `0` and `255`, otherwise
the values are scaled to be between `0` and `1`. The integer
values of the labels correspond 1-to-1 the digit that they
represent.

```julia
test_x, test_y = CIFAR10.testdata() # full datatset
test_x, test_y = CIFAR10.testdata(2) # only second observation
test_x, test_y = CIFAR10.testdata(dir="./CIFAR10") # custom folder
```

$(download_docstring("CIFAR10", DEPNAME))

Take a look at [`CIFAR10.testtensor`](@ref) and
[`CIFAR10.testlabels`](@ref) for more information.
"""
function testdata(args...; dir = nothing)
    testdata(N0f8, args...; dir = dir)
end

function testdata(::Type{T}; dir = nothing) where T
    file_path = datafile(DEPNAME, TESTSET_FILENAME, dir)
    # simply read the complete content of the testset file
    images, labels = Reader.readdata(file_path)
    # optionally transform the image array before returning
    bytes_to_type(T, images), labels
end

function testdata(::Type{T}, index::Integer; dir = nothing) where T
    @assert 1 <= index <= Reader.CHUNK_SIZE "parameter \"index\" ($index) not in 1:$(Reader.CHUNK_SIZE)"
    file_path = datafile(DEPNAME, TESTSET_FILENAME, dir)
    # read the single image+label corresponding to "index"
    image, label = Reader.readdata(file_path, index)
    # optionally transform the image array before returning
    bytes_to_type(T, image), label
end

function testdata(::Type{T}, indices::AbstractVector; dir = nothing) where T
    mi, ma = extrema(indices)
    @assert mi >= 1 && ma <= Reader.CHUNK_SIZE "not all elements in parameter \"indices\" are in 1:$(Reader.CHUNK_SIZE)"
    # preallocate a buffer we will reuse for reading individual
    # images. "buffer" is written to length(indices) times
    buffer = Array{UInt8,3}(undef, Reader.NROW, Reader.NCOL, Reader.NCHAN)
    # we know the types and dimensions of the return values,
    # so we can preallocate them
    images = Array{UInt8,4}(undef, Reader.NROW, Reader.NCOL, Reader.NCHAN, length(indices))
    labels = Array{Int,1}(undef, length(indices))
    # in contrast to the trainset, the testset only has one file
    file_path = datafile(DEPNAME, TESTSET_FILENAME, dir)
    open(file_path, "r") do io
        # iterate over the given "indices"
        @inbounds for (i, index) in enumerate(indices)
            # read the corresponding image into "buffer" and
            # the corresponding label into "y"
            _, y = Reader.readdata!(buffer, io, index)
            # write the image into the appropriate position
            # of our preallocated "images" array.
            copyto!(view(images,:,:,:,i), buffer)
            # same with labels
            labels[i] = y
        end
    end
    # optionally transform the image array before returning
    bytes_to_type(T, images::Array{UInt8,4}), labels::Vector{Int}
end

ext: jl


text:
module ThermalMonteCarlo

using Distributions: UnivariateDistribution, Dirac, Normal
using ComponentArrays: ComponentVector
using AdvancedMH: AdvancedMH
using AdvancedHMC: AdvancedHMC
using Random: Random
using UnPack: @unpack

using NQCDynamics:
    NQCDynamics,
    AbstractSimulation,
    Simulation,
    RingPolymerSimulation,
    get_temperature,
    DynamicsUtils,
    RingPolymers,
    Calculators,
    Estimators,
    DynamicsMethods,
    ndofs,
    nbeads,
    natoms,
    masses

"""
    run_advancedhmc_sampling(sim, r, steps, σ; move_ratio=0.0, internal_ratio=0.0)

Sample the configuration space for the simulation `sim` starting from `r`.

Total number of steps is given by `steps` and `σ` is the dictionary of
step sizes for each species.

`move_ratio` defaults to `0.0` and denotes the fraction of system moved each step.
If `move_ratio = 0`, every degree of freedom is moved at each step.
If `move_ratio = 1`, then nothing will happen. Experiment with this parameter to achieve
optimal sampling.

`internal_ratio` works as for `move_ratio` but for the internal modes of the ring polymer.
"""
function run_advancedmh_sampling(
    sim::AbstractSimulation,
    r,
    steps::Real,
    σ::Dict{Symbol,<:Real};
    move_ratio=0.0,
    internal_ratio=0.0
    )

    density = get_density_function(sim)

    density_model = AdvancedMH.DensityModel(density)
    proposal = get_proposal(sim, σ, move_ratio, internal_ratio)

    sampler = AdvancedMH.MetropolisHastings(proposal)

    initial_config = reshape_input(sim, copy(r))

    chain = AdvancedMH.sample(density_model, sampler, convert(Int, steps);
                              init_params=initial_config)

    return reshape_output(sim, chain)
end

function get_density_function(sim::Simulation)
    temperature = get_temperature(sim)
    function density(r)
        -DynamicsUtils.classical_potential_energy(sim, reshape(r, size(sim))) / temperature
    end
end

function get_deriv_function(sim::Simulation)
    temperature = get_temperature(sim)
    function density_deriv(r)
        lπ = -DynamicsUtils.classical_potential_energy(sim, reshape(r, size(sim))) / temperature
        Calculators.evaluate_derivative!(sim.calculator, reshape(r, size(sim)))
        lπ_grad = -sim.calculator.derivative / temperature
        return (lπ, lπ_grad[:])
    end
end

function get_density_function(sim::RingPolymerSimulation)
    temperature = get_temperature(sim)
    function density(r)
        r_local = reshape(copy(r), size(sim))
        RingPolymers.transform_from_normal_modes!(sim.beads, r_local)
        -DynamicsUtils.classical_potential_energy(sim, r_local) / temperature
    end
end

get_proposal(sim::Simulation, σ, move_ratio, internal_ratio) = ClassicalProposal(sim, σ, move_ratio)
get_proposal(sim::RingPolymerSimulation, σ, move_ratio, internal_ratio) = RingPolymerProposal(sim, σ, move_ratio, internal_ratio)

reshape_input(::Simulation, r) = r[:]
function reshape_input(sim::RingPolymerSimulation, r)
    RingPolymers.transform_to_normal_modes!(sim.beads, r)
    return r[:]
end

function reshape_output(sim::Simulation, chain)
    [reshape(config.params, size(sim)) for config in chain]
end
function reshape_output(sim::RingPolymerSimulation, chain)
    rs = [reshape(copy(config.params), size(sim)) for config in chain]
    RingPolymers.transform_from_normal_modes!.(sim.beads, rs)
    return rs
end

struct ClassicalProposal{P,T,S<:Simulation} <: AdvancedMH.Proposal{P}
    proposal::P
    move_ratio::T
    sim::S
end

struct RingPolymerProposal{P,T,S<:RingPolymerSimulation} <: AdvancedMH.Proposal{P}
    proposal::P
    move_ratio::T
    internal_ratio::T
    sim::S
end

const MolecularProposal{P,T,S} = Union{RingPolymerProposal{P,T,S}, ClassicalProposal{P,T,S}}

function ClassicalProposal(sim::Simulation, σ, move_ratio)
    proposals = Matrix{UnivariateDistribution}(undef, size(sim))
    for (i, symbol) in enumerate(sim.atoms.types) # Position proposals
        distribution = σ[symbol] == 0 ? Dirac(0) : Normal(0, σ[symbol])
        proposals[:,i] .= distribution
    end
    ClassicalProposal(proposals[:], move_ratio, sim)
end

function RingPolymerProposal(sim::RingPolymerSimulation, σ, move_ratio, internal_ratio)
    ωₖ = RingPolymers.get_matsubara_frequencies(nbeads(sim), sim.beads.ω_n)
    proposals = Array{UnivariateDistribution}(undef, size(sim))
    for i=1:nbeads(sim)
        if i == 1
            for (j, symbol) in enumerate(sim.atoms.types)
                distribution = σ[symbol] == 0 ? Dirac(0.0) : Normal(0, σ[symbol])
                proposals[:,j,i] .= distribution
            end
        else
            for j in range(sim.atoms)
                if j in sim.beads.quantum_atoms
                    Δ = sqrt(get_temperature(sim) / (sim.atoms.masses[j] * ωₖ[i]^2))
                    distribution = Normal(0, Δ)
                else
                    distribution = Dirac(0)
                end
                proposals[:,j,i] .= distribution
            end
        end
    end
    RingPolymerProposal(proposals[:], move_ratio, internal_ratio, sim)
end

function Base.rand(rng::Random.AbstractRNG, p::ClassicalProposal)
    result = map(x -> rand(rng, x), p.proposal)
    result[Random.randsubseq(eachindex(result), p.move_ratio)] .= 0
    return result
end

function Base.rand(rng::Random.AbstractRNG, p::RingPolymerProposal)
    result = map(x -> rand(rng, x), p.proposal)
    reshaped_result = reshape(result, size(p.sim))

    # Zero some of the centroid moves
    sequence = Random.randsubseq(CartesianIndices(size(p.sim)[1:2]), p.move_ratio)
    reshaped_result[sequence,1] .= 0

    # Zero some of the internal mode moves
    for i=2:nbeads(p.sim)
        sequence = Random.randsubseq(CartesianIndices((ndofs(p.sim), natoms(p.sim))), p.internal_ratio)
        reshaped_result[sequence,i] .= 0
    end

    return result
end

function AdvancedMH.propose(rng::Random.AbstractRNG, proposal::MolecularProposal,
                            ::AdvancedMH.DensityModel)
    return rand(rng, proposal)
end

function AdvancedMH.propose(rng::Random.AbstractRNG, proposal::MolecularProposal,
                            ::AdvancedMH.DensityModel, t)
    return t + rand(rng, proposal)
end

AdvancedMH.logratio_proposal_density(::MolecularProposal, state, candidiate) = 0

"""
    run_advancedhmc_sampling(sim::Simulation, r, n_samples;
        target_acceptance=0.5, kwargs...)

Perform Hamiltonian Monte Carlo sampling for the simulation `sim` using `AdvancedHMC.jl`.
"""
function run_advancedhmc_sampling(sim::Simulation, r, n_samples; target_acceptance=0.5, kwargs...)

    density = get_density_function(sim)
    deriv = get_deriv_function(sim)

    r0 = reshape_input(sim, copy(r))

    # Define a Hamiltonian system
    metric = AdvancedHMC.DiagEuclideanMetric(length(r0))
    hamiltonian = AdvancedHMC.Hamiltonian(metric, density, deriv)

    # Define a leapfrog solver, with initial step size chosen heuristically
    initial_ϵ = AdvancedHMC.find_good_stepsize(hamiltonian, r0)
    integrator = AdvancedHMC.Leapfrog(initial_ϵ)

    # Define an HMC sampler, with the following components
    #   - multinomial sampling scheme,
    #   - generalised No-U-Turn criteria, and
    #   - windowed adaption for step-size and diagonal mass matrix
    proposal = AdvancedHMC.NUTS{AdvancedHMC.MultinomialTS, AdvancedHMC.GeneralisedNoUTurn}(integrator)
    adaptor = AdvancedHMC.StanHMCAdaptor(
        AdvancedHMC.MassMatrixAdaptor(metric),
        AdvancedHMC.StepSizeAdaptor(target_acceptance, integrator)
        )

    # Run the sampler to draw samples from the specified Gaussian, where
    #   - `samples` will store the samples
    #   - `stats` will store diagnostic statistics for each sample
    samples, stats = AdvancedHMC.sample(hamiltonian, proposal, r0, convert(Int, n_samples), adaptor; kwargs...)

    return reshape.(samples, Ref(size(sim))), stats
end

end # module


text:
module Transforms
using ..Util, ..WT
using Compat
export  dwt, idwt, dwt!, idwt!,
        wpt, iwpt, wpt!, iwpt!

# TODO Use StridedArray instead of AbstractArray where writing to array.
typealias DWTArray AbstractArray
typealias WPTArray AbstractVector
typealias ValueType Union{AbstractFloat, Complex}

# DWT

"""
    dwt(x, wt[, L=maxtransformlevels(x)])

Perform a discrete wavelet transform of the array `x`.
The wavelet type `wt` determines the transform type
(filter or lifting) and the wavelet class, see `wavelet`.

The number of transformation levels `L` can be any non-negative
integer such that the size of `x` is divisible by `L`.
Performs the identity transformation if `L==0`.

# Examples
```julia
dwt(x, wavelet(WT.coif6))
```

**See also:** `idwt`, `dwt!`, `wavelet`
"""
function dwt end

"""
    idwt(x, wt[, L=maxtransformlevels(x)])

Perform an inverse discrete wavelet transform of the array `x`,
the inverse of `dwt(x, wt, L)`.

**See also:** `dwt`, `idwt!`
"""
function idwt end

"""
    dwt!(y, x, wt::OrthoFilter[, L=maxtransformlevels(x)])

    dwt!(y, wt::GLS[, L=maxtransformlevels(x)])

Same as `dwt` but without array allocation.
Perform "out of place" transform with a filter, or
a inplace transform with a lifting scheme. The difference
between the filter and lifting methods is due to the
structure of the transform algorithms.

**See also:** `idwt!`
"""
function dwt! end

"""
    idwt!(y, x, wt::OrthoFilter[, L=maxtransformlevels(x)])

    idwt!(y, wt::GLS[, L=maxtransformlevels(x)])

The inverse of `dwt!`.

**See also:** `dwt!`
"""
function idwt! end

# WPT

"""
    wpt

Perform a discrete wavelet packet transform of the array `x`.
**See also:** `dwt`, `wavelet`
"""
function wpt end

"""
    iwpt

Perform an inverse discrete wavelet packet transform of the array `x`.
**See also:** `idwt`, `wavelet`
"""
function iwpt end

"""
    wpt!

Same as `wpt` but without array allocation.
"""
function wpt! end

"""
    iwpt!

Same as `iwpt` but without array allocation.
"""
function iwpt! end

# DWT (discrete wavelet transform)

for (Xwt, Xwt!, _Xwt!, fw) in ((:dwt, :dwt!, :_dwt!, true),
                                (:idwt, :idwt!, :_dwt!, false))
@eval begin
    # filter
    function ($Xwt){T<:ValueType}(x::DWTArray{T},
                                    filter::OrthoFilter,
                                    L::Integer=maxtransformlevels(x))
        y = Array(T, size(x))
        return ($_Xwt!)(y, x, filter, L, $fw)
    end
    function ($Xwt!){T<:ValueType}(y::DWTArray{T}, x::DWTArray{T},
                                    filter::OrthoFilter,
                                    L::Integer=maxtransformlevels(x))
        return ($_Xwt!)(y, x, filter, L, $fw)
    end
    # lifting
    function ($Xwt){T<:ValueType}(x::DWTArray{T},
                                    scheme::GLS,
                                    L::Integer=maxtransformlevels(x))
        y = Array(T, size(x))
        copy!(y, x)
        return ($_Xwt!)(y, scheme, L, $fw)
    end
    function ($Xwt!){T<:ValueType}(y::DWTArray{T},
                                    scheme::GLS,
                                    L::Integer=maxtransformlevels(y))
        return ($_Xwt!)(y, scheme, L, $fw)
    end
end # begin
end # for

# WPT (wavelet packet transform)

for (Xwt, Xwt!, _Xwt!, fw) in ((:wpt, :wpt!, :_wpt!, true),
                                (:iwpt, :iwpt!, :_wpt!, false))
@eval begin
    function ($Xwt){T<:ValueType}(x::WPTArray{T},
                                    wt::DiscreteWavelet,
                                    L::Integer=maxtransformlevels(x))
        return ($Xwt)(x, wt, maketree(length(x), L, :full))
    end
    # filter
    function ($Xwt){T<:ValueType}(x::WPTArray{T},
                                    filter::OrthoFilter,
                                    tree::BitVector=maketree(x, :full))
        y = Array(T, size(x))
        return ($_Xwt!)(y, x, filter, tree, $fw)
    end
    function ($Xwt!){T<:ValueType}(y::WPTArray{T}, x::WPTArray{T},
                                    filter::OrthoFilter)
        return ($Xwt!)(y, x, filter, maketree(x, :full))
    end
    function ($Xwt!){T<:ValueType}(y::WPTArray{T}, x::WPTArray{T},
                                    filter::OrthoFilter,
                                    L::Integer)
        return ($Xwt!)(y, x, filter, maketree(length(x), L, :full))
    end
    function ($Xwt!){T<:ValueType}(y::WPTArray{T}, x::WPTArray{T},
                                    filter::OrthoFilter,
                                    tree::BitVector)
        return ($_Xwt!)(y, x, filter, tree, $fw)
    end
    # lifting
    function ($Xwt){T<:ValueType}(x::WPTArray{T},
                                    scheme::GLS,
                                    tree::BitVector=maketree(x, :full))
        y = Array(T, size(x))
        copy!(y, x)
        return ($_Xwt!)(y, scheme, tree, $fw)
    end
    function ($Xwt!){T<:ValueType}(y::WPTArray{T},
                                    scheme::GLS)
        return ($Xwt!)(y, scheme, maketree(y, :full))
    end
    function ($Xwt!){T<:ValueType}(y::WPTArray{T},
                                    scheme::GLS,
                                    L::Integer)
        return ($Xwt!)(y, scheme, maketree(length(x), L, :full))
    end
    function ($Xwt!){T<:ValueType}(y::WPTArray{T},
                                    scheme::GLS,
                                    tree::BitVector)
        return ($_Xwt!)(y, scheme, tree, $fw)
    end
end # begin
end # for

# DWTC (column-wise discrete wavelet transform)
#dwtc(::AbstractArray, ::DiscreteWavelet)
#idwtc(::AbstractArray, ::DiscreteWavelet)

# SWT (stationary wavelet transform)
#swt(::AbstractVector, ::DiscreteWavelet)
#iswt(::AbstractVector, ::DiscreteWavelet)

# CWT (continuous wavelet transform)
#cwt(::AbstractVector, ::ContinuousWavelet)
#icwt(::AbstractVector, ::ContinuousWavelet)

# CWTFT (continuous wavelet transform via FFT)
#cwtft(::AbstractVector, ::ContinuousWavelet)
#icwtft(::AbstractVector, ::ContinuousWavelet)

# Int -> Float
for Xwt in (:dwt, :idwt, :dwtc, :idwtc, :wpt, :iwpt)
    @eval $Xwt{T<:Integer}(x::AbstractArray{T}, args...) = $Xwt(float(x), args...)
end

# non-exported "out of place" functions
for (Xwt_oop!, Xwt!) in ((:dwt_oop!, :dwt!), (:idwt_oop!, :idwt!))
@eval begin
    # filter
    function ($Xwt_oop!){T<:ValueType}(y::DWTArray{T}, x::DWTArray{T},
                                    filter::OrthoFilter,
                                    L::Integer=maxtransformlevels(x))
        return ($Xwt!)(y, x, filter, L)
    end
    # lifting
    function ($Xwt_oop!){T<:ValueType}(y::DWTArray{T}, x::DWTArray{T},
                                    scheme::GLS,
                                    L::Integer=maxtransformlevels(x))
        copy!(y, x)
        return ($Xwt!)(y, scheme, L)
    end
end # begin
end # for

# Array with shared memory
function unsafe_vectorslice{T}(A::Array{T}, i::Int, n::Int)#::Vector{T}
    return unsafe_wrap(Array, pointer(A, i), n, false)::Vector{T}
end

# 2-D
row_idx(i, n) = i
col_idx(i, n) = 1 + (i-1)*n
# 3-D
row_idx(i, j, n) = row_idx(i, n) + (j-1)*n*n
col_idx(i, j, n) = col_idx(i, n) + (j-1)*n*n
hei_idx(i, j, n) = i + (j-1)*n

# filter transforms
include("transforms_filter.jl")

# lifting transforms
include("transforms_lifting.jl")

end # module


text:
using MyGraph

G = buildconnectedgraph(collect(1:6), (a,b) -> first(abs.(rand(Int8, 1))))


text:
function main(r)
    go_to_right_top_corner(r)
    global sum = 0
    Snake(r)
    print(sum)
end

function go_to_right_top_corner(r)
    while !(isborder(r,West) && isborder(r,Sud))
        if (!isborder(r,Sud))
            move!(r,Sud)
        end
        if (!isborder(r,West))
            move!(r,West)
        end
    end
end

function moves(::Robot,side::HorizonSide)::Bool
    A=[]
    global sum
    kof = false
    if !isborder(r,side)
        move!(r,side)
    else
        while (isborder(r,side) && !isborder(r,Ost))
            move!(r,Ost); push!(A,Ost)
        end
        if isborder(r,side) && isborder(r,Ost)
            kof=true
        else
            move!(r,side)
            check_line=false
            while (isborder(r,West))
                move!(r,side)
                check_line=true
            end
        end
        if !kof
            if (!check_line)
                sum+=1
            else
                sum+=2
            end
        end
        while (length(A)>0)
            side1=pop!(A); move!(r,inverse(side1))
        end
    end
    return kof
end

function Snake(r)
    side=Nord
    while !isborder(r,Ost)
        kof=false
        while (!kof)
            kof=moves(r,side)
        end
        side=inverse(side)
        move!(r,Ost)
    end
end

inverse(side::HorizonSide) = HorizonSide(mod(Int(side)+2,4))

max_stars_repo_name: VladislavBukin/mirea-progs

max_issues_repo_name: VladislavBukin/mirea-progs

max_forks_repo_name: VladislavBukin/mirea-progs


text:
# LSP indexing solution
# https://github.com/julia-vscode/julia-vscode/issues/800#issuecomment-650085983
if isdefined(@__MODULE__, :LanguageServer)
    include("../src/SpinDoctor.jl")
    using .SpinDoctor
end

using SpinDoctor
using LinearAlgebra
using GLMakie
using OrdinaryDiffEq: Rodas4

## Chose a plotting theme
set_theme!(theme_light())
set_theme!(theme_dark())
set_theme!(theme_black())

## Create model from setup recipe
# include("setups/axon.jl")
# include("setups/sphere.jl")
# include("setups/plates.jl")
# include("setups/cylinders.jl")
# include("setups/spheres.jl")
include("setups/neuron.jl")

mesh, = @time create_geometry(setup; recreate = true);
model = Model(; mesh, coeffs...);
volumes = get_cmpt_volumes(model.mesh)
D_avg = 1 / 3 * tr.(model.D)' * volumes / sum(volumes)
ncompartment = length(model.mesh.points)
@info "Number of nodes per compartment:" length.(model.mesh.points)

## Plot mesh
plot_mesh(model.mesh)

## Assemble finite element matrices
matrices = @time assemble_matrices(model);

## Magnetic field gradient
TE = 5000
φ = -π / 6
R = [cos(φ) sin(φ) 0; -sin(φ) cos(φ) 0; 0 0 1]
g⃗(t) = 1.0 * R * [sin(2π * t / TE), sin(20π * t / TE) / 5, cos(2π * t / TE)]
gradient = GeneralGradient{T,typeof(g⃗)}(; g⃗, TE)

## Solve BTPDE

# Callbacks for time stepping (plot solution, save time series)
printer = Printer(; nupdate = 1, verbosity = 2)
writer = VTKWriter(; nupdate = 5)
plotter = Plotter{T}(; nupdate = 5)
callbacks = [printer, plotter]
# callbacks = [printer, plotter, writer]

# General BTPDE for all gradients
btpde = GeneralBTPDE(;
    model,
    matrices,
    reltol = 1e-4,
    abstol = 1e-6,
    odesolver = Rodas4(autodiff = false),
)

# Solve BTPDE
ξ = @time solve(btpde, gradient; callbacks)

## Plot magnetization
plot_field(model.mesh, ξ)

## Compute signal
compute_signal(matrices.M, ξ)
compute_signal.(matrices.M_cmpts, split_field(model.mesh, ξ))

## Matrix Formalism

# Perform Laplace eigendecomposition
laplace = Laplace{T}(; model, matrices, neig_max = 400)
lap_eig = @time solve(laplace)
length_scales = eig2length.(lap_eig.values, D_avg)

# Truncate basis at minimum length scale
length_scale = 3
λ_max = length2eig(length_scale, D_avg)
lap_eig = limit_lengthscale(lap_eig, λ_max)

# Compute magnetization using the matrix formalism reduced order model
mf = MatrixFormalism(; model, matrices, lap_eig, ninterval = 500)
ξ = @time solve(mf, gradient)


text:
#=

  Global constraint among in ConstraintSolver.jl

  Among: Requires exactly 'n' variables in 'x' to take one of the values in 'v'.

  From Global Constraint Catalog:
  http://www.emn.fr/x-info/sdemasse/gccat/Camong.html
  """
  Constraint
    among(NVAR,VARIABLES,VALUES)
  ...
  Purpose
    NVAR is the number of variables of the collection VARIABLES that
    take their value in VALUES.

  Example:
  (3, <4, 5, 5, 4, 1>, <1,5,8>)

  The among constraint holds since exactly 3 values of the collection
  of values
  <4, 5, 5, 4, 1> belong to the set of values {1, 5, 8}.
  """

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function among_test(len=5,d=8,nnn=nothing,xxx=nothing,vvv=nothing,print_solutions=true,all_solutions=true,timeout=6)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>timeout,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))
    println("len=$len,d=$d,nnn=$nnn,xxx=$xxx,vvv=$vvv,print_solutions=$print_solutions,all_solutions=$all_solutions,timeout=$timeout")
    @variable(model, 1 <= x[1:len] <= d, Int)
    n = @variable(model, integer=true, lower_bound=0,upper_bound=len)

    @variable(model, 1 <= v[1:3] <= d, Int)

    if nnn !== nothing
        @constraint(model, n .== nnn)
    end

    if vvv !== nothing
        @constraint(model, v .== vvv )
    end
    increasing_strict(model, v)

    among(model,n,x,v)

    # Show all x for which n=3
    @constraint(model, n==3)

    if xxx !== nothing
        @constraint(model, x .== xxx)
    end

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    num_sols = 0
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        # println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                v_val = convert.(Integer,JuMP.value.(v; result=sol))
                n_val = convert.(Integer,JuMP.value.(n; result=sol))
                println("x:$x_val   n:$n_val  v:$v_val")

            end
        end
        println("num_sols:$num_sols\n")
    else
        println("status:$status")
    end

    return status, num_sols
end

#
# Some test cases
#
tests = Dict(
    :0 => Dict(
        :len => 5,
        :d   => 8,
        :nnn => 3,
        :xxx => nothing,
        :vvv => [1,3,8],
        :print_solutions => true,
        :all_solutions => false,
    ),
   :1 => Dict(
       :len => 5,
       :d   => 8,
       :nnn => 3,
       :xxx => nothing,
       :vvv => [1,3,8],
       :print_solutions => false,
       :all_solutions => true,
   ),
   :2 => Dict(
       :len => 5,
       :d   => 8,
       :nnn => 3,
       :xxx => [3,1,1,1,2],
       :vvv => nothing,
       :print_solutions => true,
       :all_solutions => true,
   ),
)

for p in sort(collect(keys(tests)))
    println("p:$p")
    t = tests[p]
    among_test(t[:len],t[:d],t[:nnn],t[:xxx],t[:vvv],t[:print_solutions],t[:all_solutions])
end

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
### A Pluto.jl notebook ###
# v0.15.0

using Markdown
using InteractiveUtils

# ╔═╡ c581d17a-c965-11eb-1607-bbeb44933d25
# This file imports an outdated version of PlutoPkgTestA: 0.2.1 (which is stored in the embedded Manifest file) and Artifacts, which is now a standard library (as of Julia 1.6), but it used to be a registered package (https://github.com/JuliaPackaging/Artifacts.jl). This notebook was generated on Julia 1.5, so the Manifest will be very very confusing for Julia 1.6 and up.

# It is generated on Julia 1.5 (our oldest supported Julia version, Manifest.toml is not backwards-compatible):

# 1. add our test registry:
#   pkg> registry add https://github.com/JuliaPluto/PlutoPkgTestRegistry

# 2. using Pluto, open the simple_import.jl notebook

# 3. add the `import Artifacts` cell

import PlutoPkgTestA

# ╔═╡ aef57966-ea36-478f-8724-e71430f10be9
PlutoPkgTestA.MY_VERSION |> Text

# ╔═╡ f9bdbb35-4326-4786-b308-88b6894923df
import Artifacts

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
PlutoPkgTestA = "419c6f8d-b8cd-4309-abdc-cee491252f94"

[compat]
Artifacts = "~1.3.0"
PlutoPkgTestA = "~0.2.2"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[Artifacts]]
deps = ["Pkg"]
git-tree-sha1 = "c30985d8821e0cd73870b17b0ed0ce6dc44cb744"
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
version = "1.3.0"

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[LibGit2]]
deps = ["Printf"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[Pkg]]
deps = ["Dates", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "UUIDs"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[PlutoPkgTestA]]
deps = ["Pkg"]
git-tree-sha1 = "6c9aa67135641123c559d59ba88e8cb93841773a"
uuid = "419c6f8d-b8cd-4309-abdc-cee491252f94"
version = "0.2.2"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
"""

# ╔═╡ Cell order:
# ╠═c581d17a-c965-11eb-1607-bbeb44933d25
# ╠═aef57966-ea36-478f-8724-e71430f10be9
# ╠═f9bdbb35-4326-4786-b308-88b6894923df
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


text:
# Pkg.build("QuantEcon")
import Dolo
path = Dolo.pkg_path
using AxisArrays

###############################################################################
## Markov Chain
###############################################################################
filename = joinpath(path,"examples","models","rbc_mc.yaml")
# model = Dolo.Model(joinpath(Dolo.pkg_path, "examples", "models", "rbc_mc.yaml"), print_code=true)
model = Dolo.yaml_import(filename)

@time dr_ITI  = Dolo.improved_time_iteration(model; verbose=true, tol = 1e-06, smaxit=100)

dprocess = Dolo.discretize( model.exogenous )
# init_dr = Dolo.ConstantDecisionRule(model.calibration[:controls])
# @time dr_ITI_2  = Bruteforce_module.improved_time_iteration(model, dprocess,init_dr)

@time dr_TI  = Dolo.time_iteration(model;tol_η=1e-08, maxit=1000)

################################################################################
# function profile_ITI(m)
#     dr_ITI  = Bruteforce_module.improved_time_iteration(m)
#     return dr_ITI
# end
#
#
# using ProfileView
#
# profile_ITI(model)
# Profile.clear()
# @profile profile_ITI(model)
# ProfileView.view()
#

###############################################################################
#Plotting the DR

# ITIT method
df_ITI = Dolo.tabulate(model, dr_ITI.dr, :k)

import PyPlot
plt = PyPlot;
fig = plt.figure("Decision Rule, ITI-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df_ITI[:k], df_ITI[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df_ITI[:k], df_ITI[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");
###############################################
# TI method
df = Dolo.tabulate(model, dr_TI.dr, :k)

fig = plt.figure("Decision Rule, IT-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df[:k], df[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df[:k], df[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");

################################################################################
# RBC with an AR process
################################################################################

filename = joinpath(path,"examples","models","rbc.yaml")
model_ar1 = Dolo.yaml_import(filename)
@time dr_ITI_ar1  = Dolo.improved_time_iteration(model_ar1; verbose=true, tol = 1e-06, smaxit=50)
@time dr_TI_ar1  = Dolo.time_iteration(model_ar1; tol_η=1e-08, maxit=1000)

dprocess = Dolo.discretize( model_ar1.exogenous , [3], [2])
@time dr2_ITI_ar1  = Dolo.improved_time_iteration(model_ar1, dprocess; verbose=true, tol = 1e-06, smaxit=50)
@time dr2_TI_ar1  = Dolo.time_iteration(model_ar1, dprocess; tol_η=1e-08, maxit=1000)

################################################################################
df_ITI = Dolo.tabulate(model_ar1, dr2_ITI_ar1.dr, :k)

import PyPlot
plt = PyPlot;
fig = plt.figure("Decision Rule, ITI-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df_ITI[:k], df_ITI[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df_ITI[:k], df_ITI[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");

df = Dolo.tabulate(model_ar1, dr2_TI_ar1.dr, :k)

fig = plt.figure("Decision Rule, IT-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df[:k], df[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df[:k], df[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");

################################################################################
# RBC with an IID process
################################################################################

filename = joinpath(path,"examples","models","rbc_iid.yaml")
model_iid = Dolo.yaml_import(filename)
@time dr_ITI_iid  = Bruteforce_module.improved_time_iteration(model_iid; verbose=true, tol = 1e-06, smaxit=50)
@time dr_TI_iid  = Dolo.time_iteration(model_iid; tol_η=1e-08, maxit=1000)

df_ITI = Dolo.tabulate(model_iid, dr_ITI_iid.dr, :k)

import PyPlot
plt = PyPlot;
fig = plt.figure("Decision Rule, ITI-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df_ITI[:k], df_ITI[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df_ITI[:k], df_ITI[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");

df = Dolo.tabulate(model_iid, dr_TI_iid.dr, :k)

fig = plt.figure("Decision Rule, IT-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df[:k], df[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df[:k], df[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");

################################################################################
# RBC with an AR1 with 2 shocks
################################################################################

# filename = joinpath(path,"examples","models","rbc_dtcc_ar1_2shocks.yaml")
# model_ar1_2 = Dolo.yaml_import(filename)
# @time dr_ITI_ar1_2 = Bruteforce_module.improved_time_iteration(model_ar1_2; verbose=true, tol = 1e-06, smaxit=50)
# @time dr_TI_ar_2  = Dolo.time_iteration(model_ar1_2; tol_η=1e-08, maxit=1000)
#

################################################################################
# RBC with a sudden stop. Complementarities
################################################################################

filename = joinpath(path,"examples","models","sudden_stop.yaml")
model_Sudden = Dolo.yaml_import(filename)
@time dr_ITI_s  = Dolo.improved_time_iteration(model_Sudden; complementarities = true, verbose=true, tol = 1e-06, smaxit=50)
@time dr_TI_s  = Dolo.time_iteration(model_Sudden; tol_η=1e-08, maxit=1000)

df_ITI = Dolo.tabulate(model_Sudden, dr_ITI_s.dr, :l)

import PyPlot
plt = PyPlot;
fig = plt.figure("Decision Rule, ITI-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df_ITI[:l], df_ITI[:b])
plt.ylabel("b(t-1)");
plt.xlabel("state = b");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df_ITI[:l], df_ITI[:lam])
plt.ylabel("b/c");
plt.xlabel("state = b");
plt.title("Decision Rule");

#
df = Dolo.tabulate(model_Sudden, dr_TI_s.dr, :l)

fig = plt.figure("Decision Rule, IT-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df[:l], df[:b])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df[:l], df[:lam])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");


text:
## Kullback-Leibler divergence

function datacov_generic(proc,t)
  Δt1 = t[2]-t[1]
  regularly_sampled = all(diff(t).==Δt1)
  Σ = if regularly_sampled
    N = length(t)
    τ = (0:N-1)*Δt1
    R = autocov(proc,τ)
    SymmetricToeplitz(R)
  else
    Δt = [abs(t2-t1) for t1 in t,t2 in t]
    τ = sort(unique(Δt))
    τ_idx = indexin(Δt,τ)
    R = autocov(proc,τ)
    Hermitian(R[τ_idx])
  end
  return Σ
end

"""
$SIGNATURES

Data covariance matrix `Σ` for SDE process `proc` sampled at times `t`.
"""
datacov(proc,t) = datacov_generic(proc,t)

function datacov(proc::ScalarSdeProcess,t)
  if !has_statstate(proc;rtol=1e-4) @warn "Ignoring non-stationary state" end
  return datacov_generic(proc,t)
end

datacov(proc::WhiteNoise,t) = proc.σ^2*I(length(t))

"""
$SIGNATURES

Kullback-Leibler divergence of SDE process estimate `fh` with respect to the
true SDE process `f` for a time series observed at times `t`.
The current covariance-matrix based implementation has two disadvantages:

* It has a high computational load for a large dataset, say `length(t)>1000`

* Both `f` and `fh` must be stationary.
"""
function kldiv(f,fh,t)
  N = length(t)
  Σ = datacov(f,t)
  Σh = datacov(fh,t)
  # eq:kld-covmat in continuous_discrete_bayesian_filtering_linear.lyx
  D = 1/2*(-(logdet(Σ)-logdet(Σh))+tr(Σh\Σ)-N)
  return D
end

safe_kldiv(f,fh,t,error_value=NaN) = try
  kldiv(f,fh,t)
catch e
  if !(e isa DomainError) rethrow() end
  @warn "kldiv thows error" f fh t
  error_value
end


text:
function dJv_seriesdk(k2::T,v::Int64) where {T <: Real}
# Use series expansion to compute J_v:
nmax = 100
n = 2; error = Inf; if k2 < 1; tol = eps(k2); else; tol = eps(inv(k2)); end
# Computing leading coefficient (n=0):
#coeff = 3pi/(2^(2+v)*factorial(v+2))
coeff = zero(k2)
if k2 < 1
  coeff = 3pi/(2^(2+v)*exp(lfact(v+2)))
#  println("coefficient: ",coeff)
# multiply by (2v-1)!!
  for i=2:2:2v
    coeff *= i-1
  end
# Add leading term to J_v:
  Jv = one(k2)*coeff
  dJvdk = one(k2)*coeff*(2v+1)
# Now, compute higher order terms until desired precision is reached:
  while n < nmax && abs(error) > tol
    coeff *= (n-1)*(n+2v-1)/(n*(n+2v+4))*k2
    Jv += coeff
    dJvdk += coeff*(n+2v+1)
#    error = coeff/Jv
    error = coeff
    n += 2
  end
  dJvdk *= k2^v
  Jv *= k2^v*sqrt(k2)
#  println("Jv: ",Jv," dJv/dk: ",dJvdk)
  return Jv,dJvdk
else # k^2 >= 1
  coeff = pi
  # Compute (2v-1)!!/(2^v v!):
  for i=2:2:2v
    coeff *= 1-1/i
  end
  Jv = one(k2)*coeff
  dJvdk = zero(k2)
  k2inv = inv(k2)
  while n < nmax && abs(error) > tol
    coeff *= (1-5/n)*(1-1/(n+2v))*k2inv
    Jv += coeff
    dJvdk -= n*coeff
#    error = coeff/Jv
    error = coeff
    n += 2
  end
  dJvdk /= sqrt(k2)
  return Jv,dJvdk
end
end

max_stars_repo_name: tagordon/limbdark

max_issues_repo_name: tagordon/limbdark

max_forks_repo_name: tagordon/limbdark


text:
#=
function solve!(::EulerMaruyama, Y, u, W::SamplePath, P::ProcessOrCoefficients)
    N = length(W)
    N != length(Y) && error("Y and W differ in length.")

    tt, yy = Y.tt, Y.yy
    tt .= W.tt
    ww = W.yy

    y::typeof(u) = u

    for i in 1:N-1
        dt = tt[i+1] - tt[i]
        yy[i] = y
        y = y + b(tt[i], y, P)*dt + _scale((ww[i+1]-ww[i]), σ(tt[i], y,  P))
    end
    yy[N] = y
    Y
end
=#

"""
Currently only timedependent sigma, as Ito correction is necessary
"""
function solvebackward!(::EulerMaruyama, Y, u, W::SamplePath, P::ProcessOrCoefficients)
    N = length(W)
    N != length(Y) && error("Y and W differ in length.")

    tt, yy = Y.tt, Y.yy
    tt .= W.tt
    ww = W.yy

    y::typeof(u) = u

    for i in N:-1:2
        dt = tt[i-1] - tt[i]
        yy[i] = y
        y = y + b(tt[i], y, P)*dt + _scale((ww[i-1]-ww[i]), σ(tt[i], P))
    end
    yy[1] = y
    Y
end


text:
function compute_tp(d::TruncatedMvDistribution{D,R};
                    tol::Float64 = 10e-4,
                    tol_step::Int = 10^5,
                    alg::Symbol = :hc) where {D,R}
    if alg == :mc
        i, n = 0, 0
        err = 1.0
        tp = 0.0
        while err ≥ tol
            candidate = rand(d.untruncated)
            i += intruncationregion(d.region,candidate)
            n += 1
            tp = i/n
            if n % tol_step == 0
                err = 3*√(tp*(1-tp)/n)
            end
        end
        d.state.tp = tp
    elseif alg == :hc
        d.state.tp, d.state.tp_err = hcubature((x)->pdf(d.untruncated,x),d.region.a, d.region.b)
    else
        error("Unknown algorithm $(alg)")
    end
    nothing
end

function compute_mean(d::TruncatedMvDistribution{D,R};
                    tol::Float64 = 10e-4,
                    tol_step::Int = 10^5,
                    alg::Symbol = :hc) where {D,R}
    if alg == :mc
        @error("still not implemented")
    elseif alg == :hc
        d.state.μ, d.state.μ_err = hcubature((x)->pdf(d,x)*x,d.region.a, d.region.b)
    else
        error("Unknown algorithm $(alg)")
    end
    nothing
end

function compute_cov(d::TruncatedMvDistribution{D,R};
                    tol::Float64 = 10e-4,
                    tol_step::Int = 10^5,
                    alg::Symbol = :hc) where {D,R}
    if alg == :mc
        @error("still not implemented")
    elseif alg == :hc
        μ = mean(d)
        d.state.Σ, d.state.Σ_err = hcubature((x)->pdf(d,x)*(x-μ)*(x-μ)',d.region.a, d.region.b)
    else
        error("Unknown algorithm $(alg)")
    end
    nothing
end

function compute_moment(d::TruncatedMvDistribution{D,R},k::Vector{Int};
                        tol::Float64 = 10e-4,
                        tol_step::Int = 10^5,
                        alg::Symbol = :hc) where {D,R}

    if alg == :mc
        @error("still not implemented")
    elseif alg == :hc
        return hcubature((x)->pdf(d,x)*prod(x.^k),d.region.a, d.region.b)[1]
    else
        error("Unknown algorithm $(alg)")
    end
    nothing
end


text:
struct MEM{T,S}
    Ω::T
    Φ::S
    fdspec::Matrix{Float64}
    D::Matrix{Float64}
    function MEM(S, Ω, Φ)
        D = zeros(length(Φ), length(Ω)) # transpose on return because should be frequency direction
        fdspec = similar(D)
        for (i, s) ∈ enumerate(S)
            a1, b1, a2, b2 = freqdir_fouriercoef(s)
            c1 = a1 + 1im * b1
            c2 = a2 + 1im * b2
            F₁ = (c1 - c2 * conj(c1)) / (1 - abs(c1)^2)
            F₂ = c2 - c1 * F₁

            @. D[:, i] = real(
                (1 / (2pi)) * (1 - F₁ * conj(c1) - F₂ * conj(c2)) /
                (abs(1 - F₁ * exp(-1im * Φ) - F₂ * exp(-2im * Φ))^2),
            )

            for j = 1:length(Φ)
                fdspec[j, i] = D[j, i] * real.(s[1, 1])
            end
        end
        new{typeof(Ω), typeof(Φ)}(Ω, Φ, fdspec', D')
    end
end
function freqdir_fouriercoef(S)
    denom = sqrt(abs((S[2, 2] + S[3, 3]) * S[1, 1]))
    a1 = imag(S[2,1]) / denom
    b1 = imag(S[3,1]) / denom
    a2 = real((S[2, 2] - S[3, 3]) / (S[2, 2] + S[3, 3]))
    b2 = 2 * real((S[2, 3]) / (S[2, 2] + S[3, 3]))
    return a1, b1, a2, b2
end


text:
### A Pluto.jl notebook ###
# v0.12.6

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 9bc349e8-efb0-11ea-1392-cd735828af33
# packaged loads, etc
begin
	using DelimitedFiles
	using NIfTI
	using Images
	using StatsBase
	using Plots
	using PlutoUI
	using Printf
	using LinearAlgebra
end

# ╔═╡ df5232e0-04e2-11eb-0c26-f19b9d9c7b14
begin
	using GLM
	using DataFrames
end

# ╔═╡ 8029070e-efb0-11ea-1d0a-7f402f2e9989

# ╔═╡ 4df55a88-f98a-11ea-1c9f-cd52a4b5062f
md"""
# How do we analyze data here?

## General idea

What happens if I multiply each point in a timeseries by a number (or "weight", "multiplier", ...) ?
"""

# ╔═╡ c14bf31e-f98c-11ea-33d1-f5de8f2ad0c1
md"""
**weight:** $(@bind weight Slider(range(0.25, 3.0, step=0.25), default=0.75, show_value=true))
"""

# ╔═╡ 2b2f6be4-f07f-11ea-1459-775749ca949f
md"""## Load in that CSV file """

# ╔═╡ 368e2640-f8e0-11ea-1ddf-b37069ba7588
d, h = readdlm("data/stimulus_timing.csv", ',', Float64, header=true);

# ╔═╡ cac0339a-f8e4-11ea-0876-f91dc7db1454
md"""## Stimulus timing $\rightarrow$ fMRI response timing?

"""

# ╔═╡ 1058934e-f8e1-11ea-1896-35d75c4b6a8e
begin
    md"""
*Make a slider that controls the plot*

**row:** $(@bind whichRow Slider(1:2:size(d,1), default=1, show_value=true))
"""

end

# ╔═╡ c6cab78c-f836-11ea-2b5b-23f0006f8e93
begin
	plot(view(d, 1:whichRow,1), view(d,1:whichRow,2), xlabel="Time (volumes)",
		ylabel="fMRI response", label="$(h[2])", xlims=(0, 160), ylims=(-5, 5), lw=2)
	plot!(view(d, 1:whichRow,1), view(d,1:whichRow,3),
		label="$(h[3])", xlims=(0, 160), ylims=(-5, 5), lw=3, color=:red)

end

# ╔═╡ da4ac85c-f8e4-11ea-192e-03d5b6862d5a
md"""## Scaling = multiplying"""

# ╔═╡ 1c3b6c56-f8e2-11ea-3969-f59871025c5e
begin
    md"""
How can we *scale* the idealised prediction to fit the data better?

**multiplier:** $(@bind multiplier Slider(-10:0.5:10, default=1, show_value=true))
"""
end

# ╔═╡ ddf6018c-f921-11ea-0d74-c16f4f0b7426
md"""here I am plotting the data (blue line) and the prediction for what an "ideal" face-selective response would look like... and we can scale it"""

# ╔═╡ 39d465a6-f8e2-11ea-2d71-75cbe586428c
begin
	plot(d[:,1], d[:,5],
		xlim=(0, 160), ylim=(-5, 5),
		xlabel="Time (volumes)",
	ylabel="fMRI response (%signal change)", label="data")
	plot!(d[:,1], d[:,3] * multiplier, lw=2, label="model * $(@sprintf "%.1f" multiplier )")

end

# ╔═╡ a1b031a4-fc08-11ea-2314-13e2e79c2ca3
md"""### Now: more than one component to fit"""

# ╔═╡ e6309486-fc08-11ea-3ade-c93a24b6866b
begin
	plot(d[:,1], d[:,3],
		xlim=(0, 160), ylim=(-1, 1),
		xlabel="Time (volumes)",
		ylabel="fMRI response (%signal change)",
		label="response (FACES)",
		color=:red, lw=2, layout=(2, 1), subplot=1)
	plot!(d[:,1], d[:,4],
		color=:blue, lw=2,
		label="response (OBJECTS)", subplot=2,xlim=(0, 160), ylim=(-1, 1))

end

# ╔═╡ 9c3b0258-fc08-11ea-1922-fd1bec66b4ff
begin
    md"""
How can we *scale* the idealised prediction to fit the data better?

**multiplier (faces):** $(@bind mf Slider(-10:0.5:10, default=1, show_value=true))

**multiplier (objects):** $(@bind mo Slider(-10:0.5:10, default=1, show_value=true))

"""
end

# ╔═╡ 975860d2-fc08-11ea-2def-ad1c45826911
begin
	plot(d[:,1], d[:,5],
		xlim=(0, 160), ylim=(-5, 5),
		xlabel="Time (volumes)",
	ylabel="fMRI response (%signal change)", label="data")
	plot!(d[:,1], d[:,[3,4]] * [mf; mo], lw=2, label="model * $(@sprintf "%.1f" multiplier )")

end

# ╔═╡ e7fb832e-04e0-11eb-388d-f7520b672a6e
md"""

## Calculating the weights for the best fit

*Linear regression* can be used to find/calculate the weights that will give the **best** fit of the mixture of the two curves.

**Best** here means the settings that will result in the *least squares fit*, the fit that makes the *sum of squared errors* between each measurement in the timeseries ($\text{data}_i$) and the corresponding model description ($\text{model}_i$) smallest.

$\sum_i (\text{data}_i - \text{model}_i)^2$

"""

# ╔═╡ 2b070720-04e2-11eb-3101-cbc40b17682b
# X, the design matrix
designMatrix = d[:,3:4]

# ╔═╡ 4b968440-04e2-11eb-35ce-5f050dcd4b6b
y_1 = d[:,5] #

# ╔═╡ 6464f806-04e2-11eb-37cc-7363fa76c80e
β_1 = designMatrix \ y_1

# ╔═╡ 111aea9a-04e5-11eb-066d-4d245b04dc5e
md"""
##

turn data into a `dataframe` and use `GLM` module to get a linear model fit

"""

# ╔═╡ 1e370abc-04e3-11eb-064a-21fbb9fafc7c
theData = DataFrame(d[:, 3:6], [:face_response, :object_response, :Voxel1, :Voxel2 ])

# ╔═╡ 1d60545a-04e4-11eb-1720-b5f04c5da239
ols1 = lm(@formula(Voxel1 ~ face_response + object_response), theData)

# ╔═╡ 86d10676-04e4-11eb-2121-2f0c8945cc1a
coefnames(ols1)

# ╔═╡ 979732dc-04e4-11eb-3803-7302dc91e777
coef(ols1)

# ╔═╡ a33a18f2-04e4-11eb-0587-6117c036c641
GLM.confint(ols1)

# ╔═╡ c60d9136-1d5f-11eb-0839-d760afb98985

# ╔═╡ c47f840a-1d5f-11eb-1e88-d3835fa021e3

# ╔═╡ aa954778-1d5f-11eb-00bf-034690d40e94
md"""## Plot using the `calculated` weights"""

# ╔═╡ 5d60ed36-1d5f-11eb-3e90-9315dd9bc07a
begin
	plot(d[:,1], d[:,5],
		xlim=(0, 160), ylim=(-5, 5),
		xlabel="Time (volumes)",
	ylabel="fMRI response (%signal change)", label="data")
	plot!(d[:,1], d[:,[3,4]] * coef(ols1)[2:3], lw=2, label="model * [$(@sprintf "%.2f; %.2f" coef(ols1)[2] coef(ols1)[3] )]")

end

# ╔═╡ cc5dcb6e-1d5f-11eb-0be7-57d4795af804

# ╔═╡ cb65baa0-1d5f-11eb-346f-bdb3a85d0fe7

# ╔═╡ ca5b1768-1d5f-11eb-0459-13e4d2de80c7

# ╔═╡ c3899100-04e4-11eb-37aa-394d3f60bd53
import Pkg; Pkg.add("StatsPlots")

# ╔═╡ d5c0b236-f98b-11ea-0444-6fe340642707
md"""

## Some settings and helper functions

"""

# ╔═╡ 79d08b08-fa8b-11ea-0b4b-01ceef03660a
# calculate sum of squared errors for different betas
begin

	betaVals = -10:0.5:10 ;# range(-10, 10, length=nPoints)
	e = zeros(size(betaVals))
	X = Vector(d[:,3])
	theY = Vector(d[:,5])
	for (i, theB) in enumerate(betaVals)
		evec = X * theB .- theY
		e[i] = transpose(evec) * evec
	end
	# e
	md"""calculate squared errors"""
end

# ╔═╡ 58c63ca8-fa8b-11ea-296d-0b7f613e8976
begin
	# plot error curve and current choice of M
    function plotError(betas, es, currentM, subplotNum)
	plot!(betas, es, lw=2, color=:black,
		label="error", subplot=subplotNum,)
	idx = betas .== currentM
	if es[idx]... ≈ minimum(es)
		theColor = :green
		theLabel = "min SSE"
	else
		theColor = :red
		theLabel = "SSE"

	end
	plot!(betas[idx], es[idx],
		marker=:circle,
		ms=10,
		color=theColor,
		subplot=subplotNum,
		xlabel="β value",
		ylabel="SSE",
		label=theLabel)
    end
	md"""a function that plots error for a given choice of m `plotError` """
end

# ╔═╡ 96844d74-fa87-11ea-0966-89bbbc8a4353
begin
    function plotWithBeta(xData, yData, yModel, m)
	l = @layout [ grid(2, 1) a{0.30w}
					 ]
	plot(xData, yData,
		xlim=(0, 160), ylim=(-5, 5),
		ylabel="fMRI response (%change)",
		label="data",
		layout=l,
	subplot=1)
	plot!(xData, yModel * m, lw=2,
		subplot=1,
		label="model * $(@sprintf "%.1f" m )")
	sticks!(xData, yData .- yModel * m,
			color=:gray,
			label="data-model",
				xlabel="Time (volumes)",
		subplot=2,
	xlim=(0, 160), ylim=(-5, 5))

	# betaVals from other cell
	# ditto for e
	plotError(betaVals, e, m, 3)
    end
	md"""a function that plots data + model with a
	particular choice of m `plotWithBeta` """
end

# ╔═╡ bfdd7d56-f98a-11ea-3064-a9033ba994a5
# plotly()
gr()

# ╔═╡ eda9e2dc-f98b-11ea-3eca-539ac938107c
nPoints = 20; # number of points in plot above (not used)

# ╔═╡ 778800f8-f98a-11ea-1b87-1108f500ad5f
begin
	t = 0:1:(nPoints - 1);
	y = weight .* t;

	plot(t,y, lw=2, color=:red,
		aspect_ratio=1,
		xlims=(0, 15), ylims=(0, 15),
		label="fit curve")
	plot!(t,1.25.*t, seriestype=:scatter, marker=:square, ms=5, color=:black, label=:"data points")

end

# ╔═╡ f6e4754c-f98b-11ea-2a42-034d82f52ca1
addNoise = 0.5 .* randn(nPoints);

# ╔═╡ Cell order:
# ╟─9bc349e8-efb0-11ea-1392-cd735828af33
# ╟─8029070e-efb0-11ea-1d0a-7f402f2e9989
# ╟─4df55a88-f98a-11ea-1c9f-cd52a4b5062f
# ╟─778800f8-f98a-11ea-1b87-1108f500ad5f
# ╟─c14bf31e-f98c-11ea-33d1-f5de8f2ad0c1
# ╟─2b2f6be4-f07f-11ea-1459-775749ca949f
# ╠═368e2640-f8e0-11ea-1ddf-b37069ba7588
# ╟─cac0339a-f8e4-11ea-0876-f91dc7db1454
# ╟─1058934e-f8e1-11ea-1896-35d75c4b6a8e
# ╟─c6cab78c-f836-11ea-2b5b-23f0006f8e93
# ╟─da4ac85c-f8e4-11ea-192e-03d5b6862d5a
# ╟─1c3b6c56-f8e2-11ea-3969-f59871025c5e
# ╟─ddf6018c-f921-11ea-0d74-c16f4f0b7426
# ╟─39d465a6-f8e2-11ea-2d71-75cbe586428c
# ╟─a1b031a4-fc08-11ea-2314-13e2e79c2ca3
# ╟─e6309486-fc08-11ea-3ade-c93a24b6866b
# ╟─9c3b0258-fc08-11ea-1922-fd1bec66b4ff
# ╠═975860d2-fc08-11ea-2def-ad1c45826911
# ╟─e7fb832e-04e0-11eb-388d-f7520b672a6e
# ╠═2b070720-04e2-11eb-3101-cbc40b17682b
# ╟─4b968440-04e2-11eb-35ce-5f050dcd4b6b
# ╟─6464f806-04e2-11eb-37cc-7363fa76c80e
# ╟─df5232e0-04e2-11eb-0c26-f19b9d9c7b14
# ╟─111aea9a-04e5-11eb-066d-4d245b04dc5e
# ╟─1e370abc-04e3-11eb-064a-21fbb9fafc7c
# ╠═1d60545a-04e4-11eb-1720-b5f04c5da239
# ╠═86d10676-04e4-11eb-2121-2f0c8945cc1a
# ╠═979732dc-04e4-11eb-3803-7302dc91e777
# ╟─a33a18f2-04e4-11eb-0587-6117c036c641
# ╟─c60d9136-1d5f-11eb-0839-d760afb98985
# ╟─c47f840a-1d5f-11eb-1e88-d3835fa021e3
# ╟─aa954778-1d5f-11eb-00bf-034690d40e94
# ╟─5d60ed36-1d5f-11eb-3e90-9315dd9bc07a
# ╟─cc5dcb6e-1d5f-11eb-0be7-57d4795af804
# ╟─cb65baa0-1d5f-11eb-346f-bdb3a85d0fe7
# ╟─ca5b1768-1d5f-11eb-0459-13e4d2de80c7
# ╟─c3899100-04e4-11eb-37aa-394d3f60bd53
# ╟─d5c0b236-f98b-11ea-0444-6fe340642707
# ╟─79d08b08-fa8b-11ea-0b4b-01ceef03660a
# ╟─58c63ca8-fa8b-11ea-296d-0b7f613e8976
# ╠═96844d74-fa87-11ea-0966-89bbbc8a4353
# ╟─bfdd7d56-f98a-11ea-3064-a9033ba994a5
# ╠═eda9e2dc-f98b-11ea-3eca-539ac938107c
# ╠═f6e4754c-f98b-11ea-2a42-034d82f52ca1

max_stars_repo_name: schluppeck/hands-on-brain-data

max_issues_repo_name: schluppeck/hands-on-brain-data

max_forks_repo_name: schluppeck/hands-on-brain-data


text:
### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 295178e0-ebb3-11ea-1213-b531a8ef5828
# Explicit use of own environment instead of a local one for each notebook
begin
	using Pkg
	cd(joinpath(dirname(@__FILE__),".."))
    Pkg.activate(pwd())
    using NativeSVG
end

# ╔═╡ 880bced0-ebb3-11ea-0cdf-f3f1c399de51
using Random

# ╔═╡ fb1cdbe0-ebb2-11ea-1034-a3632723963b
md"""# Game of Life

Port of [Think Complexity chapter 6](http://greenteapress.com/complexity2/html/index.html) by Allen Downey."""

# ╔═╡ 2cf79d80-ebb3-11ea-2ce7-85ca9f3edc7d
md"""Conway's Game of Life (GoL) is a 2 dimensional CA that turns out to be universal.

## Conway's GOL

- The cells in GoL are arranged in a 2-D grid, that is, an array of rows and columns. Usually the grid is considered to be infinite, but in practice it is often “wrapped”; that is, the right edge is connected to the left, and the top edge to the bottom.

- Each cell in the grid has two states — live and dead — and 8 neighbors — north, south, east, west, and the four diagonals. This set of neighbors is sometimes called a “Moore neighborhood”.

- Like the 1-D CAs in the previous chapters, GoL evolves over time according to rules, which are like simple laws of physics.

- In GoL, the next state of each cell depends on its current state and its number of live neighbors. If a cell is alive, it stays alive if it has 2 or 3 neighbors, and dies otherwise. If a cell is dead, it stays dead unless it has exactly 3 neighbors.

- This behavior is loosely analogous to real cell growth: cells that are isolated or overcrowded die; at moderate densities they flourish.

## Implementation

rule:"""

# ╔═╡ 316e5890-ebb3-11ea-247c-5104c109f9d8
function applyrulegameoflife(bits::BitArray{2})
    (nr_y, nr_x) = size(bits)
    out = falses(nr_y, nr_x)
    for y in 2:nr_y-1
        for x in 2:nr_x-1
            if bits[y, x]
                if 2 ≤ count(v->v, bits[y-1:y+1,x-1:x+1]) - 1 ≤ 3
                    out[y, x] = true
                end
            else
                if count(v->v, bits[y-1:y+1,x-1:x+1]) == 3
                    out[y, x] = true
                end
            end
        end
    end
    out
end

# ╔═╡ 66432050-ebb3-11ea-3c37-e37390422780
md"visulatisation:"

# ╔═╡ 701c7d10-ebb3-11ea-19d8-7f9587aee599
function visualize2dim(bits::BitArray{2}, dim)
    (nr_y, nr_x) = size(bits)
    width = dim * (nr_x - 1)
    height = dim * (nr_y - 1)
    Drawing(width = width, height = height) do
        for (j, y) in enumerate(2:nr_y-1)
            for (i, x) in enumerate(2:nr_x-1)
                fill = if bits[y, x] "grey" else "lightgrey" end
                rect(x = i*dim, y = j*dim, width = dim, height = dim, fill = fill)
            end
        end
    end
end

# ╔═╡ 7b710820-ebb3-11ea-3947-4548d61acafb
md"example:"

# ╔═╡ 40c8bf02-ebb4-11ea-2a2f-eff487a4a30f
mutable struct Gol
	bits :: BitArray{2}
	function Gol(xdim::Int64, ydim)
		new(falses(xdim, ydim))
	end
end

# ╔═╡ a6813e30-ebb4-11ea-2cfc-cfa8e82a0b54
function goltest()
	gol = Gol(12, 12)
	gol.bits[2:11,2:11] = reshape(bitrand(100), (10,10))
	gol
end

# ╔═╡ b725ed00-ebb7-11ea-19a0-2ff7e4729780
gol = goltest();

# ╔═╡ 69dae5b0-ebb6-11ea-0464-47c13136f76d
visualize2dim(gol.bits, 30)

# ╔═╡ 07986fe0-ebb5-11ea-3f52-2daa2bad7f22
function applyrule(gol::Gol)
	gol.bits = applyrulegameoflife(gol.bits)
end

# ╔═╡ 9506b6e0-ebb3-11ea-1922-c36f37f7fdce
visualize2dim(applyrule(gol), 30)

# ╔═╡ e6284c30-ebb5-11ea-3861-f5253bd5ace2
md"""## Life Patterns

A number of stable patterns are likely to appear.

### Beehive

A stable pattern."""

# ╔═╡ f465ed72-ebb5-11ea-2178-0928e7650f7e
function Beehive()
	beehive = Gol(5, 6)
	beehive.bits[2,3:4] = [true, true]
	beehive.bits[3,2] = true
	beehive.bits[3,5] = true
	beehive.bits[4,3:4] = [true, true]
	beehive
end

# ╔═╡ 5c8d97de-ebb6-11ea-1a01-651fceb86eb5
beehive = Beehive();

# ╔═╡ a515e600-ebc2-11ea-3691-553176d1a9cb
visualize2dim(applyrule(beehive), 30)

# ╔═╡ b0790d80-ebb6-11ea-3423-1fa5c38fd3a4
md"""### Toad

An oscillating pattern. The toad has a period of 2 timesteps."""

# ╔═╡ 8d274e00-ebb6-11ea-1a22-bdb17e0513c7
function Toad()
	toad = Gol(6, 6)
	toad.bits[3,3:5] = [true, true, true]
	toad.bits[4,2:4] = [true, true, true]
	toad
end

# ╔═╡ f7491ac0-ebb6-11ea-17d9-a99a458757d9
toad = Toad();

# ╔═╡ fbe70ce0-ebb6-11ea-2bec-8754663e7ef3
visualize2dim(applyrule(toad), 30)

# ╔═╡ 059206a0-ebb7-11ea-02e0-11ff007b95b2
md"""### Glider

Oscillation pattern that shift in space. After a period of 4 steps, the glider is back in the starting configuration, shifted one unit down and to the right."""

# ╔═╡ 584f18f0-ebb8-11ea-235c-8918cb8023bc
function Glider()
	glider = Gol(8, 8)
	glider.bits[2,3] = true
	glider.bits[3,4] = true
	glider.bits[4,2:4] = [true, true, true]
	glider
end

# ╔═╡ 83e4ef32-ebb8-11ea-3a43-7556b2cf10c7
glider = Glider();

# ╔═╡ ac886070-ebb8-11ea-1742-c9449ca01d7e
visualize2dim(glider.bits, 30)

# ╔═╡ 87ad41d0-ebb8-11ea-3e4d-5f76f10c6888
@bind toggleglider html"<input type=button value='Next'>"

# ╔═╡ 9dd17080-ebb8-11ea-11df-cfe66b93af90
if toggleglider === "Next"
	visualize2dim(applyrule(glider), 30)
else
	visualize2dim(glider.bits, 30)
end

# ╔═╡ 307f6080-ebba-11ea-210c-0182eece50c2
md"""## Methusalems

From most initial conditions, GoL quickly reaches a stable state where the number of live cells is nearly constant (possibly with some oscillation).

But there are some simple starting conditions that yield a surprising number of live cells, and take a long time to settle down. Because these patterns are so long-lived, they are called “Methusalems”.

One of the simplest Methusalems is the r-pentomino, which has only five cells, roughly in the shape of the letter 'r'."""

# ╔═╡ 90bba610-ebbb-11ea-369a-936b36d2e26f
function R_Pentomino()
	r_pentomino = Gol(66, 66)
	r_pentomino.bits[28,28:29] = [true, true]
	r_pentomino.bits[29,27:28] = [true, true]
	r_pentomino.bits[30,28] = true
	r_pentomino
end

# ╔═╡ 21cbba00-ebbc-11ea-2a6c-2f7561ea8c71
r_pentomino = R_Pentomino();

# ╔═╡ e3af2720-ebbb-11ea-0458-f51595107c9c
@bind togglerpentomino html"<input type=button value='Next'>"

# ╔═╡ 0d6ce200-ebbc-11ea-25b4-e33be459e6a2
if togglerpentomino === "Next"
	visualize2dim(applyrule(r_pentomino), 8)
else
	visualize2dim(r_pentomino.bits, 8)
end

# ╔═╡ 91aa5660-ebbc-11ea-10d6-b59fff5e0272
let r_pentomino = deepcopy(r_pentomino)
	for _ in 1:1000
		r_pentomino.bits = applyrule(r_pentomino)
	end
	visualize2dim(r_pentomino.bits, 8)
end

# ╔═╡ f7a1b350-ebbc-11ea-3951-e7e1a6c3bd7b
md"""This configuration is final in the sense that all remaining patterns are either stable, oscillators or gliders that will never collide with another pattern.

There are initial patterns that never stabilize, eg. a gun or a puffer train

The Game of Life was proved Turing complete in 1982. Since then, several people have constructed GoL patterns that implement a Turing machine or another machine known to be Turing complete."""

# ╔═╡ Cell order:
# ╟─fb1cdbe0-ebb2-11ea-1034-a3632723963b
# ╠═295178e0-ebb3-11ea-1213-b531a8ef5828
# ╟─2cf79d80-ebb3-11ea-2ce7-85ca9f3edc7d
# ╠═316e5890-ebb3-11ea-247c-5104c109f9d8
# ╟─66432050-ebb3-11ea-3c37-e37390422780
# ╠═701c7d10-ebb3-11ea-19d8-7f9587aee599
# ╟─7b710820-ebb3-11ea-3947-4548d61acafb
# ╠═880bced0-ebb3-11ea-0cdf-f3f1c399de51
# ╠═40c8bf02-ebb4-11ea-2a2f-eff487a4a30f
# ╠═a6813e30-ebb4-11ea-2cfc-cfa8e82a0b54
# ╠═b725ed00-ebb7-11ea-19a0-2ff7e4729780
# ╠═69dae5b0-ebb6-11ea-0464-47c13136f76d
# ╠═07986fe0-ebb5-11ea-3f52-2daa2bad7f22
# ╠═9506b6e0-ebb3-11ea-1922-c36f37f7fdce
# ╟─e6284c30-ebb5-11ea-3861-f5253bd5ace2
# ╠═f465ed72-ebb5-11ea-2178-0928e7650f7e
# ╠═5c8d97de-ebb6-11ea-1a01-651fceb86eb5
# ╠═a515e600-ebc2-11ea-3691-553176d1a9cb
# ╟─b0790d80-ebb6-11ea-3423-1fa5c38fd3a4
# ╠═8d274e00-ebb6-11ea-1a22-bdb17e0513c7
# ╠═f7491ac0-ebb6-11ea-17d9-a99a458757d9
# ╠═fbe70ce0-ebb6-11ea-2bec-8754663e7ef3
# ╟─059206a0-ebb7-11ea-02e0-11ff007b95b2
# ╠═584f18f0-ebb8-11ea-235c-8918cb8023bc
# ╠═83e4ef32-ebb8-11ea-3a43-7556b2cf10c7
# ╠═ac886070-ebb8-11ea-1742-c9449ca01d7e
# ╠═87ad41d0-ebb8-11ea-3e4d-5f76f10c6888
# ╠═9dd17080-ebb8-11ea-11df-cfe66b93af90
# ╟─307f6080-ebba-11ea-210c-0182eece50c2
# ╠═90bba610-ebbb-11ea-369a-936b36d2e26f
# ╠═21cbba00-ebbc-11ea-2a6c-2f7561ea8c71
# ╟─e3af2720-ebbb-11ea-0458-f51595107c9c
# ╠═0d6ce200-ebbc-11ea-25b4-e33be459e6a2
# ╠═91aa5660-ebbc-11ea-10d6-b59fff5e0272
# ╟─f7a1b350-ebbc-11ea-3951-e7e1a6c3bd7b


text:
using Documenter, SPICE

makedocs(
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true",
    ),
    sitename = "SPICE.jl",
    authors = "Helge Eichhorn",
    pages = [
        "Home" => "index.md",
        "API" => "api.md",
    ],
    doctest = false,
)

deploydocs(
    repo = "github.com/JuliaAstro/SPICE.jl.git",
    target = "build",
)


text:
#=##############################################################################################
Filename: solve_rotor_wake.jl
Author: Ryan Anderson
Contact: rymanderson@gmail.com
README: define an `Action` object to solve for the wake of a rotor system
=###############################################################################################

"""
    solve_rotor_wake(aircraft, parameters, freestream, environment, step_range, stepi, step_symbol) <: Action

# Arguments:

* `aircraft::Aircraft`: `Aircraft` system object
* `parameters<:Parameters` `Parameters` struct
* `freestream::Freestream`: `Freestream` object
* `environment::Environment` `Environment` object
* `step_range::AbstractArray`: array of times for which the simulation is run
* `stepi::Int`: index of the current step
* `step_symbol::String`: defines the step, e.g. `alpha` or `time`

`parameters <: Parameters` requires the following elements:

* `wake_function::Function` : a function describing the induced velocity behind each propeller as a function of global coordinates
* `us::Vector{Vector{Vector{Float64}}}` : each [i][j][k]th element is the axial induced velocity at ith step of the jth rotor at the kth radial section
* `vs::Vector{Vector{Vector{Float64}}}` : each [i][j][k]th element is the swirl induced velocity at ith step of the jth rotor at the kth radial section
* `wake_shape_functions::Vector{Function}` : [i]th element is a function f(Rtip, x) describing the radial distance from the rotor axis to the boundary of the wake of the ith rotor
* `axial_interpolations::Vector{Function}` : [i]th element is a function f(rs, us, r, Rtip) that returns the axial component of rotor-induced velocity at distance r from the rotor axis based on the calculated axial induced velocities output from CCBlade of the ith rotor
* `swirl_interpolations::Vector{Function}` : [i]th element is a function f(rs, us, r, Rtip) that returns the swirl component of rotor-induced velocity at distance r from the rotor axis based on the calculated axial induced velocities output from CCBlade of the ith rotor
* `axial_multipliers::Vector{Function}` : [i]th element is a function f(distance2plane, Rtip) that is multiplied by the axial induced velocity function of the ith rotor
* `swirl_multipliers::Vector{Function}` : [i]th element is a function f(distance2plane, Rtip) that is multiplied by the swirl induced velocity function of the ith rotor

"""
function solve_rotor_wake(aircraft, parameters, freestream, environment, step_range, stepi, step_symbol)

    us = parameters.us[stepi]
    vs = parameters.vs[stepi]
    wake_shape_functions = parameters.wake_shape_functions
    axial_interpolations = parameters.axial_interpolations
    swirl_interpolations = parameters.swirl_interpolations
    axial_multipliers = parameters.axial_multipliers
    swirl_multipliers = parameters.swirl_multipliers
    wakefunction = induced2wakefunction(aircraft.rotor_system, us, vs;
        wake_shape_functions,
        axial_interpolations,
        swirl_interpolations,
        axial_multipliers,
        swirl_multipliers
    )

    parameters.wake_function[stepi] = wakefunction

    return false
end

"""
    solve_rotor_wake(aircraft, step_range)

Method returns initialized elements required for the `parameters <: Parameters` struct during simulation.

# Arguments:

* `aircraft::Aircraft` : system to be simulated
* `step_range::AbstractArray` : defines each step of the simulation

# Returns:

* `wake_function::Vector{Union{Function, Nothing}}` : [i]th element is a function describing the velocity field of rotor wakes as a function of global coordinates at the ith step
* `us::Vector{Vector{Vector{Float64}}}` : each [i][j][k]th element is the axial induced velocity at ith step of the jth rotor at the kth radial section
* `vs::Vector{Vector{Vector{Float64}}}` : each [i][j][k]th element is the swirl induced velocity at ith step of the jth rotor at the kth radial section
* `wake_shape_functions::Vector{Function}` : [i]th element is a function f(Rtip, x) describing the radial distance from the rotor axis to the boundary of the wake of the ith rotor
* `axial_interpolations::Vector{Function}` : [i]th element is a function f(rs, us, r, Rtip) that returns the axial component of rotor-induced velocity at distance r from the rotor axis based on the calculated axial induced velocities output from CCBlade of the ith rotor
* `swirl_interpolations::Vector{Function}` : [i]th element is a function f(rs, us, r, Rtip) that returns the swirl component of rotor-induced velocity at distance r from the rotor axis based on the calculated axial induced velocities output from CCBlade of the ith rotor
* `axial_multipliers::Vector{Function}` : [i]th element is a function f(distance2plane, Rtip) that is multiplied by the axial induced velocity function of the ith rotor
* `swirl_multipliers::Vector{Function}` : [i]th element is a function f(distance2plane, Rtip) that is multiplied by the swirl induced velocity function of the ith rotor

"""
function solve_rotor_wake(aircraft, step_range)

    wake_function = Function[(x) -> [0.0, 0.0, 0.0] for i in 1:length(step_range)]
    _, _, _, _, us, vs = solve_rotor(aircraft, step_range)

    wake_shape_functions = Function[(Rtip, x) -> Rtip for i in 1:length(aircraft.rotor_system.index)]
    axial_interpolations = Function[(rs, us, r, Rtip) -> FM.linear(rs, us, r) for i in 1:length(aircraft.rotor_system.index)]
    swirl_interpolations = Function[(rs, vs, r, Rtip) -> FM.linear(rs, vs, r) for i in 1:length(aircraft.rotor_system.index)]
    axial_multipliers = Function[(distance2plane, Rtip) -> 2 for i in 1:length(aircraft.rotor_system.index)]
    swirl_multipliers = Function[(distance2plane, Rtip) -> 1 for i in 1:length(aircraft.rotor_system.index)]

    return wake_function, us, vs, wake_shape_functions, axial_interpolations, swirl_interpolations, axial_multipliers, swirl_multipliers
end

# # prepare solutionkeys and solutioninits for this particular system and simulation
# solutionkeys = [
#     "thrust",
#     "torque",
#     "efficiency",
#     "u",
#     "v"
# ]

max_stars_repo_name: byuflowlab/AircraftSystems

max_issues_repo_name: byuflowlab/AircraftSystems

max_forks_repo_name: byuflowlab/AircraftSystems


text:
using Test
using Metal

@testset "MTL" begin
@testset "devices" begin

devs = devices()
@test length(devs) > 0

dev = first(devs)
@test dev == devs[1]

if length(devs) > 1
    @test dev != devs[2]
end

end

@testset "buffers" begin

dev = first(devices())

buf = MtlBuffer{Int}(dev, 1)

@test sizeof(buf) == 8
@test length(buf) == 1
@test device(buf) == dev

free(buf)

end

end


text:
using Documenter, ScatteredInterpolation

makedocs(
    sitename = "ScatteredInterpolation.jl",
    pages = [
        "Home" => "index.md",
        "Supported methods" => "methods.md",
        "API" => "api.md",
    ]
)

deploydocs(
    repo = "github.com/eljungsk/ScatteredInterpolation.jl.git",
    target = "build",
)


text:
##################################################
##
## Factorization

# Pencil is made up of two factorizations
# notation in paper is
# solve V - \lambda W
# or V*inv(W) - lambda I
# but V = QR, so consider
# Q*R*inv(W)
# to avoid (introduce?) confusion,we let R be V: and have
# Q * V * inv(W) with the R Factorization store both V and W
struct RFactorizationPencil{T, S} <: AbstractRFactorization{T, S}
V::RFactorizationRankOne{T, S}
W::RFactorizationRankOne{T, S}
end

Base.copy(RF::RFactorizationPencil) = RFactorizationPencil(copy(RF.V), copy(RF.W))
Base.size(RF::RFactorizationPencil) = size(RF.V)
function Base.Matrix(RF::RFactorizationPencil)
    V = Matrix(RF.V)
    W = Matrix(RF.W)
    V[1:end-1, 1:end-1] * inv(W[1:end-1, 1:end-1])
end

## VV = [Sym("v$i$j") for i in ("i","j","k"), j in ("i","j","k")] |> triu
## WW = [Sym("w$i$j") for i in ("i","j","k"), j in ("i","j","k")] |> triu
## VV * inv(WW)[:,3]
function Base.getindex(RF::RFactorizationPencil, l, k)
    Δ = k - l
    #@assert 0 <= Δ <= 2

    V, W = RF.V, RF.W
    l <= 0 &&  return zero(RF)
    if iszero(Δ)
        return V[k,k] / W[k,k]
    elseif Δ == 1
        j = l
        return - V[j,j] * W[j,k] / (W[j,j]*W[k,k]) + V[j,k]/W[k,k]
    else
        i, j = l, l+1
        return V[i,i] * (W[i,j] * W[j,k] - W[i,k] * W[j,j]) / (W[i,i] * W[j,j] * W[k,k])  - (V[i,j] *  W[j,k]) / ( W[j,j] * W[k,k])  + V[i,k]/W[k,k]
    end

end

## Constructor for Real coefficients; with pencil

function pencil_factorization(vs, ws)

    V = r_factorization(vs)
    W = r_factorization(ws)
    RFactorizationPencil(V,W)

end

## Pass a rotator through Rfactorization with Pencil from left or right
function passthrough!(RF::RFactorizationPencil, U::AbstractRotator)

    ## Pass Ut -> W (not W^{-1} <- U
    ## Then  passthrough V
    Ut = passthrough!(U', RF.W)
    U = passthrough!(RF.V, Ut')

    U
end

function passthrough!(U::AbstractRotator, RF::RFactorizationPencil)

    U = passthrough!(U, RF.V)
    Ut = passthrough!(RF.W, U')
    U = Ut'

    U
end

## When Ct and B are identical, we can update just one and leave U,V alone
function simple_passthrough!(RF::RFactorizationPencil, U, V)
    false  # might have room for improvement here; o/w consolidate methods
end

function simple_passthrough!(RF::RFactorizationPencil, U)
    false  # might have room for improvement here
end


text:
using Test

include("FunctionTests.jl")


text:
@testset "download jlcall.m" begin
    # Download from github
    jlcall_path = tempname()
    download_jlcall(jlcall_path; latest = true)
    @test isfile(jlcall_path)

    # Copy from api folder
    download_jlcall(jlcall_path; force = true)
    jlcall_local = normpath(@__DIR__, "../api/jlcall.m")
    @test readlines(jlcall_path) == readlines(jlcall_local)
end

@testset "matlabify" begin
    for (jl, mx) in deeply_nested_pairs(; roundtrip = false)
        @test is_eq(matlabify(jl), mx)
    end
end

@testset "start/kill server" begin
    # Add two Julia workers; one for running a Julia server, and one for sending the kill signal
    port = 9876
    addprocs(2)
    @everywhere begin
        using Pkg
        Pkg.activate($(Base.active_project()); io = devnull)
        using MATDaemon
    end

    # Start a DaemonMode server on worker 2
    server_task = @spawnat 2 MATDaemon.start(port; shared = true, verbose = true)

    # Wait until server is running
    server_running = false
    while !server_running
        @test MATDaemon.DaemonMode.runexpr("@eval Main __SERVER_RUNNING__() = :SERVER_RUNNING"; port = port) === nothing
        server_running = fetch(@spawnat 2 isdefined(Main, :__SERVER_RUNNING__))
        sleep(1.0)
    end

    # Kill server and cleanup process
    @test fetch(@spawnat 3 MATDaemon.kill(port; verbose = true)) === nothing
    @test fetch(rmprocs(3)) === nothing

    # Ensure server task has completed and cleanup process
    @test fetch(server_task) === nothing
    @test fetch(rmprocs(2)) === nothing

    # Killing a nonexistent server should be a no-op
    @test MATDaemon.kill(port; verbose = true) === nothing
end

@testset "jlcall" begin
    # Test jlcall, ensuring that function definitions don't leak into Main
    for (i, (f, f_args, f_kwargs, f_output, kwargs)) in enumerate([
        ("f1(x) = LinearAlgebra.norm(x)",    mxtuple([3.0, 4.0]),    mxdict(),           mxtuple(5.0),           (; modules = ["LinearAlgebra"],)),
        ("f2(x; y) = x*y",                   mxtuple(5.0),           mxdict("y" => 3),   mxtuple(15.0),          NamedTuple()),
        ("f3() = nothing",                   mxtuple(),              mxdict(),           mxtuple(),              NamedTuple()),
        ("f4() = missing",                   mxtuple(),              mxdict(),           mxtuple(mxempty()),     NamedTuple()),
    ])
        wrap_jlcall(f, f_args, f_kwargs, f_output; kwargs...)
        @test !isdefined(Main, Symbol(:f, i))
    end

    # Test jlcall, explicitly adding methods to Main
    for (f_sym, f_str, f_args, f_kwargs, f_output, kwargs) in [
        (:f5, "@eval Main f5(x,y) = (x,y)",         mxtuple("one", 2),      mxdict(),           mxtuple("one", 2),      NamedTuple()),
        (:f6, "@eval Main f6(x,y) = x*y",           mxtuple(3.0, 2),        mxdict(),           mxtuple(6.0),           NamedTuple()),
        (:f7, "@eval Main f7(x,y) = [x*y]",         mxtuple(3.0, 2),        mxdict(),           mxtuple(6.0),           NamedTuple()),
        (:f8, "@eval Main f8(x,y) = [x,y]",         mxtuple(3.0, 2.0),      mxdict(),           mxtuple([3.0, 2.0]),    NamedTuple()),
        (:f9, "@eval Main f9(x) = Setup.mul2(x)",   mxtuple([2f0 3f0]),     mxdict(),           mxtuple([4f0 6f0]),     (; setup = joinpath(@__DIR__, "setup.jl"),)),
    ]
        wrap_jlcall(f_str, f_args, f_kwargs, f_output; kwargs...)
        @test isdefined(Main, f_sym)
    end

    # This should fail, as Base.VERSION has type VersionNumber and therefore cannot be written to .mat
    @test_throws LoadError wrap_jlcall("@eval Main f10() = Base.VERSION", mxtuple(), mxdict(), mxtuple(string(Base.VERSION)))
    @test isdefined(Main, :f10)

    # Fix `f10` by extending `matlabify` to `VersionNumber`s
    MATDaemon.matlabify(v::Base.VersionNumber) = string(v)
    wrap_jlcall("Main.f10", mxtuple(), mxdict(), mxtuple(string(Base.VERSION)))

    # Test local project
    reset_active_project() do
        wrap_jlcall("@eval Main f11 = TestProject.dot", mxtuple([1.0, 2.0, 3.0]), mxdict(), mxtuple(14.0); project = jlcall_test_project(), modules = ["TestProject"])
        @test dirname(Base.active_project()) == jlcall_test_project()
    end
    @test isdefined(Main, :f11)
end