text:
@testset "Basic Structural Model" begin
    # Univariate
    air_passengers = CSV.File(StateSpaceModels.AIR_PASSENGERS) |> DataFrame
    log_air_passengers = log.(air_passengers.passengers)

    @test has_fit_methods(BasicStructural)

    model = BasicStructural(log_air_passengers, 12)
    fit!(model)
    # Runned on Python statsmodels
    @test loglike(model) ≈ 234.33641 atol = 1e-5 rtol = 1e-5

    # forecasting
    forec = forecast(model, 10)
    @test monotone_forecast_variance(forec)
    # simualting
    scenarios = simulate_scenarios(model, 10, 10_000)
    test_scenarios_adequacy_with_forecast(forec, scenarios)

    log_ap32 = Float32.(log_air_passengers)
    model = BasicStructural(log_ap32, 12)
    @test model.system.Z == Float32[1.0; 0.0; 1.0; zeros(10)]
    @test_broken fit!(model)
    @test_broken loglike(model) ≈ 234.33641f0 atol = 1e-5 rtol = 1e-5

    model = BasicStructural(log_air_passengers, 12)
    steadystate_tol = 1e-5
    a1 = zeros(num_states(model))
    P1 = 1e6 .* Matrix{Float64}(I, num_states(model), num_states(model))
    sparse_filter = SparseUnivariateKalmanFilter(a1, P1, num_states(model), steadystate_tol)
    fit!(model; filter=sparse_filter)
    # Runned on Python statsmodels
    @test loglike(model) ≈ 234.33641 atol = 1e-5 rtol = 1e-5

    # forecasting
    forec = forecast(model, 10)
    @test monotone_forecast_variance(forec)
    # simualting
    scenarios = simulate_scenarios(model, 10, 10_000)
    test_scenarios_adequacy_with_forecast(forec, scenarios)
end


text:
"""
A [logarithmic number system](https://en.wikipedia.org/wiki/Logarithmic_number_system).

Provides the signed [`Logarithmic`](@ref) and unsigned [`ULogarithmic`](@ref) types, which represent real numbers and positive real numbers respectively.
"""
module LogarithmicNumbers

using Requires, Random

export ULogarithmic, ULogFloat16, ULogFloat32, ULogFloat64
export Logarithmic, LogFloat16, LogFloat32, LogFloat64

### Types

"""
ULogarithmic(x)

Represents the positive real number `x` by storing its logarithm.

!!! tip
    If you know `logx=log(x)` then use [`exp(ULogarithmic, logx)`](@ref exp(::Type{ULogarithmic},::Real)) instead.
"""
struct ULogarithmic{T} <: Real
    log::T
    Base.exp(::Type{ULogarithmic{T}}, x::T) where {T<:Real} = new{T}(x)
end

"""
Logarithmic(x)

Represents the real number `x` by storing its absolute value as a [`ULogarithmic`](@ref) and its sign bit.

!!! tip
    If you know `logx=log(abs(x))` then use [`exp(Logarithmic, logx)`](@ref exp(::Type{Logarithmic},::Real)) instead.
"""
struct Logarithmic{T} <: Real
    abs::ULogarithmic{T}
    signbit::Bool
    Logarithmic{T}(abs::ULogarithmic{T}, signbit::Bool=false) where {T<:Real} = new{T}(abs, signbit)
end

const ULogFloat16 = ULogarithmic{Float16}
const ULogFloat32 = ULogarithmic{Float32}
const ULogFloat64 = ULogarithmic{Float64}
const LogFloat16 = Logarithmic{Float16}
const LogFloat32 = Logarithmic{Float32}
const LogFloat64 = Logarithmic{Float64}

const AnyLogarithmic{T} = Union{ULogarithmic{T}, Logarithmic{T}}

### Constructors

function Base.exp(::Type{ULogarithmic{T}}, x::Real) where {T<:Real}
    exp(ULogarithmic{T}, convert(T, x))
end

function Base.exp(::Type{ULogarithmic}, x::T) where {T<:Real}
    exp(ULogarithmic{T}, x)
end

function Base.exp(::Type{Logarithmic{T}}, x::Real) where {T<:Real}
    Logarithmic{T}(exp(ULogarithmic{T}, x))
end

function Base.exp(::Type{Logarithmic}, x::T) where {T<:Real}
    exp(Logarithmic{T}, x)
end

uexp(x) = exp(ULogarithmic, x)
uexp(T,x) = exp(ULogarithmic{T}, x)

function ULogarithmic{T}(x::Real) where {T<:Real}
    exp(ULogarithmic{T}, log(x))
end

function ULogarithmic{T}(x::ULogarithmic{T}) where {T<:Real}
    x
end

function ULogarithmic(x::Real)
    exp(ULogarithmic, log(x))
end

function ULogarithmic(x::ULogarithmic)
    x
end

function Logarithmic{T}(x::Real) where {T<:Real}
    Logarithmic{T}(ULogarithmic{T}(abs(x)), signbit(x))
end

function Logarithmic{T}(x::Logarithmic{T}) where {T<:Real}
    x
end

function Logarithmic{T}(abs::ULogarithmic, signbit::Bool=false) where {T<:Real}
    Logarithmic{T}(ULogarithmic{T}(abs), signbit)
end

function Logarithmic(x::Real)
    Logarithmic(ULogarithmic(abs(x)), signbit(x))
end

function Logarithmic(abs::ULogarithmic{T}, signbit::Bool=false) where {T<:Real}
    Logarithmic{T}(abs, signbit)
end

function Logarithmic(x::Logarithmic)
    x
end

### float / big / signed / unsigned

function Base.float(x::AnyLogarithmic)
    AbstractFloat(x)
end

function (::Type{T})(x::ULogarithmic) where {T<:AbstractFloat}
    T(exp(float(x.log)))
end

function (::Type{T})(x::Logarithmic) where {T<:AbstractFloat}
    y = float(x.abs)
    x.signbit ? T(-y) : T(y)
end

function Base.big(x::ULogarithmic{T}) where {T}
    uexp(big(x.log))
end

function Base.big(x::Logarithmic)
    Logarithmic(big(x.abs), x.signbit)
end

function Base.unsigned(x::ULogarithmic)
    x
end

function Base.unsigned(x::Logarithmic)
    x.signbit && !iszero(x) && throw(DomainError(x))
    x.abs
end

function Base.signed(x::ULogarithmic)
    Logarithmic(x)
end

function Base.signed(x::Logarithmic)
    x
end

### Type functions

function Base.float(::Type{T}) where {T<:AnyLogarithmic}
    typeof(float(one(T)))
end

function Base.widen(::Type{ULogarithmic{T}}) where {T}
    ULogarithmic{widen(T)}
end

function Base.widen(::Type{Logarithmic{T}}) where {T}
    Logarithmic{widen(T)}
end

function Base.big(::Type{ULogarithmic{T}}) where {T}
    ULogarithmic{big(T)}
end

function Base.big(::Type{Logarithmic{T}}) where {T}
    Logarithmic{big(T)}
end

function Base.unsigned(::Type{ULogarithmic{T}}) where {T}
    ULogarithmic{T}
end

function Base.unsigned(::Type{Logarithmic{T}}) where {T}
    ULogarithmic{T}
end

function Base.signed(::Type{ULogarithmic{T}}) where {T}
    Logarithmic{T}
end

function Base.signed(::Type{Logarithmic{T}}) where {T}
    Logarithmic{T}
end

### Special values

function Base.zero(::Type{ULogarithmic{T}}) where {T}
    uexp(T, -Inf)
end

function Base.zero(::Type{ULogarithmic})
    uexp(-Inf)
end

function Base.zero(::Type{Logarithmic{T}}) where {T}
    Logarithmic(zero(ULogarithmic{T}))
end

function Base.zero(::Type{Logarithmic})
    Logarithmic(zero(ULogarithmic))
end

function Base.one(::Type{ULogarithmic{T}}) where {T}
    uexp(T, zero(T))
end

function Base.one(::Type{ULogarithmic})
    uexp(0.0)
end

function Base.one(::Type{Logarithmic{T}}) where {T}
    Logarithmic(one(ULogarithmic{T}))
end

function Base.one(::Type{Logarithmic}) where {T}
    Logarithmic(one(ULogarithmic))
end

function Base.typemin(::Type{ULogarithmic{T}}) where {T}
    uexp(typemin(T))
end

function Base.typemin(::Type{Logarithmic{T}}) where {T}
    Logarithmic{T}(typemax(ULogarithmic{T}), true)
end

function Base.typemax(::Type{ULogarithmic{T}}) where {T}
    uexp(typemax(T))
end

function Base.typemax(::Type{Logarithmic{T}}) where {T}
    Logarithmic{T}(typemax(ULogarithmic{T}))
end

### Predicates

function Base.iszero(x::ULogarithmic)
    isinf(x.log) && signbit(x.log)
end

function Base.iszero(x::Logarithmic)
    iszero(x.abs)
end

function Base.isone(x::ULogarithmic)
    iszero(x.log)
end

function Base.isone(x::Logarithmic)
    isone(x.abs) && !x.signbit
end

function Base.isinf(x::ULogarithmic)
    isinf(x.log) && !signbit(x.log)
end

function Base.isinf(x::Logarithmic)
    isinf(x.abs)
end

function Base.isfinite(x::AnyLogarithmic)
    !isinf(x)
end

function Base.isnan(x::ULogarithmic)
    isnan(x.log)
end

function Base.isnan(x::Logarithmic)
    isnan(x.abs)
end

### Ordering

function Base.sign(x::ULogarithmic)
    iszero(x) ? zero(x) : one(x)
end

function Base.sign(x::Logarithmic)
    iszero(x) ? zero(x) : x.signbit ? -one(x) : one(x)
end

function Base.signbit(x::ULogarithmic)
    false
end

function Base.signbit(x::Logarithmic)
    x.signbit
end

function Base.abs(x::ULogarithmic)
    x
end

function Base.abs(x::Logarithmic)
    x.abs
end

function Base.:(==)(x::ULogarithmic, y::ULogarithmic)
    x.log == y.log
end

function Base.:(==)(x::Logarithmic, y::Logarithmic)
    (iszero(x) && iszero(y)) || (x.abs==y.abs && x.signbit==y.signbit)
end

function Base.isequal(x::ULogarithmic, y::ULogarithmic)
    isequal(x.log, y.log)
end

function Base.isequal(x::Logarithmic, y::Logarithmic)
    isequal(x.abs, y.abs) && isequal(x.signbit, y.signbit)
end

function Base.:(<)(x::ULogarithmic, y::ULogarithmic)
    x.log < y.log
end

function Base.:(<)(x::Logarithmic, y::Logarithmic)
    if x.signbit
        if y.signbit
            y.abs < x.abs
        else
            !iszero(x) || !iszero(y)
        end
    else
        if y.signbit
            false
        else
            x.abs < y.abs
        end
    end
end

function Base.:(≤)(x::ULogarithmic, y::ULogarithmic)
    x.log ≤ y.log
end

function Base.:(≤)(x::Logarithmic, y::Logarithmic)
    if x.signbit
        if y.signbit
            y.abs ≤ x.abs
        else
            true
        end
    else
        if y.signbit
            iszero(x) && iszero(y)
        else
            x.abs ≤ y.abs
        end
    end
end

function Base.cmp(x::ULogarithmic, y::ULogarithmic)
    cmp(x.log, y.log)
end

function Base.isless(x::ULogarithmic, y::ULogarithmic)
    isless(x.log, y.log)
end

function Base.isless(x::Logarithmic, y::Logarithmic)
    if x.signbit
        if y.signbit
            isless(y.abs, x.abs)
        else
            true
        end
    else
        if y.signbit
            false
        else
            isless(x.abs, y.abs)
        end
    end
end

function Base.nextfloat(x::ULogarithmic)
    uexp(nextfloat(x.log))
end

function Base.nextfloat(x::Logarithmic)
    if x.signbit && !iszero(x)
        Logarithmic(prevfloat(x.abs), true)
    else
        Logarithmic(nextfloat(x.abs))
    end
end

function Base.prevfloat(x::ULogarithmic)
    uexp(prevfloat(x.log))
end

function Base.prevfloat(x::Logarithmic)
    if x.signbit || iszero(x)
        Logarithmic(nextfloat(x.abs), true)
    else
        Logarithmic(prevfloat(x.abs))
    end
end

### Promotion

_promote_rule(::Type, ::Type) = Union{}
for (i, A) in enumerate([ULogarithmic, Logarithmic])
    for (j, B) in enumerate([ULogarithmic, Logarithmic])
        C = i > j ? A : B
        @eval begin
            _promote_rule(::Type{$A}, ::Type{$B}) = $C
            _promote_rule(::Type{$A}, ::Type{$B{T}}) where {T} = $C{T}
            _promote_rule(::Type{$A{S}}, ::Type{$B}) where {S} = $C{S}
            _promote_rule(::Type{$A{S}}, ::Type{$B{T}}) where {S,T} = $C{promote_type(S,T)}
        end
    end
end

function Base.promote_rule(::Type{T}, ::Type{R}) where {T<:AnyLogarithmic, R<:AnyLogarithmic}
    _promote_rule(T, R)
end

@generated function Base.promote_rule(::Type{T}, ::Type{R}) where {T<:AnyLogarithmic, R<:Real}
    # TODO: Think about this some more. Always return ULogarithmic? Always Logarithmic?
    isunsigned = try
        typemin(R) ≥ 0
    catch
        false
    end
    L = isunsigned ? ULogarithmic : Logarithmic
    R2 = try
        typeof(L(one(R)))
    catch
        Union{}
    end
    promote_type(T, R2)
end

### Arithmetic

Base.:(+)(x::AnyLogarithmic) = x

Base.:(-)(x::ULogarithmic) = Logarithmic(x, true)
Base.:(-)(x::Logarithmic) = Logarithmic(x.abs, !x.signbit)

function Base.:(+)(x::T, y::T) where {T<:ULogarithmic}
    if x.log == y.log
        uexp(x.log + log1p(exp(zero(y.log) - zero(x.log))))
    elseif x.log ≥ y.log
        uexp(x.log + log1p(exp(y.log - x.log)))
    else
        uexp(y.log + log1p(exp(x.log - y.log)))
    end
end

function Base.:(+)(x::T, y::T) where {T<:Logarithmic}
    if x.signbit==y.signbit
        Logarithmic(x.abs + y.abs, x.signbit)
    elseif x.abs ≥ y.abs
        Logarithmic(x.abs - y.abs, x.signbit)
    else
        Logarithmic(y.abs - x.abs, y.signbit)
    end
