text:
using ConcolicFuzzer
using Test

import ConcolicFuzzer: assert, prove

@testset "Simple tests" begin

    function f2(x)
        b = x < 10
        y = x - 10
        c = y < 0
        return y
    end

    val, symb, trace = execute(f2, 1);

    @test val == -9
    @test symb == true
    @test length(ConcolicFuzzer.flatten(trace)) == 3

    sanity_f1(x) = x * 2
    sanity_f2(x) = x ÷ 2
    sanity_f3(x) = (x ÷ 2) * 4 + 3

    val, symb, trace = execute(sanity_f1, 10)
    @test val == 20
    @test symb

    val, symb, trace = execute(sanity_f2, 10)
    @test val == 5
    @test symb

    val, symb, trace = execute(sanity_f3, 10)
    @test val == 23
    @test symb
end

@testset "Check" begin
    function f3(x)
        assert(x<10)
        y = x - 10
        prove(y < 0)
        return y
    end

    sat, inputs = check(f3, 1)
    # Underflow possible at -9223372036854775808
    @test sat == true

    function f4(x)
        assert(0<x<10)
        y = x - 10
        prove(y < 0)
        return y
    end

    # any input works as long as we don't need to hit any branches
    # if you have branches use `fuzz_and_check`
    sat, inputs = check(f4, 1)
    @test sat == false
end

@testset "Branches" begin
    function f(x::Int)
        @assert x < 10
        y = x - 10
        @assert y < 0
        return y
    end

    val, symb, trace = execute(f, 1);

    @test val == -9
    @test symb == true
    @test length(ConcolicFuzzer.flatten(trace)) == 5

    function g(x)
        if x < 10
            return 12
        else
            return x + 2
        end
    end

    val, symb, trace = execute(g, 1);

    @test val == 12
    @test symb == false

    val, symb, trace = execute(g, 10);

    @test val == 12
    @test symb == true
end

@testset "Fuzzing!" begin
    # Examples from the PSet
    function pset_f(x)
        if x == 7
            return 100
        end
        if x*2 == x+1
            return 70
        end
        if x > 2000
            return 80
        end
        if x*2 == 1000
            return 30000
        end
        if x < 500
            return 33
        end
        if x ÷ 123 == 7
            return 1234
        end
        return 40
    end

    tested, errored = fuzz_wargs(pset_f, 2002)

    explored = collect(zip(tested...))

    outs = explored[1]
    args = explored[2]
    rands = explored[3]

    @test length(tested) == 7
    @test (7,) in args
    @test (1,) in args
    @test (500,) in args

    @test 100 in outs
    @test 70 in outs
    @test 80 in outs
    @test 30000 in outs
    @test 33 in outs
    @test 1234 in outs
    @test 40 in outs
end

@testset "Fuzzing over an abstract type" begin
    function pure_madness(x::T) where T
        if x < typemax(T)
            return 42
        end
        return 9001
    end

    tested, errored = fuzz(pure_madness, Integer)

    @test length(tested) == 21

    tested, errored = fuzz(pure_madness, Union{Int32, UInt64})
    @test length(tested) == 4
end

@testset "Randomness" begin
    function r1()
        return rand(Int)
    end
    val, symb, trace = execute(r1);
    @test symb == true

    function r2()
        x = rand(Int)
        if x < 10
            return 1337
        else
            return 42
        end
    end

    val, symb, trace = execute(r2, rands = [9]);
    @test val == 1337

    val, symb, trace = execute(r2, rands = [12]);
    @test val == 42

    function r3()
        x = rand(Int)
        y = rand(Int)

        if x + y < 10
            return 42
        end
        return 12
    end

    val, symb, trace = execute(r3, rands = [3, 4]);
    @test val == 42

    val, symb, trace = execute(r3, rands = [10, 0]);
    @test val == 12
end

@testset "Simple Loops" begin
    function hh0(x)
        acc = 0
        i = 1
        while i <= 10
            i += 1
            acc += x
        end
        acc
    end
    val, symb, trace = execute(hh0, 1);
    @test val == 10
    @test symb == true
    @test length(ConcolicFuzzer.filter(trace)) == 10

    function hh1(x)
        acc = 0
        i = 1
        while i <= x
            acc += 1
            i += 1
        end
        acc
    end

    val, symb, trace = execute(hh1, 10);
    @test val == 10
    @test symb == false
    @test length(ConcolicFuzzer.filter(trace)) == 22
end

@testset "Fuzzing of loops" begin
    # Simulate out_of_bounds
    function out_of_bounds(N)
        i = 1
        while i <= N
            @assert i <= 10
            i +=  1
        end
        return i
    end

    tested, errored = fuzz_wargs(out_of_bounds, 0)

    explored = collect(zip(tested...))
    outs = explored[1]
    args = explored[2]
    rands = explored[3]

    for i in 0:10
        @test (i,) in args
    end

    failed = collect(zip(errored...))
    outs = failed[1]
    args = failed[2]
    rands = failed[3]

    @test length(outs) >= 1
    @test first(outs) isa AssertionError
end

struct A
    x::Int
end
@testset "Structs" begin
    function propagate(x)
        a = A(x)
        return a.x
    end
    val, symb, trace = execute(propagate, 10);
    @test val == 10
    @test symb == true

    range() = 1:10
    val, symb, trace = execute(range);
    @test symb == false
end

# @testset "Arrays" begin
# function store_and_read(x)
#     A = Array{Int, 1}(10)
#     A[5] = x
#     x = 0 # Destroy symbol character of x
#     return A[5]
#  end
# end

@testset "Complex Loops" begin
    function h0()
        acc = 0
        for i in 1:10
            acc += 1
        end
        return acc
    end

    val, symb, trace = execute(h0);
    @test val == 10
    @test symb == false

    function h1(x)
        acc = 0
        for i in 1:10
            acc += x
        end
        return acc
    end

    val, symb, trace = execute(h1, 1);
    @test val == 10
    @test symb == true
    @test length(ConcolicFuzzer.filter(trace)) == 10

    # function h2(N)
    #     acc = 0
    #     for i in 1:N
    #         acc += 1
    #     end
    #     return acc
    # end

    # val, symb, trace = execute(h2, 10);
    # @test val == 10
    # @test symb == false
    # @test_broken length(ConcolicFuzzer.filter(trace)) == 25

    # function h3(N)
    #     acc = 0
    #     for i in 1:N
    #         acc += i
    #     end
    #     return acc
    # end

    # val, symb, trace = execute(h3, 10);
    # @test val == 55
    # @test_broken symb == true
end

@testset "Fuzz supported datatypes" begin
   function over_dt(::Type{T}) where T
      if T == Int64
        return true
      else
        return false
      end
    end

    tested, errored = fuzz((x)->over_dt(ConcolicFuzzer.enumerateSupportedTypes(x)), Int)
    @test length(tested) == 11

    function over_dt2(::Type{Int8})
        return false
    end

    function over_dt2(::Type{Int64})
        return true
    end

    tested, errored = fuzz((x)->over_dt2(ConcolicFuzzer.enumerateSupportedTypes(x)), Int)
    @test length(tested) == 2
    @test length(errored) == 9
end

@testset "Fuzz and Check" begin
    function fc(y)
        @assert(0 < y)
        if y < 10
            x = y - 5
        else
            x = y - 7
        end
        ConcolicFuzzer.prove(x < y)
    end
    result = ConcolicFuzzer.fuzz_and_check(fc, Int)
    @test any(map(r->r[1], result)) == false
end

# @testset "Complicated Fuzzing" begin
#     function mysum(A, N)
#         acc = 0
#         for i in 1:N
#             acc += A[i]
#         end
#         return acc
#     end

#     fuzz_wargs((N)->mysum(ones(10),N), 1)

#     function mysum_inbounds(A, N)
#         @assert N <= length(A)
#         acc = 0
#         @inbounds for i in 1:N
#             acc += A[i]
#         end
#         return acc
#     end

#     fuzz_wargs((N)->mysum_inbounds(ones(10),N), 1)
# end


text:
# TODO allow the lower and upper bounds to be changed, like uniform.

struct BetaUniformMixture <: Distribution{Float64} end

"""
    beta_uniform(theta::Real, alpha::Real, beta::Real)

Samples a `Float64` value from a mixture of a uniform distribution on [0, 1] with probability `1-theta` and a beta distribution with parameters `alpha` and `beta` with probability `theta`.
"""
const beta_uniform = BetaUniformMixture()

function logpdf(::BetaUniformMixture, x::Real, theta::Real, alpha::Real, beta::Real)
    lbeta = log(theta) + logpdf(Beta(), x, alpha, beta)
    luniform = log(1.0 - theta)
    logsumexp(lbeta, luniform)
end

function logpdf_grad(::BetaUniformMixture, x::Real, theta::Real, alpha::Real, beta::Real)
    beta_logpdf = logpdf(Beta(), x, alpha, beta)
    uniform_logpdf = logpdf(uniform_continuous, x, 0., 1.)
    beta_grad = logpdf_grad(Beta(), x, alpha, beta)
    uniform_grad = logpdf_grad(uniform_continuous, x, 0., 1.)
    w1 = 1. / (1. + exp(log(1. - theta) + uniform_logpdf - log(theta) - beta_logpdf))
    w2 = 1. - w1
    x_deriv = w1 * beta_grad[1] + w2 * uniform_grad[1]
    alpha_deriv = w1 * beta_grad[2]
    beta_deriv = w1 * beta_grad[3]
    theta_deriv = (exp(beta_logpdf) - exp(uniform_logpdf)) / (theta * exp(beta_logpdf) + (1. - theta) * exp(uniform_logpdf))
    (x_deriv, theta_deriv, alpha_deriv, beta_deriv)
end

function random(::BetaUniformMixture, theta::Real, alpha::Real, beta::Real)
    if bernoulli(theta)
        random(Beta(), alpha, beta)
    else
        random(uniform_continuous, 0., 1.)
    end
end

(::BetaUniformMixture)(theta, alpha, beta) = random(BetaUniformMixture(), theta, alpha, beta)

has_output_grad(::BetaUniformMixture) = true
has_argument_grads(::BetaUniformMixture) = (true, true, true)

export beta_uniform

ext: jl

max_stars_repo_name: dendisuhubdy/Gen

max_issues_repo_name: kumsh/Gen

max_forks_repo_name: kumsh/Gen


text:
# Module with functions for time-domain modeling and inversion using OPESCI/devito
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January, 2017
#

module TimeModeling

using JUDI, PyCall, JOLI, SegyIO, Dierckx, Distributed, LinearAlgebra, Base.Broadcast, FFTW

import Base.*, Base./, Base.+, Base.-, Base.copy!, Base.sum, Base.ndims
import Base.Broadcast.broadcasted, Base.BroadcastStyle, Base.Broadcast.DefaultArrayStyle
import Base.getindex, Base.setindex!, Base.firstindex, Base.lastindex, Base.axes, Base.ndims
import LinearAlgebra.transpose, LinearAlgebra.conj, LinearAlgebra.vcat, LinearAlgebra.adjoint
import LinearAlgebra.vec, LinearAlgebra.dot, LinearAlgebra.norm, LinearAlgebra.abs
import Base.similar, Base.isapprox, Base.isequal, Base.broadcast!
import LinearAlgebra.rmul!
#import LinearAlgebra.A_mul_B!, LinearAlgebra.Ac_mul_B!, LinearAlgebra.BLAS.axpy!

#############################################################################
# Containers
include("ModelStructure.jl")    # model container
include("InfoStructure.jl") # basic information required by all operators
include("GeometryStructure.jl") # source or receiver setup, recording time and sampling
include("OptionsStructure.jl")
include("auxiliaryFunctions.jl")

#############################################################################
# Abstract vectors
include("judiWavefield.jl") # dense RHS (wavefield)
include("judiRHS.jl")   # sparse RHS (point source(s))
include("judiVector.jl")    # Julia data container

#############################################################################
# PDE solvers
include("time_modeling_serial.jl")  # forward/adjoint linear/nonlinear modeling
include("time_modeling_parallel.jl")    # parallelization for modeling
include("fwi_objective_serial.jl")  # FWI objective function value and gradient
include("fwi_objective_parallel.jl")    # parallelization for FWI gradient

#############################################################################
# Linear operators
include("judiModeling.jl")  # nonlinear modeling operator F (no projection operators)
include("judiProjection.jl")    # source/receiver projection operator
include("judiPDEfull.jl")   # modeling operator with source and receiver projection: P*F*P'
include("judiPDE.jl")   # modeling operator with lhs projection only: P*F
include("judiJacobian.jl")  # linearized modeling operator J

#############################################################################
# Preconditioners and optimization
include("seismic_preconditioners.jl")

end


text:
immutable NoncentralF <: ContinuousUnivariateDistribution
    ndf::Float64
    ddf::Float64
    ncp::Float64
    function NoncentralF(n::Real, d::Real, nc::Real)
	n > zero(n) && d > zero(d) && nc >= zero(nc) ||
	    error("ndf and ddf must be > 0 and ncp >= 0")
	new(float64(n), float64(d), float64(nc))
    end
end

@_jl_dist_3p NoncentralF nf

@continuous_distr_support NoncentralF 0.0 Inf

mean(d::NoncentralF) = d.ddf > 2.0 ? d.ddf / (d.ddf - 2.0) * (d.ndf + d.ncp) / d.ndf : NaN

var(d::NoncentralF) = d.ddf > 4.0 ? 2.0 * d.ddf^2 *
		       ((d.ndf+d.ncp)^2 + (d.ddf - 2.0)*(d.ndf + 2.0*d.ncp)) /
		       (d.ndf * (d.ddf - 2.0)^2 * (d.ddf - 4.0)) : NaN

