text:
using DrWatson, Test
using BSON, JLD2
cd(@__DIR__)
T = 1000
N = 50 # spatial extent
Δt = 0.05 # resolution of integration
every = 10
seed = 1111

simulation = @ntuple T N Δt every seed

function f(simulation)
    @test typeof(simulation.T) <: Real
    a = rand(10); b = [rand(10) for _ in 1:10]
    return @strdict a b simulation
end

################################################################################
#                                 tagsave                                      #
################################################################################
@testset "Tagsafe ($ending)" for ending ∈ ["bson", "jld2"]
    t = f(simulation)
    tagsave(savename(simulation, ending), t, gitpath=findproject())
    file = load(savename(simulation, ending))
    @test "gitcommit" ∈ keys(file)
    @test file["gitcommit"] |> typeof == String
    rm(savename(simulation, ending))

    t = f(simulation)
    @tagsave(savename(simulation, ending), t, safe=false, gitpath=findproject())
    file = load(savename(simulation, ending))
    @test "gitcommit" ∈ keys(file)
    @test file["gitcommit"] |> typeof == String
    @test "script" ∈ keys(file)
    @test file["script"] |> typeof == String
    @test file["script"] == joinpath("test", "savefiles_tests.jl#30")

    t = f(simulation)
    @tagsave(savename(simulation, ending), t, safe=true, gitpath=findproject())
    sn = savename(simulation, ending)[1:end-5]*"_#1"*"."*ending
    @test isfile(sn)
    rm(sn)

    t = f(simulation)
    tagsave(savename(simulation, ending), t, safe=true, gitpath=findproject())
    sn = savename(simulation, ending)[1:end-5]*"_#1"*"."*ending
    @test isfile(sn)
    rm(sn)

    t = f(simulation)
    t["gitcommit"] = ""
    @test @tagsave(savename(simulation, ending), t, safe=true, gitpath=findproject())["gitcommit"] == ""
    @test isfile(sn)
    rm(sn)
    @test @tagsave(savename(simulation,ending), t, safe=true, force=true, gitpath=findproject())["gitcommit"] != ""
    @test isfile(sn)
    rm(sn)

    rm(savename(simulation, ending))
    @test !isfile(savename(simulation, ending))

    ex = @macroexpand @tagsave("testname."*ending, (@dict a b c ), storepatch=false; safe=true)
    ex2 = @macroexpand @tagsave("testname."*ending, @dict a b c; storepatch=false, safe=true)
    @test ex.args[1:end-1] == ex2.args[1:end-1]

    # Remove leftover
end

# Check if kwargs propagation works using the example of compression in JLD2.
# We need to look at the actual file sizes - JLD2 handles compression transparently
# and doesn't provide a way to check if a read dataset was compressed.
# Run twice - once without the `safe` option and once with, to implicitly test `safesave`.
@testset "Tagsafe with compression (safe=$safesave)" for safesave ∈ [false, true]
    # Create some highly compressible data
    data = Dict("data" => fill(1, 10000))
    sn_uncomp = "uncompressed.jld2"
    sn_comp = "compressed.jld2"
    # Save twice - once uncompressed and once compressed
    tagsave(sn_uncomp, data, safe=safesave, gitpath=findproject(), compress=false)
    tagsave(sn_comp, data, safe=safesave, gitpath=findproject(), compress=true)
    # Check if both files exist
    @test isfile(sn_uncomp)
    @test isfile(sn_comp)
    # Test if the compressed file is smaller
    size_uncomp = filesize(sn_uncomp)
    size_comp = filesize(sn_comp)
    @test size_uncomp > size_comp
    # Leave no trace
    rm(sn_uncomp)
    rm(sn_comp)
end

################################################################################
#                              produce or load                                 #
################################################################################

@testset "Produce or Load ($ending)" for ending ∈ ["bson", "jld2"]
    @test !isfile(savename(simulation, ending))
    sim, path = produce_or_load(simulation, f; suffix = ending)
    @test isfile(savename(simulation, ending))
    @test sim["simulation"].T == T
    @test path == savename(simulation, ending)
    sim, path = produce_or_load(simulation, f; suffix = ending)
    @test sim["simulation"].T == T
    rm(savename(simulation, ending))
    @test !isfile(savename(simulation, ending))

    @test !isfile(savename(simulation, ending))
    # Produce and save data, preserve source file name and line for test below.
    # Line needs to be saved on the same line as produce_or_load!
    sim, path = @produce_or_load(f, "", simulation, suffix = ending); fname = @__FILE__; line = @__LINE__
    @test isfile(savename(simulation, ending))
    @test sim["simulation"].T == T
    @test path == savename(simulation, ending)
    sim, path = @produce_or_load(f, "", simulation, suffix = ending)
    @test sim["simulation"].T == T
    # Test if source was included and that the file name and line number matches the first invocation
    # (and not the second!)
    @test "script" ∈ keys(sim)
    @test sim["script"] |> typeof == String
    @test sim["script"] == joinpath(relpath(fname, projectdir()) * "#$(line)")
    rm(savename(simulation, ending))
    @test !isfile(savename(simulation, ending))

    # Test if tag = true does not interfere with macro script tagging.
    # Use a semicolon before the `suffix` keyword to test that code path as well.
    sim, path = @produce_or_load(f, "", simulation, tag = true; suffix = ending); fname = @__FILE__; line = @__LINE__
    sim, path = @produce_or_load(f, "", simulation; suffix = ending)
    # Test if source was included and that the file name and line number matches the first invocation
    # (and not the second!)
    @test sim["script"] == joinpath(relpath(fname, projectdir()) * "#$(line)")
    rm(savename(simulation, ending))

    # Test that the internal function `scripttag!` properly warns if the Dict already has a `script` key.
    # This also tests the case where the `Dict` has a `Symbol` key type.
    @test_logs((:warn, "The dictionary already has a key named `script`. We won't overwrite it with the script name."),
               DrWatson.scripttag!(Dict(:script => "test"), LineNumberNode(1)))

    @test !isfile(savename(simulation, ending))
    sim, path = produce_or_load("", simulation; suffix = ending) do simulation
        @test typeof(simulation.T) <: Real
        a = rand(10); b = [rand(10) for _ in 1:10]
        return @strdict a b simulation
    end
    @test isfile(savename(simulation, ending))
    @test sim["simulation"].T == T
    @test path == savename(simulation, ending)
    sim, path = produce_or_load("", simulation; suffix = ending) do simulation
        @test typeof(simulation.T) <: Real
        a = rand(10); b = [rand(10) for _ in 1:10]
        return @strdict a b simulation
    end
    @test sim["simulation"].T == T
    rm(savename(simulation, ending))
    @test !isfile(savename(simulation, ending))

    p = String(@__DIR__)
    pre = "pre"
    expected = joinpath(p, savename(pre, simulation, ending))
    @test !isfile(expected)
    sim, path = produce_or_load(p, simulation, f; suffix = ending, prefix = pre)
    @test path == expected
    @test isfile(path)
    rm(path)
end

@testset "Produce or Load wsave keyword pass through" begin
    # Create some highly compressible data
    data = Dict("data" => fill(1, 10000))

    sn_uncomp = savename(Dict("compress" => false), "jld2")
    sn_comp = savename(Dict("compress" => true), "jld2")
    # Files cannot exist yet
    @test !isfile(sn_uncomp)
    @test !isfile(sn_comp)
    for compress in [false, true]
        wsave_kwargs = Dict(:compress => compress)
        produce_or_load("", wsave_kwargs, suffix = "jld2", wsave_kwargs=wsave_kwargs) do c
            data
        end
    end
    # Check if both files exist now
    @test isfile(sn_uncomp)
    @test isfile(sn_comp)
    # Test if the compressed file is smaller
    size_uncomp = filesize(sn_uncomp)
    size_comp = filesize(sn_comp)
    @test size_uncomp > size_comp
    # Leave no trace
    rm(sn_uncomp)
    rm(sn_comp)

    # Check the macro version
    sn_uncomp = savename(Dict("compress" => false), "jld2")
    sn_comp = savename(Dict("compress" => true), "jld2")
    # Files cannot exist yet
    @test !isfile(sn_uncomp)
    @test !isfile(sn_comp)
    for compress in [false, true]
        wsave_kwargs = Dict(:compress => compress)
        @produce_or_load("", wsave_kwargs, suffix = "jld2", wsave_kwargs=wsave_kwargs) do c
            data
        end
    end
    # Check if both files exist now
    @test isfile(sn_uncomp)
    @test isfile(sn_comp)
    # Test if the compressed file is smaller
    size_uncomp = filesize(sn_uncomp)
    size_comp = filesize(sn_comp)
    @test size_uncomp > size_comp
    # Leave no trace
    rm(sn_uncomp)
    rm(sn_comp)

end

@test produce_or_load(simulation, f; loadfile = false)[1] == nothing
rm(savename(simulation, "jld2"))
@test !isfile(savename(simulation, "jld2"))

using DataFrames
@testset "produce_or_load with dataframe" begin
    function makedf(c)
        a = c[:a]
        return DataFrame(A = a, B = rand(3))
    end
    c = (a = 0.5,)
    _, spath = produce_or_load(c, makedf; suffix = "csv")
    @test isfile(spath)
    df = DataFrame(load(spath))
    @test df.A == [0.5, 0.5, 0.5]
    rm(spath)
    @test !isfile(spath)
end

# Regression test for https://github.com/JuliaDynamics/DrWatson.jl/issues/309
@testset "@produce_or_load with value-type widening" begin
    d, spath = @produce_or_load("", Dict("a" => 5.0)) do config
        return Dict("field" => config)
    end
    @test haskey(d, "script")
    rm(spath)
end

################################################################################
#                          Backup files before saving                          #
################################################################################

@testset "Backup ($ending)" for ending ∈ ["bson", "jld2"]
    filepath = "test.#backup."*ending
    data = [Dict( "a" => i, "b" => rand(rand(1:10))) for i = 1:3]
    for i = 1:3
        safesave(filepath, data[i])
        @test data[i] == load(filepath)
    end
    @test data[2] == load("test.#backup_#1."*ending)
    @test data[1] == load("test.#backup_#2."*ending)
    @test data[3] == load("test.#backup."*ending)
    rm("test.#backup."*ending)
    rm("test.#backup_#1."*ending)
    rm("test.#backup_#2."*ending)
end

@testset "Backup (dir)" begin
    # Save contents as individual file(s) within a parent directory:
    struct Composite x end
    DrWatson._wsave(dir, data::Composite) = wsave(joinpath(dir, "x.jld2"), data.x)
    load_composite(dir) = load(joinpath(dir, "x.jld2"))

    filepath = "test.#backup.dir"
    data = [Composite(Dict( "a" => i, "b" => rand(rand(1:10)))) for i = 1:3]
    for i = 1:3
        safesave(filepath, data[i])
        @test isdir(filepath)
        @test data[i].x == load_composite(filepath)
    end
    @test data[2].x == load_composite("test.#backup_#1.dir")
    @test data[1].x == load_composite("test.#backup_#2.dir")
    @test data[3].x == load_composite("test.#backup.dir")
    rm("test.#backup.dir"; recursive=true)
    rm("test.#backup_#1.dir"; recursive=true)
    rm("test.#backup_#2.dir"; recursive=true)
end


text:
module SliceMap

export mapcols, MapCols, maprows, slicemap, tmapcols, ThreadMapCols

using MacroTools, TensorCast, JuliennedArrays

using Tracker
using Tracker: TrackedMatrix, track, @grad, data

using ZygoteRules
using ZygoteRules: pullback, @adjoint

#========== Reverse, Eachslice ==========#

"""
    mapcols(f, m) ≈ mapreduce(f, hcat, eachcol(m)) ≈ mapslices(f, m, dims=1)

This is a more efficient version of the functions on the right.
For `f(x::Vector)::Matrix` it reshapes like `mapslices(vec∘f, m, dims=1)`.

It provides a gradient for Tracker and Zygote, saving the backward function for each slice.

Any arguments after the matrix are passed to `f` as scalars, i.e.
`mapcols(f, m, args...) = reduce(hcat, f(col, args...) for col in eeachcol(m))`.
They do not get sliced/iterated (unlike `map`), nor are their gradients tracked.
"""
mapcols(f::Function, M, args...) = _mapcols(map, f, M, args...)
tmapcols(f::Function, M, args...) = _mapcols(threadmap, f, M, args...)

_mapcols(map::Function, f::Function, M::AbstractMatrix, args...) =
    reduce(hcat, map(col -> surevec(f(col, args...)), eachcol(M)))

surevec(x::Number) = [x] # to allow f vector -> scalar, as mapslices does
surevec(A) = vec(A)      # to allow f vector -> matrix, by reshaping

_mapcols(map::Function, f::Function, M::TrackedMatrix, args...) = track(_mapcols, map, f, M, args...)

@grad _mapcols(map::Function, f::Function, M::AbstractMatrix, args...) =
    ∇mapcols(map, map(col -> Tracker.forward(x -> surevec(f(x, args...)), col), eachcol(data(M))), args...)

@adjoint _mapcols(map::Function, f::Function, M::AbstractMatrix, args...) =
    ∇mapcols(map, map(col -> ZygoteRules.pullback(x -> surevec(f(x, args...)), col), eachcol(M)), args)

function ∇mapcols(bigmap, forwards, args...)
    reduce(hcat, map(data∘first, forwards)), Δ -> begin
        cols = bigmap((fwd, Δcol) -> data(last(fwd)(Δcol)[1]), forwards, eachcol(data(Δ)))
        (nothing, nothing, reduce(hcat, cols), map(_->nothing, args)...)
    end
end

"""
    maprows(f, M) ≈ mapslices(f, M, dims=2)

Like `mapcols()` but for rows.
"""
maprows(f::Function, M::AbstractMatrix, args...) =
    reduce(vcat, map(col -> transpose(surevec(f(col, args...))), eachrow(M)))

maprows(f::Function, M::TrackedMatrix, args...) = track(maprows, f, M, args...)

@grad maprows(f::Function, M::AbstractMatrix, args...) =
    ∇maprows(map(row -> Tracker.forward(x -> surevec(f(x, args...)), row), eachrow(data(M))), args)