end

function Base.:(-)(x::T, y::T) where {T<:ULogarithmic}
    if x.log == y.log
        uexp(x.log + log1p(-exp(zero(y.log) - zero(x.log))))
    elseif x.log ≥ y.log
        uexp(x.log + log1p(-exp(y.log - x.log)))
    else
        throw(DomainError((x, y), "difference is negative"))
    end
end

function Base.:(-)(x::T, y::T) where {T<:Logarithmic}
    if x.signbit == y.signbit
        if x.abs ≥ y.abs
            Logarithmic(x.abs - y.abs, x.signbit)
        else
            Logarithmic(y.abs - x.abs, !y.signbit)
        end
    else
        Logarithmic(x.abs + y.abs, x.signbit)
    end
end

function Base.:(*)(x::T, y::T) where {T<:ULogarithmic}
    uexp(x.log + y.log)
end

function Base.:(*)(x::T, y::T) where {T<:Logarithmic}
    Logarithmic(x.abs * y.abs, x.signbit ⊻ y.signbit)
end

function Base.:(/)(x::T, y::T) where {T<:ULogarithmic}
    uexp(x.log - y.log)
end

function Base.:(/)(x::T, y::T) where {T<:Logarithmic}
    Logarithmic(x.abs / y.abs, x.signbit ⊻ y.signbit)
end

function Base.:(^)(x::ULogarithmic, n::Real)
    if n == 0
        uexp(zero(x.log) * n)
    else
        uexp(x.log * n)
    end
end

function Base.:(^)(x::ULogarithmic, n::Integer)
    if n == 0
        uexp(zero(x.log) * n)
    else
        uexp(x.log * n)
    end
end

function Base.:(^)(x::Logarithmic, n::Integer)
    Logarithmic(x.abs^n, x.signbit & isodd(n))
end

function Base.:(^)(x::Logarithmic, n::Real)
    x.signbit && !iszero(x) && throw(DomainError(x))
    Logarithmic(x.abs^n)
end

function Base.inv(x::ULogarithmic)
    uexp(-x.log)
end

function Base.inv(x::Logarithmic)
    Logarithmic(inv(x.abs), x.signbit)
end

function Base.log(x::ULogarithmic)
    x.log
end

function Base.log(x::Logarithmic)
    x.signbit && !iszero(x) && throw(DomainError(x))
    log(x.abs)
end

function Base.exp(x::ULogarithmic)
    uexp(exp(x.log))
end

function Base.exp(x::Logarithmic)
    x.signbit ? inv(exp(x.abs)) : exp(x.abs)
end

### Random

function Base.rand(rng::AbstractRNG, ::Random.SamplerType{E}) where {T<:AbstractFloat, E<:AnyLogarithmic{T}}
    exp(E, -randexp(rng, T))
end

function Base.rand(rng::AbstractRNG, ::Random.SamplerType{E}) where {E<:AnyLogarithmic}
    exp(E, -randexp(rng))
end

### IO

function Base.show(io::IO, x::ULogarithmic)
    print(io, "exp(")
    show(io, x.log)
    print(io, ")")
end

function Base.show(io::IO, x::Logarithmic)
    print(io, x.signbit ? "-" : "+")
    show(io, x.abs)
end

function Base.write(io::IO, x::ULogarithmic)
    write(io, x.log)
end

function Base.write(io::IO, x::Logarithmic)
    write(io, x.abs, x.signbit)
end

function Base.read(io::IO, ::Type{ULogarithmic{T}}) where {T}
    uexp(T, read(io, T))
end

function Base.read(io::IO, ::Type{Logarithmic{T}}) where {T}
    abs = read(io, ULogarithmic{T})
    signbit = read(io, Bool)
    Logarithmic{T}(abs, signbit)
end

_overload(mod, f, logf) = quote
    if isdefined($mod, $(QuoteNode(f))) && isdefined($mod, $(QuoteNode(logf)))
        @inline function $mod.$f(::Type{T}, args...; opts...) where {T<:AnyLogarithmic}
            exp(T, $mod.$logf(args...; opts...))
        end
    end
end

_overloads(mod, fs) = :($([_overload(mod, f, logf) for (f, logf) in fs]...);)

const DISTRIBUTIONS_OVERLOADS = [f => Symbol(:log, f) for f in [:cdf, :ccdf, :pdf]]
const STATSFUNS_DISTS = [:srdist, :nchisq, :hyper, :ntdist, :tdist, :binom, :pois, :fdist, :norm, :beta, :nfdist, :chisq, :gamma, :nbeta, :nbinom]
const STATSFUNS_OVERLOADS = [Symbol(d, f) => Symbol(d, :log, f) for d in STATSFUNS_DISTS for f in [:cdf, :ccdf, :pdf]]
const SPECIALFUNCTIONS_OVERLOADS = [f => Symbol(:log, f) for f in [:gamma, :factorial, :beta, :erfc, :erfcx]]

@eval function __init__()
    @require Distributions="31c24e10-a181-5473-b8eb-7969acd0382f" $(_overloads(:Distributions, DISTRIBUTIONS_OVERLOADS))
    @require StatsFuns="4c63d2b9-4356-54db-8cca-17b64c39e42c" $(_overloads(:StatsFuns, STATSFUNS_OVERLOADS))
    @require SpecialFunctions="276daf66-3868-5448-9aa4-cd146d93841b" $(_overloads(:SpecialFunctions, SPECIALFUNCTIONS_OVERLOADS))
    # TODO: Rmath
end

end

ext: jl


text:
function gen_king_valid(source_square::UInt64)
    target_squares = zeros(UInt64, 0)

    for cms in zip(KING_MASK_FILES, KING_CLEAR_FILES, KING_SHIFTS)
        if source_square & cms[1] != EMPTY
            candidate_square = (source_square & cms[2]) << cms[3]
        else
            candidate_square = source_square << cms[3]
        end

        if candidate_square != EMPTY
            push!(target_squares, candidate_square)
        end
    end
    return target_squares
end

function gen_all_king_valid_moves()
    king_moves = Dict{UInt64, Array{UInt64,1}}()
    for i in 1:64
        king_moves[INT2UINT[i]] = gen_king_valid(INT2UINT[i])
    end
    return king_moves
end
const KING_MOVES = gen_all_king_valid_moves()
const WHITE_KING_HOME = PGN2UINT["e1"]
const BLACK_KING_HOME = PGN2UINT["e8"]


text:
# This is tools for testing ChainRulesTestUtils itself
# We might consider moving them to a package MetaTesting.jl

"""
    EncasedTestSet(desc, results) <: AbstractTestset

A custom testset that encases all test results within, not letting them out.
It doesn't let anything propagate up to the parent testset
(or to the top-level fallback testset, which throws an error on any non-passing result).
Not passes, not failures, not even errors.

This is useful for being able to observe the testsets results programatically;
without them triggering actual passes/failures/errors.
"""
struct EncasedTestSet <: Test.AbstractTestSet
    description::String
    results::Vector{Any}
end
EncasedTestSet(desc) = EncasedTestSet(desc, [])

Test.record(ts::EncasedTestSet, t) = (push!(ts.results, t); t)

function Test.finish(ts::EncasedTestSet)
    if Test.get_testset_depth() != 0
        # Attach this test set to the parent test set *if* it is also a NonPassingTestset
        # Otherwise don't as we don't want to push the errors and failures further up.
        parent_ts = Test.get_testset()
        parent_ts isa EncasedTestSet && Test.record(parent_ts, ts)
        return ts
    end
    return ts
end

"""
    nonpassing_results(f)

`f` should be a function that takes no argument, and calls some code that used `@test`.
Invoking it via `nonpassing_results(f)` will prevent those `@test` being added to the
current testset, and will return a collection of all nonpassing test results.
"""
function nonpassing_results(f)
    # Specify testset type to hijack system
    ts = @testset EncasedTestSet "nonpassing internal" begin
        f()
    end
    return _extract_nonpasses(ts)
end

"extracts as flat collection of failures from a (potential nested) testset"
_extract_nonpasses(x::Test.Result) = [x]
_extract_nonpasses(x::Test.Pass) = Test.Result[]
_extract_nonpasses(ts::EncasedTestSet) = _extract_nonpasses(ts.results)
function _extract_nonpasses(xs::Vector)
    if isempty(xs)
        return Test.Result[]
    else
        return mapreduce(_extract_nonpasses, vcat, xs)
    end
end

"""
    fails(f)

`f` should be a function that takes no argument, and calls some code that used `@test`.
`fails(f)` returns true if at least 1 `@test` fails.
If a test errors then it will display that error and throw an error of its own.
"""
function fails(f)
    results = nonpassing_results(f)
    did_fail = false
    for result in results
        did_fail |= result isa Test.Fail
        if result isa Test.Error
            # Log a error message, with original backtrace
            # Sadly we can't throw the original exception as it is only stored as a String
            error("Error occurred during `fails`")
        end
    end
    return did_fail
end

"""
    errors(f, msg_pattern="")

Returns true if at least 1 error is recorded into a testset
with a failure matching the given pattern.

`f` should be a function that takes no argument, and calls some code that uses `@testset`.
`msg_pattern` is a regex or a string, that should be contained in the error message.
If nothing is passed then it default to the empty string, which matches any error message.

If a test fails (rather than passing or erroring) then `errors` will throw an error.
"""
function errors(f, msg_pattern="")
    results = nonpassing_results(f)

    for result in results
        result isa Test.Fail && error("Test actually failed (not errored): \n $result")
        result isa Test.Error && occursin(msg_pattern, result.value) && return true
    end
    return false  # no matching error occured
end

#Meta Meta tests
@testset "meta_testing_tools.jl" begin
    @testset "Checking for non-passes" begin
        @testset "No Tests" begin
            fails = nonpassing_results(() -> nothing)
            @test length(fails) === 0
        end

        @testset "No Failures" begin
            fails = nonpassing_results(() -> @test true)
            @test length(fails) === 0
        end

        @testset "Single Test" begin
            fails = nonpassing_results(() -> @test false)
            @test length(fails) === 1
            # Julia 1.6 return a `String`, not an `Expr`.
            # Always calling  `string` on it gives gives consistency regardless of version.
            # https://github.com/JuliaLang/julia/pull/37809
            @test string(fails[1].orig_expr) == string(false)
        end

        @testset "Single Testset" begin
            fails = nonpassing_results() do
                @testset "inner" begin
                    @test false == true
                    @test true == false
                end
            end
            @test length(fails) === 2

            # Julia 1.6 return a `String`, not an `Expr`.
            # Always calling  `string` on it gives gives consistency regardless of version.
            # https://github.com/JuliaLang/julia/pull/37809
            @test string(fails[1].orig_expr) == string(:(false == true))
            @test string(fails[2].orig_expr) == string(:(true == false))
        end

        @testset "Single Error" begin
            bads = nonpassing_results(() -> error("noo"))
            @test length(bads) === 1
            @test bads[1] isa Test.Error
        end

        @testset "Single Test Erroring" begin
            bads = nonpassing_results(() -> @test error("nooo"))
            @test length(bads) === 1
            @test bads[1] isa Test.Error
        end

        @testset "Single Testset Erroring" begin
            bads = nonpassing_results() do
                @testset "inner" begin
                    error("noo")
                end
            end
            @test length(bads) === 1
            @test bads[1] isa Test.Error
        end
    end

    @testset "fails" begin
        @test !fails(() -> @test true)
        @test fails(() -> @test false)
        @test !fails(() -> @test_broken false)

        @test fails() do
            @testset "eg" begin
                @test true
                @test false
                @test true
            end
        end

        @test_throws ErrorException fails(() -> @test error("Bad"))
    end

    @testset "errors" begin
        @test !errors(() -> @test true)
        @test errors(() -> error("nooo"))
        @test errors(() -> error("nooo"), "noo")
        @test !errors(() -> error("nooo"), "ok")

        @test errors() do
            @testset "eg" begin
                @test true
                error("nooo")
                @test true
            end
        end

        @test_throws ErrorException errors(() -> @test false)
    end
end


text:
## support routines

function clopts(opts...)
    args = ["", opts...]
    API.LLVMParseCommandLineOptions(length(args), args, C_NULL)
end


text:
"""
decay_curve = EPGdecaycurve_vTE(ETL, flip_angle, TE1, TE2, nTE1, T2, T1, refcon)

Computes the normalized echo decay curve for a MR spin echo sequence with
the given parameters.

ETL:        Echo train length (number of echos)
flip_angle: Angle of refocusing pulses (degrees)
TE1:        First Interecho time (seconds)
TE2:        second Interecho time (seconds)
nTE1:       number of echoes at first echo time
T2:         Transverse relaxation time (seconds)
T1:         Longitudinal relaxation time (seconds)
refcon:     Value of Refocusing Pulse Control Angle
"""
function EPGdecaycurve_vTE_work(T, ETL)
    error("Placeholder function --- not implemented")
end
EPGdecaycurve_vTE_work(ETL) = EPGdecaycurve_vTE_work(Float64, ETL)

EPGdecaycurve_vTE(ETL::Int, flip_angle::T, TE1::T, TE2::T, nTE1::Int, T2::T, T1::T, refcon::T) where {T} =
    EPGdecaycurve_vTE!(EPGdecaycurve_vTE_work(T, ETL), ETL, flip_angle, TE1, TE2, nTE1, T2, T1, refcon)

function EPGdecaycurve_vTE!(work, ETL::Int, flip_angle::T, TE1::T, TE2::T, nTE1::Int, T2::T, T1::T, refcon::T) where {T}
    error("Placeholder function --- not implemented")
end


text:
module SegwayRobot4

    include("src4/segway.jl")
    export Segway_mutable, Segway_unmutable

    include("src4/dynamics_type.jl")
    include("src4/dynamics_vec.jl")
    #export segway_rk4, segway_rk4!, dxdt_segway
    #x = load("SegwayRobot4/Julia/src/params.jld2")

    const s_init  = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    const u_range = [1.0, 1.0]
    const params_segway = [-1.14, -5.09, -4.97, -4.59, -6.68]

    include("src4/dynmodel.jl")

    cost_func = segway_cost # set standard Environment cost-function

    function rk4_func(xt::AbstractVector, ut::AbstractVector, dt::Number)
        xt1 = copy(xt)
        dynmodel!(xt1, ut, params_segway)

        return xt1
    end

    export rk4_func, cost_func, s_init, u_range
end # module SegwayRobot47

max_stars_repo_name: NikonPic/REINFORCEMENT-JULIA

max_issues_repo_name: NikonPic/REINFORCEMENT-JULIA