function rand(d::NoncentralF)
    rn = rand(NoncentralChisq(d.ndf,d.ncp)) / d.ndf
    rd = rand(Chisq(d.ddf)) / d.ddf
    rn / rd
end


text:
using Documenter, NumCME

makedocs(sitename="𝕹𝖚𝖒ℂ𝕄𝔼",
        authors="Huy Vo",
        doctest=false,
        clean=true,
        modules=[NumCME],
        format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
        pages = [
            "Home" => "index.md",
            "Examples" => Any[
                "examples/telegraph.md",
                "examples/hog1p.md"
            ],
            "API reference" => Any[
                "api/models.md",
                "api/propensities.md",
                "api/propensity_gradients.md",
                "api/space.md",
                "api/vectors.md",
                "api/cmesolve.md",
                "api/cmesenssolve.md"
            ]
        ]
        )

deploydocs(
    repo = "github.com/voduchuy/NumCME.jl.git",
    target = "build",
    deps = Deps.pip("pygments", "mkdocs", "python-markdown-math"),
    devbranch = "main"
    # ...
)


text:
@testset "endproduct.jl" begin
    sup = Supplier(100, 10, leadtime = LeadTime(1, 3, 10))
    inv = Inventory(3, 50)
    market = Market(5, Normal, 4, 0, [30,40,20,30], 0.0)
    item = EndProduct(market, inv, sup)
    @test state(item) == [0, 30, 0.0, 40, 0.0, 20, 0.0, 30, 0.0,50,10,10,10]
    @test InventoryModels.inventory_position(item) == 80
    @test state_size(item) == 4+9
    @test action_size(item) == 2
    dest = []
    item([30])
    InventoryModels.dispatch!(item)
    @test InventoryModels.reward!(item) == -(100 + 10*30 + 30*3 + 50)
    @test state(item) == [0, 40, 0.0, 20, 0.0, 30, 0.0, 30, 0.0,30,10,10,30]
    @test market.backorder == 0
    item([0])
    InventoryModels.dispatch!(item)
    @test InventoryModels.reward!(item) == -(40)
    @test state(item) == [0.0, 20, 0.0, 30, 0.0, 30,0,40,0,0,10,30,0]
    @test market.backorder == 0
    item([0])
    InventoryModels.dispatch!(item)
    @test InventoryModels.reward!(item) == -(10*5 + 30)
    @test state(item) == [10,30, 0.0, 30,0,40,0,20,0,0,30,0,0]
    @test market.backorder == 10
    reset!(item)
    @test state(item) == [0, 30, 0.0, 40, 0.0, 20, 0.0, 30, 0.0,50,10,10,10]
    @test InventoryModels.children(item) == ()
end


text:
import Base: IteratorSize, IsInfinite, length

using TimeFrames
using TimeFrames: tonext

struct FiniteTimeFrameTrigger <: AbstractFiniteTrigger
    tf::TimeFrame
    n::Int
end

struct InfiniteTimeFrameTrigger <: AbstractInfiniteTrigger
    tf::TimeFrame
end

"""
    TimeFrameTrigger(tf::TimeFrame)

A trigger which should trigger a job at a given instant according timeframe periodicity (from [TimeFrames.jl](https://github.com/femtotrader/TimeFrames.jl))

# Example

    TimeFrameTrigger("H")

should run a job every hour
"""
function TimeFrameTrigger(tf; n=-1)
    if n < 0
        InfiniteTimeFrameTrigger(tf)
    else
        FiniteTimeFrameTrigger(tf, n)
    end
end

"""
    Trigger(tf::TimeFrame[, n=number_of_times])

Return an `TimeFrameTrigger` which should trigger a job at a given instant according timeframe periodicity. (from [TimeFrames.jl](https://github.com/femtotrader/TimeFrames.jl))
"""
Trigger(tf::TimeFrame; kwargs...) = TimeFrameTrigger(tf; kwargs...)

function get_next_dt_fire(trigger::Union{FiniteTimeFrameTrigger,InfiniteTimeFrameTrigger}, dt_previous_fire, dt_now)
    tonext(trigger.tf, dt_now)
end


text:
using RiskAdjustedLinearizations, Test, JLD2
include(joinpath(dirname(@__FILE__), "..", "..", "examples", "wachter_disaster_risk", "wachter.jl"))

# Load in guesses and true solutions
detout = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "reference", "det_ss_output.jld2"), "r")
sssout = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "reference", "homotopy_sss_output.jld2"), "r")
z = vec(detout["z"])
y = vec(detout["y"])
Ψ = zeros(length(y), length(z))

# Set up RiskAdjustedLinearization
m = WachterDisasterRisk()
ral = inplace_wachter_disaster_risk(m)
update!(ral, 1.01 .* z, 1.01 .* y, 1.01 .* Ψ)

# Solve!
@info "The following series of print statements are expected."
for i in 1:2
    try # To avoid issues arising sometimes where homotopy accidentally fails when it shouldn't
        RiskAdjustedLinearizations.homotopy!(ral, vcat(ral.z, ral.y, vec(ral.Ψ)); verbose = :high, autodiff = :central,
                                             step = .5, ftol = 1e-8) # first with finite diff NLsolve Jacobian
        break
    catch e
        update!(ral, .99 * vec(sssout["z"]), .99 * vec(sssout["y"]), .99 * sssout["Psi"])
    end
    if i == 2 # trigger error if there actually is one
        RiskAdjustedLinearizations.homotopy!(ral, vcat(ral.z, ral.y, vec(ral.Ψ)); verbose = :low, autodiff = :central,
                                             step = .5, ftol = 1e-8) # first with finite diff NLsolve Jacobian
    end
end
@test ral.z ≈ sssout["z"] atol=1e-6
@test ral.y ≈ sssout["y"] atol=1e-4
@test ral.Ψ ≈ sssout["Psi"] atol=5e-3

update!(ral, 1.01 .* z, 1.01 .* y, 1.01 .* Ψ) # now autodiff Jacobian
@test_broken RiskAdjustedLinearizations.homotopy!(ral, vcat(ral.z, ral.y, vec(ral.Ψ)); verbose = :low, autodiff = :forward,
                                                  step = .5, ftol = 1e-8) # currently can't autodiff b/c problem with chunk size selection
#=@test ral.z ≈ sssout["z"] atol=1e-6
@test ral.y ≈ sssout["y"] atol=1e-4
@test ral.Ψ ≈ sssout["Psi"] atol=5e-3=#

# Check sparse Jacobian methods work
sparsity, colorvec = compute_sparsity_pattern(ral, :homotopy; sparsity_detection = false)
jac_cache = preallocate_jac_cache(ral, :homotopy; sparsity_detection = false)
RiskAdjustedLinearizations.homotopy!(ral, 1.001 .* vcat(ral.z, ral.y, vec(ral.Ψ));
                                     verbose = :none, sparse_jacobian = true,
                                     sparsity = sparsity, colorvec = colorvec)
@test maximum(abs.(ral.z - sssout["z"])) < 1e-6
@test maximum(abs.(ral.y - sssout["y"])) < 1e-6
RiskAdjustedLinearizations.homotopy!(ral, 1.001 .* vcat(ral.z, ral.y, vec(ral.Ψ));
                                     verbose = :none, sparse_jacobian = true,
                                     jac_cache = jac_cache)
@test maximum(abs.(ral.z - sssout["z"])) < 1e-6
@test maximum(abs.(ral.y - sssout["y"])) < 1e-6
RiskAdjustedLinearizations.homotopy!(ral, 1.001 .* vcat(ral.z, ral.y, vec(ral.Ψ));
                                    verbose = :none, sparse_jacobian = true,
                                    sparsity_detection = false)
@test maximum(abs.(ral.z - sssout["z"])) < 1e-6
@test maximum(abs.(ral.y - sssout["y"])) < 1e-6

# Using SparsityDetection.jl fails
@test_broken RiskAdjustedLinearizations.homotopy!(ral, 1.001 .* vcat(ral.z, ral.y, vec(ral.Ψ));
                                    verbose = :none, sparse_jacobian = true,
                                    sparsity_detection = true)


text:
using RandomNumbers.MersenneTwisters
import Random123
using RandomNumbers.Random123
using RandomNumbers.PCG
using RandomNumbers.Xorshifts
import Random: rand
import Printf: @printf, @sprintf

include("common.jl")

function rank_timing_results(results)
    sort!(results, by = last)
    mtidx = findfirst(t -> endswith(t[1], "MersenneTwister"), results)
    timemt = results[mtidx][2]
    println("\n\nRanked:")
    for i in 1:length(results)
        rngstr, timetaken = results[i]
        @printf "%4d. %s: %6.3f ns/64 bits" i rngstr timetaken
        p = 100.0 * (timetaken-timemt)/timemt
        c = (p < 0.0 ? :green : (p > 0.0 ? :red : :normal))
        if p != 0.0
            signstr = (p > 0.0 ? "+" : "")
            printstyled("  ", signstr, @sprintf("%.2f%%", p); color=c)
        end
        println("")
    end
end

macro rngtest(rngs...)
    results = []
    for rng in rngs
        if rng.head == :call
            r = @eval $rng
            name = rng.args[1]
        else
            r = @eval $(rng.args[1])
            name = rng.args[2]
        end
        timetaken = speed_test(r)
        @printf "%20s: %6.3f ns/64 bits\n" name timetaken
        push!(results, (@sprintf("%20s", name), timetaken))
    end
    rank_timing_results(results)
end

@rngtest(
Xorshift64(123),
Xorshift64Star(123),
Xorshift128Star(123),
Xorshift128Plus(123),
Xorshift1024Star(123),
Xorshift1024Plus(123),
Xoroshiro128Plus(123),
Xoroshiro128Star(123),

MersenneTwister(123),
MT19937(123),

(Threefry2x(UInt64, (123, 321)), Threefry2x64),
(Threefry4x(UInt64, (123, 321, 456, 654)), Threefry4x64),
(Threefry2x(UInt32, (123, 321)), Threefry2x32),
(Threefry4x(UInt32, (123, 321, 456, 654)), Threefry4x32),
(Philox2x(UInt64, 123), Philox2x64),
(Philox4x(UInt64, (123, 321)), Philox4x64),
(Philox2x(UInt32, 123), Philox2x32),
(Philox4x(UInt32, (123, 321)), Philox4x32),
(AESNI1x(123), AESNI1x128),
(ARS1x(123), ARS1x128),
(AESNI4x((123, 321, 456, 654)), AESNI4x32),
(ARS4x((123, 321, 456, 654)), ARS4x32),

(PCGStateOneseq(UInt64, PCG_XSH_RS, 123), PCG_XSH_RS_128),
(PCGStateOneseq(UInt64, PCG_XSH_RR, 123), PCG_XSH_RR_128),
(PCGStateOneseq(UInt64, PCG_RXS_M_XS, 123), PCG_RXS_M_XS_64),
(PCGStateOneseq(UInt64, PCG_XSL_RR, 123), PCG_XSL_RR_128),
(PCGStateOneseq(UInt64, PCG_XSL_RR_RR, 123), PCG_XSL_RR_RR_64),

(PCGStateOneseq(UInt32, PCG_XSH_RS, 123), PCG_XSH_RS_64),
(PCGStateOneseq(UInt32, PCG_XSH_RR, 123), PCG_XSH_RR_64),
(PCGStateOneseq(UInt32, PCG_RXS_M_XS, 123), PCG_RXS_M_XS_32),
(PCGStateOneseq(UInt32, PCG_XSL_RR, 123), PCG_XSL_RR_64),

(PCGStateOneseq(UInt128, PCG_RXS_M_XS, 123), PCG_RXS_M_XS_128),
(PCGStateOneseq(UInt128, PCG_XSL_RR_RR, 123), PCG_XSL_RR_RR_128),
)


text:
# ==============================================================================

# Implicit dynamic linear maxplus system.
struct MPSysLin{T}
    A::ArrMP{T}
    B::ArrMP{T}
    C::ArrMP{T}
    D::ArrMP{T}
    x0::ArrMP{T}
end

# ==============================================================================

function Base.:(==)(x::MPSysLin, y::MPSysLin)
    (x.A == y.A) && (x.B == y.B) && (x.C == y.C) && (x.D == y.D) && (x.x0 == y.x0)
end

# ==============================================================================

function Base.show(io::IO, S::MPSysLin{T}) where T
    (@printf io "Implicit dynamic linear maxplus system:\n")
    (@printf io "  x(n) = D*x(n) + A*x(n-1) + B*u(n)\n  y(n) = C*x(n)\n  x(0) = x0\n\nwith:")
    (@printf io "\n  D  = "); Base.show(io, S.D)
    (@printf io "\n  A  = "); Base.show(io, S.A)
    (@printf io "\n  B  = "); Base.show(io, S.B)
    (@printf io "\n  C  = "); Base.show(io, S.C)
    (@printf io "\n  x0 = "); Base.show(io, S.x0)
    (@printf io "\n")
end

# ==============================================================================

function LaTeX(io::IO, S::MPSysLin)
    (@printf io "\\begin{array}{lcl}\n")
    (@printf io "x_n & = & ")
    LaTeX(io, S.D)
    (@printf io " x_n \\oplus ")
    LaTeX(io, S.A)
    (@printf io " x_{n-1} \\oplus ")
    LaTeX(io, S.B)
    (@printf io " u_n\\\\")
    (@printf io "y_n & = & ")
    LaTeX(io, S.C)
    (@printf io " x_n\\\\")
    (@printf io "x_0 & = & ")
    LaTeX(io, S.x0)
    (@printf io "\\end{array}")