@adjoint maprows(f::Function, M::AbstractMatrix, args...) =
    ∇maprows(map(row -> ZygoteRules.pullback(x -> surevec(f(x, args...)), row), eachrow(M)), args)

function ∇maprows(forwards, args)
    reduce(vcat, map(transpose∘data∘first, forwards)), Δ -> begin
        rows = map((fwd, Δrow) -> data(last(fwd)(Δrow)[1]), forwards, eachrow(data(Δ)))
        (nothing, reduce(vcat, transpose.(rows)), map(_->nothing, args)...)
    end
end

"""
    slicemap(f, A; dims) ≈ mapslices(f, A; dims)

Like `mapcols()`, but for any slice. The function `f` must preserve shape,
e.g. if `dims=(2,4)` then `f` must map matrices to matrices.

The gradient is for Zygote only.
"""
function slicemap(f::Function, A::AbstractArray{T,N}, args...; dims) where {T,N}
    code = ntuple(d -> d in dims ? (:) : (*), N)
    B = TensorCast.sliceview(A, code)
    C = [ f(slice, args...) for slice in B ]
    TensorCast.glue(C, code)
end
# TODO switch to JuliennedArrays, then rm TensorCast dep

#========== Forward, Static ==========#

using StaticArrays, ForwardDiff

struct MapCols{d} end

"""
    MapCols{d}(f, m::Matrix, args...)

Similar to `mapcols(f, m, args...)`, but slices `m` into `SVector{d}` columns.
Their length `d = size(M,1)` should ideally be provided for type-stability, but is not required.

The gradient for Tracker and Zygote uses `ForwardDiff` on each slice.
"""
MapCols(f::Function, M::AbstractMatrix, args...) =
    MapCols{size(M,1)}(f, M, args...)

MapCols{d}(f::Function, M::AbstractMatrix, args...) where {d} =
    _MapCols(map, f, M, Val(d), args...)

function _MapCols(map::Function, f::Function, M::Matrix{T}, ::Val{d}, args...) where {T,d}
    d == size(M,1) || error("expected M with $d rows")
    A = reinterpret(SArray{Tuple{d}, T, 1, d}, vec(M))
    B = map(col -> surevec(f(col, args...)), A)
    reduce(hcat, B)
end

_MapCols(map::Function, f::Function, M::TrackedMatrix, dval, args...) =
    track(_MapCols, map, f, M, dval, args...)

@grad _MapCols(map::Function, f::Function, M::TrackedMatrix, dval, args...) =
    ∇MapCols(map, f, M, dval, args...)

@adjoint _MapCols(map::Function, f::Function, M::Matrix, dval, args...) =
    ∇MapCols(map, f, M, dval, args...)

function ∇MapCols(bigmap::Function, f::Function, M::AbstractMatrix{T}, dval::Val{d}, args...) where {T,d}
    d == size(M,1) || error("expected M with $d rows")
    k = size(M,2)

    A = reinterpret(SArray{Tuple{d}, T, 1, d}, vec(data(M)))

    dualcol = SVector(ntuple(j->ForwardDiff.Dual(0, ntuple(i->i==j ? 1 : 0, dval)...), dval))
    C = bigmap(col -> surevec(f(col + dualcol, args...)), A)

    Z = reduce(hcat, map(col -> ForwardDiff.value.(col), C))

    function back(ΔZ)
        S = promote_type(T, eltype(data(ΔZ)))
        ∇M = zeros(S, size(M))
        @inbounds for c=1:k
            part = ForwardDiff.partials.(C[c])
            for r=1:d
                for i=1:size(ΔZ,1)
                    ∇M[r,c] += data(ΔZ)[i,c] * part[i].values[r]
                end
            end
        end
        (nothing, nothing, ∇M, nothing, map(_->nothing, args)...)
    end
    Z, back
end

#========== Gradients for Zygote ==========#

# @require Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"

# @init @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" include("zygote.jl")
# Now using ZygoteRules instead, mapcols etc above.

#= TensorCast =#
# These could move there, TODO

@adjoint TensorCast.sliceview(A::AbstractArray, code::Tuple) =
    TensorCast.sliceview(A, code), Δ -> (TensorCast.glue(Δ, code), nothing)

@adjoint TensorCast.red_glue(A::AbstractArray, code::Tuple) =
    TensorCast.red_glue(A, code), Δ -> (TensorCast.sliceview(Δ, code), nothing)

@adjoint TensorCast.copy_glue(A::AbstractArray, code::Tuple) =
    TensorCast.copy_glue(A, code), Δ -> (TensorCast.sliceview(Δ, code), nothing)

#= JuliennedArrays =#

@adjoint JuliennedArrays.Slices(whole, along...) =
    Slices(whole, along...), Δ -> (Align(Δ, along...), map(_->nothing, along)...)

@adjoint JuliennedArrays.Align(whole, along...) =
    Align(whole, along...), Δ -> (Slices(Δ, along...), map(_->nothing, along)...)

#= Base =#

@adjoint Base.reduce(::typeof(hcat), V::AbstractVector{<:AbstractVector}) =
    reduce(hcat, V), dV -> (nothing, collect(eachcol(dV)),)

#========== Experimenting with gradients for for eachslice / reduce ==========#

export gluecol, collecteachcol
export mapcols2, mapcols4, mapcols5, mapcols6, mapcols7

gluecol(V::AbstractVector{<:AbstractVector}) = reduce(hcat, V)

#=
gluecol(V::AbstractVector{<:TrackedVector}) = track(gluecol, V)

@grad function gluecol(V::AbstractVector)
    gluecol(data.(V)), ΔM -> (collect(eachcol(data(ΔM))),) # doesn't work
end
=#

@adjoint gluecol(V::AbstractVector) =
    gluecol(V), ΔM -> (collect(eachcol(ΔM)),) # does work!

function mapcols2(f, A)
    cols = [A[:,c] for c=1:size(A,2)]
    res = f.(cols)
    gluecol(res)
end

# Apply that straight to reduce(hcat,...)

function mapcols4(f, A)
    cols = [view(A,:,c) for c=1:size(A,2)]
    res = map(f, cols)
    reduce(hcat, res)
end

# Surprisingly dy for eachcol seems to know the answer?
# typeof(dy) = NamedTuple{(:f, :iter),Tuple{NamedTuple{(:A,),Tuple{Array{Float64,2}}},Array{Nothing,1}}}
# dy = (f = (A = [47.9325 51.3781
# Which means this works... but uses as much memory as gradient of array of views:

#=@adjoint function eachcol(x::AbstractMatrix)
    eachcol(x), dy -> (dy.f.A,) #= begin
        @show typeof(dy) dy
        dx = zero(x) .+ 0.0  # zeros(eltype(dy), size(x))
        foreach(copyto!, eachcol(dx), dy)
        (dx,)
    end =#
end=#

# @adjoint eachcol(x) = eachcol(x), dy -> (dy.f.A,)

function mapcols5(f, A)
    cols = collect(eachcol(A))
    res = map(f, cols)
    reduce(hcat, res)
end

collecteachcol(x) = collect(eachcol(x))

@adjoint function collecteachcol(x)
    collecteachcol(x), dy -> begin
        dx = _zero(x) # _zero is not in ZygoteRules, TODO
        foreach(copyto!, collecteachcol(dx), dy)
        (dx,)
    end
end

function mapcols6(f, A)
    cols = collecteachcol(A)
    res = map(f, cols)
    reduce(hcat, res)
end

# function mapcols7(f, A)
#     cols = eachcol(A) # without collect. Zygote.gradient -> StackOverflowError
#     res = map(f, cols)
#     reduce(hcat, res)
# end

# Following a suggestion? Doesn't help.
# @adjoint Base.collect(x) = collect(x), Δ -> (Δ,)

#========== Threaded Map ==========#

# What KissThreading does is much more complicated, perhaps worth investigating:
# https://github.com/mohamed82008/KissThreading.jl/blob/master/src/KissThreading.jl

# BTW I do the first one because some diffeq maps infer to ::Any,
# else you could use Core.Compiler.return_type(f, Tuple{eltype(x)})

"""
    threadmap(f, A)
    threadmap(f, A, B)

Simple version of `map` using a `Threads.@threads` loop;
only for vectors & really at most two of them, of nonzero length,
with all outputs having the same type.
"""
function threadmap(f::Function, vw::AbstractVector...)
    length(first(vw))==0 && error("can't map over empty vector, sorry")
    length(vw)==2 && (isequal(length.(vw)...) || error("lengths must be equal"))
    out1 = f(first.(vw)...)
    _threadmap(out1, f, vw...)
end
# NB barrier
function _threadmap(out1, f, vw...)
    out = Vector{typeof(out1)}(undef, length(first(vw)))
    out[1] = out1
    Threads.@threads for i=2:length(first(vw))
        @inbounds out[i] = f(getindex.(vw, i)...)
    end
    out
end

# Collect generators to allow indexing
threadmap(f::Function, v) = threadmap(f, collect(v))
threadmap(f::Function, v, w) = threadmap(f, collect(v), collect(w))
threadmap(f::Function, v, w::AbstractVector) = threadmap(f, collect(v), w)
threadmap(f::Function, v::AbstractVector, w) = threadmap(f, v, collect(w))

struct ThreadMapCols{d} end

"""
    ThreadMapCols{d}(f, m::Matrix, args...)

Like `MapCols` but with multi-threading!
"""
ThreadMapCols(f::Function, M::AbstractMatrix, args...) =
    ThreadMapCols{size(M,1)}(f, M, args...)

ThreadMapCols{d}(f::Function, M::AbstractMatrix, args...) where {d} =
    _MapCols(threadmap, f, M, Val(d), args...)

end # module

ext: jl


text:
using OrdinaryDiffEq, StochasticDiffEq, DelayDiffEq, Flux,
      Zygote, Test
println("Starting tests")
using DiffEqFlux, DiffEqSensitivity
mp = Float32[0.1,0.1]
x = Float32[2.; 0.]
xs = Float32.(hcat([0.; 0.], [1.; 0.], [2.; 0.]))
tspan = (0.0f0,1.0f0)
dudt = Chain(Dense(2,50,tanh),Dense(50,2))
fastdudt = FastChain(FastDense(2,50,tanh),FastDense(50,2))
staticdudt = FastChain(StaticDense(2,50,tanh),StaticDense(50,2))

NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)(x)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1)(x)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())(x)

NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)(xs)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1)(xs)
NeuralODE(dudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())(xs)

@info "Test some gradients"

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint())
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint(autojacvec=ZygoteVJP()))
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

## Fast

@info "Test some fast layers"

node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])

goodgrad = grads[node.p]
p = node.p

node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint(),p=p)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
@test ! iszero(grads[xs])
@test ! iszero(grads[node.p])
goodgrad2 = grads[node.p]
@test goodgrad ≈ goodgrad2 # Make sure adjoint overloads are correct

node = NeuralODE(staticdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=BacksolveAdjoint(),p=p)
grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
@test ! iszero(grads[x])
@test ! iszero(grads[node.p])

@test_throws ErrorException grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))

@info "Test some adjoints"

# Adjoint
@testset "adjoint mode" begin
    node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(dudt,tspan,Tsit5(),saveat=0.0:0.1:1.0)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(dudt,tspan,Tsit5(),saveat=0.1)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.0:0.1:1.0)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.1)
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])
end

@info "Test Tracker"

# RD
@testset "Tracker mode" begin
    node = NeuralODE(dudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    NeuralODE(dudt,tspan,Tsit5(),saveat=0.0:0.1:1.0,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(dudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node))
    @test ! iszero(grads[xs])
    @test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),save_everystep=false,save_start=false,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])

    NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.0:0.1:1.0,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])

    node = NeuralODE(fastdudt,tspan,Tsit5(),saveat=0.1,sensealg=TrackerAdjoint())
    grads = Zygote.gradient(()->sum(node(x)),Flux.params(x,node))
    @test ! iszero(grads[x])
    @test ! iszero(grads[node.p])

    @test_broken grads = Zygote.gradient(()->sum(node(xs)),Flux.params(xs,node)) isa Tuple
    #@test ! iszero(grads[xs])
    #@test ! iszero(grads[node.p])
end

@info "Test non-ODEs"

dudt2 = Chain(Dense(2,50,tanh),Dense(50,2))
fastdudt2 = FastChain(FastDense(2,50,tanh),FastDense(50,2))
NeuralDSDE(dudt,dudt2,(0.0f0,.1f0),SOSRI(),saveat=0.1)(x)
sode = NeuralDSDE(dudt,dudt2,(0.0f0,.1f0),SOSRI(),saveat=0.0:0.01:0.1)

grads = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(grads[x])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

grads = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode))
@test ! iszero(grads[xs])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

sode = NeuralDSDE(fastdudt,fastdudt2,(0.0f0,.1f0),SOSRI(),saveat=0.0:0.01:0.1)

grads = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(grads[x])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

@test_broken grads = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode)) isa Tuple
#@test ! iszero(grads[xs])
#@test ! iszero(grads[sode.p])
#@test ! iszero(grads[sode.p][end])

dudt22 = Chain(Dense(2,50,tanh),Dense(50,4),x->reshape(x,2,2))
fastdudt22 = FastChain(FastDense(2,50,tanh),FastDense(50,4),(x,p)->reshape(x,2,2))
NeuralSDE(dudt,dudt22,(0.0f0,.1f0),2,LambaEM(),saveat=0.01)(x)

sode = NeuralSDE(dudt,dudt22,(0.0f0,0.1f0),2,LambaEM(),saveat=0.0:0.01:0.1)

grads = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(grads[x])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

@test_broken grads = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode))
#@test ! iszero(grads[xs])
#@test ! iszero(grads[sode.p])
#@test ! iszero(grads[sode.p][end])

sode = NeuralSDE(fastdudt,fastdudt22,(0.0f0,0.1f0),2,LambaEM(),saveat=0.0:0.01:0.1)

grads = Zygote.gradient(()->sum(sode(x)),Flux.params(x,sode))
@test ! iszero(grads[x])
@test ! iszero(grads[sode.p])
@test ! iszero(grads[sode.p][end])

@test_broken grads = Zygote.gradient(()->sum(sode(xs)),Flux.params(xs,sode))
#@test ! iszero(grads[xs])
#@test ! iszero(grads[sode.p])
#@test ! iszero(grads[sode.p][end])