max_forks_repo_name: NikonPic/REINFORCEMENT-JULIA


text:
const LPHyperParamConfig = NamedTuple{(:M, :λ),Tuple{Float64,Float64}}

struct LPRelaxation <: Algorithm
    optimizer_factory
    regularizer::Union{Nothing,ObjectivePenalty}
    hyperparameter_set::Array{LPHyperParamConfig}
end
function LPRelaxation(optimizer_factory; regularizer=nothing)
    return LPRelaxation(
        optimizer_factory,
        regularizer,
        [(M = 2.0^M_exp_val, λ = λ_val) for λ_val in (0.0,), M_exp_val in -1:9]
    )
end

has_regularizer(algo::LPRelaxation) = algo.regularizer !== nothing
get_regularizer(algo::LPRelaxation) = algo.regularizer

function train(learner::Learner{LPRelaxation}; hypograph::Bool=false)
    formulation_fn = hypograph ? formulate_data_point_hypograph_big_m : formulate_data_point_graph_big_m
    if learner.algorithm.regularizer !== nothing
        _add_objective_penalty!(learner, learner.algorithm.regularizer)
    end
    hyperparam_results = Dict{Any,TrainingResults}()
    for (M, λ) in learner.algorithm.hyperparameter_set
        learner.param_lb .= -M
        learner.param_ub .= M
        learner.offset_lb = -M
        learner.offset_ub = M
        if λ > 0
            # TODO: Refactor so that this is less hacky
            push!(learner.objective_penalties, L2Regularizer(λ))
        end
        opt_model = formulate_big_m(learner, formulation_fn, relax_integrality=true)
        jump_model = opt_model.jump_model
        solved_opt_model = optimize!(learner, opt_model)
        @assert solved_opt_model.jump_model == jump_model
        primal_status = JuMP.primal_status(jump_model)
        if primal_status != MOI.FEASIBLE_POINT
            @error "MIP did not produce a feasible solution"
        end
        termination_status = JuMP.termination_status(jump_model)
        if termination_status != MOI.OPTIMAL
            @warn "LP relaxation not solved to optimality (status = $termination_status)"
        end
        _audit_weights(learner, solved_opt_model.linear_model.weights)
        hyperparam_results[(M, λ)] = TrainingResults(learner, solved_opt_model.linear_model, reward_bound=JuMP.objective_value(solved_opt_model.jump_model))
    end
    best_validation_reward = -Inf
    best_validation_key = nothing

    iter_results = IterationResult[]
    for (key, result) in hyperparam_results
        reward = result.final_validation_eval.reward
        if reward > best_validation_reward
            best_validation_reward = reward
            best_validation_key = key
        end
        push!(iter_results, IterationResult(result.final_is_eval, result.final_os_eval, nothing, Dict("key" => key, "validation_eval" => result.final_validation_eval)))
    end
    best_result = hyperparam_results[best_validation_key]
    best_result.iteration_evals = iter_results
    return best_result
end

max_stars_repo_name: joehuchette/reserve-price-optimization

max_issues_repo_name: joehuchette/reserve-price-optimization

max_forks_repo_name: joehuchette/reserve-price-optimization


text:
## covariance_functions.jl : utilities for Gaussian random field covariance functions

## CovarianceStructure ##
abstract type CovarianceStructure{T} end
abstract type IsotropicCovarianceStructure{T} <: CovarianceStructure{T} end
abstract type AnisotropicCovarianceStructure{T} <: CovarianceStructure{T} end

## CovarianceFunction ##
struct CovarianceFunction{d,T}
    cov::T
end

"""
	CovarianceFunction(d, cov)

Create a covariance function in `d` dimensions with covariance structure `cov`.

# Examples
```jldoctest
julia> m = Matern(0.1,1.0)
Matérn (λ=0.1, ν=1.0, σ=1.0, p=2.0)

julia> c = CovarianceFunction(2,m)
2d Matérn covariance function (λ=0.1, ν=1.0, σ=1.0, p=2.0)

```
"""
function CovarianceFunction(d::N where {N<:Integer},cov::T) where {T<:CovarianceStructure}
    d > 0 || throw(ArgumentError("dimension must be positive, got $(d)"))
    CovarianceFunction{d,T}(cov)
end

# return standard deviation of the Gaussian random field
std(cov::CovarianceFunction) = cov.cov.σ

# return number of dimension
ndims(::CovarianceFunction{d}) where {d} = d

# evaluate the covariance function
apply(cov::CovarianceFunction,x,y) = apply(cov.cov,x,y)

# apply for isotropic random fields
apply(cov::IsotropicCovarianceStructure,dx::Vector{T} where {T<:Real}) = apply(cov,sum(abs.(dx).^cov.p).^(1/cov.p))

# evaluate when pts is given as a kron product of 1d points
function apply(cov::CovarianceStructure{T}, x::Tuple, y::Tuple) where {T<:Real}
    C = zeros(T,prod(length.(x)),prod(length.(y)))
    if  size(C,1) == size(C,2)
        return apply_symmetric(cov,x,y,C)
    else
        return apply_non_symmetric(cov,x,y,C)
    end
end

function apply_symmetric(cov::CovarianceStructure{T}, x::Tuple, y::Tuple, C::Matrix{T}) where {T<:Real}
    for (j,idy) in enumerate(Base.product(y...))
        for (i,idx) in enumerate(Base.product(x...))
            if i <= j
	        @inbounds C[i,j] = apply(cov,collect(idx.-idy))
            end
        end
    end
    Symmetric(C,:U)
end

function apply_non_symmetric(cov::CovarianceStructure{T}, x::Tuple, y::Tuple, C::Matrix{T}) where {T<:Real}
    for (j,idy) in enumerate(Base.product(y...))
        for (i,idx) in enumerate(Base.product(x...))
	    @inbounds C[i,j] = apply(cov,collect(idx.-idy))
        end
    end
    C
end

# evaluate when pts is given as a Finite Element mesh
function apply(cov::CovarianceStructure{T}, tx::Tuple{T1,T2}, ty::Tuple{T1,T2}) where {T<:Real,T1<:AbstractMatrix,T2<:AbstractMatrix}
    x = first(tx) # select FE nodes
    y = first(ty)
    C = zeros(T,size(x,2),size(y,2))
    for j in 1:size(y,2), i in 1:size(x,2)
        if i <= j
            @inbounds C[i,j] = apply(cov,x[:,i].-y[:,j])
        end
    end
    return Symmetric(C,:U)
end

# evaluate for KL eigenfunctions
function apply(cov::CovarianceStructure{T}, tx::Tuple{T1,T2}, y::Tuple) where {T<:Real,T1<:AbstractMatrix,T2<:AbstractMatrix}
    x = first(tx) # select FE nodes
    C = zeros(T,size(x,2),prod(length.(y)))
    for (j,idy) in enumerate(Base.product(y...))
        for i in 1:size(x,2)
            @inbounds C[i,j] = apply(cov,x[:,i].-idy)
        end
    end
    C
end

function show(io::IO, c::CovarianceFunction{d}) where {d}
    str = split(string(c.cov)," ")
    print(io, "$(d)d $(str[1]) covariance function $(join(str[2:end]," "))")
end


text:
# %% imports

using PyCall
gym = pyimport("gym")

include("../../../Modules/Networks/net_arch.jl")
include("../../../Modules/Poisson/PoissonStateSpike.jl")
include("../../../Modules/Swingup/swingup.jl")

# %% sim functions

function count_spikes(arch::Array{Int64}, spiked::Array{Int64}, layer::Int64)
    count_layer = zeros(arch[layer])
    interval = [sum(arch[1:layer-1])+1, sum(arch[1:layer])]
    spikes_in = filter(x-> interval[1] <= x <= interval[2], spiked)
    spikes_in .-= sum(arch[1:layer-1])
    count_layer[spikes_in] .+= 1
    return count_layer
end

function left_right(count_spikes::Array{Float64}, n_left::Int64, n_right::Int64)
    sl = sum(count_spikes[1:n_left])
    sr = sum(count_spikes[n_left+1:end])
    (sl >= sr) ? action = 0 : action = 1
    return action, sl, sr
end

function borne(x,a,b)
    y = a + 0.5*(b-a)*(1+cos(pi*x/10))
    return y
end

function sig(z::Float64, alpha::Float64)
    return 1.0/(1.0+exp(-z/alpha))
end

function norm_obs(obs::Array{Float64})
    pos = obs[1]
    theta = obs[3]%pi
    vel = [obs[2], obs[4]]
    pos_scaled = (pos + 2.4)
    cone = 41.8*pi/180
    theta_scaled = (theta + cone)
    vel_scaled = (vel .+ 7.5)
    obs_scaled = [pos_scaled, vel_scaled[1], theta_scaled, vel_scaled[2]]
    return obs_scaled
end

function norm_sig(obs::Array{Float64})
    pos = obs[1] # in [-2.4,2.4]
    theta = obs[3]%pi # in [-pi, pi]
    vel = [obs[2], obs[4]] # in [-15,15]
    alpha_pos = 2.4/5
    alpha_theta = pi/5
    alpha_vel = 15/5
    pos_scaled = sig(pos,alpha_pos)
    theta_scaled = sig(theta,alpha_theta)
    vel_scaled = sig.(vel,alpha_vel)
    obs_scaled = [pos_scaled, vel_scaled[1], theta_scaled, vel_scaled[2]]
    return obs_scaled
end

function obs_input(obs::Array{Float64}, gain::Float64)
    input = Float64[]
    obs_s = norm_obs(obs)
    for idx in eachindex(obs_s)
        anti_obs = gain*(1.0-obs_s[idx])
        push!(input, obs_s[idx])
        push!(input, anti_obs)
    end
    return input
end

# %%

function play_episode(net::Network, n_steps::Int64 = 500)
    env = gym.make("CartPole-v1")
    env.seed(0)
    #obs = swingup_reset!(env)
    obs = env.reset()
    rng_action = Random.MersenneTwister(0)
    win_size = 100 # ms
    tot_reward = 0
    N_rand = 0
    N_tot = 0
    @inbounds for j in 1:n_steps
        count_tot = zeros(net.arch[end])
        input = obs_input(obs, 1.0)
        spikes = int_spikes(input) # state2spike function
        #reset voltage
        net.v .= net.param["c"]
        influence = 2
        @inbounds for msec in 1:win_size
            spiked = step_LIF!(net, spikes, false)
            if msec >= length(net.arch)-1
                count_tot .+= influence*count_spikes(net.arch, spiked, length(net.arch))
                influence *= 0.995
                #push!(spike_hist,spiked)
            end
        end
        n_left = div(net.arch[end],2)
        n_right = net.arch[end] - n_left
        action, sl, sr = left_right(count_tot, n_left, n_right)
        if sl == sr
            action = rand(rng_action, 0:1)
            tot_reward -= 0.9
            N_rand += 1
        end
        N_tot += 1
        #push!(score_hist[1], sl)
        #push!(score_hist[2], sr)
        #obs_new, reward, done, _ = swingup_step(env,action)
        obs_new, reward, done, _ = env.step(action)
        tot_reward += reward
        if done
            break
        end
        obs = obs_new
    end
    println("Rand = ", N_rand, "/", N_tot)
    env.close()
    env = nothing
    Base.GC.gc()
    return tot_reward
end

function set_weight!(net::Network, genes_w::Array{Float64})
    w = net.s.e
    for mat_idx in 1:length(w)
        for idx in eachindex(w[mat_idx])
            offset = sum([length(mat) for mat in w[1:mat_idx-1]])
            net.s.e[mat_idx][idx] = borne(genes_w[idx+offset], 0.0, 1.1)
        end
    end
end

function objective(genes::Array{Float64})
    #load the optimisation parameters
    optim_param = YAML.load(open("Julia/CartPole/LIF_base/cfg_cmaes.yml"))
    arch = optim_param["arch"]
    genes_w = genes

    params = YAML.load(open("Julia/CartPole/LIF_base/noalpha/cfglifnoalpha.yml"))
    n = Network(arch, params)
    set_weight!(n,genes_w)
    reward = play_episode(n)
    return -reward
end

max_stars_repo_name: SuReLI/RL-STDP

max_issues_repo_name: SuReLI/RL-STDP

max_forks_repo_name: SuReLI/RL-STDP


text:
immutable CatView{N,T<:Number} <: AbstractArray{T,1}
    arr::NTuple{N,SubArray{T}}
    len::NTuple{N,Integer}
    inner::NTuple{N}  # iterators for each array
end

## Constructors ##
@inline CatView(a::AbstractArray...) = CatView(a)

@generated function CatView{N,T}(arr::NTuple{N,SubArray{T}})
    quote
    len = @ntuple $N (n)->length(arr[n])
    inner = @ntuple $N (n)->eachindex(arr[n])
    CatView{N,T}(arr,len,inner)
    end
end

@generated function CatView{N,T}(arr::NTuple{N,AbstractArray{T}})
    quote
    CatView(@ntuple $N (n)->view(arr[n],:))
    end
end

## size ##
Base.size(A::CatView) = (sum(A.len),)

## get index and set index ##

function Base.getindex(A::CatView, i::Int)
    i < 1 || i > length(A) && throw(BoundsError("index out of bounds."))

    a = 0
    b = A.len[1]
    for j = 1:length(A.len)
        if i <= b
            return A.arr[j][i-a]
        else
            a = b
            b = b + A.len[j+1]
        end
    end
end

function Base.setindex!(A::CatView, val, i::Int)
    i < 1 || i > length(A) && throw(BoundsError("index out of bounds."))

    a = 0
    b = A.len[1]
    for j = 1:length(A.len)
        if i <= b
            A.arr[j][i-a] = val
            return val
        else
            a = b
            b = b + A.len[j+1]
        end
    end
end

function Base.getindex(A::CatView, idx::Tuple{Integer,Integer})
    i,j = idx
    return A.arr[i][j]
end

function Base.setindex!(A::CatView, val, idx::Tuple{Integer,Integer})
    i,j = idx
    return setindex!(A.arr[i], val, j)
end

## Fast iteration ##
@generated function Base.eachindex{N,T}(A::CatView{N,T})
    quote
    @nexprs $N (n)->(i_n = zip(repeated(n,length(A.arr[n])),eachindex(A.arr[n])))
    chain( (@ntuple $N (n)->i_n)... )
    end
end


text:
###############################################
#FEniCS tutorial demo program: Nonlinear Poisson equation.
#  -div(q(u)*grad(u)) = f   in the unit square.
#                   u = u_D on the boundary.
###############################################

module ft05

using SymPy
using FenicsPy