end

# ==============================================================================
# Utility functions against vectors! Force seeing a vector like a matrix.

# Vector 2 matrix conversion
v2m(v::Vector) = reshape(v, length(v), 1)
v2m(v::Array) = v

# Fix Vector size returning a single param
size2(A::SpaMP) = (size(A, 1), size(A, 2))
size2(A::Array) = (size(A, 1), size(A, 2))

# ==============================================================================
# Dense

function mpsyslin(A::ArrMP{T}, B::ArrMP{T}, C::ArrMP{T}, D::ArrMP{T}, x0::ArrMP{T}) where T
    (ma,na) = size2(A)
    (ma != na) && error("Matrix A shall be squared")
    (mb,nb) = size2(B)
    ((mb != na) && (mb != 0)) && error("The row dimension of B does not agree with dimensions of A")
    (mc,nc) = size2(C)
    ((nc != na) && (mc != 0)) && error("The column dimension of C $(mc) $(nc) does not agree with dimensions of A $(ma) $(na)")
    (mx0,nx0) = size2(x0)
    ((mx0 != na) || (nx0 != min(na, 1))) && error("Dimensions of x0 do not agree with dimensions of A")
    (md,nd) = size2(D)
    ((md != na) || (nd != na)) && error("The column dimension of D does not agree with dimensions of A")
    MPSysLin(A, v2m(B), v2m(C), D, v2m(x0))
end

function mpsyslin(A::ArrMP{T}, B::ArrMP{T}, C::ArrMP{T}, D::ArrMP{T}) where T
    mpsyslin(A, B, C, D, mpfzeros(T, size(A,2), 1))
end

function mpsyslin(A::ArrMP{T}, B::ArrMP{T}, C::ArrMP{T}) where T
    na = size(A,2)
    mpsyslin(A, B, C, mpeye(T, na, na), mpfzeros(T, na, 1))
end

# ==============================================================================
# Sparse

function mpsyslin(A::SpaMP{T}, B::SpaMP{T}, C::SpaMP{T}, D::SpaMP{T}, x0::SpaMP{T}) where T
    mpsyslin(full(A), full(B), full(C), full(D), full(x0))
end

function mpsyslin(A::SpaMP{T}, B::SpaMP{T}, C::SpaMP{T}, D::SpaMP{T}) where T
    mpsyslin(full(A), full(B), full(C), full(D), mpfzeros(T, size(A,2), 1))
end

function mpsyslin(A::SpaMP{T}, B::SpaMP{T}, C::SpaMP{T}) where T
    na = size(A,2)
    mpsyslin(full(A), full(B), full(C), mpeye(T, na, na), mpfzeros(T, na, 1))
end

# ==============================================================================
# %mpls_a_mpls.sci
# Parallel composition