ddudt = Chain(Dense(6,50,tanh),Dense(50,2))
NeuralCDDE(ddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.1)(x)
dode = NeuralCDDE(ddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.0:0.1:2.0)

grads = Zygote.gradient(()->sum(dode(x)),Flux.params(x,dode))
@test ! iszero(grads[x])
@test ! iszero(grads[dode.p])

@test_broken grads = Zygote.gradient(()->sum(dode(xs)),Flux.params(xs,dode)) isa Tuple
#@test ! iszero(grads[xs])
#@test ! iszero(grads[dode.p])

fastddudt = FastChain(FastDense(6,50,tanh),FastDense(50,2))
NeuralCDDE(fastddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.1)(x)
dode = NeuralCDDE(fastddudt,(0.0f0,2.0f0),(p,t)->zero(x),(1f-1,2f-1),MethodOfSteps(Tsit5()),saveat=0.0:0.1:2.0)

grads = Zygote.gradient(()->sum(dode(x)),Flux.params(x,dode))
@test ! iszero(grads[x])
@test ! iszero(grads[dode.p])

@test_broken grads = Zygote.gradient(()->sum(dode(xs)),Flux.params(xs,dode)) isa Tuple
#@test ! iszero(grads[xs])
#@test ! iszero(grads[dode.p])


text:
using ImageFiltering, Test, Statistics

@testset "gabor" begin
    σx = 8
    σy = 12
    size_x = 6*σx+1
    size_y = 6*σy+1
    γ = σx/σy
    # zero size forces default kernel width, with warnings
    @info "Four warnings are expected"
    kernel = Kernel.gabor(0,0,σx,0,5,γ,0)
    @test isequal(size(kernel[1]),(size_x,size_y))
    kernel = Kernel.gabor(0,0,σx,π,5,γ,0)
    @test isequal(size(kernel[1]),(size_x,size_y))

    for x in 0:4, y in 0:4, z in 0:4, t in 0:4
        σx1 = 2*x+1
        σy1 = 2*y+1
        λ = 2*z+1
        γ1 = σx1/σy1
        θ = 2*t+1
        kernel1 = Kernel.gabor(9,9,σx1,θ,λ,γ1,0)
        kernel2 = Kernel.gabor(9,9,σx1,θ+π,λ,γ1,0)
        @test abs(sum(kernel1[1] - kernel2[1])) < 1e-2
        @test abs(sum(kernel1[2] - kernel2[2])) < 1e-2
    end

    x1 = [j for i in 0:49,j in 0:49]
    wavelengths = (3, 10)
    images = [sin.(2*π*x1/λ) for λ in wavelengths]
    σx = 4
    σy = 5
    function match_score(image, λ)
        gabor_real = imfilter(image,centered(Kernel.gabor(6*σx+1,6*σy+1,σx,0,λ,σx/σy,0)[1]),"replicate")
        gabor_imag = imfilter(image,centered(Kernel.gabor(6*σx+1,6*σy+1,σx,0,λ,σx/σy,0)[2]),"replicate")
        gabor_result = sqrt.((gabor_real.*gabor_real) + (gabor_imag.*gabor_imag))
        return mean(gabor_result)
    end
    gabor_output = rand(Float64,2,2)
    for i = 1:2
        for j = 1:2
            gabor_output[i,j] = match_score(images[i],wavelengths[j])
        end
    end
    @test gabor_output[1,1] > gabor_output[1,2]
    @test gabor_output[2,2] > gabor_output[1,2]
    @test gabor_output[1,1] > gabor_output[2,1]
    @test gabor_output[2,2] > gabor_output[2,1]

    @test_throws ArgumentError Kernel.gabor(9,9,-2,0,5,0.1,0)
    @test_throws ArgumentError Kernel.gabor(9,9,2,0,-5,0.1,0)
    @test_throws ArgumentError Kernel.gabor(9,9,2,0,5,0,0)

end


text:
# Author: Ashwani Rathee
# Program to Find All Roots of a Quadratic Equation

println("Enter a,b,c of the quadratic equations to find roots: ")
a=readline()
a=parse(Float64,a)

b=readline()
b=parse(Float64,b)

c=readline()
c=parse(Float64,c)
det=b^2-4*a*c
println("The determinant is:",det)
root1=(b+sqrt(det))/(2*a)
root2=(b-sqrt(det))/(2*a)
println("THe Roots are: ",root1,",",root2)

max_stars_repo_name: ashwani-rathee/Learning--Julia

max_issues_repo_name: ashwani-rathee/Learning--Julia

max_forks_repo_name: ashwani-rathee/Learning--Julia


text:
# This file is a part of Julia. License is MIT: http://julialang.org/license

# TODO: make this a general purpose solution
function Base.cconvert(::Type{Ptr{DiffOptionsStruct}}, pathspecs::AbstractString)
    str_ref = Base.cconvert(Ref{Cstring}, [pathspecs])
    sa = StrArrayStruct(Base.unsafe_convert(Ref{Cstring}, str_ref), 1)
    do_ref = Ref(DiffOptionsStruct(pathspec = sa))
    do_ref, str_ref
end
function Base.unsafe_convert(::Type{Ptr{DiffOptionsStruct}}, rr::Tuple{Ref{DiffOptionsStruct}, Ref{Cstring}})
    Base.unsafe_convert(Ptr{DiffOptionsStruct}, first(rr))
end

function diff_tree(repo::GitRepo, tree::GitTree, pathspecs::AbstractString=""; cached::Bool=false)
    diff_ptr_ptr = Ref{Ptr{Void}}(C_NULL)
    if cached
        @check ccall((:git_diff_tree_to_index, :libgit2), Cint,
                     (Ptr{Ptr{Void}}, Ptr{Void}, Ptr{Void}, Ptr{Void}, Ptr{DiffOptionsStruct}),
                     diff_ptr_ptr, repo.ptr, tree.ptr, C_NULL, isempty(pathspecs) ? C_NULL : pathspecs)
    else
        @check ccall((:git_diff_tree_to_workdir_with_index, :libgit2), Cint,
                     (Ptr{Ptr{Void}}, Ptr{Void}, Ptr{Void}, Ptr{DiffOptionsStruct}),
                     diff_ptr_ptr, repo.ptr, tree.ptr, isempty(pathspecs) ? C_NULL : pathspecs)
    end
    return GitDiff(repo, diff_ptr_ptr[])
end

function diff_tree(repo::GitRepo, oldtree::GitTree, newtree::GitTree)
    diff_ptr_ptr = Ref{Ptr{Void}}(C_NULL)
    @check ccall((:git_diff_tree_to_tree, :libgit2), Cint,
                  (Ptr{Ptr{Void}}, Ptr{Void}, Ptr{Void}, Ptr{Void}, Ptr{DiffOptionsStruct}),
                   diff_ptr_ptr, repo.ptr, oldtree.ptr, newtree.ptr, C_NULL)
    return GitDiff(repo, diff_ptr_ptr[])
end

function Base.count(diff::GitDiff)
    return ccall((:git_diff_num_deltas, :libgit2), Cint, (Ptr{Void},), diff.ptr)
end

function Base.getindex(diff::GitDiff, i::Integer)
    delta_ptr = ccall((:git_diff_get_delta, :libgit2),
                      Ptr{DiffDelta},
                      (Ptr{Void}, Csize_t), diff.ptr, i-1)
    delta_ptr == C_NULL && throw(BoundsError(diff, (i,)))
    return unsafe_load(delta_ptr)
end


text:
include("cantorPairing.jl");

function cantorPairingRecursive(arrayTuple)
	len = length(arrayTuple);
	if (len == 2)
		return cantorPairing(arrayTuple[1], arrayTuple[2]);
	else
		return cantorPairing(arrayTuple[1], cantorPairingRecursive(arrayTuple[2:len]));
	end
end

max_stars_repo_name: driesverhees/dobble

max_issues_repo_name: driesverhees/dobble

max_forks_repo_name: driesverhees/dobble


text:
α = 0.025

# Hepers

@testset "correlations" begin
  @testset "copula functions" begin
    @test Ccl([1., 0.5], 3.) == 0.5
    @test Cg([1., 0.5], 3.) == 0.5
  end
  @testset "Kendall cor" begin
    @test Debye(0.5, 1) ≈ 0.8819271567906056
    @test τ2θ(0.6, "frank") ≈ 7.929642284264058
    @test τ2θ(0.5, "gumbel") ≈ 2.
    @test τ2θ(1/3, "clayton") ≈ 1.
    @test τ2θ(1/4, "amh") ≈ 0.8384520912688538
    @test_throws AssertionError τ2θ(1/3, "frak")
    @test_throws DomainError τ2θ(0., "frank")
    @test frankτ2θ(0.6) ≈ 7.929642284264058
    @test AMHτ2θ(1/4) ≈ 0.8384520912688538
    @test AMHτ2θ(0.) ≈ 0. atol=1.0e-6
    @test AMHτ2θ(0.28) ≈ 0.9999
    @test AMHτ2θ(-2/11) ≈ -0.9999
  end
  @testset "Spearman cor" begin
    @test dilog(0.5) ≈ 0.5822405264650125
    @test claytonρ2θ(1/3) ≈ 0.58754 atol=1.0e-5
    @test claytonρ2θ(0.01) ≈ 0. atol=1.0e-1
    @test gumbelρ2θ(0.5) ≈ 1.5410704204332681
    @test gumbelρ2θ(0.0001) == 1.
    @test ρ2θ(1/3, "frank") ≈ 2.116497 atol=1.0e-5
    @test ρ2θ(0.2, "amh") ≈ 0.5168580913147318
    @test_throws AssertionError ρ2θ(1/3, "frak")
    @test_throws DomainError ρ2θ(0., "frank")
    @test frankρ2θ(1/3) ≈ 2.116497 atol=1.0e-5
    @test AMHρ2θ(0.2) ≈ 0.5168580913147318
    @test claytonθ2ρ(3.) ≈ 0.78645 atol=1.0e-4
    @test claytonθ2ρ(0.001) ≈ 0. atol=1.0e-2
    @test gumbelθ2ρ(3.) ≈ 0.8489 atol=1.0e-4
    @test AMHρ2θ(0.) ≈ 0. atol=1.0e-4
    @test AMHρ2θ(0.49) ≈ 1 atol=1.0e-4
    @test AMHρ2θ(-0.273) ≈ -1 atol=1.0e-4
  end
  @testset "negative cor" begin
    @test ρ2θ(-0.2246, "amh") ≈ -0.8 atol=1.0e-3
    @test τ2θ(-0.1505, "amh") ≈ -0.8 atol=1.0e-3
    @test ρ2θ(-0.5572, "frank") ≈ -4. atol=1.0e-3
    @test τ2θ(-0.3881, "frank") ≈ -4. atol=1.0e-3
    @test τ2θ(-1/3, "clayton") ≈ -.5 atol=1.0e-5
    @test ρ2θ(-0.4668, "clayton") ≈ -.5 atol=1.0e-3
  end
end
@testset "archimedean copulas axiliary functions" begin
  @test getV0(2., [0.2, 0.4, 0.6, 0.8], "clayton") ≈ [0.0320924, 0.137498, 0.354163, 0.821187] atol=1.0e-4
  @test_throws AssertionError getV0(2., [0.2, 0.4, 0.6, 0.8], "clayto")
  @test phi([0.2 0.6; 0.4 0.8], 2., "clayton") ≈ [0.912871  0.790569; 0.845154  0.745356] atol=1.0e-4
  @test_throws AssertionError phi([0.2 0.6; 0.4 0.8], 2., "clayto")
  c = copulagen("clayton", [0.2 0.4 0.8; 0.2 0.8 0.6; 0.3 0.9 0.6], 1.)
  @test c ≈ [0.5 0.637217; 0.362783 0.804163; 0.432159 0.896872] atol=1.0e-5
  @test useτ(0.5, "clayton") == 2.
  @test useρ(0.75, "gumbel") ≈ 2.285220798876495
end

@testset "archimedean copulas exceptions" begin
  @test_throws DomainError testθ(0.5, "gumbel")
  @test_throws DomainError useρ(0.6, "amh")
  @test_throws DomainError useτ(0.45, "amh")
  @test_throws DomainError archcopulagen(100000, 4, -0.6, "frank")
end

@testset "gumbel copula" begin
  Random.seed!(43)
  x = archcopulagen(500000, 3, 2., "gumbel");
  @test pvalue(ExactOneSampleKSTest(x[:,1], Uniform(0,1))) > α
  @test pvalue(ExactOneSampleKSTest(x[:,2], Uniform(0,1))) > α
  @test pvalue(ExactOneSampleKSTest(x[:,3], Uniform(0,1))) > α
  @test tail(x[:,1], x[:,2], "r") ≈ 2-2^(1/2) atol=1.0e-1
  @test tail(x[:,1], x[:,2], "r") ≈ 2-2^(1/2) atol=1.0e-1
  @test tail(x[:,1], x[:,2], "l", 0.00001) ≈ 0.
  @test tail(x[:,1], x[:,3], "l", 0.00001) ≈ 0.
  @test corkendall(x) ≈ [1. 1/2 1/2; 1/2 1. 1/2; 1/2 1/2 1.] atol=1.0e-2
  Random.seed!(43)
  x = archcopulagen(500000, 2, 1.5, "gumbel"; rev = true)
  @test tail(x[:,1], x[:,2], "l") ≈ 2-2^(1/1.5) atol=1.0e-1
  @test tail(x[:,1], x[:,2], "r", 0.00001) ≈ 0.
  @test pvalue(ExactOneSampleKSTest(x[:,1], Uniform(0,1))) > α
  Random.seed!(43)
  x = archcopulagen(500000, 2, 0.5, "gumbel"; cor = "Kendall")
  @test corkendall(x) ≈ [1. 0.5; 0.5 1.] atol=1.0e-4
  Random.seed!(43)
  x = archcopulagen(500000, 2, 0.5, "gumbel"; cor = "Spearman")
  @test corspearman(x) ≈ [1. 0.5; 0.5 1.] atol=1.0e-2