# nonlinear coefficient
q(u) = 1 + u^2

# Use SymPy to compute f from the manufactured solution u
x, y = symbols("x[0], x[1]")
u = 1 + x + 2*y
f = - diff(q(u)*diff(u, x), x) - diff(q(u)*diff(u, y), y)
f = simplify(f)
u_code = sympy.printing.ccode(u)
f_code = sympy.printing.ccode(f)
println("u =", u_code)
println("f =", f_code)

# Create mesh and define function space
mesh = UnitSquareMesh(8, 8)
V = FunctionSpace(mesh, "P", 1)

# Define boundary condition
u_D = Expression(u_code, degree=2)

bc = DirichletBC(V, u_D, "on_boundary")

# Define variational problem
u = FeFunction(V)  # Note: not TrialFunction!
v = TestFunction(V)
f = Expression(f_code, degree=2)
F = q(u)*dot(grad(u), grad(v))*dx - f*v*dx

# Compute solution
# both SymPy and FenicsPy export "solve"; uses of it in module Main must be qualified
FenicsPy.solve(F == 0, u, bc)

# Plot solution
#plot(u)

# Compute maximum error at vertices. This computation illustrates
# an alternative to using compute_vertex_values as in poisson.py
# both SymPy and FenicsPy export "interpolate"; uses of it in module Main must be qualified.
u_e = FenicsPy.interpolate(u_D, V)
error_max = max(abs.(array(u_e) - array(u))...)
print("error_max = ", error_max)

end # module ft05


text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã 2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

module SortByLikelihood

export sortbylikelihood

using ..DefineSave
using ..SaveHelpers

"""
Input directory indir of
"""
function sortbylikelihood{T<:AbstractString}(files::Vector{T};
    outfile::T="sort_results.txt",
    descending::Bool=true)
    scores = map(files) do f
        d = trajLoad(f)
        likelihood = sum(sv_simlog_tdata_vid(d, "logProb", "logProb"))
    end
    perm = sortperm(scores, rev=descending) #descending order
    fout = open(outfile, "w")
    for p in perm
        println(fout, files[p], ", ", scores[p])
    end
    close(fout)
    outfile
end

end

max_stars_repo_name: sisl/RLESCAS

max_issues_repo_name: sisl/RLESCAS

max_forks_repo_name: sisl/RLESCAS


text:
##############################################################################
#                                Tensors                                     #
##############################################################################

"""
    from_nnlib(x)

Convert argument from NNlib-friendly to ONNX-friendly format.
The reverse operation is available as [`from_onnx`](@ref).

See also: [`from_onnx_conv`](@ref), [`from_onnx_spatial`](@ref),
[`from_onnx_conv`](@ref) and similar.
"""
from_nnlib(x::AbstractArray) = permutedims(x, ndims(x):-1:1)

"""
    from_onnx(x)

The reverse of [`from_nnlib`](@ref).
"""
function from_onnx(x::AbstractArray)
    if ndims(x) > 0
        permutedims(x, ndims(x):-1:1)
    else
        x
    end
end

##############################################################################
#                             Conv Attributes                                #
##############################################################################

"""
    from_nnlib_spatial(x, d)

Convert spatial attributes such as Conv's stride or dilation
from Julia to ONNX format.

`x` is a Julia values of that attribute
`d` is the number of spatial dimensions
"""
from_nnlib_spatial(x::Int, d::Int) = [x for i=1:d]
from_nnlib_spatial(x::Tuple, d::Int) = collect(reverse(x))

from_onnx_spatial(x) = x
from_onnx_spatial(x::AbstractVector) = Tuple(reverse(x))

function from_nnlib_pad(pad::Int, N::Int)
    return Tuple([pad for i=1:2N])
end

function from_nnlib_pad(pad::NTuple{N, T}, d::Int) where {N, T}
    @assert(N == d || N == 2d,
        "Padding should be a tuple of either `N` or 2N elements where N is the number " *
        "of spatial dimensions, but got `pad = $pad` and N = $d")
    pad = collect(pad)
    if N != 2d
        # e.g. [1, 2] -> [1, 1, 2, 2]
        pad = repeat(pad, inner=2)
    end
    # notation: (1, 2, 3) - dimensions; b - beginning, e - end of dimension
    # our goal (for 3D case): [b1, e1, b2, e2, b3, e3] -> [b3, b2, b1, e3, e2, e1]
    d = length(pad) ÷ 2
    return [pad[2d-1:-2:1]; pad[2d:-2:2]]
end

function from_onnx_pad(pad::Vector{Int})
    # notation: (1, 2, 3) - dimensions; b - beginning, e - end of dimension
    # our goal (for 3D case): [b3, b2, b1, e3, e2, e1] -> [b1, e1, b2, e2, b3, e3]
    d = length(pad) ÷ 2
    # step 1: [b3, b2, b1, e3, e2, e1] -> [b1, b2, b3, e1, e2, e3]
    out = [pad[d:-1:1]; pad[2d:-1:d+1]]
    # step 2: [b1, b2, b3, e1, e2, e3] -> [(b1, e1), (b2, e2), (b3, e3)]
    out = [(out[i], out[d + i]) for i=1:d]
    # step 3: [(b1, e1), (b2, e2), (b3, e3)] -> [b1, e1, b2, e2, b3, e3]
    out = reduce((x, y) -> [x...; y...], out)
    return Tuple(out)
end

function from_nnlib_conv(attrs::Dict, d::Int)
    out = Dict{Symbol, Any}()
    if haskey(attrs, :stride)
        out[:strides] = from_nnlib_spatial(attrs[:stride], d)
    end
    if haskey(attrs, :dilation)
        out[:dilations] = from_nnlib_spatial(attrs[:dilation], d)
    end
    if haskey(attrs, :kernel)
        out[:kernel_shape] = from_nnlib_spatial(attrs[:kernel], d)
    end
    if haskey(attrs, :groups)
        out[:group] = attrs[:groups]
    end
    if haskey(attrs, :pad)
        out[:pads] = from_nnlib_pad(attrs[:pad], d)
    end
    return out
end

function from_onnx_conv(attrs::Dict)
    out = Dict{Symbol, Any}()
    haskey(attrs, :auto_pad) && error("auto_pad attribute is currently not supported")
    if haskey(attrs, :strides)
        out[:stride] = from_onnx_spatial(attrs[:strides])
    end
    if haskey(attrs, :dilations)
        out[:dilation] = from_onnx_spatial(attrs[:dilations])
    end
    if haskey(attrs, :kernel_shape)
        out[:kernel] = from_onnx_spatial(attrs[:kernel_shape])
    end
    if haskey(attrs, :group)
        out[:groups] = attrs[:group]
    end
    if haskey(attrs, :pads)
        out[:pad] = from_onnx_pad(attrs[:pads])
    end
    return out
end


text:
module BasisDefinition
using ..DynamicDefinition

import Base

export Basis, DualComposedWithDynamic, ProjectDualElement, AverageZero, assemble, integral_covector, one_vector, is_integral_preserving, strong_norm, weak_norm, aux_norm

abstract type Basis end

length(B::Basis) = @error "Not Implemented"

struct DualComposedWithDynamic{B<:Basis, D<:Dynamic}
	basis::B
	dynamic::D
	ϵ::Float64
end

Base.iterate(S::DualComposedWithDynamic{B, D}, state) where {B<:Basis, D<:Dynamic} = @error "Not implemented"

struct ProjectDualElement{B<:Basis}
	basis::B
	j_min
	j_max
	dual_element
end
ProjectDualElement(basis::B, j_min, j_max, y) where {B} = ProjectDualElement{B}(basis, j_min, j_max, y)

nonzero_on(B::Basis, I) = @error "Not Implemented"

function ProjectDualElement(B::Basis, y)
 	j_min, j_max = nonzero_on(B, y)
 	return ProjectDualElement(B, j_min, j_max, y)
end

Base.iterate(S::ProjectDualElement{B}, state) where {B} = @error "Not Implemented"

evaluate(B::Basis, i, x) = @error "Not Implemented"
evaluate_integral(B::Basis, i; T = Float64) = @error "Not Implemented"

strong_norm(B::Basis) = @error "Must be specialized"
weak_norm(B::Basis) = @error "Must be specialized"
aux_norm(B::Basis) = @error "Must be specialized"

"""
Check if Bfine is a refinement of Bcoarse
"""
is_refinement(Bfine::Basis, Bcoarse::Basis) = @error "Not implemented"

"""
	Covector that represents the integral
"""
integral_covector(B::Basis) = @error "Must be specialized"

"""
	Vector that represents the function 1
"""
one_vector(B::Basis) = @error "Must be specialized"

"""
	Integral-preserving discretizations may specialize this to "true"
"""
is_integral_preserving(B::Basis) = false

"""
Integral of a function in U_h

Args:
	v (any type of vector):

Returns:
	the integral, computed with the arithmetic of v.
"""
function integral(B::Basis, v; T = Float64)
	return sum([T(v[i])*evaluate_integral(B, i, T) for i in 1:length(B)])
end

"""
	Yields a basis of the space of average zero vectors
"""
struct AverageZero{B<:Basis}
	basis::B
end

Base.iterate(S::AverageZero{B}, state) where {B} = @error "Not Implemented"
Base.length(S::AverageZero) = length(S.basis)-1
#Base.zeros(B::Basis) = zeros(length(B))

"""
	Return a constant Kh (typically scales as h ~ 1/n) such that `||P_h f-f||\\leq Kh ||f||_s`
	Must be rounded up correctly!

	Arg:
		B::Basis
"""
weak_projection_error(B::Basis) = @error "Not implemented"

"""
	Return a constant Eh (typically scales as h ~ 1/n) such that `|||P_h f|||\\leq |||f|||+ Eh * ||f||_s`
	Must be rounded up correctly!

	Arg:
		B::Basis
"""
aux_normalized_projection_error(B::Basis) = @error "Not implemented"

"""
	Return a constant M₁n such that for a vector v in Uₕ `||v||_s\\leq M1n*||v||`
	Must be rounded up correctly!
"""
strong_weak_bound(B::Basis) = @error "Not implemented"

"""
	Return a constant M₂ such that for a vector v in Uₕ `|||v|||\\leq M_2||v||`
	Must be rounded up correctly!
"""
aux_weak_bound(B::Basis) = @error "Not implemented"

"""
	Return constants S₁, S₂ such that for a vector v in Uₕ `||v||\\leq S_1||v||_s+S_2|||v|||`
	Must be rounded up correctly!
"""
weak_by_strong_and_aux_bound(B::Basis) = @error "Not implemented"

"""
	Return constants W₁, W₂ such that for a vector v in Uₕ `||v||\\leq W_1||v||_1+W_2||v||_{\\infty}`
	Must be rounded up correctly!
"""
bound_weak_norm_from_linalg_norm(B::Basis) = @error "Not implemented"

"""
	Return a constant A such that for a vector v in Uₕ `||v||_1\\leq A||v||`
	Must be rounded up correctly!
"""
bound_linalg_norm_L1_from_weak(B::Basis) = @error "Not implemented"

"""
	Return a constant A such that for a vector v in Uₕ `||v||_\\infty \\leq A||v||`
	Must be rounded up correctly!
"""
bound_linalg_norm_L∞_from_weak(B::Basis) = @error "Not implemented"

"""
	Bounds ||u||_s, where ||u|| is the invariant measure normalized with integral(u)=1.
"""
invariant_measure_strong_norm_bound(B::Basis, D::Dynamic) = @error "Must be specialized"

end

max_stars_repo_name: orkolorko/ResponseManneville

max_issues_repo_name: orkolorko/ResponseManneville

max_forks_repo_name: orkolorko/ResponseManneville


text:
using RDatasets
using LowRankModels

# pick a data set
df = RDatasets.dataset("psych", "msq")

# initialize
glrm, labels = GLRM(df,2)

println("Fitting model with random initialization")
X, Y, ch = fit!(glrm)
println("final objective is ", objective(glrm))

println("Fitting model with svd initialization; fit should be faster and final solution (slightly) better")
init_svd!(glrm)
X, Y, ch = fit!(glrm)
println("final objective is ", objective(glrm))


text:
# pogm_restart.jl

using MIRT: pogm_restart

using LinearAlgebra: norm, opnorm, I
using Random: seed!
using Test: @test, @test_throws