function Base.:(+)(x::MPSysLin{T}, y::MPSysLin{T}) where T
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    MPSysLin([x.A mpfzeros(T, n1, n2); mpfzeros(T, n2, n1) y.A],
             [x.B; y.B],
             [x.C y.C],
             [x.D mpfzeros(T, n1, n2); mpfzeros(T, n2, n1) y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# %mpls_m_mpls.sci
# Series composition.

function Base.:(*)(y::MPSysLin{T}, x::MPSysLin{T}) where T
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    nb1 = size(x.B, 2)
    nc2 = size(y.C, 1)
    MPSysLin([x.A mpfzeros(T, n1, n2); mpfzeros(T, n2, n1) y.A],
             [x.B; mpfzeros(T, n2, nb1)],
             [mpfzeros(T, nc2, n1) y.C],
             [x.D mpfzeros(T, n1, n2); y.B * x.C y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# %mpls_g_mpls.sci
# Diagonal composition

function Base.:(|)(x::MPSysLin{T}, y::MPSysLin{T}) where T
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    nb1 = size(x.B, 2)
    nb2 = size(y.B, 2)
    nc1 = size(x.C, 1)
    nc2 = size(y.C, 1)
    MPSysLin([x.A mpfzeros(T, n1, n2); mpfzeros(T, n2, n1) y.A],
             [x.B mpfzeros(T, n1, nb2); mpfzeros(T, n2, nb1) y.B],
             [x.C mpfzeros(T, nc1, n2); mpfzeros(T, nc2, n1) y.C],
             [x.D mpfzeros(T, n1, n2); mpfzeros(T, n2, n1) y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# %mpls_m_mpls.sci

# ==============================================================================
# %mpls_f_mpls.sci
# computes [S1;S2] that is :  inputs in common,  concatenation of outputs

function Base.:vcat(x::MPSysLin{T}, y::MPSysLin{T}) where T
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    nc1 = size(x.C, 1)
    nc2 = size(y.C, 1)
    MPSysLin([x.A mpfzeros(T, n1, n2); mpfzeros(T, n2, n1) y.A],
             [x.B; y.B],
             [x.C mpfzeros(T, nc1, n2); mpfzeros(T, nc2, n1) y.C],
             [x.D mpfzeros(T, n1, n2); mpfzeros(T, n2, n1) y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# %mpls_c_mpls.sci
# computes [S1,S2] that is  concatenation of inputs, addition of outputs
# TODO [S1 S2] works but not [S1, S2]

function Base.:hcat(x::MPSysLin{T}, y::MPSysLin{T}) where T
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    nb1 = size(x.B, 2)
    nb2 = size(y.B, 2)
    MPSysLin([x.A mpfzeros(T, n1, n2); mpfzeros(T, n2, n1) y.A],
             [x.B mpfzeros(T, n1, nb2); mpfzeros(T, n2, nb1) y.B],
             [x.C y.C],
             [x.D mpfzeros(T, n1, n2); mpfzeros(T, n2, n1) y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# %mpls_v_mpls.sci
# Feedback composition: computes star(S1*S2)*S1 in state-space form.

function Base.:(/)(x::MPSysLin{T}, y::MPSysLin{T}) where T
    n1 = size(x.A, 1)
    n2 = size(y.A, 1)
    nb1 = size(x.B, 2)
    nc1 = size(y.C, 1)
    MPSysLin([x.A mpfzeros(T, n1, n2); mpfzeros(T, n2, n1) y.A],
             [x.B; mpfzeros(T, n2, nb1)],
             [x.C mpfzeros(T, nc1, n2)],
             [x.D x.B * y.C; y.B * x.C y.D],
             [x.x0; y.x0])
end

# ==============================================================================
# %mpls_m_talg.sci
# S1 * k

function Base.:(*)(S::MPSysLin{T}, k::U) where {T,U}
    MPSysLin(S.A, S.B * k, S.C, S.D, S.x0)
end

function Base.:(*)(k::U, S::MPSysLin{T}) where {T,U}
    MPSysLin(S.A, S.B, S.C * k, S.D, S.x0)
end

# %mpls_m_talg.sci
function Base.:(*)(S::MPSysLin{T}, M::SpaMP{T}) where T
    MPSysLin(S.A, S.B * M, S.C, S.D, S.x0)
end

function Base.:(*)(M::SpaMP{T}, S::MPSysLin{T}) where T
    MPSysLin(S.A, M * S.B, S.C, S.D, S.x0)
end

# ==============================================================================
# mplssize.sci
function Base.:size(S::MPSysLin{T}) where T
    [size(S.B, 1), size(S.B, 2), size(S.C, 1)]
end

# ==============================================================================
# explicit.sci
#mpstar ko pour identity

function mpexplicit(S::MPSysLin{T}) where T
    ds = mpstar(S.D)
    bs = ds * S.B
    ac = [ds * S.A; S.C]
    zerocol = map(x -> Bool(x.λ), mpones(T, 1, size(ac, 1)) * (ac .!= mpzero(T)))
    keep = findall(zerocol[1,:])
    c = [S.C[i] for i in keep]
    mpsyslin([ac[i, j] for i in keep, j in keep],
             [bs[i] for i in keep],
             reshape(c, 1, length(c)))
end

# ==============================================================================

#TODO
"""
    mpsyslin(A::SpaMP{T}, x0::Vector{MP{T}}, k::Int64; history=false)

Compute states X of an autonomous linear max-plus system:
`x(n+1) = A x(n)`  for n = 0 .. k
where:
- `A` is a system matrix
- `x0` is a initial state vector,
- `k` is the number of iterations
- when history is set to true save all computed states, else return the last one.

# Examples
```julia-repl
julia> S1 = mpsyslin(MP([1.0 2; 3 4]), MP([0.0; 0]), MP([0.0 0]), mpeye(Float64, 2,2));
julia> mpsimul(S1, MP(1:10))
1×10 Array{MP{Float64},2}:
 1  5  9  13  17  21  25  29  33  37
julia> mpsimul(S1, MP(1:10), history=false)
1×1 Array{MP{Float64},2}:
 37
```
"""
function mpsimul(S::MPSysLin{T}, u::ArrMP{T}, history::Bool) where T
    x = S.x0
    k = size(u, 1)
    if history
        Y = mpones(T, size(S.C, 1), k)
        for i = 1:k
            x = S.A * x + S.B * u[i,:]
            Y[:,i] = S.C * x
        end
        Y
    else
        for i = 1:k
            x = S.A * x + S.B * u[i,:]
        end
        Y = S.C * x
        Y[:,1]
    end
end

function mpsimul(S::MPSysLin{T}, u::VecMP{U}, history::Bool) where {T,U}
    mpsimul(S, map(x -> MP(T(x.λ)), reshape(u, length(u), 1)), history)
end


text:
module ExtendedVariantDictionary

export Variants

const Variants = String[
	"AddSeekers",
	"AffectExistingChasers",
	"AirFriction",
	"AllStrawberriesAreGoldens",
	"AllowLeavingTheoBehind",
	"AllowThrowingTheoOffscreen",
	"AnxietyEffect",
	"BackgroundBlurLevel",
	"BackgroundBrightness",
	"BadelineAttackPattern",
	"BadelineBossCount",
	"BadelineBossNodeCount",
	"BadelineBossesEverywhere",
	"BadelineChasersEverywhere",
	"BadelineLag",
	"BlurLevel",
	"BoostMultiplier",
	"BounceEverywhere",
	"ChangePatternsOfExistingBosses",
	"ChangeVariantsInterval",
	"ChangeVariantsRandomly",
	"ChaserCount",
	"ColorGrading",
	"CoyoteTime",
	"DashCount",
	"DashDirection",
	"DashLength",
	"DashSpeed",
	"DashTrailAllTheTime",
	"DelayBetweenBadelines",
	"DisableClimbJumping",
	"DisableClimbingUpOrDown",
	"DisableMadelineSpotlight",
	"DisableNeutralJumping",
	"DisableOshiroSlowdown",
	"DisableRefillsOnScreenTransition",
	"DisableSeekerSlowdown",
	"DisableSuperBoosts",
	"DisableWallJumping",
	"DisplayDashCount",
	"DontRefillDashOnGround",
	"DontRefillStaminaOnGround",
	"EveryJumpIsUltra",
	"EverythingIsUnderwater",
	"ExplodeLaunchSpeed",
	"FallSpeed",
	"FirstBadelineSpawnRandom",
	"ForceDuckOnGround",
	"ForegroundEffectOpacity",
	"Friction",
	"FriendlyBadelineFollower",
	"GameSpeed",
	"GlitchEffect",
	"Gravity",
	"HeldDash",
	"HiccupStrength",
	"HyperdashSpeed",
	"InvertDashes",
	"InvertGrab",
	"InvertHorizontalControls",
	"InvertVerticalControls",
	"JellyfishEverywhere",
	"JumpCount",
	"JumpHeight",
	"NoFreezeFrames",
	"OshiroCount",
	"OshiroEverywhere",
	"PreserveExtraDashesUnderwater",
	"RefillJumpsOnDashRefill",
	"RegularHiccups",
	"RestoreDashesOnRespawn",
	"ReverseOshiroCount",
	"RisingLavaEverywhere",
	"RisingLavaSpeed",
	"ScreenShakeIntensity",
	"SnowballDelay",
	"SnowballsEverywhere",
	"SpeedX",
	"Stamina",
	"SuperdashSteeringSpeed",
	"SwimmingSpeed",
	"TheoCrystalsEverywhere",
	"UpsideDown",
	"WallBouncingSpeed",
	"WindEverywhere",
	"ZoomLevel"
]

end

max_stars_repo_name: Kokasgui/Everest-ExtendedVariants

max_issues_repo_name: Kokasgui/Everest-ExtendedVariants

max_forks_repo_name: Kokasgui/Everest-ExtendedVariants


text:
module TestConnectorComp

using Mimi
using Test

import Mimi:
    reset_compdefs, compdef

reset_compdefs()

@defcomp Long begin
    x = Parameter(index=[time])
end

@defcomp Short begin
    a = Parameter()
    b = Variable(index=[time])

    function run_timestep(p, v, d, t)
        v.b[t] = p.a * t.t
    end
end

years = 2000:2010
late_start = 2005
dim = Mimi.Dimension(years)

#------------------------------------------------------------------------------
#  1. Use the connect_param! method with backup data (ConnectorComp gets added
#       under the hood during build)
#------------------------------------------------------------------------------

model1 = Model()
set_dimension!(model1, :time, years)
add_comp!(model1, Short; first=late_start)
add_comp!(model1, Long)
set_param!(model1, :Short, :a, 2.)
connect_param!(model1, :Long, :x, :Short, :b, zeros(length(years)))

run(model1)

@test length(components(model1.mi)) == 3    # ConnectorComp is added to the list of components in the model isntance
@test length(model1.md.comp_defs) == 2      # The ConnectorComp shows up in the model instance but not the model definition

b = model1[:Short, :b]
x = model1[:Long, :x]

# Test that all allocated datum arrays are the full length of the time dimension
@test length(b) == length(years)
@test length(x) == length(years)

@test all(ismissing, b[1:dim[late_start]-1])
@test all(iszero, x[1:dim[late_start]-1])

# Test the values are right after the late start
@test b[dim[late_start]:end] ==
    x[dim[late_start]:end] ==
    [2 * i for i in 1:(years[end]-late_start + 1)]

@test Mimi.datum_size(model1.md, Mimi.compdef(model1.md, :Long), :x) == (length(years),)

# Test the dataframe size
b = getdataframe(model1, :Short, :b)
@test size(b) == (length(years), 2)

#------------------------------------------------------------------------------
#  2. Test with a short component that ends early (and test Variable timesteps)
#------------------------------------------------------------------------------

years_variable = [2000:2004..., 2005:5:2030...]
dim_variable = Mimi.Dimension(years_variable)

early_last = 2010

model2 = Model()
set_dimension!(model2, :time, years_variable)
add_comp!(model2, Short; last=early_last)
add_comp!(model2, Long)
set_param!(model2, :Short, :a, 2.)
connect_param!(model2, :Long, :x, :Short, :b, zeros(length(years_variable)))

run(model2)

@test length(components(model2.mi)) == 3
@test length(model2.md.comp_defs) == 2      # The ConnectorComp shows up in the model instance but not the model definition

b = model2[:Short, :b]
x = model2[:Long, :x]

# Test that all allocated datum arrays are the full length of the time dimension
@test length(b) == length(years_variable)
@test length(x) == length(years_variable)

@test all(ismissing, b[dim_variable[early_last]+1 : end])
@test all(iszero, x[dim_variable[early_last]+1 : end])

# Test the values are right after the late start
@test b[1 : dim_variable[early_last]] ==
    x[1 : dim_variable[early_last]] ==
    [2 * i for i in 1:dim_variable[early_last]]

#------------------------------------------------------------------------------
#  3. A model that requires multiregional ConnectorComps
#------------------------------------------------------------------------------

@defcomp Long_multi begin
    regions = Index()

    x = Parameter(index = [time, regions])
end

@defcomp Short_multi begin
    regions = Index()

    a = Parameter(index=[regions])
    b = Variable(index=[time, regions])

    function run_timestep(p, v, d, ts)
        for r in d.regions
            v.b[ts, r] = ts.t + p.a[r]
        end
    end
end

regions = [:A, :B]

model3 = Model()
set_dimension!(model3, :time, years)
set_dimension!(model3, :regions, regions)
add_comp!(model3, Short_multi; first=late_start)
add_comp!(model3, Long_multi)
set_param!(model3, :Short_multi, :a, [1,2])
connect_param!(model3, :Long_multi, :x, :Short_multi, :b, zeros(length(years), length(regions)))

run(model3)

@test length(components(model3.mi)) == 3
@test length(model3.md.comp_defs) == 2      # The ConnectorComp shows up in the model instance but not the model definition

b = model3[:Short_multi, :b]
x = model3[:Long_multi, :x]

# Test that all allocated datum arrays are the full length of the time dimension
@test size(b) == (length(years), length(regions))
@test size(x) == (length(years), length(regions))

@test all(ismissing, b[1:dim[late_start]-1, :])
@test all(iszero, x[1:dim[late_start]-1, :])

# Test the values are right after the late start
@test b[dim[late_start]:end, :] ==
    x[dim[late_start]:end, :] ==
    [[i + 1 for i in 1:(years[end]-late_start + 1)] [i + 2 for i in 1:(years[end]-late_start + 1)]]

#------------------------------------------------------------------------------
#  4. Test where the short component starts late and ends early
#------------------------------------------------------------------------------

first, last = 2002, 2007

model4 = Model()
set_dimension!(model4, :time, years)
set_dimension!(model4, :regions, regions)
add_comp!(model4, Short_multi; first=first, last=last)
add_comp!(model4, Long_multi)

set_param!(model4, :Short_multi, :a, [1,2])
connect_param!(model4, :Long_multi=>:x, :Short_multi=>:b, zeros(length(years), length(regions)))

run(model4)

@test length(components(model4.mi)) == 3
@test length(model4.md.comp_defs) == 2      # The ConnectorComp shows up in the model instance but not the model definition

b = model4[:Short_multi, :b]
x = model4[:Long_multi, :x]

# Test that all allocated datum arrays are the full length of the time dimension
@test size(b) == (length(years), length(regions))
@test size(x) == (length(years), length(regions))

@test all(ismissing, b[1:dim[first]-1, :])
@test all(ismissing, b[dim[last]+1:end, :])
@test all(iszero, x[1:dim[first]-1, :])
@test all(iszero, x[dim[last]+1:end, :])

# Test the values are right after the late start
@test b[dim[first]:dim[last], :] ==
    x[dim[first]:dim[last], :] ==
    [[i + 1 for i in 1:(years[end]-late_start + 1)] [i + 2 for i in 1:(years[end]-late_start + 1)]]

#------------------------------------------------------------------------------
#  5. Test errors with backup data
#------------------------------------------------------------------------------

late_start_long = 2002

model5 = Model()
set_dimension!(model5, :time, years)
add_comp!(model5, Short; first = late_start)
add_comp!(model5, Long; first = late_start_long)    # starts later as well, so backup data needs to match this size
set_param!(model5, :Short, :a, 2)

# A. test wrong size (needs to be length of component, not length of model)
@test_throws ErrorException connect_param!(model5, :Long=>:x, :Short=>:b, zeros(length(years)))
@test_throws ErrorException connect_param!(model4, :Long_multi=>:x, :Short_multi=>:b, zeros(length(years), length(regions)+1)) # test case with >1 dimension

# B. test no backup data provided
@test_throws ErrorException connect_param!(model5, :Long=>:x, :Short=>:b)   # Error because no backup data provided

#------------------------------------------------------------------------------
#  6. Test connecting Short component to Long component (does not add a
#       connector component)
#------------------------------------------------------------------------------

@defcomp foo begin
    par = Parameter(index=[time])
    var = Variable(index=[time])
    function run_timestep(p, v, d, ts)
        v.var[ts] = p.par[ts]
    end
end

model6 = Model()
set_dimension!(model6, :time, years)
add_comp!(model6, foo, :Long)
add_comp!(model6, foo, :Short; first=late_start)
connect_param!(model6, :Short=>:par, :Long=>:var)
set_param!(model6, :Long, :par, years)

run(model6)

@test length(components(model6.mi)) == 2

short_par = model6[:Short, :par]
short_var = model6[:Short, :var]

@test short_par == years    # The parameter has values instead of `missing` for years when this component doesn't run,
                            # because they are coming from the longer component that did run

@test all(ismissing, short_var[1:dim[late_start]-1])
@test short_var[dim[late_start]:end] == years[dim[late_start]:end]

end #module


text:
function _datasetregistry(; name = "Datasets")
    registry = Registry(
        (
            id = Field(
                String,
                name = "ID",
                formatfn=FeatureRegistries.string_format),
           description = Field(
                String;
                name = "Description",
                optional = true,
                description = "More information about the dataset",
                formatfn=FeatureRegistries.md_format),
            size = Field(
                String;
                name = "Size",
                description = "Download size of the dataset",
                optional = true),
            tags = Field(
                Vector{String};
                name = "Tags",
                defaultfn = (row, key) -> String[]),
            package = Field(
                Module;
                name = "Package",
                formatfn = FeatureRegistries.code_format),
            downloaded = Field(
                Bool;
                name = "Is downloaded",
                description = """
                    Whether the dataset has been downloaded and is available locally.
                    Updates after session restart.
                """,
                computefn = (row, key) -> isavailable(row.loader)),
            loader = Field(
                DatasetLoader;
                name = "Dataset loader",
                formatfn = x -> FeatureRegistries.type_format),
        );
        name,
        loadfn = function (row)
            if row.loader isa DataDepLoader && startswith(row.loader.datadep, "fastai-")
                # Change download format for fastai datasets without having to redownload
                # The download process was slightly changed; this saves having to
                # redownload to update.
                dir = Datasets.loaddata(row.loader)
                if isdir(joinpath(dir, row.id))
                    cd(dir) do
                        temp = mktempdir()
                        mv(joinpath(dir, row.id), temp, force=true)
                        mv(temp, pwd(), force=true)
                    end
                end
            end
            loaddata(row.loader)
        end,
        description = """
        A registry for datasets. `load`ing an entry will download a dataset (if it
        hasn't been already), and return a path to where the files were downloaded.

        ```julia
        path = load(datasets()[id])
        ```

        See `datarecipes` to load these datasets in a format compatible with learning
        tasks.
        """
    )
    return registry
end

const DATASETS = _datasetregistry()

"""
    datasets(; filters...)

Show a registry of available datasets. Pass in filters as keyword
arguments to look at a subset.

See also [finding functionality](/documents/docs/discovery.md), [`learningtasks`](#),
and [`datarecipes`](#). For more information on registries, see
[FeatureRegistries.jl](https://github.com/lorenzoh/FeatureRegistries.jl).

## Examples

Show all available learning tasks:

{cell}
```julia
using FastAI
datasets()
```

Download a dataset:

{cell}
```julia
path = load(datasets()["imagenette2-160"])
```

Get an explanation of fields in the dataset registry:

{cell}
```julia
info(datasets())
```

Show all datasets with `"image"` in their name:

{cell}
```julia
datasets(id="image")
```
"""
datasets(; kwargs...) = isempty(kwargs) ? DATASETS : filter(DATASETS; kwargs...)

function _registerdatasets(registry::Registry)
    for config in FastAI.Datasets.DATASETCONFIGS
        id = config.datadepname
        haskey(registry, id) && continue
        push!(registry, (
            id,
            description = config.description == "" ? missing : config.description,
            loader = DataDepLoader("fastai-$(config.datadepname)"),
            size=config.size === "???" ? missing : config.size,
            package = FastAI,
        ))
    end
end


text:
using Documenter, Literate
using TimeSeriesCausality

EXAMPLE_DIR = joinpath(@__DIR__, "literate")
OUT_DIR = joinpath(@__DIR__, "src/generated")

## Use Literate.jl to generate docs and notebooks of examples
for example in readdir(EXAMPLE_DIR)
    EXAMPLE = joinpath(EXAMPLE_DIR, example)

    Literate.markdown(EXAMPLE, OUT_DIR; documenter=true) # markdown for Documenter.jl
    Literate.notebook(EXAMPLE, OUT_DIR) # .ipynb notebook
    Literate.script(EXAMPLE, OUT_DIR) # .jl script
end

## Build docs
makedocs(;
    sitename="TimeSeriesCausality.jl",
    format=Documenter.HTML(),
    modules=[TimeSeriesCausality],
    pages=["Home" => "index.md",
           "PSI Examples" => "generated/psi_examples.md",
           "Granger Examples" => "generated/granger_examples.md",
           "Benchmark" => "generated/benchmark.md"],
)

deploydocs(;
    repo="github.com/ssnio/TimeSeriesCausality.jl.git", devbranch="main", branch="gh-pages"
)


text:
"""
    pgba_heuristic(inst::Instance, maxbreaks::Int)

Finds a heuristic start solution satisfying all hard contraints of a given
instance making use of a pattern generating decomposition scheme.

# Arguments
- `inst::Instance` - the instance to be solved
- `maxbreaks::Int` - limits the number of patterns generated by setting an upper
    bound on the number of breaks a pattern is allowed to have

# Returns
- `Matrix{Int}` - a feasible schedule
"""
function pgba_heuristic(
    inst::Instance,
    maxbreaks::Int,
    schedbuf::Matrix{Int};
    gamtimelimit = GRB_INFINITY,
)

    nteams, nslots = inst.rr2.nteams, inst.rr2.nslots

    # pattern data
    pid = 0
    patternidx = Dict{Int,Int}()
    team_to_patidx = [Vector{Int}() for i = 1:nteams]
    patternCons, patternSetCons, genCons = group_constraints(inst; f = ishard)
    lazycutbuf = Vector{PGBACut}()
    for nbreaks = 0:maxbreaks
        #adding new patterns
        for i = 1:nteams
            pat = generate_patterns(inst.rr2, nbreaks, patternCons[i])
            for p in pat
                idx = get!(patternidx, p) do
                    pid += 1
                end
                push!(team_to_patidx[i], patternidx[p])
            end
        end
        # model is always infeasible if no patterns with breaks are included
        nbreaks == 0 && continue
        # building model
        patterns = sort!(collect(keys(patternidx)), by = k -> patternidx[k])
        model = pgba_model(inst.rr2, patterns, team_to_patidx, patternSetCons)
        x = model[:x]
        # adding previously generated lazy cuts to model
        for cut in lazycutbuf
            @constraint(
                model,
                sum(x[a[1], a[2]] for a in cut.assignments) <= cut.ub
            )
        end
        # solve with callback
        MOI.set(model, MOI.RawParameter("LazyConstraints"), 1)
        MOI.set(
            model,
            Gurobi.CallbackFunction(),
            (cb_data, cb_where) -> pgba_callback(
                inst.rr2,
                model,
                patterns,
                genCons,
                lazycutbuf,
                schedbuf,
                gamtimelimit,
                cb_data,
                cb_where,
            ),
        )
        optimize!(model)

        # retrieve solution if feasible solution found
        if termination_status(model) == MOI.OPTIMAL
            return schedbuf
        end
    end
end

"""
    pgba_model(
        rr2::RR2,
        patterns::Vector{Int64},
        teamToPatidx::Vector{Vector{Int64}},
        cons::Vector{T},
    ) where {T<:Constraint}

Builds the master problem of the PGBA decomposition.

# Arguments
- `rr2::RR2` - tournament format data
- `patterns::Vector{Int64}` - each pattern is represented as a `Int64` where the
    `s`-th bit indicates if the pattern has a home game in slot `s`.
- `teamToPatidx::Vector{Vector{Int64}}` - `teamToPatidx[i]` is a vector of
    indices of patterns feasible for team `i` in `patterns`.
- `cons::Vector{T}` - Vector of patternset constraints

# Model Description

The model uses binary variables `x[i,p]` indicating if team `i` gets
assigned pattern `patterns[p]`.
"""
function pgba_model(
    rr2::RR2,
    patterns::Vector{Int64},
    teamToPatidx::Vector{Vector{Int64}},
    cons::Vector{T},
) where {T<:Constraint}

    model = direct_model(Gurobi.Optimizer(GRB_ENV))
    npatterns = length(patterns)
    nteams, nslots = rr2.nteams, rr2.nslots

    # 1 if pattern at index p has home game in slot s, 0 otw.
    h(p, s) = patterns[p] >>> (s - 1) & 1

    @variable(model, x[1:nteams, 1:npatterns], Bin)

    # each team gets assigned exactly one of its feasible patterns
    @constraint(model, [i = 1:nteams], sum(x[i, teamToPatidx[i]]) == 1)
    @constraint(model, [i = 1:nteams], sum(x[i, :]) == 1)

    # same number of home and away games in each slot
    @constraint(
        model,
        [s = 1:nslots],
        sum(sum(h(p, s) * x[i, p] for p in 1:npatterns) for i in 1:nteams) ==
        nteams ÷ 2
    )

    for con in cons
        add_patternset_con!(model, con, nteams, patterns)
    end

    return model
end

"""
    pgba_callback(
        rr2::RR2,
        model::Model,
        patterns::Vector{Int},
        gencons::Vector{Constraint},
        lazycuts::Vector{PGBACut},
        schedbuf::Matrix{Int},
        useiis::Bool,
        cb_data,
        cb_where::Cint,
    )

Callback function used to generate the Gurobi callback for the master problem.

# Arguments

- `rr2::RR2`- tournament format data
- `model::Model`- the direct Gurobi model
- `patterns::Vector{Int64}` - each pattern is represented as a `Int64` where the
    `s`-th bit indicates if the pattern has a home game in slot `s`.
- `gencons::Vector{Constraint}`- generic constraints of the instance
- `lazycuts::Vector{PGBACut}` - stores all cuts generated in callbacks
- `schedbuf::Matrix{Int}` - solution schedule buffer
- `cb_data`, `cb_where::Cint`- Gurobi callback arguments
"""
function pgba_callback(
    rr2::RR2,
    model::Model,
    patterns::Vector{Int},
    gencons::Vector{Constraint},
    lazycuts::Vector{PGBACut},
    schedbuf::Matrix{Int},
    gamtimelimit::Float64,
    cb_data::Gurobi.CallbackData,
    cb_where::Cint,
)
    cb_where != GRB_CB_MIPSOL && return
    x = model[:x]
    patternset, patidx = cb_load_solution(model, patterns, cb_data, cb_where)

    # check meeting ranges and add cuts accordingly. if cuts are added return.
    infeasPairs = meeting_ranges_check(model, rr2, patternset)
    if !isempty(infeasPairs)
        for ij in infeasPairs
            addcut!(model, lazycuts, cb_data, ij, patidx, true)
        end
        return
    end

    # check pattern diversity with incr. subsetsize. Once cut is added return.
    for subsetsize = 3:rr2.nteams
        infeasTeams =
            pattern_diversity_check(model, rr2, subsetsize, patternset)
        if infeasTeams !== nothing
            addcut!(model, lazycuts, cb_data, infeasTeams, patidx, false)
            @show subsetsize
            return
        end
    end

    val = assign_games!(
        schedbuf,
        rr2,
        patternset,
        gencons;
        timelimit = gamtimelimit,
        usestart = false,
    )

    if val == GRB_INFINITY
        addcut!(model, lazycuts, cb_data, 1:rr2.nteams, patidx, false)
    end
end

"""
    get_patterns_from_breaks(nslots::Int, breaks::Vector{Int})

Computes the two potential patterns of length `nslots` with breaks in slots
given by `breaks`.
"""
function get_patterns_from_breaks(nslots::Int, breaks::Vector{Int})
    pat1 = sum(2^i for i = 0:2:nslots-1)
    pat2 = sum(2^i for i = 1:2:nslots-1)
    for bridx = 1:2:length(breaks)
        idx1 = breaks[bridx]
        idx2 = bridx == length(breaks) ? nslots : breaks[bridx+1] - 1
        pat1 = flipbits(pat1, idx1, idx2)
        pat2 = flipbits(pat2, idx1, idx2)
    end
    return pat1, pat2
end

"""
    generate_patterns(rr2::RR2, numbreaks::Int, patterncons::Vector{Constraint})

Generate all feasible patterns with a given number of breaks wrt. to some
pattern constraints.

# Arguments
- `rr2::RR2` - tournament format data
- `numbreaks::Int` - number of breaks each generated pattern must have
- `patterncons::Vector{Constraint}` - pattern constraints to be satisfied

# Returns
- `Set{Int64}` - the set of patterns satisfying `patterncons` where each pattern
    is represented as a 64-bit integer
"""
function generate_patterns(
    rr2::RR2,
    numbreaks::Int,
    patterncons::Vector{T},
) where {T<:Constraint}

    patterns = Set{Int64}()
    for breaks in combinations(2:rr2.nslots, numbreaks)
        pat1, pat2 = get_patterns_from_breaks(rr2.nslots, breaks)
        # only check pat1 since pat2 is simply bitwise complement
        if isPattern(rr2.nslots, pat1)
            # validating pattern constraints for pat1
            if all(c -> validate_pattern_constraint(rr2, pat1, c), patterncons)
                push!(patterns, pat1)
            end
            # validating pattern constraints for pat2
            if all(c -> validate_pattern_constraint(rr2, pat2, c), patterncons)
                push!(patterns, pat2)
            end
        end
    end
    return patterns
end

"""
    add_patternset_con!(
        model::Model,
        con::T,
        nteams::Int,
        patterns::Vector{Int}
    ) where {T <: Union{CAcon, BRcon}

Adds a patternset constraint `con` to `model`.
"""
function add_patternset_con! end

function add_patternset_con!(
    model::Model,
    con::CAcon,
    nteams::Int,
    patterns::Vector{Int},
)
    @assert con.mode != MODE_HA
    x = model[:x]
    if con.mode == MODE_H
        @constraint(
            model,
            sum(
                count_ones(patterns[p] & con.slots.binrepr) * x[i, p] for
                i in con.teams1, p in 1:length(patterns)
            ) <= con.ub
        )
    else # mode_A
        @constraint(
            model,
            sum(
                count_ones(~patterns[p] & con.slots.binrepr) * x[i, p] for
                i in con.teams1, p in 1:length(patterns)
            ) <= con.ub
        )
    end
end

function add_patternset_con!(
    model::Model,
    con::BRcon,
    nteams::Int,
    patterns::Vector{Int},
)
    @assert con.mode == MODE_HA

    x = model[:x]
    @constraint(
        model,
        sum(
            (
                count_ones(
                    patterns[p] & (patterns[p] << 1) & con.slots.binrepr,
                ) + count_ones(
                    ~patterns[p] & (~patterns[p] << 1) & con.slots.binrepr,
                )
            ) * x[i, p] for i in con.teams, p = 1:length(patterns)
        ) <= con.ub
    )
end

function pattern_diversity_model(
    rr2::RR2,
    model::Model,
    patternset::Vector{Int64},
    subsetsize::Int,
    slotrange::UnitRange{Int64},
)

    ps(j, s) = patternset[j] >>> (s - 1) & 1
    ubm = direct_model(Gurobi.Optimizer(GRB_ENV_2))
    MOI.set(ubm, MOI.RawParameter("OutputFlag"), 0)
    @variable(ubm, b[slotrange] >= 0)
    @variable(ubm, a[1:rr2.nteams], Bin)
    @variable(ubm, d[slotrange], Bin)
    @constraint(ubm, sum(a) == subsetsize)

    @constraint(
        ubm,
        [s in slotrange],
        b[s] - sum(ps(j, s) * a[j] for j = 1:rr2.nteams) + subsetsize -
        subsetsize * d[s] >= 0
    )
    @constraint(
        ubm,
        [s in slotrange],
        b[s] - sum((1 - ps(j, s)) * a[j] for j = 1:rr2.nteams) +
        subsetsize * d[s] >= 0
    )
    @objective(ubm, Min, sum(b))
    return ubm
end

function pattern_diversity_check(model, rr2, subsetsize, patternset)
    if rr2.isphased
        for phase = 1:2
            slots = phase == 1 ? (1:rr2.nslots÷2) : (rr2.nslots÷2+1:rr2.nslots)
            ubm = pattern_diversity_model(
                rr2,
                model,
                patternset,
                subsetsize,
                1:rr2.nslots÷2,
            )
            optimize!(ubm)
        end
        if objective_value(ubm) < (subsetsize * (subsetsize - 1)) / 2
            a = ubm[:a]
            return filter(i -> value(a[i]) ≈ 1, 1:rr2.nteams)

        end
    else
        ubm = pattern_diversity_model(
            rr2,
            model,
            patternset,
            subsetsize,
            1:rr2.nslots,
        )
        optimize!(ubm)
        if objective_value(ubm) < (subsetsize * (subsetsize - 1))
            a = ubm[:a]
            return filter(i -> value(a[i]) ≈ 1, 1:rr2.nteams)
        end
    end
    return nothing
end

function meeting_ranges_check(model, rr2, patternset)
    infeasPairs = Vector{Tuple{Int,Int}}()
    for (i, j) in combinations(1:rr2.nteams, 2)
        hg = patternset[i] & ~patternset[j] # slots for game (i,j)
        ag = ~patternset[i] & patternset[j] # slots for game (j,i)
        phase1 = 2^(rr2.nslots ÷ 2) - 1
        phase2 = phase1 << (rr2.nslots ÷ 2)
        if hg == 0 || ag == 0
            push!(infeasPairs, (i, j))
        elseif rr2.isphased
            if hg & phase1 == ag & phase1 == 0 ||
               hg & phase2 == ag & phase2 == 0
                push!(infeasPairs, (i, j))
            end
        end
    end
    return infeasPairs
end

"""
    cb_load_solution(model::Model, cb_data, cb_where::Cint)

Loads solution values for the PGBA master `model` during a callback.

# Returns
`Vector{Int}`, `Vector{Int}` - patternset and pattern index vectors
"""
function cb_load_solution(model::Model, patterns::Vector{Int}, cb_data, cb_where::Cint)
    @assert cb_where == GRB_CB_MIPSOL

    Gurobi.load_callback_variable_primal(cb_data, cb_where)
    x_val = callback_value.(Ref(cb_data), model[:x])
    patternset = Vector{Int64}(undef, size(x_val, 1))
    idx = Vector{Int64}(undef, size(x_val, 1))
    for p = 1:size(x_val, 2)
        for i = 1:size(x_val, 1)
            if x_val[i, p] ≈ 1
                patternset[i] = patterns[p]
                idx[i] = round(Int, p)
            end
            continue
        end
    end
    return patternset, idx
end

function addcut!(model, lazycuts, cut::PGBACut, cb_data)
    x = model[:x]
    lzcon = @build_constraint(sum(x[a[1], a[2]] for a in cut.assignments) <= cut.ub)
    push!(lazycuts, cut)
    MOI.submit(model, MOI.LazyConstraint(cb_data), lzcon)
end

"""
Adds a cut to model requiring at least one team in teams to have a different
pattern assigned to them
"""
function addcut!(model, lazycuts, cb_data, teams, patidx, team_specific::Bool)
    x = model[:x]
    if team_specific
        cut = PGBACut(length(teams)-1, collect((i,patidx[i]) for i in teams))
    else
        cut = PGBACut(length(teams)-1, vcat([(j, patidx[i]) for i in teams, j=1:size(x,1)]...))
    end
    addcut!(model, lazycuts, cut, cb_data)
end


text:
using Optim
using BITEModel
const BM=BITEModel

######
# Choose glacier & options
######
@def plotyes = false

# ITMIX glacier with IV data:
glaciers_iv = [:Brewster, :Hellstugubreen, :NorthGlacier,
               :SouthGlacier, :Tasman, :Unteraar, :Synthetic1,
               :Synthetic2, :Synthetic3]
glacier = glaciers_iv[1]
glacier = :Unteraar
gid = [ITMIXGlacier(glacier), SyntheticGlacier(:bench)][1]
pl_kws = Dict{Symbol,Any}()
if isa(gid, ITMIXGlacier)
    pl_kws[:use_glogem] = true
end

gl,gb,pp,pm,pn,pl = BM.init_forward(gid; pl_kws...)

####################
# Do it
####################

(theta0, logposterior, logposterior1d, logprior, logprior_debug,
 loglikelihood, loglikelihood1d, fwdm_fn, fwdm1d_fn,
 pmcmc_defaults, fit_target) =
     BM.init_inverse(gb, pp, pm, pn, n_1d_theta=n_1d_theta, run=run,
                     theta0_dict = theta0_dict     )
# some tests:
fwdm_fn(theta0.th0);
print("Time to run forward model:")
@time fwdsol = fwdm_fn(theta0.th0);
print("Posterior value:")
@show logposterior(theta0.th0)[1]
print("Time to calculate posterior:")
@time logposterior(theta0.th0)[1]

########
# max posterior estimation
#######
if false
    println("Optimizing for maximum posterior...")
    theta_max2 = optimize(theta_vec -> -logposterior(theta_vec)[1], float(theta0.th0), iterations=200, f_tol=1.0)
    # these don't work so well:
    #   theta_max2 = optimize(theta_vec -> -logposterior(theta_vec)[1], float(theta0.th0), GradientDescent())
    #theta_max2 = optimize(theta_vec -> -logposterior(theta_vec)[1], float(theta0.th0),  LBFGS())

    @show theta_max2
    error("asdf")
    theta0.th0[:] = theta_max2.minimizer
end

###############
## MCMC fitting
###############
pmcmc = BM.MCMCNum(pmcmc,
                   niter=0.1*10^6,
                   nthin = 50)

# sample posterior
res = mcmc(logposterior, theta0, pmcmc; verbose=true)
varnames = BM.get_varnames(theta0)
using KissMCMC
print_results(res.thetas, res.accept_ratio, names=varnames)

# sample the prior to get that distribution too:
pmcmc_prior = BM.MCMCNum(pmcmc, niter=0.1*10^6,
                         nthin = 50)
res_prior = mcmc(logprior, theta0, pmcmc_prior; verbose=true)

# BM.savemcmc(thetas, blobs, gb,
#                   pp, pm, pn, pmcmc,
#                   theta0, run, sigmas, error_on_dirty=false)

#end
# ### plots

if plotyes
    display(BM.plotinv1d(gb, blobs, reuse=false))
    display(BM.plotinv1d_err(gb, blobs, reuse=false))

    display(BM.plotinv2d(gl, blobs, reuse=false))
    display(BM.plotinv2d_h(gl, blobs, reuse=false))
    display(BM.plotinv2d_iverr(gl, blobs, reuse=false))

    display(BM.plottheta(thetas, theta0))
    display(BM.plottheta(thetas, theta0, toplot=:btilde, reuse=false))
    display(BM.plottheta(thetas, theta0, toplot=:fsl, reuse=false))
    display(BM.plottheta(thetas, theta0, toplot=:temp, reuse=false))

    display(BM.plottheta_violin((thetas, thetas_prior), theta0, :btilde, reuse=false, width=1))
    display(BM.plottheta_violin((thetas, thetas_prior), theta0, :fsl, reuse=false, width=1))
    display(BM.plottheta_violin((thetas, thetas_prior), theta0, :temp, reuse=false, width=1))
end


text:
function resample(table, index_col, new_inds, methods)
    @assert Tables.istable(table)
    ct = Tables.columntable(table)
    cols = Tables.columnnames(ct)
    index_col = get_index_col(cols, index_col)
    methods = get_table_methods(methods, cols, index_col)

    ct = Tables.columntable(sort(Tables.rowtable(ct), by=r->getproperty(r, index_col)))
    new_inds = get_new_indices(getproperty(ct, index_col), new_inds)

    if any(m -> isa(m, AbstractWeightedSampleMethod), values(methods))
        weighted_indices = resample_indices(getproperty(ct, index_col), new_inds, Mean())
        weights = Dict{AbstractWeightedSampleMethod, Vector{Vector{Float64}}}()
        for method in WEIGHTED_SAMPLE_METHODS
            if method in values(methods)
                weights[method] = resample_weights(
                    getproperty(ct, index_col),
                    new_inds,
                    weighted_indices,
                    method,
                )
            end
        end
    end

    single_indices = Dict{AbstractSingleSampleMethod, Vector{Int}}()
    for method in SINGLE_SAMPLE_METHODS
        if method in values(methods)
            single_indices[method] = resample_indices(getproperty(ct, index_col), new_inds, method)
        end
    end

    nt = (;
        (
            col => if col == index_col
                new_inds
            elseif methods[col] in WEIGHTED_SAMPLE_METHODS
                [dot(getproperty(ct, col)[weighted_indices[i]], weights[methods[col]][i]) for i in eachindex(new_inds)]
            else # if methods[col] in SINGLE_SAMPLE_METHODS
                [getproperty(ct, col)[single_indices[methods[col]][i]] for i in eachindex(new_inds)]
            end for col in cols
        )...
    )
end

resample(table, index_col, new_inds; methods...) = resample(table, index_col, new_inds, methods)

get_index_col(cols, index_col::Symbol) = index_col
get_index_col(cols, index_col::Int) = cols[index_col]

function get_table_methods(d, cols, index_col)
    col_methods = Dict{Symbol, AbstractSampleMethod}()

    for col in keys(d)
        if col in cols
            col_methods[col] = d[col]
        else
            @warn "The column name `:$col` could not be found and has been ignored."
        end
    end

    for col in cols
        if col != index_col && !haskey(col_methods, col)
            col_methods[col] = Mean()
        end
    end

    return col_methods
end

function get_table_methods(a::AbstractVector{<:AbstractSampleMethod}, cols, index_col)
    n = length(cols) - 1

    if length(a) < n
        @warn "The provided methods vector is not long enough and will be repeated until the length $n is reached."
        a = a[mod1.(1:n), length(a)]
    elseif length(a) > n
        @warn "The provided methods vector is too long. Only the first $n elements will be used"
    end

    col_methods = Dict{Symbol, AbstractSampleMethod}()
    i = 1

    for col in cols
        if col != index_col
            col_methods[col] = a[i]
            i += 1
        end
    end

    return col_methods
end

function get_table_methods(m::AbstractSampleMethod, cols, index_col)
    return Dict(col => m for col in cols if col != index_col)
end


text:
module TestExamplesTreeMesh2DPart3

using Test
using Trixi

include("test_trixi.jl")

# Start with a clean environment: remove Trixi output directory if it exists
outdir = "out"
isdir(outdir) && rm(outdir, recursive=true)

@testset "TreeMesh2D Part 3" begin

# Run basic tests
@testset "Examples 2D" begin
  # MHD
  include("test_tree_2d_mhd.jl")

  # MHD Multicomponent
  include("test_tree_2d_mhdmulti.jl")

  # Lattice-Boltzmann
  include("test_tree_2d_lbm.jl")

  # Shallow water
  include("test_tree_2d_shallowwater.jl")

  # FDSBP methods on the TreeMesh
  include("test_tree_2d_fdsbp.jl")
end

# Clean up afterwards: delete Trixi output directory
@test_nowarn rm(outdir, recursive=true)

end # TreeMesh2D Part 3

end #module


text:
#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

@testset "graph.bipartite" begin

@testset "maxcardmatch" begin
    relation1 = Dict(
        1 => Set([1]), 2 => Set([1, 2, 3]),
        3 => Set([1, 2]), 4 => Set([2, 3, 4])
    )
    @test maxcard(Set(1:4), Set(1:4), relation1) == 4
    relation2 = Dict(
        1 => Set([2]), 2 => Set([2, 3]),
        3 => Set(1:5), 4 => Set([3]), 5 => Set([4])
    )
    @test maxcard(Set(1:5), Set(1:5), relation2) == 4
    relation3 = Dict(
        1 => Set([1, 2]), 2 => Set([3]),
        3 => Set([4, 5]), 4 => Set([3, 4]), 5 => Set([2, 5])
    )
    @test maxcard(Set(1:5), Set(1:5), relation3) == 5
    relation4 = Dict(
        1 => Set(1:6), 2 => Set([4]), 3 => Set([4]), 4 => Set([4])
    )
    @test maxcard(Set(1:4), Set(1:6), relation4) == 2
end

@testset "twocoloring" begin
    noedge = SubgraphView(pathgraph(5), Set(1:5), Set{Int}())
    a, b = twocoloring(noedge)
    @test length(a) == 5
    @test isempty(b)

    K5_5 = completebipartite(5, 5)
    a, b = twocoloring(K5_5)
    @test length(a) == 5
    @test length(b) == 5

    CL5 = circularladder(5)
    n = twocoloring(CL5)
    @test n === nothing

    CL6 = circularladder(6)
    a, b = twocoloring(CL6)
    @test length(a) == 6
    @test length(b) == 6
end

end # graph.bipartite


text:
using DelimitedFiles

function solve(input)
    raw = readdlm(input,'\n', Int16)
    serial_number=raw[1]
    @show search(serial_number,300,3:3)
    @show search(serial_number,300,1:300)
end

function compute_power_level(x,y,serial_number)
    rack_id = x + 10
    power_level = rack_id * y
    power_level += serial_number
    power_level *= rack_id
    power_level = floor(power_level%1000/100)
    power_level -= 5
end

@assert compute_power_level(3,5,8) == 4
@assert compute_power_level(122,79,57) == -5
@assert compute_power_level(217,196,39) == 0
@assert compute_power_level(101,153,71) == 4

function search(serial_number, grid_size, square_sizes)
    power_sum = zeros(Int32,grid_size,grid_size)
    x_max, y_max, square_size_max = 0,0,0
    for x in 2:grid_size
        for y in 2:grid_size
            power=compute_power_level(x,y,serial_number)
            power_sum[x,y] = power + power_sum[x,y-1] + power_sum[x-1,y] - power_sum[x-1,y-1];
        end
    end
    largest_total_power = 0
    for s in square_sizes
        for x in s+1:grid_size
            for y in s+1:grid_size
                #current = sum(grid[x:x+s-1,y:y+s-1])
                current = power_sum[x,y] - power_sum[x,y-s] - power_sum[x-s,y] + power_sum[x-s,y-s]
                if current > largest_total_power
                    largest_total_power = current
                    x_max = x
                    y_max = y
                    square_size_max = s
                end
            end
        end
    end
    return (x_max-square_size_max+1,y_max-square_size_max+1,square_size_max)
end
@assert search(18,300,3:3) == (33,45,3)
@assert search(42,300,3:3) == (21,61,3)
@assert search(18,300,1:300) == (90,269,16)
@assert search(42,300,1:300) == (232,251,12)

solve("input.txt")


text:
struct Brendel{T} <: OptProp
    dl :: Model
    σ  :: T
end

struct SerieBrendel <: OptProp
   serie :: Vector{Brendel}
end

permittivity(material :: Brendel,w) = permittivity(material.dl, w)
susceptibility(material :: Model,w) = permittivity(material, w) - material.eps0

function susceptibility(material :: Brendel, w)
    @unpack eps0,w0,wp,gamma0 = material.dl
    @unpack σ = material
    integrand(x) = exp(-(x-w0)^2/(2.0*σ^2))*susceptibility(Model(eps0,wp,x,gamma0),w)/sqrt(2.0*pi)/σ
    val, err = quadgk(integrand,10*(σ-w0),10*(σ+w0))
    return val
end

function susceptibility(material :: SerieBrendel, w)
    χ = 0.0 + 0.0*im
    for i in material.serie
        χ += susceptibility(i, w)
    end
    return χ
end

permittivity(material :: SerieBrendel , w) = material.serie[1].dl.eps0 + susceptibility(material,w)

b1 = Brendel(Model(6.12, 1240*c0*2*pi, 826*c0*pi, 2*c0*pi), 283*c0*2*pi)
b2 = Brendel(Model(6.12,    1*c0*2*pi, 925*c0*pi, 4*c0*pi),  44*c0*2*pi)
b3 = Brendel(Model(6.12,   63*c0*2*pi,1063*c0*pi,36*c0*pi),3468*c0*2*pi)
b4 = Brendel(Model(6.12,  702*c0*2*pi,1185*c0*pi,57*c0*pi), 592*c0*2*pi)
b5 = Brendel(Model(6.12,    0*c0*2*pi,7931*c0*pi, 1*c0*pi),2987*c0*2*pi)

SN20_300 = SerieBrendel([b1 , b2, b3, b4 ,b5])


text:
using DiffEqFlux
using OrdinaryDiffEq
using Optimisers
using Flux

using PyCall
using PyCallChainRules.Torch: TorchModuleWrapper, torch

u0 = Float32[2.; 0.]
datasize = 30
tspan = (0.0f0, 1.5f0)

function trueODEfunc(du,u,p,t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end
t = range(tspan[1],tspan[2],length=datasize)
prob = ODEProblem(trueODEfunc,u0,tspan)
ode_data = Array(solve(prob,Tsit5(),saveat=t))

torch_module = torch.nn.Sequential(
    torch.nn.Linear(2, 50), torch.nn.Tanh(),
    torch.nn.Linear(50, 2), torch.nn.Tanh(),
)
# Mix of Flux layers and Torch layers
jlmod = Chain(Dense(2, 2, tanh), TorchModuleWrapper(torch_module), Dense(2, 2,))
p, re = Optimisers.destructure(jlmod)

dudt(u, p, t) = re(p)(u)
prob = ODEProblem(dudt, u0, tspan)

function predict_n_ode(p)
    Array(solve(prob,Tsit5(),u0=u0,p=p,saveat=t))
end

function loss_n_ode(p)
    pred = predict_n_ode(p)
    loss = sum(abs2,ode_data .- pred)
    loss
end

loss_n_ode(p)

data = Iterators.repeated((), 1000)

@info "before" loss_n_ode(p)

function train(p;nsteps=100)
    opt = Optimisers.ADAM(0.01)
    state = Optimisers.setup(opt, p)

    for i in 1:nsteps
        gs, = Flux.gradient(p) do ps
            loss_n_ode(ps)
        end
        state, p = Optimisers.update(state, p, gs)
    end
    return p
end

newp = train(p)

@info "after" loss_n_ode(newp)


text:
using BinaryBuilder

name = "Qemu"
version = v"4.1.0"

# Collection of sources required to build libffi
sources = [
    GitSource("https://gitlab.com/virtio-fs/qemu.git",
              "bf5775237ee563b4baa1c7f3c1a65b7c93b93fca"),
    DirectorySource("./bundled"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/qemu

# Patch out usage of MADV_NOHUGEPAGE which does not exist in glibc 2.12.X
atomic_patch -p1 "${WORKSPACE}/srcdir/patches/qemu_madv_nohugepage.patch"

# Patch in adapter for `clock_gettime()` on macOS 10.12-
atomic_patch -p1 "${WORKSPACE}/srcdir/patches/qemu_clock_gettime.patch"

# Patch to fix pointer mismatch between `size_t` and `uint64_t`
atomic_patch -p1 "${WORKSPACE}/srcdir/patches/qemu_size_uint64.patch"

# Configure, ignoring some warnings that we don't need, etc...
./configure --host-cc="${HOSTCC}" --extra-cflags="-I${prefix}/include -Wno-unused-result" --disable-cocoa --prefix=$prefix

echo '#!/bin/true ' > /usr/bin/Rez
echo '#!/bin/true ' > /usr/bin/SetFile
chmod +x /usr/bin/Rez
chmod +x /usr/bin/SetFile
make -j${nproc}
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = [
    Platform("x86_64", "linux"; libc="glibc"),
    Platform("x86_64", "macos"),
]

# The products that we will ensure are always built
products = [
    ExecutableProduct("qemu-system-x86_64", :qemu_system_x86_64),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency("Pixman_jll"),
    Dependency("Glib_jll", v"2.59.0"; compat="2.59.0"),
    Dependency("PCRE_jll"),
    # TOOD: verify Gettext is actually needed at runtime
    Dependency("Gettext_jll", v"0.20.1"; compat="=0.20.1"),
    Dependency("libcap_jll"),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
"""A simple and fast module to solve the harmonic power flow problem for given distribution grids."""
module HarmonicPowerFlow

include("infrastructure.jl")
include("settings.jl")
include("model.jl")
include("evaluation.jl")
include("plots.jl")

export init_power_grid, init_settings, import_nodes_from_csv, import_lines_from_csv, pf, hpf, barplot_THD, heatmap_NE

end # module


text:
using MyPackage
f([3,4])


text:
include("Include.jl")

# load the data -
path_data_file = joinpath(_PATH_TO_DATA, "Thrombin-TF.csv")
experimental_data_table = load(path_data_file)

# get input and output data -
input_data = convert.(Float32, Matrix(experimental_data_table[!, 3:13]))
output_data = convert.(Float32, Matrix(experimental_data_table[!, 14:18]))

# load the model -
@load "deep_coag_model.bson" deep_coag_model

max_stars_repo_name: varnerlab/UVM-TopDown-LegacyModel

max_issues_repo_name: varnerlab/UVM-TopDown-LegacyModel

max_forks_repo_name: varnerlab/UVM-TopDown-LegacyModel


text:
using ManifoldLearning

export ManifoldKnnIndex,
    ExactEuclidean,
    ExactManhattan,
    ExactChebyshev,
    ExactCosine,
    ExactAngle,
    ApproxEuclidean,
    ApproxManhattan,
    ApproxChebyshev,
    ApproxCosine,
    ApproxAngle

"""
    ManifoldKnnIndex{DistType,MinRecall}

Implements the `ManifoldLearning.AbstractNearestNeighbors` interface to interoperate with
the non-linear dimensionality reduction methods of the `ManifoldLearning` package.

It should be passed to the `fit` method as a type, e.g.,

```julia
fit(ManifoldKnnIndex{L2Distance,0.9})  # will use an approximate index with an expected recall of 0.9
```

`DistType` should be any in [`SimilaritySearch`](https://sadit.github.io/SimilaritySearch.jl/dev/api/#Distance-functions) package, or `Distances.jl`
or any other following the `SemiMetric`.

The second argument of the composite type indicates the quality and therefore the type of index to use:
- It takes values between `0` and `1`.
- `0` means for a `SearchGraph` index using `ParetoRecall` optimization for the construction and searching, this will try to achieve a competitive structure in both quality and search speed
- `1` means for a `ExhaustiveSearch` index, this will compute the exact solution (exact knns) but at cost speed. Can work pretty well on small datasets and very high dimensionality. Really high dimensions suffer from the _curse of dimensionality_ such that an index like `SearchGraph` degrades to ExhaustiveSearch.
- `0 < value < 1`: Uses a `SearchGraph` and is the minimum recall-score quality that the index should perform. In particular, it constructs the index using `ParetoRecall` and the use a final optimization with `MinRecall`. It takes values from 0 to 1, small values produce faster searches with lower qualities and high values slower searches with higher quality. Values 0.8 or 0.9 should work pretty well.

Note: The minimum performance is evaluated in a small training set took from the database, this could yield to some kind of overfitting in the parameters, and therefore, perform not so good in an unseen query set. If you note this effect, please see `SimilaritySearch` documentation [function `optimize!`](https://sadit.github.io/SimilaritySearch.jl/dev/api/#SimilaritySearch.optimize!-Tuple{SearchGraph,%20OptimizeParameters}).

"""
struct ManifoldKnnIndex{DistType,MinRecall} <: ManifoldLearning.AbstractNearestNeighbors
    index::AbstractSearchContext  # A simple wrapper
end

"""
    ExactEuclidean

`ManifoldKnnIndex`'s type specialization for exact search with the Euclidean distance.
"""
const ExactEuclidean = ManifoldKnnIndex{L2Distance,1}

"""
    ExactManhattan

`ManifoldKnnIndex`'s type specialization for exact search with the Manhattan distance.
"""
const ExactManhattan = ManifoldKnnIndex{L1Distance,1}

"""
    ExactChebyshev

`ManifoldKnnIndex`'s type specialization for exact search with the Chebyshev distance.
"""
const ExactChebyshev = ManifoldKnnIndex{LInftyDistance,1}

"""
    ExactCosine

`ManifoldKnnIndex`'s type specialization for exact search with the cosine distance.
"""
const ExactCosine = ManifoldKnnIndex{CosineDistance,1}

"""
    ExactAngle

`ManifoldKnnIndex`'s type specialization for exact search with the angle distance.
"""
const ExactAngle = ManifoldKnnIndex{AngleDistance,1}

"""
    ApproxEuclidean

`ManifoldKnnIndex`'s type specialization for approximate search with the Euclidean distance (expected recall of 0.9)
"""
const ApproxEuclidean = ManifoldKnnIndex{L2Distance,0.9}

"""
    ApproxManhattan

`ManifoldKnnIndex`'s type specialization for approximate search with the Manhattan distance (expected recall of 0.9)
"""
const ApproxManhattan = ManifoldKnnIndex{L1Distance,0.9}

"""
    ApproxChebyshev

`ManifoldKnnIndex`'s type specialization for approximate search with the Chebyshev distance (expected recall of 0.9)
"""
const ApproxChebyshev = ManifoldKnnIndex{LInftyDistance,0.9}

"""
    ApproxCosine

`ManifoldKnnIndex`'s type specialization for approximate search with the Cosine distance (expected recall of 0.9)
"""
const ApproxCosine = ManifoldKnnIndex{CosineDistance,0.9}

"""
    ApproxAngle

`ManifoldKnnIndex`'s type specialization for approximate search with the angle distance (expected recall of 0.9)
"""
const ApproxAngle = ManifoldKnnIndex{AngleDistance,0.9}

## ManifoldLearning api for nearest neighbor algorithms

Base.size(G::ManifoldKnnIndex) = (length(G.index[1]), length(G.index))

function fit(::Type{ManifoldKnnIndex{DistType,MinRecall_}}, X) where {DistType,MinRecall_}
    db = MatrixDatabase(X)
    dist = DistType()

    index = if MinRecall_ == 1
        ExhaustiveSearch(; dist, db)
    else
        G = SearchGraph(; dist, db)
        parallel_block = length(db) < 512 || Threads.nthreads() == 1 ? 1 : 4 * Threads.nthreads()
        index!(G; parallel_block)
        if MinRecall_ > 0
            minrecall = (MinRecall_ isa AbstractFloat) ? MinRecall_ : 0.9
            optimize!(G, OptimizeParameters(; kind=MinRecall(minrecall)))
        end
        G
    end

    ManifoldKnnIndex{DistType,MinRecall_}(index)
end

"""
    knn(G::ManifoldKnnIndex, Q::AbstractMatrix{T}, k::Integer; self::Bool=false, weights::Bool=true, kwargs...) where {T<:Real}

Solves `k` nearest neighbors queries of `Q` using `G` (a SimilaritySearch index).
"""
function ManifoldLearning.knn(G::ManifoldKnnIndex, Q::AbstractMatrix{T}, k::Integer; self::Bool=false, weights::Bool=true, kwargs...) where {T<:Real}
    m, n = size(Q)
    self && throw(ArgumentError("SimilaritySearch `"))
    n > k || throw(ArgumentError("Number of observations must be more than $(k)"))
    Q = MatrixDatabase(Q)
    KNNS = [KnnResult(k+self) for _ in 1:n]  # we can't use `allknn` efficiently from ManifoldLearning
    @time searchbatch(G.index, Q, KNNS; parallel=true)

    E = [res.id for res in KNNS] # reusing the internal structure
    W = [res.dist for res in KNNS] # `KnnResult` distances are always Float32

    if self  # removes zero/close-to-zero neighbors
        for i in eachindex(E)
            while length(W[i]) > 0
                if W[1] < eps(W)
                    popfirst!(E)
                    popfirst!(W)
                end
            end
        end
    end

    E, W
end


text:
##############################################################################
##
## Reshaping
##
## Also, see issue # ??
##
##############################################################################

##############################################################################
##
## stack()
## melt()
##
##############################################################################

function stack(df::DataFrame, measure_vars::Vector{Int}, id_vars::Vector{Int})
    res = DataFrame[insert!(df[[i, id_vars]], 1, names(df)[i], :variable) for i in measure_vars]
    # fix column names
    map(x -> names!(x, [:variable, :value, names(df[id_vars])]), res)
    res = vcat(res)
    res
end
stack(df::DataFrame, measure_vars, id_vars) = stack(df, [df.colindex[measure_vars]], [df.colindex[id_vars]])
stack(df::DataFrame, measure_vars) = stack(df, [df.colindex[measure_vars]], _setdiff([1:ncol(df)], [df.colindex[measure_vars]]))

melt(df::DataFrame, id_vars) = stack(df, _setdiff([1:ncol(df)], [df.colindex[id_vars]]))
melt(df::DataFrame, id_vars, measure_vars) = stack(df, measure_vars, id_vars)

##############################################################################
##
## unstack()
##
##############################################################################

function unstack(df::AbstractDataFrame, rowkey::Int, colkey::Int, value::Int)
    # `rowkey` integer indicating which column to place along rows
    # `colkey` integer indicating which column to place along column headers
    # `value` integer indicating which column has values
    keycol = PooledDataArray(df[rowkey])
    valuecol = df[value]
    # TODO make a version with a default refkeycol
    refkeycol = PooledDataArray(df[colkey])
    remainingcols = _setdiff([1:ncol(df)], [rowkey, value])
    Nrow = length(refkeycol.pool)
    Ncol = length(keycol.pool)
    # TODO make fillNA(type, length)
    payload = DataFrame({DataArray([fill(valuecol[1],Nrow)], fill(true, Nrow))  for i in 1:Ncol}, map(symbol, keycol.pool))
    nowarning = true
    for k in 1:nrow(df)
        j = int(keycol.refs[k])
        i = int(refkeycol.refs[k])
        if i > 0 && j > 0
            if nowarning && !isna(payload[j][i])
                println("Warning: duplicate entries in unstack.")
                nowarning = false
            end
            payload[j][i]  = valuecol[k]
        end
    end
    insert!(payload, 1, refkeycol.pool, names(df)[colkey])
end
unstack(df::AbstractDataFrame, rowkey, value, colkey) =
    unstack(df, index(df)[rowkey], index(df)[value], index(df)[colkey])

##############################################################################
##
## pivottable()
##
##############################################################################

# Limitations:
#  - only one `value` column is allowed (same as dcast)
#  - `fun` must reduce to one value
#  - no margins
#  - can't have zero rows or zero columns
#  - the resulting data part is Float64 (`payload` below)

function pivottable(df::AbstractDataFrame, rows::Vector{Int}, cols::Vector{Int}, value::Int, fun::Function)
    # `rows` vector indicating which columns are keys placed in rows
    # `cols` vector indicating which columns are keys placed as column headers
    # `value` integer indicating which column has values
    # `fun` function applied to the value column during aggregation
    cmb_df = by(df, [rows, cols], d->fun(d[value]))
    row_pdv = PooledDataArray(paste_columns(cmb_df[[length(rows):-1:1]]))  # the :-1: is to reverse the columns for sorting
    row_idxs = int(row_pdv.refs)
    Nrow = length(row_pdv.pool)
    col_pdv = PooledDataArray(paste_columns(cmb_df[[length(rows) + (1:length(cols))]]))
    col_idxs = int(col_pdv.refs)
    Ncol = length(col_pdv.pool)
    # `payload` is the main "data holding" part of the resulting DataFrame
    payload = DataFrame(Nrow, Ncol)
    names!(payload, col_pdv.pool)
    for i in 1:length(row_idxs)
        payload[row_idxs[i], col_idxs[i]] = cmb_df[i,"x1"]
    end
    # find the "row" key DataFrame
    g = groupby(cmb_df[[1:length(rows)]], [1:length(rows)])
    row_key_df = g.parent[g.idx[g.starts],:]
    hcat(row_key_df, payload)
end
# `mean` is the default aggregation function:
pivottable(df::AbstractDataFrame, rows, cols, value) = pivottable(df, rows, cols, value, mean)
pivottable(df::AbstractDataFrame, rows, cols, value, fun) = pivottable(df, [index(df)[rows]], [index(df)[cols]], index(df)[value], fun)

##############################################################################
##
## Reshaping using referencing (issue #145)
## New AbstractVector types (all read only):
##     StackedVector
##     RepeatedVector
##     EachRepeatedVector
##
##############################################################################

## StackedVector({[1,2], [9,10], [11,12]}) is equivalent to [1,2,9,10,11,12]
type StackedVector <: AbstractVector{Any}
    components::Vector{Any}
end

## RepeatedVector([1,2], 3) is equivalent to [1,2,1,2,1,2]
type RepeatedVector{T} <: AbstractVector{T}
    parent::AbstractVector{T}
    n::Int
end

## EachRepeatedVector([1,2], 3) is equivalent to [1,1,1,2,2,2]
type EachRepeatedVector{T} <: AbstractVector{T}
    parent::AbstractVector{T}
    n::Int
end

function Base.getindex(v::StackedVector,i::Real)
    lengths = [length(x)::Int for x in v.components]
    cumlengths = [0, cumsum(lengths)]
    j = searchsortedlast(cumlengths .+ 1, i)
    if j > length(cumlengths)
        error("indexing bounds error")
    end
    k = i - cumlengths[j]
    if k < 1 || k > length(v.components[j])
        error("indexing bounds error")
    end
    v.components[j][k]
end

function Base.getindex{I<:Real}(v::StackedVector,i::AbstractVector{I})
    result = similar(v.components[1], length(i))
    for idx in 1:length(i)
        result[idx] = v[i[idx]]
    end
    result
end
Base.getindex(v::StackedVector,i::Union(Ranges, Vector{Bool}, BitVector)) = getindex(v, [i])

Base.size(v::StackedVector) = (length(v),)
Base.length(v::StackedVector) = sum(map(length, v.components))
Base.ndims(v::StackedVector) = 1
Base.eltype(v::StackedVector) = promote_type(map(eltype, v.components)...)
Base.similar(v::StackedVector, T, dims::Dims) = similar(v.components[1], T, dims)

Base.show(io::IO, v::StackedVector) = internal_show_vector(io, v)
Base.repl_show(io::IO, v::StackedVector) = internal_repl_show_vector(io, v)

DataArrays.PooledDataArray(v::StackedVector) = PooledDataArray(v[:]) # could be more efficient

function Base.getindex{T,I<:Real}(v::RepeatedVector{T},i::AbstractVector{I})
    j = mod(i .- 1, length(v.parent)) .+ 1
    v.parent[j]
end
function Base.getindex{T}(v::RepeatedVector{T},i::Real)
    j = mod(i - 1, length(v.parent)) + 1
    v.parent[j]
end
Base.getindex(v::RepeatedVector,i::Union(Ranges, Vector{Bool}, BitVector)) = getindex(v, [i])

Base.size(v::RepeatedVector) = (length(v),)
Base.length(v::RepeatedVector) = v.n * length(v.parent)
Base.ndims(v::RepeatedVector) = 1
Base.eltype{T}(v::RepeatedVector{T}) = T
Base.reverse(v::RepeatedVector) = RepeatedVector(reverse(v.parent), v.n)
Base.similar(v::RepeatedVector, T, dims::Dims) = similar(v.parent, T, dims)

Base.show(io::IO, v::RepeatedVector) = internal_show_vector(io, v)
Base.repl_show(io::IO, v::RepeatedVector) = internal_repl_show_vector(io, v)

Base.unique(v::RepeatedVector) = unique(v.parent)

function DataArrays.PooledDataArray(v::RepeatedVector)
    res = PooledDataArray(v.parent)
    res.refs = rep(res.refs, v.n)
    res
end

function Base.getindex{T}(v::EachRepeatedVector{T},i::Real)
    j = div(i - 1, v.n) + 1
    v.parent[j]
end
function Base.getindex{T,I<:Real}(v::EachRepeatedVector{T},i::AbstractVector{I})
    j = div(i .- 1, v.n) .+ 1
    v.parent[j]
end
Base.getindex(v::EachRepeatedVector,i::Union(Ranges, Vector{Bool}, BitVector)) = getindex(v, [i])

Base.size(v::EachRepeatedVector) = (length(v),)
Base.length(v::EachRepeatedVector) = v.n * length(v.parent)
Base.ndims(v::EachRepeatedVector) = 1
Base.eltype{T}(v::EachRepeatedVector{T}) = T
Base.reverse(v::EachRepeatedVector) = EachRepeatedVector(reverse(v.parent), v.n)
Base.similar(v::EachRepeatedVector, T, dims::Dims) = similar(v.parent, T, dims)

Base.show(io::IO, v::EachRepeatedVector) = internal_show_vector(io, v)
Base.repl_show(io::IO, v::EachRepeatedVector) = internal_repl_show_vector(io, v)

Base.unique(v::EachRepeatedVector) = unique(v.parent)

DataArrays.PooledDataArray(v::EachRepeatedVector) = PooledDataArray(v[:], dropna(unique(v.parent)))

function DataArrays.PooledDataArray(v::EachRepeatedVector)
    res = PooledDataArray(v.parent)
    res.refs = rep(res.refs, rep(v.n,length(res.refs)))
    res
end

# The default values of show and repl_show don't work because
# both try to reshape the vector into a matrix, and these
# types do not support that.

function internal_show_vector(io::IO, v::AbstractVector)
    Base.show_delim_array(io, v, '[', ',', ']', true)
end

function internal_repl_show_vector(io::IO, v::AbstractVector)
    print(io, summary(v))
    if length(v) < 21
        print_matrix(io, v[:]'')    # the double transpose ('') reshapes to a matrix
    else
        println(io)
        print_matrix(io, v[1:10]'')
        println(io)
        println(io, "  \u22ee")
        print_matrix(io, v[end - 9:end]'')
    end
end

##############################################################################
##
## stackdf()
## meltdf()
## Reshaping using referencing (issue #145), using the above vector types
##
##############################################################################

# Same as `stack`, but uses references
# I'm not sure the name is very good
function stackdf(df::AbstractDataFrame, measure_vars::Vector{Int}, id_vars::Vector{Int})
    N = length(measure_vars)
    remainingcols = _setdiff([1:ncol(df)], measure_vars)
    cnames = names(df)[id_vars]
    insert!(cnames, 1, "value")
    insert!(cnames, 1, "variable")
    DataFrame({EachRepeatedVector(names(df)[measure_vars], nrow(df)), # variable
               StackedVector({df[:,c] for c in 1:N}),                    # value
               ## RepeatedVector([1:nrow(df)], N),                       # idx - do we want this?
               [RepeatedVector(df[:,c], N) for c in id_vars]...},        # id_var columns
              cnames)
end
stackdf(df::AbstractDataFrame, measure_vars) = stackdf(df, [index(df)[measure_vars]])

stackdf(df::AbstractDataFrame, measure_vars, id_vars) = stackdf(df, [index(df)[measure_vars]], [index(df)[id_vars]])
stackdf(df::AbstractDataFrame, measure_vars) = stackdf(df, [index(df)[measure_vars]], _setdiff([1:ncol(df)], [index(df)[measure_vars]]))

meltdf(df::AbstractDataFrame, id_vars) = stackdf(df, _setdiff([1:ncol(df)], index(df)[id_vars]))
meltdf(df::AbstractDataFrame, id_vars, measure_vars) = stackdf(df, measure_vars, id_vars)


text:
@doc raw"""
    BPR(dim::Int64, reg_coeff::Float64)

Bayesian personalized ranking model. The model evaluates user-item triplet ``(u ,i ,j)``, which expresses "the user ``u`` prefers item ``i`` to item ``j``. Here the following matrix factoriazation model is adopted to model this relation:

```math
p_{uij} = \bm u_u \cdot \bm v_i - \bm u_u \cdot \bm v_j
```

# Constructor arguments
- `dim`: dimension of user/item vectors.
- `reg_coeff`: ``L_2`` regularization coefficients for model parameters.
"""
mutable struct BPR <: AbstractRecommender
    dim::Int64
    loss::LossFunction
    reg_coeff::Float64

    user_embedding::Union{Nothing,Matrix{Float64}}
    item_embedding::Union{Nothing,Matrix{Float64}}

    user2uidx::Union{Dict,Nothing}
    item2iidx::Union{Dict,Nothing}
    iidx2item::Union{Dict,Nothing}
    user_history::Union{Dict,Nothing}

    BPR(dim::Int64, reg_coeff::Float64) = new(
        dim,
        BPRLoss(), # default
        reg_coeff,
        nothing,
        nothing,
        nothing,
        nothing,
        nothing,
        nothing,
    )
end

function predict(model::BPR, uidx, iidx)::Float64
    return model.user_embedding[:, uidx]' * model.item_embedding[:, iidx]
end

function predict(model::BPR, uidx, iidx, jidx)::Float64
    return predict(model, uidx, iidx) - predict(model, uidx, jidx)
end

function sgd!(model::BPR, uidx, iidx, jidx, grad_value, lr)
    reg = model.reg_coeff

    uvec = @view model.user_embedding[:, uidx]
    ivec = @view model.item_embedding[:, iidx]
    jvec = @view model.item_embedding[:, jidx]

    model.user_embedding[:, uidx] -= lr * (grad_value * (ivec - jvec) + reg * uvec)
    model.item_embedding[:, iidx] -= lr * (grad_value * uvec + reg * ivec)
    model.item_embedding[:, jidx] -= lr * (grad_value * (-uvec) + reg * ivec)
end

"""
    fit!(model::BPR, table; callbacks = Any[], col_user = :userid, col_item = :item_id, n_epochs = 2, learning_rate = 0.01, n_negatives = 1, verbose = -1)

Fit the `BPR` model by stochastic grandient descent. Instead the learnBPR algorithm proposed by the original paper, the simple SGD with negative sampling is implemented.

# Model-specific arguments
- `n_epochs`: number of epochs. During one epoch, all the row in `table` is read once.
- `learning_rate`: Learing rate of SGD.
- `n_negatives`: Number of negative item samples per positive (user, item) pair.
- `verbose`: If set to positive integer, the training info is printed once per `verbose`.
- `callbacks`: Additional callback functions during SGD. One can implement, for instance, monitoring the validation metrics and the early stopping. See [Callbacks](@ref).

# References
Rendel et. al. (2012), [BPR: Bayesian Personalized Ranking from Implicit Feedback](http://arxiv.org/abs/1205.2618)
"""
function fit!(
    model::BPR,
    table;
    callbacks = Any[],
    col_user = :userid,
    col_item = :item_id,
    n_epochs = 2,
    n_negatives = 1,
    learning_rate = 0.01,
    verbose = -1,
    kwargs...,
)
    model.user_history = Dict()
    for (userid, history) in
        zip(make_u2i_dataset(table, col_user = col_user, col_item = col_item)...)
        model.user_history[userid] = history
    end

    table, model.user2uidx, model.item2iidx, model.iidx2item =
        make_idmap(table, col_user = col_user, col_item = col_item)

    n_user = length(keys(model.user2uidx))

    unique_items = collect(keys(model.iidx2item))
    n_item = length(unique_items)

    model.user_embedding = rand(model.dim, n_user)
    model.item_embedding = rand(model.dim, n_item)

    # callback is any callbale with same interface
    callbacks = append!(Any[LogTrainLoss()], callbacks)
    for cb in callbacks
        if typeof(cb) <: AbstractCallback
            initialize!(cb, model, col_user = col_user, col_item = col_item)
        end
    end

    for epoch = 1:n_epochs
        train_loss = 0
        n_sample = 0
        for row in Tables.rows(table)
            uidx = row[col_user]
            iidx = row[col_item]

            # sample negative
            for _ = 1:n_negatives
                jidx = rand(unique_items)
                while jidx in model.user_history[uidx]
                    jidx = rand(unique_items)
                end

                pred = predict(model, uidx, iidx, jidx)
                train_loss = (model.loss(pred) + train_loss * n_sample) / (n_sample + 1)
                n_sample += 1

                grad_value = grad(model.loss, pred)
                sgd!(model, uidx, iidx, jidx, grad_value, learning_rate)
            end
        end

        try
            for cb in callbacks
                cb(model, train_loss, epoch, verbose)
            end
        catch StopTrain
            break
        end
    end

end

@doc raw"""
    predict_u2i(model::BPR, userid::Union{AbstractString,Int}, n::Int64; drop_history = false)

Make predictions by using ``\bm u_u \cdot \bm v_i``.
"""
function predict_u2i(
    model::BPR,
    userid::Union{AbstractString,Int},
    n::Int64;
    drop_history = false,
    kwargs...,
)
    if !(userid in keys(model.user2uidx))
        return []
    end
    uidx = model.user2uidx[userid]
    unique_iidx = collect(keys(model.iidx2item))
    preds = [predict(model, uidx, iidx) for iidx in unique_iidx]
    pred_iidx = unique_iidx[sortperm(preds, rev = true)]
    pred_items = [model.iidx2item[iidx] for iidx in pred_iidx]
    if drop_history
        filter!(e -> !(e in model.user_history[userid]), pred_items)
    end
    n = min(n, length(pred_items))
    return pred_items[1:n]
end