end
@testset "clayton copula" begin
  Random.seed!(43)
  x = archcopulagen(500000, 3, 1., "clayton");
  @test pvalue(ExactOneSampleKSTest(x[:,1], Uniform(0,1))) > α
  @test pvalue(ExactOneSampleKSTest(x[:,2], Uniform(0,1))) > α
  @test pvalue(ExactOneSampleKSTest(x[:,3], Uniform(0,1))) > α
  @test tail(x[:,1], x[:,2], "l") ≈ 2.0^(-1) atol=1.0e-1
  @test tail(x[:,1], x[:,3], "l") ≈ 2.0^(-1) atol=1.0e-1
  @test tail(x[:,1], x[:,2], "r", 0.0001) ≈ 0
  @test corkendall(x) ≈ [1. 1/3 1/3; 1/3 1. 1/3; 1/3 1/3 1.] atol=1.0e-2
  Random.seed!(43)
  x = archcopulagen(500000, 2, 0.5, "clayton"; cor = "Kendall")
  @test corkendall(x) ≈ [1. 0.5; 0.5 1.] atol=1.0e-3
  Random.seed!(43)
  x = archcopulagen(500000, 2, -0.9, "clayton")
  @test pvalue(ExactOneSampleKSTest(x[:,2], Uniform(0,1))) > α
  @test pvalue(ExactOneSampleKSTest(x[:,1], Uniform(0,1))) > α
  @test corkendall(x)[1,2] ≈ -0.9/(2-0.9) atol=1.0e-3
end
@testset "frank copula" begin
  Random.seed!(43)
  x = archcopulagen(500000, 3, 0.8, "frank")
  @test pvalue(ExactOneSampleKSTest(x[:,1], Uniform(0,1))) > α
  @test pvalue(ExactOneSampleKSTest(x[:,2], Uniform(0,1))) > α
  @test pvalue(ExactOneSampleKSTest(x[:,3], Uniform(0,1))) > α
  @test tail(x[:,1], x[:,2], "l", 0.0001) ≈ 0
  @test tail(x[:,2], x[:,3], "r", 0.0001) ≈ 0
  Random.seed!(43)
  x = archcopulagen(500000, 2, 0.2, "frank"; cor = "Kendall")
  @test corkendall(x) ≈ [1. 0.2; 0.2 1.] atol=1.0e-3
end
@testset "Ali-Mikhail-Haq copula" begin
  Random.seed!(43)
  x = archcopulagen(500000, 3, 0.8, "amh")
  @test pvalue(ExactOneSampleKSTest(x[:,1], Uniform(0,1))) > α
  @test pvalue(ExactOneSampleKSTest(x[:,2], Uniform(0,1))) > α
  @test pvalue(ExactOneSampleKSTest(x[:,3], Uniform(0,1))) > α
  @test tail(x[:,1], x[:,2], "l", 0.0001) ≈ 0
  @test tail(x[:,1], x[:,2], "r", 0.0001) ≈ 0
  @test corkendall(x)[1:2, 1:2] ≈ [1. 0.23373; 0.23373 1.] atol=1.0e-3
  x = archcopulagen(500000, 2, 0.25, "amh"; cor = "Kendall")
  @test corkendall(x) ≈ [1. 0.25; 0.25 1.] atol=1.0e-3
end


text:
# Defines the Vrev object which holds a value and an reverse function...
"""
    Vrev{T<:Number,F}

Structure used to wrap the value and reverse function
"""
struct Vrev{T<:Number,F}
  val::T
  rev::F
end
Vrev(x::T) where T = Vrev{T,typeof(identity)}(x, identity)
@inline _val(d::Vrev{T,F}) where {T,F} = d.val

zero(::Vrev{T,F}) where {T<:Number,F} = Vrev(zero(T))

macro norev_scalar(c, f)
  esc(quote
    function Cassette.overdub(::$c, ::typeof($f), x::Vrev{T,F}) where {T<:Number,F}
        r = y -> x.rev(_val(x))
        return Vrev{T,typeof(r)}(($f)(x), r)
    end
  end)
end

macro no_prop(f)
  esc(quote
    function Cassette.overdub(::$c, ::typeof($f), x::Vrev{T,F}) where {T<:Number,F}
        r = y -> x.rev(_val(x))
        return Vrev{T,typeof(r)}(($f)(x), r)
    end
  end)
end

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
function return_kwargs(; kwargs...)
    return kwargs
end

function kwargs_test()
    @testset "kwargs" begin
        defaultVals = KwArgs(Dict([
            :a => "a",
            :b => 2,
            :c => [1, 2, 3]
        ]));

        @test has_default(:b, defaultVals)
        @test !has_default(:bb, defaultVals)
        @test default_value(:c, defaultVals) == [1,2,3]

        kwV = return_kwargs(; :a => "aa", :z => "zz");
        # In both
        @test kw_arg(:a, defaultVals; kwV...) == kwV[:a]
        # in defaults only
        @test kw_arg(:b, defaultVals; kwV...) == 2
        # in kwV only
        @test kw_arg(:z, defaultVals; kwV...) == kwV[:z]

        # Convenience function
        get_arg(name :: Symbol) = kw_arg(name, defaultVals; kwV...);
        @test get_arg(:a) == kwV[:a]
    end
end

@testset "KwArgs" begin
    kwargs_test()
end

# ---------


text:
using Random, Statistics
Random.seed!(1)
allData = [rand(24), rand(15), rand(73)]

xBarArray = mean.(allData)
nArray = length.(allData)
xBarTotal = mean(vcat(allData...))
L = length(nArray)

ssBetween=sum([nArray[i]*(xBarArray[i] - xBarTotal)^2 for i in 1:L])
ssWithin=sum([sum([(ob - xBarArray[i])^2 for ob in allData[i]]) for i in 1:L])
ssTotal=sum([sum([(ob - xBarTotal)^2 for ob in allData[i]]) for i in 1:L])

println("Sum of squares between groups: ", ssBetween)
println("Sum of squares within groups: ", ssWithin)
println("Sum of squares total: ", ssTotal)

max_stars_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_issues_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_forks_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook


text:
"""
    UnweightedParticleRepresentation{S} <: ParticleRepresentation{S} <: AbstractFilterRepresentation{S}

Abstract type for the representation of the conditional distribution over the hidden state by unweighted (i.e. equally weighted) particles (samples).
"""
abstract type UnweightedParticleRepresentation{S} <: ParticleRepresentation{S} end

max_stars_repo_name: simsurace/FeedbackParticleFilters

max_issues_repo_name: simsurace/FeedbackParticleFilters

max_forks_repo_name: simsurace/FeedbackParticleFilters


text:
using Plasmo
using KaHyPar
using Ipopt
using SchwarzSolver

T = 3000          #number of time points
d = sin.(1:T)    #disturbance vector
overlap = 5
imbalance = 0.1

graph = OptiGraph()
@optinode(graph,state[1:T])
@optinode(graph,control[1:T-1])

for (i,node) in enumerate(state)
    @variable(node,x)
    @constraint(node, x >= 0)
    @objective(node,Min,0.001*x^2) #- 2*x*d[i])
end
for node in control
    @variable(node,u)
    @constraint(node, u >= -1000)
    @objective(node,Min,u^2)
end
n1 = state[1]
@constraint(n1,n1[:x] == 0)

@linkconstraint(graph,links[i = 1:T-1], state[i][:x] + control[i][:u] + d[i] == state[i+1][:x])

#Partition the problem
hypergraph,hyper_map = gethypergraph(graph) #create hypergraph object based on graph
partition_vector = KaHyPar.partition(hypergraph,8,configuration = :connectivity,imbalance = imbalance)
partition = Partition(hypergraph,partition_vector,hyper_map)
make_subgraphs!(graph,partition)

#Solve directly with expanded subgraphs
subgraphs = getsubgraphs(graph)
expanded_subs = expand.(Ref(graph),subgraphs,Ref(5))
schwarz_solve(graph,expanded_subs;sub_optimizer = optimizer_with_attributes(Ipopt.Optimizer,"tol" => 1e-12,"print_level" => 0),max_iterations = 100,tolerance = 1e-10,
dual_links = [],primal_links = [])


text:
__precompile__()

module DataStructures

    import Base: <, <=, ==, length, isempty, start, next, done, delete!,
                 show, dump, empty!, getindex, setindex!, get, get!,
                 in, haskey, keys, merge, copy, cat,
                 push!, pop!, shift!, unshift!, insert!,
                 union!, delete!, similar, sizehint!,
                 isequal, hash,
                 map, reverse,
                 first, last, eltype, getkey, values, sum,
                 merge, merge!, lt, Ordering, ForwardOrdering, Forward,
                 ReverseOrdering, Reverse, Lt,
                 isless,
                 union, intersect, symdiff, setdiff, issubset,
                 find, searchsortedfirst, searchsortedlast, endof, in

    using Compat: uninitialized, Nothing, Cvoid, AbstractDict

    export Deque, Stack, Queue, CircularDeque
    export deque, enqueue!, dequeue!, dequeue_pair!, update!, reverse_iter
    export capacity, num_blocks, front, back, top, top_with_handle, sizehint!

    export Accumulator, counter, reset!, inc!, dec!

    export ClassifiedCollections
    export classified_lists, classified_sets, classified_counters

    export IntDisjointSets, DisjointSets, num_groups, find_root, in_same_set, root_union!

    export AbstractHeap, compare, extract_all!
    export BinaryHeap, binary_minheap, binary_maxheap, nlargest, nsmallest
    export MutableBinaryHeap, mutable_binary_minheap, mutable_binary_maxheap
    export heapify!, heapify, heappop!, heappush!, isheap

    export OrderedDict, OrderedSet
    export DefaultDict, DefaultOrderedDict
    export Trie, subtrie, keys_with_prefix, path

    export LinkedList, Nil, Cons, nil, cons, head, tail, list, filter, cat,
           reverse
    export SortedDict, SortedMultiDict, SortedSet
    export SDToken, SDSemiToken, SMDToken, SMDSemiToken
    export SetToken, SetSemiToken
    export startof
    export pastendsemitoken, beforestartsemitoken
    export searchsortedafter, searchequalrange
    export packcopy, packdeepcopy
    export exclusive, inclusive, semitokens
    export orderobject, ordtype, Lt, compare, onlysemitokens

    export MultiDict, enumerateall

    import Base: eachindex, keytype, valtype

    include("delegate.jl")

    include("deque.jl")
    include("circ_deque.jl")
    include("stack.jl")
    include("queue.jl")
    include("accumulator.jl")
    include("classified_collections.jl")
    include("disjoint_set.jl")
    include("heaps.jl")

    include("dict_support.jl")
    include("ordered_dict.jl")
    include("ordered_set.jl")
    include("default_dict.jl")
    include("trie.jl")

    include("int_set.jl")

    include("list.jl")
    include("balanced_tree.jl")
    include("tokens.jl")

    import .Tokens: IntSemiToken

    include("multi_dict.jl")
    include("sorted_dict.jl")
    include("sorted_multi_dict.jl")
    include("sorted_set.jl")
    include("tokens2.jl")
    include("container_loops.jl")

    include("dict_sorting.jl")

    export
        CircularBuffer,
        capacity,
        isfull
    include("circular_buffer.jl")

    export status
    export deref_key, deref_value, deref, advance, regress

    export PriorityQueue, peek

    include("priorityqueue.jl")

    # Deprecations

    # Remove when Julia 0.7 (or whatever version is after v0.6) is released
    @deprecate DefaultDictBase(default, ks::AbstractArray, vs::AbstractArray) DefaultDictBase(default, zip(ks, vs))
    @deprecate DefaultDictBase(default, ks, vs) DefaultDictBase(default, zip(ks, vs))
    @deprecate DefaultDictBase(::Type{K}, ::Type{V}, default) where {K,V} DefaultDictBase{K,V}(default)

    @deprecate DefaultDict(default, ks, vs) DefaultDict(default, zip(ks, vs))
    @deprecate DefaultDict(::Type{K}, ::Type{V}, default) where {K,V} DefaultDict{K,V}(default)

    @deprecate DefaultOrderedDict(default, ks, vs) DefaultOrderedDict(default, zip(ks, vs))
    @deprecate DefaultOrderedDict(::Type{K}, ::Type{V}, default) where {K,V} DefaultOrderedDict{K,V}(default)

    function SortedMultiDict(ks::AbstractVector{K},
                             vs::AbstractVector{V},
                             o::Ordering=Forward) where {K,V}
        Base.depwarn("SortedMultiDict(ks, vs, o::Ordering=Forward) is deprecated.\n" * "Use SortedMultiDict(o, zip(ks,vs)) or SortedMultiDict(zip(ks, vs))", :SortedMultiDict)
        if length(ks) != length(vs)
            throw(ArgumentError("SortedMultiDict(ks,vs,o): ks and vs arrays must be the same length"))
        end
        SortedMultiDict(o, zip(ks,vs))
    end

    @deprecate PriorityQueue(::Type{K}, ::Type{V}) where {K,V} PriorityQueue{K,V}()
    @deprecate PriorityQueue(::Type{K}, ::Type{V}, o::Ordering) where {K,V} PriorityQueue{K,V,typeof(o)}(o)
    @deprecate (PriorityQueue{K,V,ForwardOrdering}() where {K,V}) PriorityQueue{K,V}()

    function PriorityQueue(ks::AbstractVector{K},
                           vs::AbstractVector{V},
                           o::Ordering=Forward) where {K,V}
        Base.depwarn("PriorityQueue(ks, vs, o::Ordering=Forward) is deprecated.\n" * "Use PriorityQueue(o, zip(ks,vs)) or PriorityQueue(zip(ks, vs))", :PriorityQueue)
        if length(ks) != length(vs)
            throw(ArgumentError("PriorityQueue(ks,vs,o): ks and vs arrays must be the same length"))
        end
        PriorityQueue(o, zip(ks,vs))
    end

end


text:
module ContinuedFractions
"""
Calculate the continued fraction until convergence against specified tolerance
∑ᵢ uᵢ/vᵢ = v(0) + u(1) / (v(1) + u(2) / (v(2) + u(3)/(v(3) + ...)))
"""

export continuedfraction