# Tikhonov regularized LS problem
seed!(0); M = 30; N = 6; A = randn(M,N); y = randn(M)
a2 = opnorm(A)^2
reg = 0.1 * a2
xh = (A'A + reg*I) \ A'y
Fcost = (x) -> 1/2 * norm(A * x - y)^2
cost = (x) -> Fcost(x) + reg/2 * norm(x)^2
fun = (iter, x, y, is_restart) ->
	(cost(x) - cost(xh), norm(x - xh) / norm(xh), time())
f_grad = (x) -> A' * (A * x - y)
g_prox = (z,c) -> z / (1 + reg * c) # proximal operator for 2-norm
grad = (x) -> A' * (A * x - y) + reg * x
x0 = A \ y

@test_throws ArgumentError pogm_restart(x0, Fcost, f_grad, a2; mom=:bad)

x, _ = pogm_restart(x0, Fcost, f_grad, a2;
	f_mu=0, mom=:fpgm, niter=100, g_prox=g_prox, fun=fun)
pogm_restart(x0, Fcost, f_grad, a2;
	f_mu=0, mom=:pgm, niter=2, g_prox=g_prox)
pogm_restart(x0, Fcost, f_grad, a2;
	f_mu=reg, mom=:fpgm, niter=2, g_prox=g_prox) # with mu
pogm_restart(x0, Fcost, f_grad, a2;
	f_mu=reg, mom=:pgm, niter=2, g_prox=g_prox) # with mu

# todo: @inferred

x, out = pogm_restart(x0, Fcost, f_grad, a2;
	f_mu=0, mom=:pogm, restart=:gr, restart_cutoff=0.,
	bsig=1, niter=100, g_prox=g_prox, fun=fun)
@test x ≈ xh

x, out = pogm_restart(x0, Fcost, f_grad, a2;
	f_mu=reg, mom=:pogm, restart=:gr, restart_cutoff=0., # with mu
	bsig=1, niter=100, g_prox=g_prox, fun=fun)
@test x ≈ xh


text:
#---------------------------------------------------------
# # [2D dynamic](@id 3-2d-t)
#---------------------------------------------------------

#src 2019-06-13, Jeff Fessler
#src 2019-06-23 update to use more realistic simulated sensitivity maps
#src 2020-06-20 update
#src 2021-12-08 Literate

#=
This example illustrates 2D dynamic MRI image reconstruction
from golden angle (GA) radial sampled k-space data
collected with multiple coils
(parallel MRI),
with temporal "TV" regularizer (corner-rounded)
using the Julia language.

This entire page was generated using a single Julia file:
[3-2d-t.jl](@__REPO_ROOT_URL__/mri/3-2d-t.jl).
=#
#md # In any such Julia documentation,
#md # you can access the source code
#md # using the "Edit on GitHub" link in the top right.

#md # The corresponding notebook can be viewed in
#md # [nbviewer](http://nbviewer.jupyter.org/) here:
#md # [`3-2d-t.ipynb`](@__NBVIEWER_ROOT_URL__/mri/3-2d-t.ipynb),
#md # and opened in [binder](https://mybinder.org/) here:
#md # [`3-2d-t.ipynb`](@__BINDER_ROOT_URL__/mri/3-2d-t.ipynb).

using ImageGeoms: ImageGeom
using ImagePhantoms: shepp_logan, SouthPark, phantom, Ellipse
using MIRTjim: jim, prompt
using MIRT: Anufft, diffl_map, ncg
using MIRT: ir_mri_sensemap_sim, ir_mri_kspace_ga_radial
using Plots: gui, plot, scatter, default; default(markerstrokecolor=:auto, label="")
using Plots: gif, @animate, Plots
using LinearAlgebra: norm, dot, Diagonal
using LinearMapsAA: LinearMapAA, block_diag
using Random: seed!
jim(:abswarn, false); # suppress warnings about display of |complex| images

# The following line is helpful when running this jl-file as a script;
# this way it will prompt user to hit a key after each image is displayed.

isinteractive() && jim(:prompt, true);

# ## Create (synthetic) data

# Generate dynamic image sequence:
N = (60,64)
fov = 220
nt = 8 # frames
ig = ImageGeom(; dims = N, deltas=(fov,fov) ./ N)

object0 = shepp_logan(SouthPark(); fovs = (fov,fov))
objects = Array{typeof(object0)}(undef, nt)
xtrue = Array{ComplexF32}(undef, N..., nt)
for it=1:nt
    tmp = copy(object0)
    width2 = 15 + 5 * sin(2*pi*it/nt) # mouth open/close
    mouth = tmp[2]
    tmp[2] = Ellipse(mouth.center, (mouth.width[1], width2), mouth.angle[1], mouth.value)
    objects[it] = tmp
    xtrue[:,:,it] = phantom(axes(ig)..., tmp, 4)
end
jimxy = (args...; kwargs...) -> jim(axes(ig)..., args...; kwargs...)
jimxy(xtrue, "True images")

# Animate true image:
anim1 = @animate for it in 1:nt
    jimxy(xtrue[:,:,it], title="Frame $it")
end
gif(anim1; fps = 6)

# Plot one time course to see temporal change:
ix,iy = 30,14
plot(1:nt, abs.(xtrue[ix,iy,:]), label="ix=$ix, iy=$iy",
    marker=:o, xlabel="frame")

#
isinteractive() && jim(:prompt, true);

# Generate k-space sampling and data:
accelerate = 3
nspf = round(Int, maximum(N)/accelerate) # spokes per frame

Nro = maximum(N)
Nspoke = nspf * nt
kspace = ir_mri_kspace_ga_radial(Nro = Nro, Nspoke = Nspoke)
fovs = (fov, fov)
kspace[:,:,1] ./= fovs[1]
kspace[:,:,2] ./= fovs[2]
kspace = reshape(kspace, Nro, nspf, nt, 2)
(size(kspace), extrema(kspace))

# Plot sampling (in units of cycles/pixel):
ps = Array{Any}(undef, nt)
for it=1:nt
    ps[it] = scatter(kspace[:,:,it,1] * fovs[1], kspace[:,:,it,2] * fovs[2],
        xtick=(-1:1)*0.5, ytick=(-1:1)*0.5, xlim=(-1,1).*0.52, ylim=(-1,1).*0.52,
        aspect_ratio=1, markersize=1, title="Frame $it")
#src plot(ps[it]); gui()
end
plot(ps..., layout=(2,4))

#
isinteractive() && jim(:prompt, true);

anim2 = @animate for it in 1:nt
    plot(ps[it])
end
gif(anim2; fps = 6)

# Make sensitivity maps, normalized so SSoS = 1:
ncoil = 2
smap = ir_mri_sensemap_sim(dims=N, ncoil=ncoil, orbit_start=[90])
p1 = jim(smap, "Sensitivity maps raw");

ssos = sqrt.(sum(abs.(smap).^2, dims=ndims(smap))) # SSoS
ssos = selectdim(ssos, ndims(smap), 1)
p2 = jim(ssos, "SSoS for ncoil=$ncoil");

for ic=1:ncoil
    selectdim(smap, ndims(smap), ic) ./= ssos
end
p3 = jim(smap, "Sensitivity maps");

ssos = sqrt.(sum(abs.(smap).^2, dims=ndims(smap))) # SSoS
@assert all(isapprox.(ssos,1))
jim(p1, p2, p3)

# Make system matrix for dynamic non-Cartesian parallel MRI:
Fs = Array{Any}(undef, nt)
for it=1:nt # a NUFFT object for each frame
    Ω = [kspace[:,:,it,1][:] kspace[:,:,it,2][:]] * fov * 2pi
    Fs[it] = Anufft(Ω, N, n_shift = [N...]/2)
end

# Block diagonal system matrix, with one NUFFT per frame
S = [Diagonal(vec(selectdim(smap, ndims(smap), ic))) for ic=1:ncoil]
SO = s -> LinearMapAA(s ; idim=N, odim=N) # LinearMapAO for coil maps
AS1 = F -> vcat([F * SO(s) for s in S]...); # [A1*S1; ... ; A1*Sncoil]

# Input is [N... nt], output is [nspf*Nro Ncoil nt]
A = block_diag([AS1(F) for F in Fs]...) # todo: refine show()
(size(A), A._odim, A._idim)

# Simulate k-space data via an inverse crime (todo):
ytrue = A * xtrue
snr2sigma = (db, yb) -> # compute noise sigma from SNR (no sqrt(2) needed)
    10^(-db/20) * norm(yb) / sqrt(length(yb))
sig = Float32(snr2sigma(50, ytrue))
seed!(0)
y = ytrue + sig * randn(ComplexF32, size(ytrue))
20*log10(norm(ytrue) / norm(y - ytrue)) # verify SNR

# Initial image via zero-fill and scaling:
# todo: should use density compensation, perhaps via
# [VoronoiDelaunay.jl](https://github.com/JuliaGeometry/VoronoiDelaunay.jl)
x0 = A' * y # zero-filled recon (for each frame)
tmp = A * x0 # Nkspace × Ncoil × Nframe
x0 = (dot(tmp,y) / norm(tmp)^2) * x0 # scale sensibly
jimxy(x0, "initial image")

# Temporal finite differences:
Dt = diffl_map((N..., nt), length(N)+1 ; T=eltype(A))
tmp = Dt' * (Dt * xtrue)
jimxy(tmp, "time diff")

# Run nonlinear CG on "temporal edge-preserving" regularized LS cost function
niter = 90
delta = Float32(0.1) # small relative to temporal differences
reg = Float32(2^20) # trial and error here
ffair = (t,d) -> d^2 * (abs(t)/d - log(1 + abs(t)/d))
pot = z -> ffair(z, delta)
dpot = z -> z / (Float32(1) + abs(z/delta))
cost = x -> 0.5 * norm(A*x - y)^2 + reg * sum(pot.(Dt * x))
fun = (x,iter) -> cost(x)
gradf = [v -> v - y, u -> reg * dpot.(u)]
curvf = [v -> Float32(1), u -> reg]
(xh, out) = ncg([A, Dt], gradf, curvf, x0 ; niter, fun)
costs = [out[i+1][1] for i=0:niter];

# Show results
plot(
    jimxy(xtrue, "xtrue"),
    jimxy(xh, "recon"),
    jimxy(xh-xtrue, "error"),
    scatter(0:niter, log.(costs), label="cost", xlabel="iteration"),
)

#
isinteractive() && jim(:prompt, true);

# Animate true, recon, error
anim3 = @animate for it in 1:nt
    plot(
        jimxy(xtrue[:,:,it], clim=(0,120), title="True"),
        jimxy(xh[:,:,it], clim=(0,120), title="|Recon|"),
        jimxy(xh[:,:,it] - xtrue[:,:,it], clim=(0,30), title="|Error|"),
        plot_title = "Frame $it",
        layout = (1,3),
    )
end
gif(anim3; fps = 6)

max_stars_repo_name: JuliaImageRecon/Examples

max_issues_repo_name: JuliaImageRecon/Examples

max_forks_repo_name: JuliaImageRecon/Examples


text:
function convergence_factor!(vector)
    length = argmin(vector)[1]
    return round(((vector[length] / vector[1])^(1.0 / length)), digits=3)
end

# graphs methods:

function unet_vs_vcycle_graph!(title, vc_unet_res, unet_res, vc_res, j_unet_res; after_vcycle=false, e_vcycle_input=false)
    iterations = length(vc_unet_res)
    iter = range(1, length=iterations)

    # Unet
    factor = convergence_factor!(unet_res)
    factor_text = "u=$(factor)"
    p = plot(iter,unet_res,label="UNet(r) $(factor)")

    # Jacobi Unet
    factor = convergence_factor!(j_unet_res)
    factor_text = "$(factor_text) ju=$(factor)"
    plot!(iter,j_unet_res,label="J(UNet(r)) $(factor)")

    # Vcycle
    factor = convergence_factor!(vc_res)
    factor_text = "$(factor_text) v=$(factor)"
    plot!(iter,vc_res,label="V(r,0) $(factor)")

    # Vcycle Unet
    factor = convergence_factor!(vc_unet_res)
    factor_text = "$(factor_text) vu=$(factor)"
    plot!(iter,vc_unet_res,label="V(r,UNet(r)) $(factor)")

    yaxis!(L"\Vert b - Hx \Vert_2", :log10)
    xlabel!("iterations")

    savefig("test/unet/results/$(title) $(factor_text)")
    @info "$(Dates.format(now(), "HH:MM:SS.sss")) - Convergence Factors : $(factor_text)"
    return "$(title) $(factor_text)"
end

function unet_vs_vcycle_graph_times!(title, vc_unet_res, vc_unet_times, unet_res, unet_times, vc_res, vc_times, j_unet_res, j_unet_times; after_vcycle=false, e_vcycle_input=false)
    iterations = length(vc_unet_res)
    iter = range(1, length=iterations)

    # Unet
    factor = convergence_factor!(unet_res)
    factor_text = "u=$(factor)"
    p = plot(unet_times,unet_res,label="UNet(r) $(factor)")

    # Jacobi Unet
    factor = convergence_factor!(j_unet_res)
    factor_text = "$(factor_text) ju=$(factor)"
    plot!(j_unet_times,j_unet_res,label="J(UNet(r)) $(factor)")

    # Vcycle
    factor = convergence_factor!(vc_res)
    factor_text = "$(factor_text) v=$(factor)"
    plot!(vc_times,vc_res,label="V(r,0) $(factor)")

    # Vcycle Unet
    factor = convergence_factor!(vc_unet_res)
    factor_text = "$(factor_text) vu=$(factor)"
    plot!(vc_unet_times,vc_unet_res,label="V(r,UNet(r)) $(factor)")

    yaxis!(L"\Vert b - Hx \Vert_2", :log10)
    xlabel!("milliseconds")

    savefig("test/unet/results/$(title) times $(factor_text)")
    @info "$(Dates.format(now(), "HH:MM:SS.sss")) - Convergence Factors : $(factor_text)"
    return "$(title) times $(factor_text)"
end

function load_model!(test_name, e_vcycle_input, kappa_input, gamma_input;kernel=(3,3), model_type=SUnet, k_type=NaN, resnet_type=SResidualBlock, k_chs=-1, indexes=3, σ=elu, arch=0)
    model = create_model!(e_vcycle_input, kappa_input, gamma_input; kernel=kernel, type=model_type, k_type=k_type, resnet_type=resnet_type, k_chs=k_chs, indexes=indexes, σ=σ, arch=arch)
    model = model|>pu
    @load test_name model #"../../models/$(test_name).bson" model
    @info "$(Dates.format(now(), "HH:MM:SS.sss")) - Load Model"
    model = model|>cgpu

    return model
end

function model_tuning!(model, test_name, kappa, omega, gamma, n, m, f, dataset_size, bs, iter, lr, threshold, axb, jac, k_kernel, restrt, kappa_type; residual_loss=false)
    name = "$(test_name) $(n) $(m) $(dataset_size) $(iter)"
    model, train_loss, test_loss = train_residual_unet!(model, name, n, m, f, kappa, omega, gamma,
                                                        dataset_size, 10, bs, iter, lr;
                                                        e_vcycle_input=false, v2_iter=10, level=3, data_augmentetion=false,
                                                        kappa_type=kappa_type, threshold=threshold, kappa_input=true, kappa_smooth=true, k_kernel=k_kernel,
                                                        gamma_input=true, kernel=3, smaller_lr=24, axb=axb, jac=jac, norm_input=norm_input, model_type=model_type, k_type=k_type, k_chs=k_chs, indexes=indexes,
                                                        data_path="", full_loss=false, residual_loss=residual_loss, gmres_restrt=restrt,σ=σ)
    return model, name
end

function loss!(x, y)
    return norm(x .- y) / norm(y)
end

# check model methods:

function check_model_times!(test_name, model, n, m, f, kappa, omega, gamma, e_vcycle_input, kappa_type, kappa_input, gamma_input, kernel, dataset_size, restrt, max_iter; v2_iter=10, level=3, smooth=false, k_kernel=3, threshold=50, axb=false, norm_input=false, before_jacobi=false, log_error=false, unet_in_vcycle=false, indexes=3, arch=1) # arch - Standart (0) SplitUNet (1) FeaturesUNet (2)
    unet_results = zeros(dataset_size,6,restrt*max_iter+1)
    vcycle_results = zeros(dataset_size,4,restrt*max_iter+1)
    mkpath("test/unet/results")
    for i=1:dataset_size
        x_true = randn(c_type,n-1,m-1, blocks, 1)|>pu
        kappa = r_type.(generate_kappa!(n, m; type=kappa_type, smooth=smooth, threshold=threshold, kernel=k_kernel)|>pu)
        heatmap(kappa, color=:blues)
        savefig("test/unet/results/kappa $(kappa_type) $(threshold) $(k_kernel) $(i)")
        CSV.write("test/unet/results/kappa $(kappa_type) $(threshold) $(k_kernel) $(i).csv", DataFrame(K=vec(kappa)), delim = ';')
        kappa_features = NaN
        if arch != 0
            kappa_input = reshape(kappa, n-1, m-1, 1, 1)
            if indexes != 3
                kappa_input = cat(kappa_input, reshape(gamma, n-1, m-1, 1, 1), dims=3)
            end
            kappa_features = model.kappa_subnet(kappa_input|>cgpu)|>pu
            # @info "$(Dates.format(now(), "HH:MM:SS.sss")) - after kappa features $(size(kappa_features))"
        end
        resvec1, resvec2, times1, times2 = unet_vs_vcycle_gpu_gmres!(model, n, m, kappa, kappa_features, omega, gamma, x_true, true, e_vcycle_input, kappa_input, gamma_input, restrt, max_iter, false; v2_iter=v2_iter, level=level, axb=axb, norm_input=norm_input, log_error=log_error, test_name=test_name, before_jacobi=before_jacobi, unet_in_vcycle=unet_in_vcycle, arch=arch)
        unet_results[i,1,:] = resvec1
        unet_results[i,2,:] = times1
        vcycle_results[i,1,:] = resvec2
        vcycle_results[i,2,:] = times2
        resvec1, resvec2, times1, times2 = unet_vs_vcycle_gpu_gmres!(model, n, m, kappa, kappa_features, omega, gamma, x_true, false, e_vcycle_input, kappa_input, gamma_input, restrt, max_iter, false; v2_iter=v2_iter, level=level, axb=axb, norm_input=norm_input, log_error=log_error, test_name=test_name, before_jacobi=before_jacobi, unet_in_vcycle=unet_in_vcycle, arch=arch)
        unet_results[i,3,:] = resvec1
        unet_results[i,4,:] = times1
        vcycle_results[i,3,:] = resvec2
        vcycle_results[i,4,:] = times2
        resvec1, resvec2, times1, times2 = unet_vs_vcycle_gpu_gmres!(model, n, m, kappa, kappa_features, omega, gamma, x_true, false, e_vcycle_input, kappa_input, gamma_input, restrt, max_iter, true; v2_iter=v2_iter, level=level, axb=axb, norm_input=norm_input, log_error=log_error, test_name=test_name, before_jacobi=before_jacobi, unet_in_vcycle=unet_in_vcycle, arch=arch)
        unet_results[i,5,:] = resvec1
        unet_results[i,6,:] = times1
    end
    graph_title = unet_vs_vcycle_graph_times!("$(test_name) t_n=$(n) t_m=$(m) t_axb=$("$(axb)"[1]) t_norm=$("$(norm_input)"[1])",
                    mean(unet_results[:,1,:],dims=1)', mean(unet_results[:,2,:],dims=1)', mean(unet_results[:,3,:],dims=1)', mean(unet_results[:,4,:],dims=1)',
                    mean(vcycle_results[:,3,:],dims=1)', mean(vcycle_results[:,4,:],dims=1)', mean(unet_results[:,5,:],dims=1)', mean(unet_results[:,6,:],dims=1)')

    CSV.write("test/unet/results/$(graph_title).csv", DataFrame(VU=vec(mean(unet_results[:,1,:],dims=1)'),
                                                        VU_T=vec(mean(unet_results[:,2,:],dims=1)'),
                                                        U=vec(mean(unet_results[:,3,:],dims=1)'),
                                                        U_T=vec(mean(unet_results[:,4,:],dims=1)'),
                                                        V=vec(mean(vcycle_results[:,3,:],dims=1)'),
                                                        V_T=vec(mean(vcycle_results[:,4,:],dims=1)'),
                                                        JU=vec(mean(unet_results[:,5,:],dims=1)'),
                                                        JU_T=vec(mean(unet_results[:,6,:],dims=1)')), delim = ';')
end

function check_model!(test_name, model, n, m, f, kappa, omega, gamma, e_vcycle_input, kappa_type, kappa_input, gamma_input, kernel, dataset_size, restrt, max_iter; v2_iter=10, level=3, smooth=false, k_kernel=3, threshold=50, axb=false, norm_input=false, before_jacobi=false, log_error=false, unet_in_vcycle=false, indexes=3, arch=1) # arch - Standart (0) SplitUNet (1) FeaturesUNet (2)
    unet_results = zeros(dataset_size,3,restrt*max_iter+1)
    vcycle_results = zeros(dataset_size,2,restrt*max_iter+1)
    mkpath("test/unet/results")
    for i=1:dataset_size
        x_true = randn(c_type,n-1,m-1, blocks, 1)|>pu
        kappa = r_type.(generate_kappa!(n, m; type=kappa_type, smooth=smooth, threshold=threshold, kernel=k_kernel)|>pu)
        heatmap(kappa, color=:blues)
        savefig("test/unet/results/kappa $(kappa_type) $(threshold) $(k_kernel) $(i)")
        CSV.write("test/unet/results/kappa $(kappa_type) $(threshold) $(k_kernel) $(i).csv", DataFrame(K=vec(kappa)), delim = ';')
        kappa_features = NaN
        if arch != 0
            kappa_input = reshape(kappa, n-1, m-1, 1, 1)
            if indexes != 3
                kappa_input = cat(kappa_input, reshape(gamma, n-1, m-1, 1, 1), dims=3)
            end
            kappa_features = model.kappa_subnet(kappa_input|>cgpu)|>pu
        end
        resvec1, resvec2 = unet_vs_vcycle_cpu_gmres!(model, n, m, kappa, kappa_features, omega, gamma, x_true, true, e_vcycle_input, kappa_input, gamma_input, restrt, max_iter, false; v2_iter=v2_iter, level=level, axb=axb, norm_input=norm_input, log_error=log_error, test_name=test_name, before_jacobi=before_jacobi, unet_in_vcycle=unet_in_vcycle, arch=arch)
        unet_results[i,1,:] = resvec1
        vcycle_results[i,1,:] = resvec2
        resvec1, resvec2 = unet_vs_vcycle_cpu_gmres!(model, n, m, kappa, kappa_features, omega, gamma, x_true, false, e_vcycle_input, kappa_input, gamma_input, restrt, max_iter, false; v2_iter=v2_iter, level=level, axb=axb, norm_input=norm_input, log_error=log_error, test_name=test_name, before_jacobi=before_jacobi, unet_in_vcycle=unet_in_vcycle, arch=arch)
        unet_results[i,2,:] = resvec1
        vcycle_results[i,2,:] = resvec2
        resvec1, resvec2 = unet_vs_vcycle_cpu_gmres!(model, n, m, kappa, kappa_features, omega, gamma, x_true, false, e_vcycle_input, kappa_input, gamma_input, restrt, max_iter, true; v2_iter=v2_iter, level=level, axb=axb, norm_input=norm_input, log_error=log_error, test_name=test_name, before_jacobi=before_jacobi, unet_in_vcycle=unet_in_vcycle, arch=arch)
        unet_results[i,3,:] = resvec1
    end
    graph_title = unet_vs_vcycle_graph!("$(test_name) t_n=$(n) t_m=$(m) t_axb=$("$(axb)"[1]) t_norm=$("$(norm_input)"[1])",
            mean(unet_results[:,1,:],dims=1)', mean(unet_results[:,2,:],dims=1)',
            mean(vcycle_results[:,2,:],dims=1)', mean(unet_results[:,3,:],dims=1)')
    # graph_title = "$(Dates.format(now(), "HH_MM_SS_sss"))_$(test_name)_t_n=$(n)_t_m=$(m)"
    @info "$(Dates.format(now(), "HH:MM:SS.sss")) - $(graph_title)"
    CSV.write("test/unet/results/$(graph_title).csv", DataFrame(VU=vec(mean(unet_results[:,1,:],dims=1)'),
                                                        U=vec(mean(unet_results[:,2,:],dims=1)'),
                                                        V=vec(mean(vcycle_results[:,2,:],dims=1)'),
                                                        JU=vec(mean(unet_results[:,3,:],dims=1)')), delim = ';')
end

function check_point_source_problem!(test_name, model, n, m, kappa, omega, gamma, e_vcycle_input, kappa_input, gamma_input; v2_iter=10, level=3)
    b = zeros(c_type,n-1,m-1)|>pu
    h = 2.0 / (n+m)
    b[floor(Int32,n / 2.0),floor(Int32,m / 2.0)] = 1.0 ./ (h.^2);
    full_sol = fgmres_v_cycle_helmholtz!(n, m, h, b, kappa|>pu, omega, gamma|>pu; restrt=10, maxIter=5)
    heatmap(real(reshape(full_sol, n-1, m-1)), color=:blues)
    savefig("test/unet/results/$(test_name) $(n) e true")
    CSV.write("test/unet/results/$(test_name) $(n) e true.csv", DataFrame(E=vec(real(full_sol))), delim = ';')

    shifted_laplacian_matrix, helmholtz_matrix = get_helmholtz_matrices!(kappa, omega, gamma; alpha=0.5)
    coefficient = h^2

    input = complex_grid_to_channels!(reshape(coefficient .* b , n-1, m-1, 1, 1))
    if arch == 1
        kappa_features = model.kappa_subnet(reshape(kappa, n-1, m-1, 1, 1)|>cgpu)|>pu
        input = cat(input, reshape(kappa, n-1, m-1, 1, 1), reshape(gamma, n-1, m-1, 1, 1), kappa_features, dims=3)
        e_unet = (model.solve_subnet(u_type.(input)|>cgpu)|>cpu)
    elseif arch == 2
        kappa_features = model.kappa_subnet(reshape(kappa, n-1, m-1, 1, 1)|>cgpu)|>pu
        input = cat(input, reshape(kappa, n-1, m-1, 1, 1), reshape(gamma, n-1, m-1, 1, 1), dims=3)
        e_unet = (model.solve_subnet(u_type.(input)|>cgpu, kappa_features|>cgpu)|>cpu)
    else
        input = cat(input, reshape(kappa, n-1, m-1, 1, 1), reshape(gamma, n-1, m-1, 1, 1), dims=3)
        e_unet = (model(u_type.(input)|>cgpu)|>cpu)
    end

    heatmap(e_unet[:,:,1,1]|>cpu, color=:blues)
    savefig("test/unet/results/$(test_name) $(n) e unet")
    CSV.write("test/unet/results/$(test_name) $(n) e unet.csv", DataFrame(E=vec(real(e_unet[:,:,1,1]))), delim = ';')

    approx = zeros(c_type,n-1,m-1)|>pu
    approx, = v_cycle_helmholtz!(n, m, h, approx, b, kappa, omega, gamma; v2_iter=v2_iter, level=level)
    heatmap(real(approx), color=:blues)
    savefig("test/unet/results/$(test_name) $(n) e vcycle")
    CSV.write("test/unet/results/$(test_name) $(n) e vcycle.csv", DataFrame(E=vec(real(approx))), delim = ';')
    ax = helmholtz_chain!(reshape(approx,n-1,m-1,1,1), helmholtz_matrix; h=h)
    @info "Vcycle error=$(norm_diff!(approx|>cpu,full_sol|>cpu)) residual=$(norm_diff!(ax|>cpu,b|>cpu))"

    approx, = v_cycle_helmholtz!(n, m, h, approx, b, kappa, omega, gamma; v2_iter=v2_iter, level=level)
    heatmap(real(approx), color=:blues)
    savefig("test/unet/results/$(test_name) $(n) e vcycle vcycle")
    CSV.write("test/unet/results/$(test_name) $(n) e vcycle vcycle.csv", DataFrame(E=vec(real(approx))), delim = ';')
    ax = helmholtz_chain!(reshape(approx,n-1,m-1,1,1), helmholtz_matrix; h=h)
    @info "Vcycle Vcycle error=$(norm_diff!(approx|>cpu,full_sol|>cpu)) residual=$(norm_diff!(ax|>cpu,b|>cpu))"

    approx = e_unet[:,:,1,1]+im*e_unet[:,:,2,1]
    ax = helmholtz_chain!(reshape(approx|>pu,n-1,m-1,1,1), helmholtz_matrix; h=h)
    @info "Unet error=$(norm_diff!(approx|>cpu,full_sol|>cpu)) residual=$(norm_diff!(ax|>cpu,b|>cpu))"
    approx, = v_cycle_helmholtz!(n, m, h, approx|>pu, b, kappa, omega, gamma; v2_iter=v2_iter, level=level)
    heatmap(real(approx), color=:blues)
    savefig("test/unet/results/$(test_name) $(n) e vcycle unet")
    CSV.write("test/unet/results/$(test_name) $(n) e vcycle unet.csv", DataFrame(E=vec(real(approx))), delim = ';')
    ax = helmholtz_chain!(reshape(approx,n-1,m-1,1,1), helmholtz_matrix; h=h)
    @info "Vcycle Unet error=$(norm_diff!(approx|>cpu,full_sol|>cpu)) residual=$(norm_diff!(ax|>cpu,b|>cpu))"

    approx = (e_unet[:,:,1,1]+im*e_unet[:,:,2,1])
    approx = jacobi_helmholtz_method!(n, m, h, approx|>pu, b, helmholtz_matrix)
    heatmap(real(approx), color=:blues)
    savefig("test/unet/results/$(test_name) $(n) e jacobi unet")
    CSV.write("test/unet/results/$(test_name) $(n) e jacobi unet.csv", DataFrame(E=vec(real(approx))), delim = ';')
    ax = helmholtz_chain!(reshape(approx,n-1,m-1,1,1), helmholtz_matrix; h=h)
    @info "Jacobi Unet error=$(norm_diff!(approx|>cpu,full_sol|>cpu)) residual=$(norm_diff!(ax|>cpu,b|>cpu))"
end

function unet_vs_vcycle_cpu_gmres!(model, n, m, kappa, kappa_features, omega, gamma, x_true, after_vcycle, e_vcycle_input, kappa_input, gamma_input, restrt, max_iter, relax_jacobi; v2_iter=10, level=3, axb=false, norm_input=false, log_error=true, test_name="", before_jacobi=false, unet_in_vcycle=false, arch=1)

    shifted_laplacian_matrix, helmholtz_matrix = get_helmholtz_matrices!(kappa, omega, gamma; alpha=r_type(0.5))
    h = 2.0./(m+n)

    if axb == true
        r_vcycle = zeros(c_type,n-1,m-1,1,1)
        r_vcycle[floor(Int32,n / 2.0),floor(Int32,m / 2.0),1,1] = r_type(1.0 ./mean(h.^2))
        r_vcycle = vec(r_vcycle)
        for i = 2:blocks
            r_vcycle1 = zeros(c_type,n-1,m-1,1,1)
            r_vcycle1[floor(Int32,(n / blocks)*(i-1)),floor(Int32,(m / blocks)*(i-1)),1,1] = r_type(1.0 ./mean(h.^2))
            r_vcycle = cat(r_vcycle, vec(r_vcycle1), dims=2)
        end
    else
        x_true = randn(c_type,n-1,m-1, 1, 1)
        r_vcycle, _ = generate_r_vcycle!(n, m, kappa, omega, gamma, reshape(x_true,n-1,m-1,1,1))
        r_vcycle = vec(r_vcycle)
        for i = 2:blocks
            x_true = randn(c_type,n-1,m-1, 1, 1)
            r_vcycle1, _ = generate_r_vcycle!(n, m, kappa, omega, gamma, reshape(x_true,n-1,m-1,1,1))
            r_vcycle = cat(r_vcycle, vec(r_vcycle1), dims=2)
        end
    end

    coefficient = r_type(h^2)

    A(v) = vec(helmholtz_chain!(reshape(v, n-1, m-1, 1, 1), helmholtz_matrix; h=h))
    function As(v)
        res = vec(A(v[:,1]))
        for i = 2:blocks
            res = cat(res, vec(A(v[:,i])), dims=2)
        end

        return res
    end

    function M_Unet(r)
        r = reshape(r, n-1, m-1)
        rj = reshape(r, n-1, m-1)
        e = zeros(c_type, n-1, m-1)
        ej = zeros(c_type, n-1, m-1)

        if after_vcycle != true
            ej = jacobi_helmholtz_method!(n, m, h, e, r, helmholtz_matrix)
            rj = r - reshape(A(ej), n-1, m-1)
        end

        input = complex_grid_to_channels!(reshape(rj , n-1, m-1, 1, 1))
        if arch == 1
            input = cat(input, reshape(kappa, n-1, m-1, 1, 1), reshape(gamma, n-1, m-1, 1, 1), kappa_features, dims=3)
            e_unet = (model.solve_subnet(u_type.(input)|>cgpu)|>cpu)
        elseif arch == 2
            input = cat(input, reshape(kappa, n-1, m-1, 1, 1), reshape(gamma, n-1, m-1, 1, 1), dims=3)
            e_unet = (model.solve_subnet(u_type.(input)|>cgpu, kappa_features|>cgpu)|>cpu)
        else
            input = cat(input, reshape(kappa, n-1, m-1, 1, 1), reshape(gamma, n-1, m-1, 1, 1), dims=3)
            e_unet = (model(u_type.(input)|>cgpu)|>cpu)
        end

        e_unet = reshape((e_unet[:,:,1,1] + im*e_unet[:,:,2,1]),n-1,m-1,1,1) .* coefficient
        e = reshape(e_unet, n-1,m-1)

        e = ej + e

        if relax_jacobi == true
            e = jacobi_helmholtz_method!(n, m, h, e, r, helmholtz_matrix)
        end

        if after_vcycle == true
            e, = v_cycle_helmholtz!(n, m, h, e, r, kappa, omega, gamma; v2_iter = v2_iter, level = level)
        end
        return vec(e)
    end

    function M_Unets(rs)
        res = M_Unet(rs[:,1])
        for i = 2:blocks
            res = cat(res, M_Unet(rs[:,i]), dims=2)
        end

        return res
    end

    function M(r)
        e_vcycle = zeros(c_type,n-1,m-1)
        e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[:,1], n-1, m-1), kappa, omega, gamma; v2_iter = v2_iter, level = 3)
        if after_vcycle == true
            e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[:,1], n-1, m-1), kappa, omega, gamma; v2_iter = v2_iter, level = 3)
        end
        res = vec(e_vcycle)
        for i = 2:blocks
            e_vcycle = zeros(c_type,n-1,m-1)
            e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[:,i], n-1, m-1), kappa, omega, gamma; v2_iter = v2_iter, level = 3)
            if after_vcycle == true
                e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[:,i], n-1, m-1), kappa, omega, gamma; v2_iter = v2_iter, level = 3)
            end
            res = cat(res, vec(e_vcycle), dims=2)
        end

        return res
    end

    function SM(r)
        e_vcycle = zeros(c_type,n-1,m-1)
        e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[:,1], n-1, m-1), kappa, omega, gamma; v2_iter = v2_iter, level = 3)
        res = vec(e_vcycle)
        for i = 2:blocks
            e_vcycle = zeros(c_type,n-1,m-1)
            e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[:,i], n-1, m-1), kappa, omega, gamma; v2_iter = v2_iter, level = 3)
            res = cat(res, vec(e_vcycle), dims=2)
        end

        return res
    end

    x_init = zeros(c_type,(n-1)*(m-1),blocks)
    x3,flag3,err3,iter3,resvec3= KrylovMethods.blockFGMRES(As, r_vcycle, 3, tol=1e-30, maxIter=1,
                                                    M=SM, X=x_init, out=-1,flexible=true)
    i = 1
    x1 = x3
    x1,flag1,err1,iter1,resvec1 = KrylovMethods.blockFGMRES(As, r_vcycle, restrt, tol=1e-30, maxIter=max_iter,
                                                            M=M_Unets, X =x1, out=-1,flexible=true)

    x_init = zeros(c_type,(n-1)*(m-1),blocks)
    x3,flag3,err3,iter3,resvec3= KrylovMethods.blockFGMRES(As, r_vcycle, 3, tol=1e-30, maxIter=1,
                                                    M=SM, X=x_init, out=-1,flexible=true)
    i = 1
    x2 = x3
    x2,flag2,err2,iter2,resvec2 = KrylovMethods.blockFGMRES(As, r_vcycle, restrt, tol=1e-30, maxIter=max_iter,
                                                    M=M, X=x2, out=-1,flexible=true)

    return cat([resvec3[end]],resvec1, dims=1), cat([resvec3[end]],resvec2,dims=1)
end

function unet_vs_vcycle_gpu_gmres!(model, n, m, kappa, kappa_features, omega, gamma, x_true, after_vcycle, e_vcycle_input, kappa_input, gamma_input, restrt, max_iter, relax_jacobi; v2_iter=10, level=3, axb=false, norm_input=false, log_error=true, test_name="", before_jacobi=false, unet_in_vcycle=false, arch=1)

    sl_matrix_level3, h_matrix_level3 = get_helmholtz_matrices!(kappa, omega, gamma; alpha=r_type(0.5))
    kappa_coarse = down(reshape(kappa, n-1, m-1, 1, 1)|>pu)[:,:,1,1]
    gamma_coarse = down(reshape(gamma, n-1, m-1, 1, 1)|>pu)[:,:,1,1]
    sl_matrix_level2, h_matrix_level2 = get_helmholtz_matrices!(kappa_coarse, omega, gamma_coarse; alpha=r_type(0.5))
    kappa_coarse = down(reshape(kappa_coarse, Int64((n/2)-1),  Int64((m/2)-1), 1, 1)|>pu)[:,:,1,1]
    gamma_coarse = down(reshape(gamma_coarse,  Int64((n/2)-1),  Int64((m/2)-1), 1, 1)|>pu)[:,:,1,1]
    sl_matrix_level1, h_matrix_level1 = get_helmholtz_matrices!(kappa_coarse, omega, gamma_coarse; alpha=r_type(0.5))

    e_zeros = zeros(c_type, n-1, m-1)|>pu
    shifted_laplacian_matrix, helmholtz_matrix = get_helmholtz_matrices!(kappa, omega, gamma; alpha=r_type(0.5))
    h = r_type(2.0./(n+m))

    if axb == true
        src = [64,64]

        r_vcycle = zeros(c_type,n-1,m-1,1,1)|>pu
        r_vcycle[floor(Int32,n / 2.0),floor(Int32,m / 2.0),1,1] = r_type(1.0 ./mean(h.^2))
    else
        x_true = randn(c_type,n-1,m-1, 1, 1)|>pu
        r_vcycle, _ = generate_r_vcycle!(n, m, kappa, omega, gamma, reshape(x_true,n-1,m-1,1,1))
        r_vcycle = vec(r_vcycle)
        for i = 2:blocks
            x_true = randn(c_type,n-1,m-1, 1, 1)|>pu
            r_vcycle1, _ = generate_r_vcycle!(n, m, kappa, omega, gamma, reshape(x_true,n-1,m-1,1,1))
            r_vcycle = cat(r_vcycle, vec(r_vcycle1), dims=2)
        end
    end

    coefficient = r_type(h^2)

    A(v::a_type) = vec(helmholtz_chain!(reshape(real(v), n-1, m-1, 1, 1), helmholtz_matrix; h=h) + im*helmholtz_chain!(reshape(imag(v), n-1, m-1, 1, 1), helmholtz_matrix; h=h))

    function As(v)
        res = vec(A(v[1:(n-1)*(m-1)]))
        for i = 2:blocks
            res = cat(res, vec(A(v[(i-1)*((n-1)*(m-1))+1:i*((n-1)*(m-1))])), dims=2)
        end

        return vec(res)
    end

    function M_Unet(r::a_type)
        r = reshape(r, n-1, m-1)
        rj = reshape(r, n-1, m-1)
        e = ej = e_zeros

        if after_vcycle != true
            ej = jacobi_helmholtz_method!(n, m, h, e, r, helmholtz_matrix)
            rj = r - reshape(A(ej), n-1, m-1)
        end
        input = complex_grid_to_channels!(reshape(rj , n-1, m-1, 1, 1))
        if arch == 1
            input = cat(input, reshape(kappa, n-1, m-1, 1, 1), reshape(gamma, n-1, m-1, 1, 1), kappa_features, dims=3)
            e_unet = (model.solve_subnet(u_type.(input)|>cgpu)|>pu)
        elseif arch == 2
            input = cat(input, reshape(kappa, n-1, m-1, 1, 1), reshape(gamma, n-1, m-1, 1, 1), dims=3)
            e_unet = r_type.((model.solve_subnet(u_type.(input), kappa_features)))|>pu
        else
            input = cat(input, reshape(kappa, n-1, m-1, 1, 1), reshape(gamma, n-1, m-1, 1, 1), dims=3)
            e_unet = model(input)
        end

        e_unet = reshape((e_unet[:,:,1,1] + im*e_unet[:,:,2,1]),n-1,m-1,1,1) .* coefficient
        e = reshape(e_unet, n-1,m-1)

        e = ej + e

        if relax_jacobi == true
            e = jacobi_helmholtz_method!(n, m, h, e, r, helmholtz_matrix)
        end

        if after_vcycle == true
            e, = v_cycle_helmholtz!(n, m, h, e, r, h_matrix_level1, sl_matrix_level1, h_matrix_level2, sl_matrix_level2, h_matrix_level3, sl_matrix_level3; v2_iter = v2_iter, level = level)
        end
        return vec(e)
    end

    function M_Unets(rs)
        res = M_Unet(rs[1:((n-1)*(m-1))])
        for i = 2:blocks
            res = cat(res, M_Unet(rs[(i-1)*((n-1)*(m-1))+1:i*((n-1)*(m-1))]), dims=2)
        end

        return vec(res)
    end

    function M(r::a_type)
        e_vcycle = e_zeros
        e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[1:((n-1)*(m-1))], n-1, m-1), h_matrix_level1, sl_matrix_level1, h_matrix_level2, sl_matrix_level2, h_matrix_level3, sl_matrix_level3; v2_iter = v2_iter, level = 3)
        if after_vcycle == true
            e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[1:((n-1)*(m-1))], n-1, m-1), h_matrix_level1, sl_matrix_level1, h_matrix_level2, sl_matrix_level2, h_matrix_level3, sl_matrix_level3; v2_iter = v2_iter, level = 3)
        end
        res = vec(e_vcycle)
        for i = 2:blocks
            e_vcycle = e_zeros
            e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[(i-1)*((n-1)*(m-1))+1:i*((n-1)*(m-1))], n-1, m-1), h_matrix_level1, sl_matrix_level1, h_matrix_level2, sl_matrix_level2, h_matrix_level3, sl_matrix_level3; v2_iter = v2_iter, level = 3)
            if after_vcycle == true
                e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[(i-1)*((n-1)*(m-1))+1:i*((n-1)*(m-1))], n-1, m-1), h_matrix_level1, sl_matrix_level1, h_matrix_level2, sl_matrix_level2, h_matrix_level3, sl_matrix_level3; v2_iter = v2_iter, level = 3)
            end
            res = cat(res, vec(e_vcycle), dims=2)
        end

        return vec(res)
    end

    function SM(r::a_type)
        e_vcycle = zeros(c_type,n-1,m-1)|>pu
        e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[1:((n-1)*(m-1))], n-1, m-1), h_matrix_level1, sl_matrix_level1, h_matrix_level2, sl_matrix_level2, h_matrix_level3, sl_matrix_level3; v2_iter = v2_iter, level = 3)
        res = vec(e_vcycle)
        for i = 2:blocks
            e_vcycle = zeros(c_type,n-1,m-1)|>pu
            e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, reshape(r[(i-1)*((n-1)*(m-1))+1:i*((n-1)*(m-1))], n-1, m-1), h_matrix_level1, sl_matrix_level1, h_matrix_level2, sl_matrix_level2, h_matrix_level3, sl_matrix_level3; v2_iter = v2_iter, level = 3)
            res = cat(res, vec(e_vcycle), dims=2)
        end

        return vec(res)
    end

    x_init = zeros(c_type,(n-1)*(m-1),blocks)|>pu
    x3,flag3,err3,iter3,resvec3,times3= fgmres_func(As, vec(r_vcycle), 3, tol=1e-30, maxIter=1,
                                                    M=SM, x=vec(x_init), out=-1,flexible=true)
    i = 1
    x1 = x3
    x1,flag1,err1,iter1,resvec1,times1 = fgmres_func(As, vec(r_vcycle), restrt, tol=1e-30, maxIter=max_iter,
                                                            M=M_Unets, x=vec(x1), out=-1,flexible=true)

    x_init = zeros(c_type,(n-1)*(m-1),blocks)|>pu
    x3,flag3,err3,iter3,resvec3,times3= fgmres_func(As, vec(r_vcycle), 3, tol=1e-30, maxIter=1,
                                                    M=SM, x=vec(x_init), out=-1,flexible=true)
    i = 1
    x2 = x3
    x2,flag2,err2,iter2,resvec2,times2 = fgmres_func(As, vec(r_vcycle), restrt, tol=1e-30, maxIter=max_iter,
                                                    M=M, x=vec(x2), out=-1,flexible=true)

    return cat([resvec3[end]],resvec1, dims=1), cat([resvec3[end]],resvec2,dims=1), cat([0],times1,dims=1), cat([0],times2,dims=1)