"""
    continuedfraction(v, u; atol=0, rtol=sqrt(eps(typeof(real(v(1) / u(1)))))
Calculate the continued fraction until convergence against specified tolerance
∑ᵢ uᵢ/vᵢ = v(0) + u(1) / (v(1) + u(2) / (v(2) + u(3)/(v(3) + ...)))
"""
function continuedfraction(v::V, u::U; atol=0,
    rtol=sqrt(eps(typeof(real(v(1) / u(1)))))
    ) where {V<:Function, U<:Function}
  n = 8
  a, b = continuedfraction(v, u, n, 2n)
  n *= 2
  converged = isapprox(a, b, rtol=rtol, atol=atol)
  while !converged && n <= 2^10
    a = deepcopy(b)
    n *= 2
    b = continuedfraction(v, u, n)
    converged = isapprox(a, b, rtol=rtol, atol=atol)
  end
  return b, converged
end
function continuedfraction(v::V, u::U, n::Int, m::Int
    ) where {V<:Function, U<:Function}
  @assert m > n
  output_m = u(m) / v(m)
  for i ∈ reverse(n+1:m-1)
    output_m = u(i) / (v(i) + output_m)
  end
  ui, vi = u(n), v(n)
  output_m = ui / (vi + output_m)
  output_n = ui / vi
  for i ∈ reverse(1:n-1)
    ui, vi = u(i), v(i)
    output_m = ui / (vi + output_m)
    output_n = ui / (vi + output_n)
  end
  vi = v(0)
  return vi + output_n, vi + output_m
end
function continuedfraction(v::V, u::U, n::Int) where {V<:Function, U<:Function}
  output = u(n) / v(n)
  for i ∈ reverse(1:n-1)
    output = u(i) / (v(i) + output)
  end
  return v(0) + output
end

end


text:
# Static satellite analysis
##
# 1. static_location_plot
# 2. multi_satellite_plot
# 3. multi_satellite_contour
# 4. wave analysis
# 5. check_wave_type
##
# Galileo G8 In/Out locations in GPhiO cooridinates:
# In : -1.525      -2.269       0.694
# Out: -1.212       2.771       0.785
#
# Eventually this should become a module!
#
# Hongyang Zhou, hyzhou@umich.edu 12/28/2019

using DelimitedFiles, Statistics, PyPlot, LinearAlgebra, Distributions, PyCall
gridspec = pyimport("matplotlib.gridspec")

include("utility.jl")
if !isdefined(Main, :Ganymede)
   include("Ganymede.jl")
   using .Ganymede
end

"Import satellite data. Require one line header presented."
function read_satellite_data(fname)
   f = readdlm(fname, ',', Float32, '\n'; header=true)

   header = f[2][1:end-1]
   data   = f[1]

   satelliteNo = unique(data[:,1]) # number of static satellites

   if minimum(satelliteNo) > 0.1
      @error "Some thing is wrong with the input data!"
   end

   header, data, satelliteNo
end

"""
    static_location_plot(filename, dir, nShift)

Plot all MHD quantities for a static location during the simulation.
# Arguments
- `filename::String`: Input file generated by ParaView.
- `dir::String`: file directory.
- `nShift::Integer`: shift w.r.t. the satellite starting location.
"""
function static_location_plot(filename="satellites_PIC.txt",
   dir="/Users/hyzhou/Documents/Computer/ParaView/data/", nShift=91)

   @assert nShift≥0 "nShift must be non-negative!"

   header, data, satelliteNo = read_satellite_data(dir*filename)

   index_ = findall(x->x==satelliteNo[1], data[:,1])

   id = getIndex(header)
   Rho_, Ux_, Uy_, Uz_, Bx_, By_, Bz_, Pe_, P_, U_, B_ =
   id.Rho_+1, id.Ux_+1, id.Uy_+1, id.Uz_+1, id.Bx_+1, id.By_+1, id.Bz_+1,
   id.Pe_+1, id.P_+1, id.U_.+1, id.B_.+1

   # magnetic pressure, [nPa]
   PB = sum(data[index_ .+ nShift,B_].^2; dims=2) ./ (2*μ₀) ./ (1e9)

   # dynamic pressure, [nPa]
   Pd = sum(data[index_ .+ nShift,Rho_] .* data[index_ .+ nShift,U_].^2; dims=2) * amu * 1e21

   # All quantities at one location
   figure(figsize=(12.5,6))
   plt.rc("font", family="serif", size=14)
   gs = gridspec.GridSpec(4, 1)
   gs.update(wspace=0.0, hspace=0.3) # set the spacing between axes.
   f1 = subplot(get(gs, (0, 0))) # rho

   plot(data[index_ .+ nShift,Rho_], label=L"\rho\ [amu/cc]")
   xlim(0,1200)
   plt.axhline(y=upstream_value[Rho_-1], color="k", linestyle="--", alpha=0.5)
   grid(true)
   f1.axes.xaxis.set_ticklabels([])
   legend(loc="lower left", bbox_to_anchor=(0.0, 0.8), ncol=4, frameon=false)

   f2 = subplot(get(gs, (1, 0))) # p
   plot(data[index_ .+ nShift,P_], alpha=0.9, label=L"P_i")
   plot(data[index_ .+ nShift,Pe_], alpha=0.9, label=L"P_e")
   plot(Pd, alpha=0.9, label=L"P_d")
   plot(PB, alpha=0.9, label=L"P_B\ [nPa]")
   xlim(0,1200)
   ylim(0,maximum(Pd[60:end])+0.02)
   legend(loc="lower left", bbox_to_anchor=(0.0, 0.8), ncol=4, frameon=false)
   plt.axhline(y=upstream_value[P_-1], color="k", linestyle="--", alpha=0.5)
   grid(true)
   f2.axes.xaxis.set_ticklabels([])

   f3 = subplot(get(gs, (2, 0))) # U
   plot(data[index_ .+ nShift,Ux_], alpha=0.9, label="Ux")
   plot(data[index_ .+ nShift,Uy_], alpha=0.9, label="Uy")
   plot(data[index_ .+ nShift,Uz_], alpha=0.7, label="Uz  [km/s]")
   xlim(0,1200)
   #ylim(-100,200)#ylim(-70,180)
   legend(loc="lower left", bbox_to_anchor=(0.0, 0.8), ncol=3, frameon=false)
   plt.axhline(y=upstream_value[Ux_-1], color="k", linestyle="--", alpha=0.5)
   plt.axhline(y=upstream_value[Uy_-1], color="k", linestyle="--", alpha=0.5)
   plt.axhline(y=upstream_value[Uz_-1], color="k", linestyle="--", alpha=0.5)
   grid(true)
   f3.axes.xaxis.set_ticklabels([])

   f4 = subplot(get(gs, (3, 0))) # B
   plot(data[index_ .+ nShift,Bx_], alpha=0.9, label="Bx")
   plot(data[index_ .+ nShift,By_], alpha=0.9, label="By")
   plot(data[index_ .+ nShift,Bz_], alpha=0.7, label="Bz  [nT]")
   xlim(0,1200)
   legend(loc="lower left", bbox_to_anchor=(0.0, 0.8), ncol=3, frameon=false)
   plt.axhline(y=upstream_value[Bx_-1], color="k", linestyle="--", alpha=0.5)
   plt.axhline(y=upstream_value[By_-1], color="k", linestyle="--", alpha=0.5)
   #plt.axhline(y=upstream_value[Bz_], color="k", linestyle="--", alpha=0.5)
   grid(true)
   xlabel("simulation time [s]")

   suptitle("Near magnetopause, location: "*
      string(data[nShift+1,11:end])[8:end]*L"R_G", fontsize=14)
   #plt.tight_layout(rect=[0, -0.02, 1, 0.98])
end

"""
    multi_satellite_plot(filename, dir)

One variable `var` plotted at 6 locations.
"""
function multi_satellite_plot(filename="satellites_PIC.txt",
   dir="/Users/hyzhou/Documents/Computer/ParaView/data/", var="Rho")

   header, data, satelliteNo = read_satellite_data(dir*filename)

   var_ = findfirst(x->x==var, header) + 1

   index_ = findall(x->x==satelliteNo[1], data[:,1])

   shift = [i*46 for i in 0:5]
   #shift = [i + 88 for i in 0:5]
   satelliteLoc = data[1 .+ shift,11:13]

   # one variable at 6 locations
   figure(figsize=(10,6))
   for iplot in 1:6
      plt.subplot(3,2,iplot)
      plot(data[index_ .+ shift[iplot],var_])
      plt.axhline(y=upstream_value[var_-1], color="r",linestyle="--", alpha=0.5)
      plt.title(string(satelliteLoc[iplot,:])[8:end])
      grid(true)
      f1 = plt.gca()
      if iplot ∉ [5,6]
         f1.axes.xaxis.set_ticklabels([])
         #f1.axes.yaxis.set_ticklabels([])
      end
   end

   suptitle("$var at 6 satellite locations", fontsize=14)
   plt.tight_layout()

end

"""
    multi_satellite_contour(filename, dir; DoSave=false,
       DoSubtractMean = true, nLead=10, nTrail=10)

Static satellite analysis, contour plots of all the variables as a function of
location and time.
`x` is the satellite position, `y` is the time, and `z` are values.
# Arguments
- `filename::String`: Input file generated by ParaView.
- `dir::String`: file directory.
- `plane::Char`: {'y','z'} cut plane.
- `DoSave::Bool`: Save the plots.
- `DoSubtractMean::Bool`: Subtract the average state for each variable.
- `nLead::Int`: the number of points removing from the start.
- `nTrail::Int`: the number of points removing from the end.

Issue: nLead must be ≥ 1, which should not be the case!
"""
function multi_satellite_contour(filename="satellites_PIC.txt",
   dir="/Users/hyzhou/Documents/Computer/ParaView/data/"; plane='y',
   DoSave=false, DoSubtractMean = true, nLead=10, nTrail=10)

   header, data, satelliteNo = read_satellite_data(dir*filename)

   # Remove the leading and trailing satellites
   satelliteNo = satelliteNo[1+nLead:end-nTrail]

   index_ = findall(x->x==0.0f0, data[:,1])

   if length(index_) ≤ 2
      @error "number of snapshots must be equal or larger than 2!"
   end

   c = Array{Float32, 2}(undef, length(index_), length(satelliteNo))

   crange = (-9.5,9.5)
   clength= 40
   vdisp = range(crange..., length=11)
   vplot = range(crange..., length=clength)

   # Subtract the average
   for var in header[1:end-3]
      var_ = findfirst(x->x==var, header) + 1
      @show var, var_

      for i in eachindex(satelliteNo)
         c[:,i] = data[index_ .+ Int(satelliteNo[i]),var_]
      end

      figure(figsize=(4,8))
      if DoSubtractMean
         cmean = mean(c, dims=1)
         if plane == 'y'
            contourf(data[Int.(satelliteNo),end], 1:length(index_), c .- cmean, vplot)
         elseif plane == 'z'
            contourf(data[Int.(satelliteNo),end-1], 1:length(index_), c .- cmean, vplot, extend="both")
         end
         #contourf(data[Int.(satelliteNo),end], 1:length(index_), c .- cmean, 50)
         #contourf(c .- cmean)
      else
         contourf(data[Int.(satelliteNo),end], 1:length(index_),c,50)
      end

      if plane == 'y'
         xlabel(L"z\ [R_G]")
      elseif plane == 'z'
         xlabel(L"y\ [R_G]")
      end
      ylabel("simulation time [s]")
      plt.set_cmap("seismic")
      #colorbar()
      colorbar(boundaries=vplot, ticks=vdisp)
      title(var)
      tight_layout()

      if DoSave
         if occursin("pic",lowercase(filename))
            subname = "pic"
         elseif occursin("hall",lowercase(filename))
            subname = "hall"
         end
         if DoSubtractMean
            savefig(var*"_traj_contour_subtractMean_"*subname*".png")
         else
            savefig(var*"_traj_contour_"*subname*".png")
         end
         close()
      end
   end

end

function satellite_p_contour(filename="satellites_y0_PIC.txt",
   dir="/Users/hyzhou/Documents/Computer/ParaView/data/"; plane='y',
   DoSubtractMean = true, nLead=10, nTrail=10, No=1)

   header, data, satelliteNo = read_satellite_data(dir*filename)

   # Remove the leading and trailing satellites
   satelliteNo = satelliteNo[1+nLead:end-nTrail]

   index_ = findall(x->x==0.0f0, data[:,1])

   p = Array{Float32, 2}(undef, length(index_), length(satelliteNo))
   pmean = similar(p)

   crange = (-9.5,9.5)
   clength= 40
   vdisp = range(crange..., length=11)
   vplot = range(crange..., length=clength)

   pe_ = 9  # index for electron pressure
   pi_ = 10 # index for ion pressure

   for i in eachindex(satelliteNo)
      p[:,i] = data[index_ .+ Int(satelliteNo[i]),pe_] .+
               data[index_ .+ Int(satelliteNo[i]),pi_]
      pmean[:,i] = sma(p[:,i],100) # box average over +-100s for each location
   end

   t = 1:length(index_)

   if plane == 'y'
      z = data[Int.(satelliteNo),end]
   elseif plane == 'z'
      y = data[Int.(satelliteNo),end-1]
   end

   DN = matplotlib.colors.DivergingNorm

   fig, ax = subplots(figsize=(3.5,7))
   plt.rc("font", family="serif", size=14)

   if DoSubtractMean
      if plane == 'y'
         z = data[Int.(satelliteNo),end]
         contourf(z, t, p .- pmean, norm=DN(0), vplot)
      elseif plane == 'z'
         y = data[Int.(satelliteNo),end-1]
         contourf(y, 1:length(index_), p .- pmean, norm=DN(0), vplot,
         extend="both")
      end

   else
      contourf(data[Int.(satelliteNo),end], t,c,50)
   end

   if plane == 'y'
      #zpickUp_ = [index for (index, value) in enumerate(z) if value > 0.2]
      #zUpMean = mean((p.-pmean)[:,zpickUp_], dims=2)

      #zpickDn_ = [index for (index, value) in enumerate(z) if value < -0.2]
      #zDnMean = mean((p.-pmean)[:,zpickDn_], dims=2)

      peakUp_index = Int[]
      peakDn_index = Int[]
      σUp = std(p[:,59]) # at z=0.5
      σDn = std(p[:,8]) # at z=-0.5
      @info σUp, σDn
      tGap = 9 # peaks must be differed by a time range to be picked
      pmean = mean(p, dims=1) # Averaged pressure at each location over time

      for i = 1:size(p,1)
         if p[i,59] - pmean[59] > 1.0σUp
         #if zUpMean[i] > 1.5σUp
            if isempty(peakUp_index)
               append!(peakUp_index, i)
            elseif i - peakUp_index[end] > tGap
               append!(peakUp_index, i)
            end
         end

         if p[i,8] - pmean[8] > 1.0σDn
         #if zDnMean[i] > 1.5σDn
            if isempty(peakDn_index)
               append!(peakDn_index, i)
            elseif i - peakDn_index[end] > tGap
               append!(peakDn_index, i)
            end
         end
      end

      @info "number of FTE = $(length(peakUp_index)+length(peakDn_index))"

      zUp = fill(data[Int.(satelliteNo),end][59], size(peakUp_index))
      zDn = fill(data[Int.(satelliteNo),end][8], size(peakDn_index))
      plot(zUp, peakUp_index, linestyle="", marker="P", markersize=6,
      markerfacecolor="gray", markeredgecolor="None")

      plot(zDn, peakDn_index, linestyle="", marker="P", markersize=6,
      markerfacecolor="k", markeredgecolor="None")

      xlabel(L"z\ [R_G]")
   elseif plane == 'z'
      xlabel(L"y\ [R_G]")
   end

   plt.set_cmap("seismic")
   # create an axes on the right side of ax. The width of cax will be 5%
   # of ax and the padding between cax and ax will be fixed at 0.05 inch.
   axes_grid1 = pyimport("mpl_toolkits.axes_grid1")
   divider = axes_grid1.make_axes_locatable(ax)
   cax = divider.append_axes("right", size="5%", pad=0.05)

   plt.colorbar(cax=cax, boundaries=vplot, ticks=vdisp)
   ax.set_title(L"\Delta P_t\ [nPa]")

   ax.annotate("($('a'+No-1))", xy=(-0.17, 1.0), xycoords="axes fraction")

   if occursin("pic",lowercase(filename))
      ax.set_ylabel("MHD-EPIC, simulation time [s]")
   else
      ax.set_ylabel("Hall MHD, simulation time [s]")
   end

   tight_layout()

   if plane == 'y'
      return peakUp_index, peakDn_index
   else
      return nothing
   end