end

# not in use:
# function check_gmres_alternatively_model!(test_name, model, n, f, kappa, omega, gamma, e_vcycle_input, kappa_type, kappa_input, gamma_input, kernel, m, restrt, max_iter; v2_iter=10, level=3, smooth=false, threshold=50, axb=false, norm_input=false, before_jacobi=false)
#     vu_results = zeros(m,max_iter*5)
#     vv_results = zeros(m,max_iter*5)
#     v_results = zeros(m,max_iter*5)
#     vuv_results = zeros(m,max_iter*5)
#     vuj_results = zeros(m,max_iter*5)
#
#     h = r_type(2.0 ./ (n+m))
#
#     for i=1:m
#         x_true = randn(c_type, n-1, m-1, 1, 1)|>pu
#         kappa = generate_kappa!(n; type=kappa_type, smooth=smooth, threshold=threshold)
#         s_matrix, h_matrix = get_helmholtz_matrices!(kappa, omega, gamma; alpha=0.5)
#         if axb == true
#             b_true = helmholtz_chain!(x_true, h_matrix; h=h)
#         else
#             b_true, _ = generate_r_vcycle!(n, m, kappa, omega, gamma, x_true)
#         end
#
#         A(v) = vec(helmholtz_chain!(reshape(v, n-1, m-1, 1, 1), h_matrix; h=h))
#
#         function M(r)
#             r = reshape(r, n-1, m-1)
#             e_vcycle = zeros(c_type, n-1, m-1)
#             e_vcycle, = v_cycle_helmholtz!(n, m, h, e_vcycle, r, kappa, omega, gamma; v2_iter = v2_iter, level = level)
#             return vec(e_vcycle)
#         end
#
#         x_vcycle = zeros(c_type,n-1,m-1)
#         for j = 1:max_iter
#             x_vcycle,flag,err,iter,resvec = fgmres_func(A, vec(b_true), 5, tol=1e-30, maxIter=1,
#                                                             M=M, x=vec(x_vcycle), out=-1, flexible=true)
#
#             r = b_true - helmholtz_chain!(reshape(x_vcycle,n-1,m-1,1,1), h_matrix; h=h) |> pu
#             x_unet = (model(cat(complex_grid_to_channels!(r, n), kappa, gamma, dims=3)|> cgpu) .* (h^2))|> pu
#             @info "VU $(i) $(j) x_unet $(norm(x_unet)) x_vcycle $(norm(x_vcycle)) r $(norm(r)) b_true $(norm(b_true))"
#             x_vcycle = reshape(x_vcycle,n-1,m-1) + (x_unet[:,:,1,1] + im*x_unet[:,:,2,1])
#             vu_results[i,5*(j-1)+1:5*(j-1)+5] = resvec
#         end
#         x_vcycle = zeros(c_type,n-1,m-1)
#         for j = 1:max_iter
#             x_vcycle,flag,err,iter,resvec = fgmres_func(A, vec(b_true), 5, tol=1e-30, maxIter=1,
#                                                             M=M, x=vec(x_vcycle), out=-1, flexible=true)
#
#             v_results[i,5*(j-1)+1:5*(j-1)+5] = resvec
#         end
#         x_vcycle = zeros(c_type,n-1,m-1)
#         for j = 1:max_iter
#             x_vcycle,flag,err,iter,resvec = fgmres_func(A, vec(b_true), 5, tol=1e-30, maxIter=1,
#                                                             M=M, x=vec(x_vcycle), out=-1, flexible=true)
#             x_vcycle, = v_cycle_helmholtz!(n, m, h, reshape(x_vcycle,n-1,m-1), reshape(b_true,n-1,m-1), kappa, omega, gamma; v2_iter = v2_iter, level = level)
#             vv_results[i,5*(j-1)+1:5*(j-1)+5] = resvec
#         end
#         x_vcycle = zeros(c_type,n-1,m-1)
#         for j = 1:max_iter
#
#             x_vcycle,flag,err,iter,resvec = fgmres_func(A, vec(b_true), 5, tol=1e-30, maxIter=1,
#                                                             M=M, x=vec(x_vcycle), out=-1, flexible=true)
#             r = b_true - helmholtz_chain!(reshape(x_vcycle,n-1,m-1,1,1), h_matrix; h=h) |> pu
#             x_unet = (model(cat(complex_grid_to_channels!(r, n), kappa, gamma, dims=3)|> cgpu) .* (h^2))|> pu
#             @info "VUV $(i) $(j) x_unet $(norm(x_unet)) x_vcycle $(norm(x_vcycle)) r $(norm(r)) b_true $(norm(b_true))"
#             x_vcycle = reshape(x_vcycle,n-1,m-1) + (x_unet[:,:,1,1] + im*x_unet[:,:,2,1])
#             x_vcycle, = v_cycle_helmholtz!(n, m, h, x_vcycle, reshape(b_true,n-1,m-1), kappa, omega, gamma; v2_iter = v2_iter, level = level)
#             vuv_results[i,5*(j-1)+1:5*(j-1)+5] = resvec
#         end
#         x_vcycle = zeros(c_type,n-1,m-1)
#         for j = 1:max_iter
#             x_vcycle,flag,err,iter,resvec = fgmres_func(A, vec(b_true), 5, tol=1e-30, maxIter=1,
#                                                             M=M, x=vec(x_vcycle), out=-1, flexible=true)
#             r = b_true - helmholtz_chain!(reshape(x_vcycle,n-1,m-1,1,1), h_matrix; h=h) |> pu
#             x_unet = (model(cat(complex_grid_to_channels!(r, n), kappa, gamma, dims=3)|> cgpu) .* (h^2))|> pu
#             @info "VUJ $(i) $(j) x_unet $(norm(x_unet)) x_vcycle $(norm(x_vcycle)) r $(norm(r)) b_true $(norm(b_true))"
#             x_vcycle = reshape(x_vcycle,n-1,m-1) + (x_unet[:,:,1,1] + im*x_unet[:,:,2,1])
#             x_vcycle = jacobi_helmholtz_method!(n, m, h, x_vcycle, reshape(b_true,n-1,m-1), h_matrix)
#             vuj_results[i,5*(j-1)+1:5*(j-1)+5] = resvec
#         end
#     end
# end
# function check_full_solution_model!(test_name, model, n, f, kappa, omega, gamma, e_vcycle_input, kappa_type, kappa_input, gamma_input, kernel, m, restrt, max_iter; v2_iter=10, level=3, smooth=false, threshold=50, axb=false, norm_input=false)
#     unet_results = zeros(m,restrt*max_iter-1)
#     h = r_type(2.0 ./ (n+m))
#     for i=1:m
#         x_true = randn(c_type,n-1,m-1, 1, 1)|>pu
#         kappa = generate_kappa!(n; type=kappa_type, smooth=smooth, threshold=threshold)
#         s_matrix, h_matrix = get_helmholtz_matrices!(kappa, omega, gamma; alpha=0.5)
#
#         if axb == true
#             r = helmholtz_chain!(x_true, h_matrix; h=h)
#         else
#             r, _ = generate_r_vcycle!(n, m, kappa, omega, gamma, x_true)
#         end
#
#         e = zeros(c_type, n-1, m-1, 1, 1)
#         ae = zeros(c_type, n-1, m-1, 1, 1)
#         for j = 1:restrt*max_iter-1
#             # e, = v_cycle_helmholtz!(n, m, h, reshape(e,n-1,m-1), reshape(r,n-1,m-1), kappa, omega, gamma; v2_iter = v2_iter, level = level)
#             # e = reshape(e, n-1, m-1, 1, 1)
#             # e = jacobi_helmholtz_method!(n, m, h, reshape(e,n-1,m-1), reshape(r,n-1,m-1), h_matrix)
#             ae = helmholtz_chain!(e, h_matrix; h=h)|> pu
#             correct_r = r - ae
#
#             unet_results[i,j] = norm(correct_r) / norm(r)
#             @info "$(i) $(j) : $(unet_results[i,j]) r=$(norm(r)) e=$(norm(e)) ae=$(norm(ae)) correct_r=$(norm(correct_r))"
#
#             #
#             correct_e = model(cat(complex_grid_to_channels!(reshape(correct_r,n-1,m-1,1,1).* (h^2), n), kappa, gamma, dims=3)|> cgpu)|> pu
#             e = e + (reshape((correct_e[:,:,1,1] + im*correct_e[:,:,2,1]),n-1,m-1,1,1))
#             @info "$(i) $(j) : e=$(norm(e)) correct_e=$(norm(correct_e))"
#             # ae = helmholtz_chain!(e, h_matrix; h=h)|> pu
#             # e = jacobi_helmholtz_method!(n, m, h, e, r, h_matrix)
#
#         end
#     end
#     CSV.write("test/unet/results/$(test_name) only unet.csv", DataFrame(U=vec(mean(unet_results[:,:],dims=1)')), delim = ';')
# end