end

"""
    satellite_p_contour_test(filename, dir; plane='y', nLead=10, nTrail=10,
       No=1, DoSubtractMean=true)

Test on different methods of obtaining the FTEs.
"""
function satellite_p_contour_test(filename="satellites_y0_PIC.txt",
   dir="/Users/hyzhou/Documents/Computer/ParaView/data/"; plane='y',
   DoSubtractMean=true, nLead=10, nTrail=10, No=1)

   header, data, satelliteNo = read_satellite_data(dir*filename)

   # Remove the leading and trailing satellites
   satelliteNo = satelliteNo[1+nLead:end-nTrail]
   # Starting index for the first satellite
   index_ = findall(x->x==0.0f0, data[:,1])
   # Pressure at each time for each satellite location
   p = Array{Float32, 2}(undef, length(index_), length(satelliteNo))
   #pmean = similar(p)

   crange = (-9.5,9.5)
   clength= 40
   vdisp = range(crange..., length=11)
   vplot = range(crange..., length=clength)

   # Subtract the average
   var_ = 10

   for i in eachindex(satelliteNo)
      p[:,i] = data[index_ .+ Int(satelliteNo[i]),var_]
      #pmean[:,i] = sma(p[:,i],100) # box average over +-100s for each location
   end
   pmean = mean(p, dims=1) # Averaged pressure at each location over time

   z = data[Int.(satelliteNo),end] # z coordinates of satellites

   zpickUp_ = [index for (index, value) in enumerate(z) if value > 0.2]
   zUpMean = mean((p.-pmean)[:,zpickUp_], dims=2)

   zpickDn_ = [index for (index, value) in enumerate(z) if value < -0.2]
   zDnMean = mean((p.-pmean)[:,zpickDn_], dims=2)

   peakUp_index = Int[]
   peakDn_index = Int[]
   σUp = std(p[:,59]) # may be changed
   σDn = std(p[:,8])  # may be changed
   tGap = 10 # peaks must be differed by a time range to be picked
   for i in eachindex(index_)
      #if c[i,59] - cmean[59] > 1.5σUp
      if zUpMean[i] > 1.5σUp
         if isempty(peakUp_index)
            append!(peakUp_index, i)
         elseif i - peakUp_index[end] > tGap
            append!(peakUp_index, i)
         end
      end

      #if c[i,8] - cmean[8] > 1.5σDn
      if zDnMean[i] > 1.5σDn
         if isempty(peakDn_index)
            append!(peakDn_index, i)
         elseif i - peakDn_index[end] > tGap
            append!(peakDn_index, i)
         end
      end
   end

   zMean = 0.5.*(zUpMean.+zDnMean)
   figure()
   plot(zMean)
   tight_layout()
   #[plt.axvline(x=i, color="black", linestyle="--") for i in peakUp_index]

   #figure()
   #plot(zDnMean)
   #[plt.axvline(x=i, color="black", linestyle="--") for i in peakDn_index]

   t = 1:length(index_)
   Δp = p[:,59].-pmean[59]

   fig, ax = subplots(figsize=(8,3))
   plot(t,Δp)
   fill_between(t, 0, 1, where=Δp .> 1.2σUp,
      color="green", alpha=0.5, transform=ax.get_xaxis_transform())
   title("ΔP at z=0.5")
   tight_layout()

   #return zMean
   return Δp
end

"""
    wave_analysis(nShift; DoPlot=false, filename, dir, iPlot=1, verbose=false)

Static satellite wave analysis.
`x` is the satellite position, `y` is the time, and `z` are values.
# Arguments
- `nShift::Integer`: shift w.r.t. the satellite starting location.
- `DoPlot::Bool`: Plot the perturbations.
- `filename::String`: Input file generated by ParaView.
- `dir::String`: file directory.
- `iPlot::Integer=1`: label tag index in alphabetical order starting from (a).
- `verbose::Bool=false`: Display information.
"""
function wave_analysis(nShift=115; DoPlot=false, filename="satellites_PIC.txt",
   dir="/Users/hyzhou/Documents/Computer/ParaView/data/", iPlot=1, verbose=true)

   header, data, satelliteNo = read_satellite_data(dir*filename)

   index_ = findall(x->x==satelliteNo[1], data[:,1])

   verbose && @info "satellite location: ",nShift, " ",data[nShift+1,11:end]

   id = getIndex(header)
   Rho_, Ux_, Uy_, Uz_, Bx_, By_, Bz_, Pe_, P_, U_, B_ =
   id.Rho_, id.Ux_, id.Uy_, id.Uz_, id.Bx_, id.By_, id.Bz_, id.Pe_, id.P_,
   id.U_, id.B_

   ρ = data[index_ .+ nShift,Rho_]
   P = data[index_ .+ nShift,P_]
   Pe= data[index_ .+ nShift,Pe_]
   Ux= data[index_ .+ nShift,Ux_]
   Uy= data[index_ .+ nShift,Uy_]
   Uz= data[index_ .+ nShift,Uz_]
   Bx= data[index_ .+ nShift,Bx_]
   By= data[index_ .+ nShift,By_]
   Bz= data[index_ .+ nShift,Bz_]
   B = @. √(Bx^2 + By^2 + Bz^2)

   # Subtract the background
   ρ̄ = sma(ρ)
   P̄ = sma(P)
   P̄e= sma(Pe)
   Ūx= sma(Ux)
   Ūy= sma(Uy)
   Ūz= sma(Uz)
   B̄x= sma(Bx)
   B̄y= sma(By)
   B̄z= sma(Bz)
   B̄ = sma(B)

   Δρ = ρ .- ρ̄
   ΔUx= Ux .- Ūx
   ΔUy= Uy .- Ūy
   ΔUz= Uz .- Ūz
   ΔBx= Bx .- B̄x
   ΔBy= By .- B̄y
   ΔBz= Bz .- B̄z
   ΔB = @. hypot(Bx, By, Bz) - hypot(B̄x, B̄y, B̄z)
   ΔP = @. P - P̄ + (Pe - P̄e)

   VA = @. B/√(μ₀*ρ*amu)*1e-15 # [km/s]

   start_ = 100 # remove the starting points in correlation analysis

   # Determine if ΔB and ΔU are correlated --> Alfven wave?
   rUBx = cor((ΔUx./VA)[start_:end],(ΔBx./B)[start_:end])
   rUBy = cor((ΔUy./VA)[start_:end],(ΔBy./B)[start_:end])
   rUBz = cor((ΔUz./VA)[start_:end],(ΔBz./B)[start_:end])

   rAlfven = sort([rUBx, rUBy, rUBz])[2] # Take the intermediate value

   if verbose
      @info "correlation between ΔUx and ΔBx = ", round(rUBx, digits=2)
      @info "correlation between ΔUy and ΔBy = ", round(rUBy, digits=2)
      @info "correlation between ΔUz and ΔBz = ", round(rUBz, digits=2)
   end

   n = length(ρ)
   α = 0.05 # acceptance quantile

   t = abs(rUBx) / √(1-rUBx^2)*√(n-start_-2)
   if t ≤ quantile(TDist(n-start_-2), 1-α)
      verbose && @info "ΔUx and ΔBx hypothesis test rejected!"*
         " No strong evidence of correlation!"
   end

   t = abs(rUBy) / √(1-rUBy^2)*√(n-start_-2)

   if t ≤ quantile(TDist(n-start_-2), 1-α)
      verbose && @info "ΔUy and ΔBy hypothesis test rejected!"*
         " No strong evidence of correlation!"
   end

   t = abs(rUBz) / √(1-rUBz^2)*√(n-start_-2)
   if t ≤ quantile(TDist(n-start_-2), 1-α)
      verbose && @info "ΔUz and ΔBz hypothesis test rejected!"*
         " No strong evidence of correlation!"
   end

   # Determine if ΔB and Δρ are correlated --> fast/slow wave?
   rUρ = cor(ΔB[start_:end], Δρ[start_:end]) # this one is not actually used

   verbose && @info "correlation between ΔB and Δρ = ", round(rUρ, digits=2)

   t = abs(rUρ) / √(1-rUρ^2)*√(n-start_-2)

   if t ≤ quantile(TDist(n-start_-2), 1-α)
      verbose && @info "ΔB and Δρ hypothesis test rejected!"*
         " No strong evidence of correlation!"
   end

   # thermal pressure perturbation vs. magnetic pressure perturbation
   ΔPB = @. (Bx^2 + By^2 + Bz^2 - B̄x^2 - B̄y^2 - B̄z^2)/(2*μ₀)*1e-9
   rPBPt = cor(ΔPB[start_:end], ΔP[start_:end])

   verbose && @info "correlation between ΔPB and ΔPt = ", round(rPBPt, digits=2)

   if rPBPt > 0.5
      verbose && @info "possibly fast magnetosonic wave"
   elseif rPBPt < -0.5
      verbose && @info "possibly slow magnetosonic wave"
   end

   if DoPlot
      figure(figsize=(10.0,2.0))
      ax = subplot(111)
      plt.rc("font", family="serif", size=12)
      plot(ΔP, label=L"\Delta P_t")
      plot(ΔPB, label=L"\Delta P_B")
      xlim(100,1200)
      #ylim(-0.16,0.16)
      ylim(-3.5,3.5)
      legend(loc=(0.02,-0.05), ncol=2, frameon=false)
      xlabel("simulation time [s]")
      plt.axhline(y=0.0, color="k", linestyle="--", alpha=0.5)
      plt.annotate("cc=$(round(rPBPt,digits=2))", xy=(0.9, 0.05),
         xycoords="axes fraction")
      if occursin("PIC", filename)
         titletext = "MHD-EPIC, Location: "*
            string(data[nShift+1,11:end])[8:end]*L"R_G"
      else
         titletext = "Hall MHD, Location: "*
            string(data[nShift+1,11:end])[8:end]*L"R_G"
      end
      plt.title(titletext)
      ax.annotate("($('a'+iPlot-1))", xy=(0.0, 1.1), xycoords="axes fraction")
      ax.minorticks_on()
      ax.tick_params(which="both",top=true, right=true)
      ax.tick_params(which="both", direction="in")
      tight_layout()
   end

   # Energy perturbations
   Ek = @. (((ΔUx*Ux+ΔUy*Uy+ΔUz*Uz)+0.5*(ΔUx^2 + ΔUy^2 + ΔUz^2))*ρ̄ +
      Δρ*(Ūx^2 + Ūy^2 + Ūz^2))*amu*1e12
   Et = @. ΔP*1e-9/(γ-1)
   Eb = @. ((ΔBx*Bx+ΔBy*By+ΔBz*Bz)+0.5*(ΔBx^2 + ΔBy^2 + ΔBz^2))/μ₀*1e-18

   verbose && @info "correlation between Ek and Et = ",
      round(cor(Ek[start_:end], Et[start_:end]), digits=2)
   verbose && @info "correlation between Ek and Eb = ",
      round(cor(Ek[start_:end], Eb[start_:end]), digits=2)
   verbose && @info "correlation between Et and Eb = ",
      round(cor(Et[start_:end], Eb[start_:end]), digits=2)

   #println("Average kinetic energy = ",mean(Ek))
   #println("Average magnetic energy = ",mean(Eb))

   if DoPlot
      figure(figsize=(12.0,3.5))
      plt.rc("font", family="serif", size=12)
      gs = gridspec.GridSpec(3, 1)
      gs.update(wspace=0.0, hspace=0.0) # set the spacing between axes.
      f1 = subplot(get(gs, (0, 0)))
      plot(ΔUx ./ VA, label=L"\Delta U_x / V_A")
      plot(ΔBx ./ B, label=L"\Delta B_x / B_0")
      plt.axhline(y=0.0, color="k", linestyle="--", alpha=0.5)
      xlim(100,1200)
      #legend(loc=(0.36,-0.06), ncol=2, frameon=false)
      legend(loc=(0.00,0.6), ncol=2, frameon=false)
      f1.axes.xaxis.set_ticklabels([])
      f1.minorticks_on()
      f1.tick_params(which="both",top=true, right=true)
      f1.tick_params(which="both", direction="in")
      plt.annotate("cc=$(round(rUBx,digits=2))", xy=(0.885, 0.05),
         xycoords="axes fraction")

      f2 = subplot(get(gs, (1, 0)))
      plot(ΔUy ./ VA, label=L"\Delta U_y / V_A")
      plot(ΔBy ./ B, label=L"\Delta B_y / B_0")
      plt.axhline(y=0.0, color="k", linestyle="--", alpha=0.5)
      xlim(100,1200)
      #legend(loc=(0.36,-0.06), ncol=2, frameon=false)
      legend(loc=(0.00,0.6), ncol=2, frameon=false)
      f2.axes.xaxis.set_ticklabels([])
      f2.minorticks_on()
      f2.tick_params(which="both",top=true, right=true)
      f2.tick_params(which="both", direction="in")
      plt.annotate("cc=$(round(rUBy,digits=2))", xy=(0.885, 0.05),
         xycoords="axes fraction")

      f3 = subplot(get(gs, (2, 0)))
      plot(ΔUz ./ VA, label=L"\Delta U_z / V_A")
      plot(ΔBz ./ B, label=L"\Delta B_z / B_0")
      plt.axhline(y=0.0, color="k", linestyle="--", alpha=0.5)
      xlim(100,1200)
      #legend(loc=(0.36,-0.06), ncol=2, frameon=false)
      legend(loc=(0.00,0.06), ncol=2, frameon=false)

      #xlabel("simulation time [s]", x=0.5, y=0.5)
      plt.annotate("simulation time [s]", xy=(0.43, -0.36),
         xycoords="axes fraction")
      f3.minorticks_on()
      f3.tick_params(which="both",top=true, right=true)
      f3.tick_params(which="both", direction="in")
      plt.annotate("cc=$(round(rUBz,digits=2))", xy=(0.885, 0.05),
         xycoords="axes fraction")

      if occursin("PIC", filename)
         titletext = "MHD-EPIC, Location: "*
            string(data[nShift+1,11:end])[8:end]*L"R_G"
      else
         titletext = "Hall MHD, Location: "*
            string(data[nShift+1,11:end])[8:end]*L"R_G"
      end
      suptitle(titletext, x=0.5, y=0.94)
      f1.annotate("($('a'+iPlot-1))", xy=(0.0, 1.1), xycoords="axes fraction")
   end

   if DoPlot
      figure(figsize=(12.0,2.5))
      plt.rc("font", family="serif", size=12)
      plot(ΔB, alpha=0.95, label=L"\Delta B")
      plot(Δρ, alpha=0.95, label=L"\Delta \rho")
      xlim(0,1200)
      #ylim(-3,3)
      legend(loc="upper left", ncol=2, frameon=false)
      xlabel("simulation time [s]")
      plt.title("Location: "*string(data[nShift+1,11:end])[8:end]*L"R_G")
      tight_layout()
   end

   return rAlfven, rPBPt