text:
# Solve a trivial problem
A = reshape([1.0],(1,1))
solution = SCS_solve(SCS.Direct, 1, 1, A, [1.0], [1.0], 1, 0, Int[], Int[], 0, 0, Float64[]);
@test solution.ret_val == 1

feasible_basic_conic(SCS.Direct)

feasible_exponential_conic(SCS.Direct)

feasible_sdp_conic(SCS.Direct)

feasible_pow_conic(SCS.Direct)


text:
const P = Dict{Symbol, Any}
mutable struct LatticeParameter
    dict :: P
    function LatticeParameter(p::P)
        ld = new(P())
        for key in keys(p)
            val = p[key]
            if isa(val,AbstractArray)
                if length(val)>0 && isa(first(val), P)
                    setproperty!(ld, key, LatticeParameter.(val))
                else
                    setproperty!(ld, key, val)
                end
            elseif isa(val,P)
                setproperty!(ld, key, LatticeParameter(val))
            else
                setproperty!(ld, key, val)
            end
        end
        return ld
    end
end

import Base: getindex, setindex!, getproperty, setproperty!, propertynames, convert

getproperty(ld::LatticeParameter, name::Symbol) = getfield(ld, :dict)[name]
setproperty!(ld::LatticeParameter, name::Symbol, value) = getfield(ld, :dict)[name] = value
getindex(ld::LatticeParameter, name::Symbol) = ld.name
setindex!(ld::LatticeParameter, name::Symbol, value) = ld.name = value
propertynames(ld::LatticeParameter) = keys(getfield(ld,:dict))

generator(ld::LatticeParameter) = get(getfield(ld, :dict), :generator, generatelattice_std)

convert(::Type{LatticeParameter}, p::P) = LatticeParameter(p)


text:
"""
    fit_adc(b, S)

Fit Apparent Diffusion Coefficient (ADC) using a polynomial logfit of the normalized signal `S` against the b-values `b`.
"""
function fit_adc(b::Array{T,1}, S::Array{T,1}) where {T<:Real}

    # Fit log of signal
    p = fit(b, log.(S))

    # ADC is the negative linear term
    -p[1]
end


text:
"""
`rowcoord(df::DataFrame)`

Returns the :latitude and :longitude columns of a dataframe
as a tuple of (latitude, longitude). If the dataframe has more
than 1 row a vector of these tuples is returned.
"""
rowcoord(df::DataFrame) = [singlerowcoord(df[i, :]) for i=1:nrow(df)]
rowcoord(row::DataFrameRow) = singlerowcoord(row)

# Returns Coordinate from a single dataframe row
singlerowcoord(df::DataFrameRow) = df[:latitude][1], df[:longitude][1]

# Convert distance from default meters
function dist_convert(distance::AbstractFloat,
                      newunit::Unitful.FreeUnits,
                      returntype::Type=AbstractFloat)
   v = 1.0u"m" * distance |> newunit
   return v.val |> returntype
end


text:
# Functions for preforming gradient descent on the likelihood of an embedding for a given graph

""" Likelihood of an edge given the embedding coordinate """
function edge_probability(f1, i::Int, j::Int, coordinates)
    yi, yj = coordinates[[i,j], :]
    return f1((norm(yi - yj))^2)
end

""" A likelihood function """
f1(x) = 1/(1+exp(x))

""" Derivative of the likelihood function """
f1prime(x) = -exp(x)/((exp(x)+1)^2)

""" Gradient of the log-likelihood """
function diff_log_sig(x,y)
    z = (norm(x - y))^2
    return (f1prime(z)/f1(z))*2*(x-y)
end

""" Gradient of the log-likelihood with negative edge samples """
function diff_f(x,y,neg_sample)
    x1 = diff_log_sig(x,y)
    x2 = diff_log_sig(y,x)
    for sample in neg_sample
        x2 = x2 + diff_log_sig(y,sample)
    end
    return [x1,x2]
end

"""
    diff_likelihood(coordinates,g,nodes, γ, weight2,num_negative_edges,negsource,negdestin) -> coordinates_revised
Gradient descent step on the embedding `coordinates` given the graph `g`.
"""
function diff_likelihood(coordinates,g,nodes, γ, weight2,num_negative_edges,negsource,negdestin)
    coordinates_revised = [[0.0,0.0] for x in nodes]
    for ed in edges(g)
        ei, ej = src(ed), dst(ed)

        indi = findall(x->x==ei, nodes)[1]
        indj = findall(x->x==ej, nodes)[1]

        yi = coordinates[indi, :]
        yj = coordinates[indj, :]

        x1 = diff_log_sig(yi,yj)
        x2 = diff_log_sig(yj,yi)

        neg_ind = findall(x->x==ej,negdestin)

        num_of_neg_ej = length(neg_ind)
        neg_ind_weight = weight2[neg_ind]
        neg_ind_weight_norm = neg_ind_weight./(sum(neg_ind_weight)+0.000001)

        for i in 1:num_negative_edges
            edge_ind = sample(neg_ind, Weights(neg_ind_weight_norm))

            indii = findall(x->x==negsource[edge_ind][1], nodes)[1]
            indjj = findall(x->x==negdestin[edge_ind][1], nodes)[1]

            yii = coordinates[indii, :]
            yjj = coordinates[indjj, :]

            x2 = x2 .+ diff_log_sig(yj,yii)
        end

        coordinates_revised[indi] = coordinates_revised[indi] .+ x1
        coordinates_revised[indj] = coordinates_revised[indj] .+ x2

    end

    coordinates_revised = [[x[1]/ne(g),x[2]/ne(g)] for x in coordinates_revised]

    return coordinates_revised
end