end

function check_wave_type(filename="satellites_PIC.txt",
   dir="/Users/hyzhou/Documents/Computer/ParaView/data/")

   f = readdlm("/Users/hyzhou/Documents/Computer/ParaView/data/G8Traj.csv",
      ',', Float32, '\n'; header=true)
   nS = size(f[1])[1]

   # X, Y, Z, rAlfven, rFastSlow
   fnew = Array{Float32,2}(undef, nS, 5)

   header, data, satelliteNo = read_satellite_data(dir*filename)

   if nS != length(satelliteNo)
      error("Number of satellites not right! Check the input files!")
   end

   for iShift = 1:nS
      r1,r2 = wave_analysis(iShift-1; DoPlot=false, filename=filename,
         dir=dir, verbose=false)
      fnew[iShift,1:3] = f[1][iShift,:]
      fnew[iShift,3] = 0.75 # For plotting purpose
      fnew[iShift,4:5] .= r1, r2
      println("iShift = ", iShift)
   end

   # Write to file
   open(dir*"waveAlongTrajG8.csv", "w") do io
      write(io, "\"X\",\"Y\",\"Z\",\"rA\",\"rM\"\n")
      writedlm(io, fnew, ',')
   end

   return fnew
end

## Execution
#static_location_plot("satellites_PIC.txt",
#   "/Users/hyzhou/Documents/Computer/ParaView/data/", 185)
#multi_satellite_plot()
#multi_satellite_contour("satellites_y0_PIC.txt", DoSubtractMean=true)
#multi_satellite_contour("satellites_boundary_PIC.txt", plane='z', DoSubtractMean=true)

peakUp, peakDn = satellite_p_contour("satellites_y0_PIC.txt"; No=2, plane='y')
#satellite_p_contour("satellites_boundary_Hall.txt"; No=3, plane='z')

#pMeanHall = satellite_p_contour_test("satellites_y0_Hall.txt"; No=1)
#pMeanPIC = satellite_p_contour_test("satellites_y0_PIC.txt"; No=1)

#nShift = 130
#static_location_plot("satellites_Hall.txt",
#   "/Users/hyzhou/Documents/Computer/ParaView/data/", nShift)
#wave_analysis(nShift; DoPlot=true, filename="satellites_Hall.txt", iPlot=1,
#   verbose=true)
#fnew = check_wave_type()
#check_wave_type()

#=
# Cross correlation check
using DSP

figure()
subplot(411)
plot(xcorr(ΔUx, ΔBx))
subplot(412)
plot(xcorr(ΔUy, ΔBy))
subplot(413)
plot(xcorr(ΔUz, ΔBz))
subplot(414)
plot(xcorr(ΔB,Δρ))
=#

#=
using StatsBase
figure()
subplot(411)
plot(crosscor(ΔUx, ΔBx))
subplot(412)
plot(crosscor(ΔUy, ΔBy))
subplot(413)
plot(crosscor(ΔUz, ΔBz))
subplot(414)
plot(crosscor(ΔB,Δρ))
=#


text:
function create_buffer_resource(allocate_buffer, device::Device, vdata)
    resource = allocate_buffer(device, buffer_size(vdata))
    upload_data(resource.memory, vdata)
    resource
end

function allocate_buffer(device::Device, size, usage, memory_properties; sharing_mode = SHARING_MODE_EXCLUSIVE, queue_indices = [0])
    vbuffer = Buffer(device, size, usage, sharing_mode, queue_indices)
    vmemory = DeviceMemory(vbuffer, memory_properties)
    GPUResource(vbuffer, vmemory)
end

allocate_vertex_buffer(device::Device, size) = allocate_buffer(device, size, BUFFER_USAGE_VERTEX_BUFFER_BIT, MEMORY_PROPERTY_HOST_VISIBLE_BIT | MEMORY_PROPERTY_HOST_COHERENT_BIT)
allocate_index_buffer(device::Device, size) = allocate_buffer(device, size, BUFFER_USAGE_INDEX_BUFFER_BIT, MEMORY_PROPERTY_HOST_VISIBLE_BIT | MEMORY_PROPERTY_HOST_COHERENT_BIT)

vertex_data_type(w::Widget) = vertex_data_type(typeof(w))
mesh_encoding_type(w::Widget) = mesh_encoding_type(typeof(w))
nvertices(T::Type{<:Widget}) = not_implemented_for(T)
nvertices(w::Widget) = nvertices(typeof(w))
resource_types(w::Widget) = resource_types(typeof(w))
push_data(w::Widget) = nothing

indices(encoding::IndexEncoding) = convert(Vector{UInt32}, collect(encoding.indices))

function indices(encoding::IndexList)
    class = GeometryExperiments.topology_class(typeof(encoding))
    n = GeometryExperiments.parametric_dimension(class)
    convert(Vector{Point{n,UInt32}}, encoding.indices)
end

function Rhyolite.ShaderDependencies(w::Widget)
    mesh = MeshVertexEncoding(w)
    ShaderDependencies(mesh.vertex_data, indices(mesh.encoding), [])
end

include("widgets/image.jl")
include("widgets/box.jl")
include("widgets/text.jl")


text:
#=
The tests here merely check that the interface functions work.
Testing of the implementation is done in ldlt.jl.
=#

using Test
using Random
using SparseArrays
using LinearAlgebra
using SparseFactorizations

@testset "ldlt API" begin
    @testset "ldlt_structure" begin
        Random.seed!(1)
        A = sprand(10,10,0.1)
        A = I + A + A'
        Fref = lu(Matrix(A),Val(false))
        P = ldlt_structure(A)
        @test P == (Fref.L .!= 0)
    end

    @testset "level_of_fill" begin
        Random.seed!(1)
        A = sprand(10,10,0.1)
        A = I + A + A'
        Fref = lu(Matrix(A),Val(false))
        P = level_of_fill(A)
        nonzeros(P) .= 1
        @test P == (Fref.L .!= 0)
    end

    @testset "simplicial_ldlt" begin
        @testset for (conj,adj) in ((conj,adjoint),(identity,transpose))
            Random.seed!(1)
            A = sprand(ComplexF64,10,10,0.1)
            A = I + A + adj(A)
            Fref = lu(Matrix(A),Val(false))
            F = simplicial_ldlt(A; conj=conj)

            @test F.L ≈ Fref.L
            @test F.d ≈ diag(Fref.U)
            @test F.Lt == adj(F.L)
            @test F.D == Diagonal(F.d)
        end
    end

    @testset "ildlt" begin
        @testset for (conj,adj) in ((conj,adjoint),(identity,transpose))
            Random.seed!(1)
            A = sprand(ComplexF64,10,10,0.1)
            A = I + A + adj(A)
            Fref = lu(Matrix(A),Val(false))
            P = ldlt_structure(A)
            F = ildlt(A,P; conj=conj)

            @test F.L ≈ Fref.L
            @test F.d ≈ diag(Fref.U)
            @test F.Lt == adj(F.L)
            @test F.D == Diagonal(F.d)
        end
    end

    @testset "ildlt_lof" begin
        @testset for (conj,adj) in ((conj,adjoint),(identity,transpose))
            Random.seed!(1)
            A = sprand(ComplexF64,10,10,0.1)
            A = I + A + adj(A)
            Fref = lu(Matrix(A),Val(false))
            F = ildlt_lof(A,size(A,2); conj=conj)

            @test F.L ≈ Fref.L
            @test F.d ≈ diag(Fref.U)
            @test F.Lt == adj(F.L)
            @test F.D == Diagonal(F.d)
        end
    end

    @testset "ildlt_tol" begin
        @testset for (conj,adj) in ((conj,adjoint),(identity,transpose))
            Random.seed!(1)
            A = sprand(ComplexF64,10,10,0.1)
            A = I + A + adj(A)
            Fref = lu(Matrix(A),Val(false))
            F = ildlt_tol(A,0.0; conj=conj)

            @test F.L ≈ Fref.L
            @test F.d ≈ diag(Fref.U)
            @test F.Lt == adj(F.L)
            @test F.D == Diagonal(F.d)
        end
    end
end

max_stars_repo_name: ettersi/SparseFactorizations

max_issues_repo_name: ettersi/SparseFactorizations

max_forks_repo_name: ettersi/SparseFactorizations


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule CUTENSOR_CUDA111_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("CUTENSOR_CUDA111")
JLLWrappers.@generate_main_file("CUTENSOR_CUDA111", UUID("0e40f886-cfff-5d56-8a11-88482f468a43"))
end  # module CUTENSOR_CUDA111_jll


text:
"variable: `0 <= active_bus[l] <= 1` for `l` in `bus`es"
function variable_bus_active_indicator(pm::_PM.AbstractPowerModel; nw::Int=nw_id_default, relax = false,  report::Bool=true)
    if relax == false
        z_bus = _PM.var(pm, nw)[:z_bus] = JuMP.@variable(pm.model,
            [l in _PM.ids(pm, nw, :bus)],
            base_name="$(nw)_active_bus",
            binary = true,
            start = _PM.comp_start_value(_PM.ref(pm, nw, :bus, l), "bus_active_start")
        )
    else
        z_bus = _PM.var(pm, nw)[:z_bus] = JuMP.@variable(pm.model,
            [l in _PM.ids(pm, nw, :bus)],
            base_name="$(nw)_active_bus",
            lower_bound = 0,
            upper_bound = 1,
            start = _PM.comp_start_value(_PM.ref(pm, nw, :bus, l), "bus_active_start")
        )
    end

    report && _IM.sol_component_value(pm, _PM.pm_it_sym, nw, :bus, :status, _PM.ids(pm, nw, :bus), z_bus)
end


text:
#
#--------------------------------------#
export DSS
#--------------------------------------#
"""
Q*Q'*u where Q: local -> global operator
"""
function DSS(u,l2g,g2l)

    Qu   = NNlib.scatter(+,u,l2g)
    QQtu = NNlib.gather(Qu,g2l)

    return v
end


text:
module Fuji

export
FujiRequest,
FujiServer,
after,
before,
delete,
get,
patch,
post,
put,
render_template,
route,
start,
unroute

include("log.jl")
include("route.jl")
include("request.jl")
include("server.jl")
include("template.jl")

server = FujiServer(Nullable{Function}(), Nullable{Function}(), Array{Route,1}())

end


text:
"""
See [1, "Online Implementation of IRL: A Hybrid Optimal Adaptive Controller"].

# Refs
[1] “Reinforcement Learning and Feedback Control: Using Natural Decision Methods to Design Optimal Adaptive Controllers,” IEEE Control Syst., vol. 32, no. 6, pp. 76–105, Dec. 2012, doi: 10.1109/MCS.2012.2214134.

# Notes
- T: Data stack period
- N: The maximum length of stacked data
"""
mutable struct LinearIRL <: AbstractIRL
    Q::AbstractMatrix
    R_inv::AbstractMatrix
    B::AbstractMatrix
    T::Real
    N::Int
    i::Int
    i_init::Int
    terminated::Bool
    function LinearIRL(Q::AbstractMatrix, R::AbstractMatrix, B::AbstractMatrix;
            T=0.04,
            N=nothing,
            i_init=0,  # iteration number
        )
        n1, n2 = size(Q)
        @assert n1 == n2
        N_min = Int(n1*(n1 + 1)/2)
        if N == nothing
            N = N_min
        else
            @assert N >= N_min
        end
        @assert T > 0 && N > 0
        n1, n2 = size(Q)
        @assert n1 == n2
        N_min = Int(n1*(n1 + 1)/2)
        if N == nothing
            N = N_min
        else
            @assert N >= N_min
        end
        @assert T > 0 && N > 0
        R_inv = inv(R)
        i = i_init
        terminated = false
        new(Q, R_inv, B, T, N, i, i_init, terminated)
    end
end

function reset!(irl::LinearIRL)
    irl.i = irl.i_init
end

"""
Value iteration [1, Eq. 99]; updated in least-square sense
# Notes
w: critic parameter (vectorised)
- ϕs_prev: the vector of bases (evaluated)
- V̂: the vector of approximate values (evaluated)
"""
function value_iteration!(irl::LinearIRL, buffer::DataBuffer, w;
        sc::AbstractStopCondition=DistanceStopCondition(),
    )
    @unpack i, N = irl
    @unpack data_array = buffer
    data_filtered = filter(x -> x.i == i, data_array)  # data from the current policy
    if length(data_filtered) >= N + 1
        data_sorted = sort(data_filtered, by=x -> x.t)  # sort by time index
        ϕs_prev = data_sorted[end-N:end-1] |> Map(datum -> datum.ϕ) |> collect
        # V̂s = data_sorted[end-(N-1):end] |> Map(datum -> datum.V̂) |> collect
        xs = data_sorted[end-(N-1):end] |> Map(datum -> datum.x) |> collect
        ∫rs = data_sorted[end-(N-1):end] |> Map(datum -> datum.∫r) |> collect
        P = convert_to_matrix(w)
        V̂s_with_prev_P = xs |> Map(x -> value(irl, P, x)) |> collect
        V̂s = ∫rs .+ V̂s_with_prev_P
        # update the critic parameter
        w_new = ( hcat(V̂s...) * pinv(hcat(ϕs_prev...)) )'[:]  # to reduce the computation time; [:] for vectorisation
        if !irl.terminated
            if is_terminated(sc, w, w_new)
                irl.terminated = true
            else
                w .= w_new
                update_index!(irl)
            end
        end
    end
end

"""
Policy iteration [1, Eq. 98]; updated in least-square sense
# Notes
w: critic parameter (vectorised)
- Δϕs: the vector of (ϕ - ϕ_prev) (evaluated)
- ∫rs: the vector of integral running cost by numerical integration (evaluated)
"""
function policy_iteration!(irl::LinearIRL, buffer::DataBuffer, w;
        sc::AbstractStopCondition=DistanceStopCondition(),
    )
    @unpack i, N = irl
    @unpack data_array = buffer
    data_filtered = filter(x -> x.i == i, data_array)  # data from the current policy
    if length(data_filtered) >= N + 1
        data_sorted = sort(data_filtered, by=x -> x.t)  # sort by time index
        ϕs_prev_and_present = data_sorted[end-N:end] |> Map(datum -> datum.ϕ) |> collect
        Δϕs = diff(ϕs_prev_and_present)
        ∫rs = data_sorted[end-(N-1):end] |> Map(datum -> datum.∫r) |> collect
        # update the critic parameter
        w_new = ( hcat(∫rs...) * pinv(hcat(-Δϕs...)) )'[:]  # to reduce the computation time; [:] for vectorisation
        if !irl.terminated
            if is_terminated(sc, w, w_new)
                irl.terminated = true
            else
                w .= w_new
                update_index!(irl)
            end
        end
    end
end

function update_index!(irl::LinearIRL)
    irl.i += 1
end

"""
Policy improvement [1, Eq. 96].
"""
function _optimal_input(R_inv, B, P, x)
    -0.5 * R_inv * B' * (2 * P * x)
end

function optimal_input(irl::LinearIRL, x, w::AbstractVector)
    @unpack R_inv, B = irl
    P = convert_to_matrix(w)
    _optimal_input(R_inv, B, P, x)
end

function optimal_input(irl::LinearIRL, x, P::AbstractMatrix)
    @unpack R_inv, B = irl
    _optimal_input(R_inv, B, P, x)
end

function value(irl::LinearIRL, P::AbstractMatrix, x)
    x' * P * x
end

function value(irl::LinearIRL, w::AbstractVector, x)
    P = convert_to_matrix(w)
    value(irl, P, x)
end

"""
w: critic parameter (vectorised)
"""
function Base.push!(buffer::DataBuffer, irl::LinearIRL, cost::AbstractCost;
        t, x, u, w,
    )
    # P = convert_to_matrix(w)
    @unpack data_array = buffer
    @unpack i = irl
    data_sorted = sort(data_array, by = x -> x.t)  # sorted by t
    # V̂_with_prev_P = value(irl, P, x)
    # prev data
    if length(data_sorted) != 0
        t_prev = data_sorted[end].t
        x_prev = data_sorted[end].x
        u_prev = data_sorted[end].u
        # numerical integration
        Δt = t - t_prev
        r = cost(x, u)
        r_prev = cost(x_prev, u_prev)
        ∫r = 0.5 * (r + r_prev) * Δt # trapezoidal
        # V̂ = ∫r + V̂_with_prev_P
    else
        # V̂ = missing
        ∫r = missing
    end
    ϕ = convert_quadratic_to_linear_basis(x)  # x'Px = w'ϕ(x)
    datum = (;
             t=t,
             x=x,
             u=u,
             w=w,  # logging
             ϕ=ϕ,
             ∫r=∫r,
             # V̂=V̂,
             i=i,  # iteration number
            )
    push!(buffer.data_array, datum)
end


text:
module GmshSDK

using Artifacts
using StaticArrays
using RecipesBase
using OrderedCollections
using Printf

using WavePropBase
using WavePropBase.Geometry
using WavePropBase.Interpolation
using WavePropBase.Mesh

const GMSH_VERSION   = "4.8.4"

const GMSH_PATH = artifact"gmsh4.8.4"

const dirs      = readdir(GMSH_PATH,join=true)

@assert length(dirs)==1 "there should be only one directory for the untared gmsh file. Got $dirs"

include(joinpath(first(dirs),"lib","gmsh.jl"))

include("gmshIO.jl")

export
    gmsh,
    # re-export useful modules from WavePropBase
    Geometry,
    Mesh,
    # macros
    @gmsh
    # methods

end # module

max_stars_repo_name: WaveProp/GmshSDK

max_issues_repo_name: WaveProp/GmshSDK

max_forks_repo_name: WaveProp/GmshSDK


text:
using seisimu

vp = [5000.0]
vs = [3500.0]
rho = 2.8
depth = 6000
Horizon = 6000
pkf = 30
T = 2
ext = 10
iflag = 2
dx = nothing
dz = nothing
dt = nothing
nT = nothing

model = initmodel(vp,rho,depth,Horizon,pkf,T)

sn = 1
loc = [2500 2500]
ot = [0.001]
mt = [1 5 1]
wavelet = Ricker(pkf,model.medium.dt)

sou = initsource(sn,loc,ot,mt,wavelet,model)

@time runsimu(model,sou; showevery = 20000)


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

using Markdown
using InteractiveUtils

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

# ╔═╡ 4e9478cb-11d3-4d00-8456-8e1e1e20178e
begin
	using Plots
	using PlutoUI
	using StaticArrays
	using StatsBase

	gr()

	src_dir = joinpath(@__DIR__, "../../build")
	work_dir = @__DIR__
end

# ╔═╡ 1ef4e7be-b0ca-4d37-9b87-416495802473
struct VertexData
    cell_area :: Float64

    coord :: SVector{3, Float64}
    species_count :: SVector{3, Float64}
end

# ╔═╡ 78e4be62-f030-4e3b-8e5c-a1bff309e735
struct FrameData
    time :: Float64
    vertices :: Vector{VertexData}
end

# ╔═╡ 00f87b79-d447-4732-9ebf-4378dd804905
function read_vertices(filepath::AbstractString)
    frames = Vector{FrameData}(undef, 0)
    open(filepath) do f
        for line ∈ eachline(f)
            splitted = split(line)
            if splitted[1] == "FRAME"
                push!(frames, FrameData(parse(Float64, splitted[2]), Vector{VertexData}(undef, 0)))
            else
                vertices = frames[end].vertices
                push!(
                    vertices,
                    VertexData(
                        parse(Float64, splitted[2]),
                        SVector{3}([parse(Float64, splitted[3]), parse(Float64, splitted[4]), parse(Float64, splitted[5])]),
                        SVector{3}([parse(Int, splitted[6]), parse(Int, splitted[7]), parse(Int, splitted[8])])
                    )
                )
            end
        end
    end
    frames
end

# ╔═╡ 18500426-5217-4784-b6e3-9e04030270bc
function read_triangles(filepath::AbstractString)
	open(filepath) do f
		num_triangles = parse(Int, readline(f))
		# The file uses 0-based indexing, and must be converted to 1-based indexing here.
		inds = map(x->parse(Int, x) + 1, f |> readline |> split)
		vertex_indices = reshape(inds, 3, num_triangles)
	end
end

# ╔═╡ 1b83f02a-dc6a-4124-9b7f-eaaef5edb95d
begin
	frames = read_vertices(joinpath(src_dir, "side-surface-mesh-diffusion.txt"))
	triangle_vertex_indices = read_triangles(joinpath(src_dir, "side-surface-mesh-diffusion-triangle.txt"))

    max_copy = 10000
    β = 1 / 4.1

    total_area = 480 * 480
    areas = [v.cell_area for v ∈ frames[1].vertices]
    xs = [v.coord[1] for v ∈ frames[1].vertices]
    ys = [v.coord[2] for v ∈ frames[1].vertices]

	energy_at_coord(coord) = -5 * sin(2π / 1000 * coord[1])

    # calculate expected copy numbers
    expected_uniform = [max_copy * v.cell_area / total_area for v ∈ frames[1].vertices]
    expected_sin = [v.cell_area * exp(-β * energy_at_coord(v.coord)) for v ∈ frames[1].vertices]
    expected_sin *= max_copy / sum(expected_sin)

	expected_conc = [expected_uniform, expected_uniform, expected_sin]
end

# ╔═╡ f7e0efbd-2370-48af-bb4a-929d7fcb9d5d
md"
Species index $(@bind species_index Slider(1:3; show_value=true))
"

# ╔═╡ 2e4b60fd-e59c-42e0-b68e-9da50f73c09c

function make_avg_plot()
	frame_range = 10:length(frames)
	conc_avg = zeros(length(frames[1].vertices))

	for frame_index ∈ frame_range
		conc_avg += [v.species_count[species_index] / v.cell_area for v ∈ frames[frame_index].vertices]
	end
	conc_avg /= length(frame_range)

	# avg copy numbers
	conc_plot = scatter3d(
		xs, ys, conc_avg;
		label = "average frame $frame_range",
		xlabel = "x (nm)",
		ylabel = "y (nm)",
		zlabel = "concentration (nm^-2)",
	)

	scatter3d!(conc_plot, xs, ys, expected_conc[species_index] ./ areas; label = "expected")
	conc_plot
end

# ╔═╡ 5b4da84b-f765-4e49-9251-0c2426985ff0
begin
	avg_plot = make_avg_plot()
	savefig(avg_plot, joinpath(work_dir, "surface-mesh-diffusion-avg.png"))
	avg_plot
end

# ╔═╡ 61623a45-1616-44d0-b468-0658a13fb49c
begin
	anim = @animate for (fi, frame) ∈ enumerate(frames[1:20])
		conc_plot = scatter3d(
			xs, ys, [v.species_count[species_index] / v.cell_area for v ∈ frame.vertices];
			label = "frame $fi",
			xlabel = "x (nm)",
			ylabel = "y (nm)",
			zlabel = "concentration (nm^-2)",
		)

		scatter3d!(conc_plot, xs, ys, expected_conc[species_index] ./ areas; label = "expected")
	end
	gif(anim, joinpath(work_dir, "surface-mesh-diffusion-anim.gif"); fps=4)
end

# ╔═╡ 5e0f3cd3-b0b6-46c5-a812-f1edc0c56e9e
begin
	# Make mesh graph with energy
	p = plot(;
		xlabel = "x (nm)",
		ylabel = "y (nm)",
		xlim = (0, 500),
		ylim = (0, 500),
		aspect_ratio = :equal,
	)

	# Energy part
	heat_xrange = 0:5:500
	heat_yrange = 0:100:500
	heatmap!(p,
		heat_xrange, heat_yrange,
		(x, y) -> energy_at_coord(@SVector [x, y, 0]);
		color = :copper,
	)

	# Mesh part
	for ti ∈ 1:size(triangle_vertex_indices)[2]
		vs = triangle_vertex_indices[[1,2,3,1], ti]
		xs = map(vi -> frames[1].vertices[vi].coord[1], vs)
		ys = map(vi -> frames[1].vertices[vi].coord[2], vs)
		plot!(p, xs, ys; legend=:none, color=:white, colorbar_title="energy (pN nm)")
	end

	savefig(p, joinpath(work_dir, "surface-mesh-diffusion-setup.png"))
	p
end

# ╔═╡ Cell order:
# ╠═4e9478cb-11d3-4d00-8456-8e1e1e20178e
# ╠═1ef4e7be-b0ca-4d37-9b87-416495802473
# ╠═78e4be62-f030-4e3b-8e5c-a1bff309e735
# ╟─00f87b79-d447-4732-9ebf-4378dd804905
# ╟─18500426-5217-4784-b6e3-9e04030270bc
# ╠═1b83f02a-dc6a-4124-9b7f-eaaef5edb95d
# ╟─f7e0efbd-2370-48af-bb4a-929d7fcb9d5d
# ╟─2e4b60fd-e59c-42e0-b68e-9da50f73c09c
# ╠═5b4da84b-f765-4e49-9251-0c2426985ff0
# ╠═61623a45-1616-44d0-b468-0658a13fb49c
# ╠═5e0f3cd3-b0b6-46c5-a812-f1edc0c56e9e

max_stars_repo_name: allen-cell-animated/medyan

max_issues_repo_name: allen-cell-animated/medyan

max_forks_repo_name: allen-cell-animated/medyan