text:
struct CPUEvent <: Event
    task::Core.Task
end
isdone(ev::CPUEvent) = Base.istaskdone(ev.task)
failed(ev::CPUEvent) = Base.istaskfailed(ev.task)

function Event(::CPU)
    return NoneEvent()
end

"""
    Event(f, args...; dependencies, progress, sticky=true)

Run function `f` with `args` in a Julia task. If `sticky` is `true` the task
is run on the thread that launched it.
"""
function Event(f, args...; dependencies=nothing, progress=nothing, sticky=true)
    T = Task() do
        wait(MultiEvent(dependencies), progress)
        f(args...)
    end
    T.sticky = sticky
    Base.schedule(T)
    return CPUEvent(T)
end

wait(ev::Union{CPUEvent, NoneEvent, MultiEvent}, progress=nothing) = wait(CPU(), ev, progress)
wait(::CPU, ev::NoneEvent, progress=nothing) = nothing

function wait(cpu::CPU, ev::MultiEvent, progress=nothing)
    events = ev.events
    N = length(events)
    alldone = ntuple(i->false, N)

    while !all(alldone)
        alldone = ntuple(N) do i
            if alldone[i]
                true
            else
                isdone(events[i])
            end
        end
        if progress === nothing
            yield()
        else
            progress()
        end
    end

    if any(failed, events)
        ex = CompositeException()
        for event in events
            if failed(event) && event isa CPUEvent
                push!(ex, Base.TaskFailedException(event.task))
            end
        end
        throw(ex)
    end
end

function wait(::CPU, ev::CPUEvent, progress=nothing)
    if progress === nothing
        wait(ev.task)
    else
        while !Base.istaskdone(ev.task)
            progress()
        end
        if Base.istaskfailed(ev.task)
            throw(Base.TaskFailedException(ev.task))
        end
    end
end

function async_copy!(::CPU, A, B; dependencies=nothing, progress=nothing)
    Event(copyto!, A, B, dependencies=dependencies, progress=progress)
end

function (obj::Kernel{CPU})(args...; ndrange=nothing, workgroupsize=nothing, dependencies=nothing, progress=nothing)
    if ndrange isa Integer
        ndrange = (ndrange,)
    end
    if workgroupsize isa Integer
        workgroupsize = (workgroupsize, )
    end
    if dependencies isa Event
        dependencies = (dependencies,)
    end

    if KernelAbstractions.workgroupsize(obj) <: DynamicSize && workgroupsize === nothing
        workgroupsize = (1024,) # Vectorization, 4x unrolling, minimal grain size
    end
    iterspace, dynamic = partition(obj, ndrange, workgroupsize)
    # partition checked that the ndrange's agreed
    if KernelAbstractions.ndrange(obj) <: StaticSize
        ndrange = nothing
    end

    if length(blocks(iterspace)) == 0
        return MultiEvent(dependencies)
    end

    Event(__run, obj, ndrange, iterspace, args, dynamic,
          dependencies=dependencies, progress=progress)
end

# Inference barriers
function __run(obj, ndrange, iterspace, args, dynamic)
    N = length(iterspace)
    Nthreads = Threads.nthreads()
    if Nthreads == 1
        len, rem = N, 0
    else
        len, rem = divrem(N, Nthreads)
    end
    # not enough iterations for all the threads?
    if len == 0
        Nthreads = N
        len, rem = 1, 0
    end
    if Nthreads == 1
        __thread_run(1, len, rem, obj, ndrange, iterspace, args, dynamic)
    else
        @sync for tid in 1:Nthreads
            Threads.@spawn __thread_run(tid, len, rem, obj, ndrange, iterspace, args, dynamic)
        end
    end
    return nothing
end

function __thread_run(tid, len, rem, obj, ndrange, iterspace, args, dynamic)
    # compute this thread's iterations
    f = 1 + ((tid-1) * len)
    l = f + len - 1
    # distribute remaining iterations evenly
    if rem > 0
        if tid <= rem
            f = f + (tid-1)
            l = l + tid
        else
            f = f + rem
            l = l + rem
        end
    end
    # run this thread's iterations
    for i = f:l
        block = @inbounds blocks(iterspace)[i]
        ctx = mkcontext(obj, block, ndrange, iterspace, dynamic)
        Cassette.overdub(ctx, obj.f, args...)
    end
    return nothing
end

Cassette.@context CPUCtx

function mkcontext(kernel::Kernel{CPU}, I, _ndrange, iterspace, ::Dynamic) where Dynamic
    metadata = CompilerMetadata{ndrange(kernel), Dynamic}(I, _ndrange, iterspace)
    Cassette.disablehooks(CPUCtx(pass = CompilerPass, metadata=metadata))
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Local_Linear), idx::CartesianIndex)
    indices = workitems(__iterspace(ctx.metadata))
    return @inbounds LinearIndices(indices)[idx]
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Group_Linear), idx::CartesianIndex)
    indices = blocks(__iterspace(ctx.metadata))
    return @inbounds LinearIndices(indices)[__groupindex(ctx.metadata)]
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Global_Linear), idx::CartesianIndex)
    I = @inbounds expand(__iterspace(ctx.metadata), __groupindex(ctx.metadata), idx)
    @inbounds LinearIndices(__ndrange(ctx.metadata))[I]
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Local_Cartesian), idx::CartesianIndex)
    return idx
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Group_Cartesian), idx::CartesianIndex)
    __groupindex(ctx.metadata)
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__index_Global_Cartesian), idx::CartesianIndex)
    return @inbounds expand(__iterspace(ctx.metadata), __groupindex(ctx.metadata), idx)
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__validindex), idx::CartesianIndex)
    # Turns this into a noop for code where we can turn of checkbounds of
    if __dynamic_checkbounds(ctx.metadata)
        I = @inbounds expand(__iterspace(ctx.metadata), __groupindex(ctx.metadata), idx)
        return I in __ndrange(ctx.metadata)
    else
        return true
    end
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(__print), items...)
    __print(items...)
end

generate_overdubs(CPUCtx)

# Don't recurse into these functions
const cpufuns = (:cos, :cospi, :sin, :sinpi, :tan,
          :acos, :asin, :atan,
          :cosh, :sinh, :tanh,
          :acosh, :asinh, :atanh,
          :log, :log10, :log1p, :log2,
          :exp, :exp2, :exp10, :expm1, :ldexp,
          :isfinite, :isinf, :isnan, :signbit,
          :abs,
          :sqrt, :cbrt,
          :ceil, :floor,)
for f in cpufuns
    @eval function Cassette.overdub(ctx::CPUCtx, ::typeof(Base.$f), x::Union{Float32, Float64})
        @Base._inline_meta
        return Base.$f(x)
    end
end

###
# CPU implementation of shared memory
###
@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(SharedMemory), ::Type{T}, ::Val{Dims}, ::Val) where {T, Dims}
    MArray{__size(Dims), T}(undef)
end

###
# CPU implementation of scratch memory
# - private memory for each workitem
# - memory allocated as a MArray with size `Dims + WorkgroupSize`
###
struct ScratchArray{N, D}
    data::D
    ScratchArray{N}(data::D) where {N, D} = new{N, D}(data)
end

@inline function Cassette.overdub(ctx::CPUCtx, ::typeof(Scratchpad), ::Type{T}, ::Val{Dims}) where {T, Dims}
    return ScratchArray{length(Dims)}(MArray{__size((Dims..., __groupsize(ctx.metadata)...)), T}(undef))
end

Base.@propagate_inbounds function Base.getindex(A::ScratchArray, I...)
    return A.data[I...]
end

Base.@propagate_inbounds function Base.setindex!(A::ScratchArray, val, I...)
    A.data[I...] = val
end


text:
using GlobalSensitivity, Test, OrdinaryDiffEq

@info "Morris Method"

A = reshape([1,0,2,3],2,2)
function f_morris(p)
    A*p
end

function linear_batch(X)
    A= 7
    B= 0.1
    @. A*X[1,:]+B*X[2,:]
end

function neg_linear_batch(X)
    A= -7
    B= 0.1
    @. A*X[1,:]+B*X[2,:]
end

function ishi(X)
    A= 7
    B= 0.1
    sin(X[1]) + A*sin(X[2])^2+ B*X[3]^4 *sin(X[1])
end

function ishi_linear(X)
    A= 7
    B= 0.1
    [sin(X[1]) + A*sin(X[2])^2+ B*X[3]^4 *sin(X[1]),A*X[1]+B*X[2]]
end

function ishi_linear_batch(X)
    A= 7
    B= 0.1
    X1 = @. sin(X[1,:]) + A*sin(X[2,:])^2+ B*X[3,:]^4 *sin(X[1,:])
    X2 = @. A*X[1,:]+B*X[2,:]
    vcat(X1',X2')
end

lb = -ones(4)*π
ub = ones(4)*π

m = gsa(f_morris,Morris(p_steps=[10,10],total_num_trajectory=1000,num_trajectory=150),[[1,5],[1,5]])
@test m.means[:,1] ≈ A[:,1] atol=1e-12
@test m.means[:,2] ≈ A[:,2] atol=1e-12
@test m.means_star[:,1] ≈ A[:,1] atol=1e-12
@test m.means_star[:,2] ≈ A[:,2] atol=1e-12
@test m.variances ≈ reshape([0,0,0,0],2,2) atol=1e-12

m = gsa(f_morris,Morris(p_steps=[10,10],relative_scale=true,total_num_trajectory=1000,num_trajectory=150),[[1,5],[1,5]])
@test m.means[2,1] ≈ 0 atol=1e-12
@test m.means_star[2,1] ≈ 0 atol=1e-12
@test m.variances[2,1] ≈ 0 atol=1e-12
@test m.means[1,2] < m.means[2,2]
@test m.means_star[1,2] < m.means_star[2,2]

m = gsa(linear_batch,Morris(p_steps=[10,10],num_trajectory=50000),[[1,5],[1,5]],batch=true)
@test m.means ≈ [7.0  0.1] atol = 1e-2
@test m.means_star ≈ [7.0  0.1] atol = 1e-2
@test m.variances ≈ reshape([0,0], 1, 2) atol=1e-12

m = gsa(neg_linear_batch,Morris(p_steps=[10,10],num_trajectory=50000),[[1,5],[1,5]],batch=true)
@test m.means ≈ [-7.0  0.1] atol = 1e-2
@test m.means_star ≈ [7.0  0.1] atol = 1e-2
@test m.variances ≈ reshape([0,0], 1, 2) atol=1e-12

m = gsa(ishi, Morris(num_trajectory=500000), [[lb[i],ub[i]] for i in 1:4])
@test m.means_star[1,:] ≈ [2.25341,4.40246,2.5049,0.0] atol = 5e-2
@test m.means[1, :] ≈ [-0.416876, -0.0077712, -0.015714,  0.0] atol = 5e-2
@test m.means_star[1,:] ≈ [2.25341,4.40246,2.5049,0.0] atol = 5e-2

m = gsa(ishi_linear,Morris(num_trajectory=500000),[[lb[i],ub[i]] for i in 1:4])
@test m.means_star[1,:] ≈ [2.25341,4.40246,2.5049,0.0] atol = 5e-2
@test m.means_star[2,:] ≈ [7.0,0.1,0.0,0.0] rtol = 1e-12

m = gsa(ishi_linear_batch,Morris(num_trajectory=500000),[[lb[i],ub[i]] for i in 1:4],batch=true)
@test m.means_star[1,:] ≈ [2.25341,4.40246,2.5049,0.0] atol = 5e-2
@test m.means_star[2,:] ≈ [7.0,0.1,0.0,0.0] rtol = 1e-12

function f(du,u,p,t)
  du[1] = p[1]*u[1] - p[2]*u[1]*u[2] #prey
  du[2] = -p[3]*u[2] + p[4]*u[1]*u[2] #predator
end

u0 = [1.0;1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0]
prob = ODEProblem(f,u0,tspan,p)
t = collect(range(0, stop=10, length=200))

f1 = function (p)
  prob1 = remake(prob;p=p)
  sol = solve(prob1,Tsit5();saveat=t)
  return sol
end

m = gsa(f1,Morris(total_num_trajectory=1000,num_trajectory=150),[[1,5],[1,5],[1,5],[1,5]])

ext: jl


text:
# ParamSet types
#
# The variable names are:
# u       = Location in world coordinates (formerly mu)
# e_dev   = Weight given to a galaxy of type 1 (formerly theta)
# e_axis  = Galaxy minor/major ratio (formerly rho)
# e_angle = Galaxy angle (formerly phi)
# e_scale = Galaxy scale (sigma)
# For r1 and r2, the first row is stars, and the second is galaxies.
# r1      = Iax1 lognormal mean parameter for r_s.
# r2      = Iax1 lognormal variance parameter for r_s.
# c1      = C_s means (formerly beta)
# c2      = C_s variances (formerly lambda)
# a       = probability of being a star or galaxy.  a[1, 1] is the
#           probability of being a star and a[2, 1] of being a galaxy.
#           (formerly chi)
# k       = {D|D-1}xIa matrix of color prior component indicators.
#           (formerly kappa)

import ..SensitiveFloats: ParamSet

type StarPosParams <: ParamSet
    u::Vector{Int}
    StarPosParams() = new([1, 2])
end
const star_ids = StarPosParams()
getids(::Type{StarPosParams}) = star_ids
length(::Type{StarPosParams}) = 2

type GalaxyShapeParams <: ParamSet
    e_axis::Int
    e_angle::Int
    e_scale::Int
    GalaxyShapeParams() = new(1, 2, 3)
end
const gal_shape_ids = GalaxyShapeParams()
getids(::Type{GalaxyShapeParams}) = gal_shape_ids
length(::Type{GalaxyShapeParams}) = 3

type GalaxyPosParams <: ParamSet
    u::Vector{Int}
    e_dev::Int
    e_axis::Int
    e_angle::Int
    e_scale::Int
    GalaxyPosParams() = new([1, 2], 3, 4, 5, 6)
end
const gal_ids = GalaxyPosParams()
getids(::Type{GalaxyPosParams}) = gal_ids
length(::Type{GalaxyPosParams}) = 6

type BrightnessParams <: ParamSet
    r1::Int
    r2::Int
    c1::Vector{Int}
    c2::Vector{Int}
    BrightnessParams() = new(1, 2,
                             collect(3:(3+(B-1)-1)),
                             collect((3+B-1):(3+2*(B-1)-1)))
end
const bids = BrightnessParams()
getids(::Type{BrightnessParams}) = bids
length(::Type{BrightnessParams}) = 2 + 2 * (B-1)

type CanonicalParams <: ParamSet
    u::Vector{Int}
    e_dev::Int
    e_axis::Int
    e_angle::Int
    e_scale::Int
    r1::Vector{Int}
    r2::Vector{Int}
    c1::Matrix{Int}
    c2::Matrix{Int}
    a::Matrix{Int}
    k::Matrix{Int}
    CanonicalParams() =
        new([1, 2], 3, 4, 5, 6,
            collect(7:(7+Ia-1)),  # r1
            collect((7+Ia):(7+2Ia-1)), # r2
            reshape((7+2Ia):(7+2Ia+(B-1)*Ia-1), (B-1, Ia)),  # c1
            reshape((7+2Ia+(B-1)*Ia):(7+2Ia+2*(B-1)*Ia-1), (B-1, Ia)),  # c2
            reshape((7+2Ia+2*(B-1)*Ia):(7+3Ia+2*(B-1)*Ia-1), (Ia, 1)),  # a
            reshape((7+3Ia+2*(B-1)*Ia):(7+3Ia+2*(B-1)*Ia+D*Ia-1), (D, Ia))) # k
end
const ids = CanonicalParams()
getids(::Type{CanonicalParams}) = ids
length(::Type{CanonicalParams}) = 6 + 3*Ia + 2*(B-1)*Ia + D*Ia

type UnconstrainedParams <: ParamSet
    u::Vector{Int}
    e_dev::Int
    e_axis::Int
    e_angle::Int
    e_scale::Int
    r1::Vector{Int}
    r2::Vector{Int}
    c1::Matrix{Int}
    c2::Matrix{Int}
    a::Matrix{Int}
    k::Matrix{Int}
    UnconstrainedParams() =
        new([1, 2], 3, 4, 5, 6,
            collect(7:(7+Ia-1)),  # r1
            collect((7+Ia):(7+2Ia-1)), # r2
            reshape((7+2Ia):(7+2Ia+(B-1)*Ia-1), (B-1, Ia)),  # c1
            reshape((7+2Ia+(B-1)*Ia):(7+2Ia+2*(B-1)*Ia-1), (B-1, Ia)),  # c2
            reshape((7+2Ia+2*(B-1)*Ia):
                    (7+2Ia+2*(B-1)*Ia+(Ia-1)-1), (Ia - 1, 1)),  # a
            reshape((7+2Ia+2*(B-1)*Ia+(Ia-1)):
                    (7+2Ia+2*(B-1)*Ia+(Ia-1)+(D-1)*Ia-1), (D-1, Ia))) # k
end
const ids_free = UnconstrainedParams()
getids(::Type{UnconstrainedParams}) = ids_free
length(::Type{UnconstrainedParams}) =  6 + 2*Ia + 2*(B-1)*Ia + (D-1)*Ia + Ia-1

# Parameters for a representation of the PSF
type PsfParams <: ParamSet
    mu::Vector{Int}
    e_axis::Int
    e_angle::Int
    e_scale::Int
    weight::Int

    function PsfParams()
      new([1, 2], 3, 4, 5, 6)
    end
end
const psf_ids = PsfParams()
getids(::Type{PsfParams}) = psf_ids
length(::Type{PsfParams}) = 6

# define length(value) in addition to length(Type) for ParamSets
length{T<:ParamSet}(::T) = length(T)

#TODO: build these from ue_align, etc., here.
align(::StarPosParams, CanonicalParams) = ids.u
align(::GalaxyPosParams, CanonicalParams) =
   [ids.u; ids.e_dev; ids.e_axis; ids.e_angle; ids.e_scale]
align(::CanonicalParams, CanonicalParams) = collect(1:length(CanonicalParams))
align(::GalaxyShapeParams, GalaxyPosParams) =
  [gal_ids.e_axis; gal_ids.e_angle; gal_ids.e_scale]

# The shape and brightness parameters for stars and galaxies respectively.
const shape_standard_alignment = (ids.u,
   [ids.u; ids.e_dev; ids.e_axis; ids.e_angle; ids.e_scale])
bright_ids(i) = [ids.r1[i]; ids.r2[i]; ids.c1[:, i]; ids.c2[:, i]]
const brightness_standard_alignment = (bright_ids(1), bright_ids(2))

# Note that gal_shape_alignment aligns the shape ids with the GalaxyPosParams,
# not the CanonicalParams.
const gal_shape_alignment = align(gal_shape_ids, gal_ids)

function get_id_names(
  ids::Union{CanonicalParams, UnconstrainedParams})
  ids_names = Array(String, length(ids))
  for name in fieldnames(ids)
    inds = getfield(ids, name)
    if length(size(inds)) == 0
      ids_names[inds] = "$(name)"
    elseif length(size(inds)) == 1
      for i = 1:size(inds)[1]
          ids_names[inds[i]] = "$(name)_$(i)"
      end
    elseif length(size(inds)) == 2
      for i = 1:size(inds)[1], j = 1:size(inds)[2]
          ids_names[inds[i, j]] = "$(name)_$(i)_$(j)"
      end
    else
      error("Names of 3d parameters not supported ($(name))")
    end
  end
  return ids_names
end

const ids_names = get_id_names(ids)
const ids_free_names = get_id_names(ids_free)


text:
@testset "cloud 1" begin
    bc1 = benchmarkcloud1()
    p = RANSACParameters()
    vs = bc1.vertices
    ns = bc1.normals

    for i in eachindex(bc1.shapes)
        s = bc1.shapes[i]
        good_indexes = findall(x->x==i, bc1.indexes)
        goodsize = size(good_indexes, 1)
        curr_v = vs[good_indexes]
        curr_n = ns[good_indexes]

        if s isa FittedPlane
            sel, _ = RANSAC.compatiblesPlane(s, curr_v, curr_n, p)
            @test count(sel) == goodsize
        elseif s isa FittedCylinder
            sel, _ = RANSAC.compatiblesCylinder(s, curr_v, curr_n, p)
            @test count(sel) == goodsize
        elseif s isa FittedCone
            sel = RANSAC.compatiblesCone(s, curr_v, curr_n, p)
            @test count(sel) == goodsize
        elseif s isa FittedSphere
            sel, _, _ = RANSAC.compatiblesSphere(s, curr_v, curr_n, p)
            @test count(sel) == goodsize
        end
    end
end

## cloud 2
@testset "cloud 2" begin
    bc2 = benchmarkcloud2()
    p = RANSACParameters()
    vs2 = bc2.vertices
    ns2 = bc2.normals

    for i in eachindex(bc2.shapes)
        s = bc2.shapes[i]
        good_indexes = findall(x->x==i, bc2.indexes)
        goodsize = size(good_indexes, 1)
        curr_v = vs2[good_indexes]
        curr_n = ns2[good_indexes]

        if s isa FittedPlane
            sel, _ = RANSAC.compatiblesPlane(s, curr_v, curr_n, p)
            @test count(sel) == goodsize
        elseif s isa FittedCylinder
            sel, _ = RANSAC.compatiblesCylinder(s, curr_v, curr_n, p)
            @test count(sel) == goodsize
        elseif s isa FittedCone
            sel = RANSAC.compatiblesCone(s, curr_v, curr_n, p)
            @test count(sel) == goodsize
        elseif s isa FittedSphere
            sel, _, _ = RANSAC.compatiblesSphere(s, curr_v, curr_n, p)
            @test count(sel) == goodsize
        end
    end
end


text:
type Container
  nid::Int
  parent::Nullable{Container}
  subcontainers::Dict{Symbol,Container}
end

# Simplified constructor for usual cases (i.e. not the root container)
function Container(parent::Container, opts::Dict=Dict())
  nid = getnid()
  nct = Container(nid, Nullable(parent), Dict{Symbol,Container}())

  args = Dict(:newnid   => nid,
              :compname => "html-node",
              :params   => Dict(:html=>""),
              :deco     => opts)

  send(currentSession, parent.nid, "append", args)
  nct
end

# walk the container tree, creating containers along the way
# if they do not exist
function findorcreate(ct::Container, path::Vector{Symbol},
                      index::Int, opts::Dict)
  nn = path[index]
  if length(path) == index # we are at the leaf container
    if haskey(ct.subcontainers, nn) # container exists => clear
      nct = ct.subcontainers[nn]
      send(currentSession, nct.nid, "clear", Dict(:deco=>opts))
      nct.subcontainers = Dict{Symbol,Container}()
    else  # no container with this name => create
      nct = Container(ct, opts) # apply opts only on leaf
      ct.subcontainers[nn] = nct
    end
    return nct
  else
    if ! haskey(ct.subcontainers, nn)  # no container with this name => create
      ct.subcontainers[nn] = Container(ct) # no opts, opts apply only to leaf
    end
    return findorcreate(ct.subcontainers[nn], path, index+1, opts)
  end
end

macro container(args...)
  objects, opts = parseargs(args...)

  name = [Symbol(randstring())]
  if length(objects) > 0
    nt = unfold(objects[1]) # try to interpret as container name path
    length(nt) > 0 && (name = nt)
  end

  # create or clear container identified by its path in `name`
  nct = findorcreate(currentSession.root_container, name, 1, opts)
  currentSession.active_container = nct
  nothing
end


text:
@testset "Base functions" begin

    @testset "Distributed data transfers -- local" begin
        data = rand(5)
        save_at(1, :test, data)

        @test fetch(get_from(1, :test)) == data
        @test get_val_from(1, :test) == data

        remove_from(1, :test)

        @test sizeof(get_val_from(1, :test)) == 0 #should be "nothing" but this is more generic and generally describes what we want
    end

    addprocs(3)
    @everywhere using DistributedData
    W = workers()

    @testset "Distributed data transfers -- with workers" begin
        data = [rand(5) for i in W]
        for (i, w) in enumerate(W)
            save_at(w, :test, data[i])
        end

        @test [fetch(get_from(w, :test)) for w in W] == data
        @test [get_val_from(w, :test) for w in W] == data

        unscatter(:test, W)

        @test sum([sizeof(get_val_from(w, :test)) for w in W]) == 0
    end

    @testset "Data distribution" begin
        d = rand(100, 5)

        di = scatter_array(:test, d, W)

        @test di.val == :test
        @test Set(di.workers) == Set(W)
        @test begin
            d1 = get_val_from(di.workers[1], :test)
            d1 == d[1:size(d1, 1), :]
        end

        #TODO test actual sizes of the distributed pieces

        @test gather_array(di, free = false) == d #TODO test with true
        @test sum([sizeof(get_val_from(w, :test)) for w in W]) > 0
        unscatter(di)
        @test sum([sizeof(get_val_from(w, :test)) for w in W]) == 0
    end

    @testset "Distributed computation" begin
        di = dtransform(:(), x -> rand(5), W, :test)

        @test get_val_from(W[1], :test) == gather_array(di)[1:5]

        orig = gather_array(di)

        @test isapprox(
            dmapreduce(:test, d -> sum(d .^ 2), (a, b) -> a + b, W),
            sum(orig .^ 2),
        )

        dtransform(di, d -> d .* 2)

        @test orig .* 2 == gather_array(:test, W)

        @test isapprox(
            dmapreduce(di, d -> sum(d .^ 2), (a, b) -> a + b),
            sum((orig .* 2) .^ 2),
        )

        t = zeros(length(W))
        exp = zeros(length(W))

        t[1] = 2
        exp[1] = sum(2 .* get_val_from(W[1], :test))

        @test dmap(t, (i) -> eval(:(sum($i .* $(di.val)))), W) == exp

        unscatter(di)

        @test dmapreduce(:noname, x -> x, (a, b) -> a + b, []) == nothing
    end

    @testset "`pmap` on distributed data" begin
        map(fetch, save_at.(W, :test, 1234321))
        di = Dinfo(:test, W) # also test the example in docs
        @test dpmap(
            x -> :($(di.val) + $x),
            WorkerPool(di.workers),
            [4321234, 1234, 4321],
        ) == [5555555, 1235555, 1238642]
        map(fetch, remove_from.(W, :test))
    end

    @testset "Internal utilities" begin
        @test DistributedData.tmp_symbol(:test) != :test
        @test DistributedData.tmp_symbol(:test, prefix = "abc", suffix = "def") ==
              :abctestdef
        @test DistributedData.tmp_symbol(Dinfo(:test, W)) != :test
    end

    @testset "Persistent distributed data" begin
        di = dtransform(:(), x -> rand(5), W, :test)

        files = DistributedData.defaultFiles(di.val, di.workers)
        @test allunique(files)

        orig = gather_array(di)
        dstore(di, files)
        dtransform(di, x -> "erased")
        dload(di, files)

        @test orig == gather_array(di)

        dstore(di.val, di.workers, files)
        di2 = dload(:test2, di.workers, files)

        @test orig == gather_array(di2)

        unscatter(di)
        unscatter(di2)

        dunlink(di)

        @test all([!isfile(f) for f in files])
    end

    @testset "@remote macro" begin
        di = dtransform(:(), _ -> myid(), W, :test)

        test = 333

        for pid in W
            @test remotecall_fetch(() -> test + (@remote test), pid) == test + pid
        end
    end

    rmprocs(W)
    W = nothing

end


text:
module StandardPatterns
# This module is designed for creating complex patterns from the primtive ones.

using MLStyle
using MLStyle.toolz.List: cons, nil
using MLStyle.Infras
using MLStyle.MatchCore

struct TypeVar
    t :: Symbol
end

struct Relation
    l :: Symbol
    op :: Symbol
    r
end

function any_constraint(t, forall)

    function is_rel(::Relation)
        true
    end

    function is_rel(::TypeVar)
        false
    end

    !(t isa Symbol) || any(is_rel, collect(extract_tvars(forall)))
end

macro type_matching(t, forall)
    quote
        NAME = mangle(mod)
        __T__ = $t
        __FORALL__ = $forall
        if !($any_constraint(__T__, __FORALL__))
            function (body)
                @format [body, tag, NAME, TARGET, __T__] quote
                    @inline __L__ function NAME(TARGET :: __T__) where {$(__FORALL__...)}
                        __T__ # if not put this here, an error would be raised : "local variable XXX cannot be used in closure declaration"
                        body
                    end
                    NAME(tag)
                end
            end
        else
            function (body)
                @format [body, tag, NAME, TARGET, __T__] quote
                    @inline __L__ function NAME(TARGET :: __T__) where {$(__FORALL__...)}
                        __T__
                        body
                    end
                    @inline __L__ function NAME(_)
                        failed
                    end
                    NAME(tag)
                end
            end
        end
    end |> esc
end

function extract_tvars(t :: AbstractArray)
    @match t begin
        [] => nil()
        [hd && if hd isa Symbol end, tl...] => cons(TypeVar(hd), extract_tvars(tl))
        [:($hd <: $r), tl...] =>  cons(Relation(hd, :<:, r), extract_tvars(tl))
        [:($hd >: $(r)), tl...] =>  cons(Relation(hd, Symbol(">:"), r), extract_tvars(tl))
        _ => @error "invalid tvars"
    end
end

def_pattern(StandardPatterns,
        predicate = x -> x isa Expr && x.head == :(::),
        rewrite = (tag, case, mod) ->
                let args   = (case.args..., ),
                    TARGET = mangle(mod)
                    function for_type(t)
                        @match t begin
                            :($typ where {$(tvars...)}) => (@type_matching typ tvars)
                            _ => @typed_as t

                        end
                    end

                    function f(args :: NTuple{2, Any})
                        pat, t = args
                        for_type(t) ∘ mk_pattern(TARGET, pat, mod)
                    end

                    function f(args :: NTuple{1, Any})
                        t = args[1]
                        for_type(t)
                    end
                    f(args)
                end
)

export @active
"""
simple active pattern.

@active F(x) begin
    if x > 0
        nothing
    else
        :ok
    end
end
@match -1 begin
    F(:ok) => false
    _ => true
end
"""

macro active(case, active_body)
    mod = __module__
    (case_name, IDENTS, param) = @match case begin
        :($(case_name :: Symbol)($param)) => (case_name, nothing, param)
        :($(case_name :: Symbol){$(idents...)}($param)) => (case_name, idents, param)
    end
    TARGET = mangle(mod)
    mod.eval(quote struct $case_name end end)
    case_obj = getfield(mod, case_name)
    if IDENTS === nothing
        def_app_pattern(mod,
            predicate = (hd_obj, args) -> hd_obj === case_obj,
            rewrite = (tag, hd_obj, args, mod) -> begin
                arg = length(args) == 1 ? args[1] : Expr(:tuple, args...)
                function (body)
                    @format [tag, param, TARGET, active_body, body] quote
                        let  TARGET =
                            let param = tag
                                active_body
                            end
                            TARGET === nothing ?  failed : body
                        end
                    end
                end ∘ mk_pattern(TARGET, arg, mod)
        end)
    else
        n_idents = length(IDENTS)
        def_gapp_pattern(mod,
            predicate = (spec_vars, hd_obj, args) -> hd_obj === case_obj && length(spec_vars) === n_idents,
            rewrite   = (tag, forall, spec_vars, hd_obj, args, mod) -> begin
                arg = length(args) == 1 ? args[1] : Expr(:tuple, args...)
                assign_elts_and_active_body =
                    let arr = [:($IDENT = $(spec_vars[i]))
                                for (i, IDENT)
                                in enumerate(IDENTS)]
                        Expr(:let, Expr(:block, arr...), Expr(:block, active_body))
                    end
                function (body)
                    @format [tag, param, TARGET, assign_elts_and_active_body, body] quote
                        let TARGET =
                            let param = tag
                                assign_elts_and_active_body
                            end
                            TARGET === nothing ?  failed : body
                        end
                    end
                end ∘ mk_pattern(TARGET, arg, mod)
        end)
    end
    nothing
end

end


text:
HistoricAnalysis.unstack_tasks(; options...) =
    HistoricAnalysis.unstack_tasks(Historic.events(); options...)

HistoricAnalysis.unstack_tasks(recordmodule::Module; options...) =
    HistoricAnalysis.unstack_tasks(Historic.events(recordmodule); options...)

function HistoricAnalysis.unstack_tasks(table; options...)
    table = HistoricAnalysis.simplified(table; options...)
    long = DataFrame(table)
    long[!, :taskname], mapping = tasknames(long.taskid)
    wide = unstack(long, :taskname, :name)

    left = ["task$n" for n in 1:length(mapping)]
    right = ["taskid"]
    cols = union(left, setdiff!(names(wide), intersect(right, names(wide))))
    append!(cols, right)
    return select!(wide, cols)
end

function tasknames(taskids)
    mapping = Dict{eltype(taskids),String}()
    ntasks = Ref(0)
    ys = if Missing <: eltype(taskids)
        Union{String,Missing}[]
    else
        String[]
    end
    for id in taskids
        if id === missing
            push!(ys, missing)
            continue
        end
        y = get!(mapping, id) do
            n = ntasks[] += 1
            "task$n"
        end
        push!(ys, y)
    end
    @assert length(taskids) == length(ys)
    return ys, mapping
end


text:
using GeometricFlux
using Flux
using Flux: onehotbatch, onecold, crossentropy, throttle
using JLD2  # use v0.1.2
using Statistics: mean
using SparseArrays
using LightGraphs.SimpleGraphs
using CuArrays

@load "data/cora_features.jld2" features
@load "data/cora_labels.jld2" labels
@load "data/cora_graph.jld2" g

num_nodes = 2708
num_features = 1433

heads  = 8
hidden = 8
target_catg = 7
epochs = 10

## Preprocessing data
train_X = features |> gpu  # dim: num_features * num_nodes
train_y = labels |> gpu  # dim: target_catg * num_nodes

## Model
model = Chain(GATConv(g, num_features=>hidden, heads=heads),
              Dropout(0.6),
              GATConv(g, hidden=>target_catg, heads=heads),
              softmax) |> gpu
# test model
# model(train_X)

## Loss
loss(x, y) = crossentropy(model(x), y)
accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))

## Training
ps = Flux.params(model)
train_data = [(train_X, train_y)]
opt = ADAM(0.01)
evalcb() = @show(accuracy(train_X, train_y))

for i = 1:epochs
    Flux.train!(loss, ps, train_data, opt, cb=throttle(evalcb, 10))
end


text:
import ClimateMachine.BalanceLaws:
# declaration
    vars_state,
# initialization
    init_state_prognostic!,
    init_state_auxiliary!,
    nodal_init_state_auxiliary!,
# rhs computation
    compute_gradient_argument!,
    compute_gradient_flux!,
    flux_first_order!,
    flux_second_order!,
    source!,
# boundary conditions
    boundary_conditions,
    boundary_state!

import ClimateMachine.DGMethods.NumericalFluxes:
    CentralNumericalFluxGradient,
    CentralNumericalFluxSecondOrder,
    NumericalFluxFirstOrder,
    NumericalFluxSecondOrder,
    RusanovNumericalFlux,
    numerical_boundary_flux_second_order!,
    numerical_flux_second_order!,
    numerical_boundary_flux_first_order!,
    numerical_flux_first_order!,
    normal_boundary_flux_second_order!

"""
    Boundary condition modifications for the coupler
"""
@inline boundary_conditions(model::ModelSetup) = model.boundary_conditions

@inline function boundary_state!(numerical_flux, bc::FluidBC, model::ModelSetup, args...)
    # We need to apply boundary conditions for state variables.
    # This depends on first, second, and high-order
    # operations, hence the diffusion model dependence.
    # TODO!: make work for higher-order diffusion
    diffusion = model.physics.diffusion # defaults to `nothing`
    calc_boundary_state!(numerical_flux, bc.ρu, model, diffusion, args...)
    calc_boundary_state!(numerical_flux, bc.ρθ, model, diffusion, args...)
end

#calc_boundary_state!(numerical_flux, bc, model, diffusion, args...) = nothing

# first-order flux and gradient flux follow the original Impenetrable() / Freeslip() conditions.

# second order fluxes require additional treatment

function numerical_boundary_flux_second_order!(
    numerical_flux::Union{PenaltyNumFluxDiffusive},
    bctype::FluidBC,
    balance_law::ModelSetup,
    args...,
) where {S, D, A, HD}

    numerical_boundary_flux_second_order!( # this overwrites the total flux for ρθ (using methods below)
        numerical_flux,
        bctype.ρθ,
        balance_law,
        args...,
    )

    numerical_boundary_flux_second_order!( # this overwrites the total flux for ρu (using methods below)
        numerical_flux,
        bctype.ρu,
        balance_law,
        args...,
    )

end

numerical_boundary_flux_second_order!(::PenaltyNumFluxDiffusive, ::Impenetrable{FreeSlip}, _...) = nothing # no contribution to fluxᵀn.ρu from nF_diffusive

function numerical_boundary_flux_second_order!(
    numerical_flux::Union{PenaltyNumFluxDiffusive},
    bctype::CoupledPrimaryBoundary,
    balance_law::ModelSetup,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_gradient_flux⁻::Vars{D},
    state_hyperdiffusive⁻::Vars{HD},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_gradient_flux⁺::Vars{D},
    state_hyperdiffusive⁺::Vars{HD},
    state_auxiliary⁺::Vars{A},
    t,
    state1⁻::Vars{S},
    diff1⁻::Vars{D},
    aux1⁻::Vars{A},
) where {S, D, A, HD}
    F_tot = calculate_land_sfc_fluxes(balance_law, state_prognostic⁻, state_auxiliary⁻, t)  # W/m^2
    fluxᵀn.ρθ += -F_tot / balance_law.parameters.cp_d

end

function numerical_boundary_flux_second_order!(
    numerical_flux::Union{PenaltyNumFluxDiffusive},
    bctype::Insulating,
    balance_law::ModelSetup,
    fluxᵀn::Vars{S},
    _...,
) where {S}
    FT = eltype(fluxᵀn)
    fluxᵀn.ρθ += FT(0)
end

"""
function numerical_flux_second_order!(::PenaltyNumFluxDiffusive,
  Penalty flux formulation of second order numerical flux. This formulation
  computes the CentralNumericalFluxSecondOrder term first (which is just the average
  of the + and - fluxes and an edge), and then adds a "penalty" flux that relaxes
  the edge state + and - toward each other.
"""
function numerical_flux_second_order!(
    ::PenaltyNumFluxDiffusive,
    bl::ModelSetup,
    fluxᵀn::Vars{S},
    n::SVector,
    state⁻::Vars{S},
    diff⁻::Vars{D},
    hyperdiff⁻::Vars{HD},
    aux⁻::Vars{A},
    state⁺::Vars{S},
    diff⁺::Vars{D},
    hyperdiff⁺::Vars{HD},
    aux⁺::Vars{A},
    t,
) where {S, HD, D, A}

    numerical_flux_second_order!(
        CentralNumericalFluxSecondOrder(),
        bl,
        fluxᵀn,
        n,
        state⁻,
        diff⁻,
        hyperdiff⁻,
        aux⁻,
        state⁺,
        diff⁺,
        hyperdiff⁺,
        aux⁺,
        t,
    )

    Fᵀn = parent(fluxᵀn)
    FT = eltype(Fᵀn)
    tau = FT(0.0) # decide if want to use this penalty
    Fᵀn .+= tau * (parent(state⁻) - parent(state⁺))
end

numerical_flux_second_order!(::PenaltyNumFluxDiffusive, bl::SlabLandModelSetup, _...) = nothing

max_stars_repo_name: CliMA/CouplerMachine

max_issues_repo_name: CliMA/CouplerMachine

max_forks_repo_name: CliMA/CouplerMachine


text:
module classopath_test

using Test, ConstrainedLasso, ECOS, Random

@info("Test lsq_classopath: sum-to-zero constraint")

# set up
Random.seed!(456)
n, p = 100, 20
# truth with sum constraint sum(β) = 0
β = zeros(p)
β[1:round(Int, p / 4)] .= 0
β[(round(Int, p / 4) + 1):round(Int, p / 2)] .= 1
β[(round(Int, p / 2) + 1):round(Int, 3p / 4)] .= 0
β[(round(Int, 3p / 4) + 1):p] .= -1
# generate data
X = randn(n, p)
y = X * β + randn(n)
# equality constraints
Aeq = ones(1, p)
beq = [0.0]
penwt  = ones(p)

β̂path1, ρpath1, objpath, = lsq_classopath(X, y; Aeq = Aeq, beq = beq)
# print(β̂path1)
# print(size(β̂path1))
# print(sum(β̂path1,1))
@test all(abs.(sum(β̂path1, 1)) .< 1.0e-6)

@info("Test lsq_classopath: non-negativity constraint")

# set up
n, p = 20, 100
# truth with a few non-zero coefficients
β = zeros(p)
β[1:10] = 1:10
# generate data
Random.seed!(123)
X = randn(n, p)
y = X * β + randn(n)
# inequality constraints
Aineq = - eye(p)
bineq = zeros(p)

solver = ECOSSolver(verbose=0, maxit=1e8)

β̂path2, ρpath2, = lsq_classopath(X, y; Aineq = Aineq, bineq = bineq, solver=solver)

@testset "non-negativity" begin
for i in reshape(β̂path2, p * length(ρpath2), 1)
  @test i >= 0
end
end

end # classopath_test module


text:
using VLConstraintBasedModelGenerationUtilities

# setup path to protein sequence file -
path_to_vff_file = "/Users/jeffreyvarner/Desktop/julia_work/VLConstraintBasedModelGenerationUtilities.jl/test/data/Test.vff"

# let's build the reaction table -
metabolic_reaction_table = build_metabolic_reaction_table(path_to_vff_file) |> check

# get list of species -
list_of_species = build_species_symbol_array(metabolic_reaction_table) |> check

# build the stm -
stm = build_stoichiometric_matrix(metabolic_reaction_table) |> check


text:
struct SemObsCommon{
        A <: Union{AbstractArray, Nothing},
        B <: AbstractArray,
        C <: Union{AbstractArray, Nothing},
        D <: AbstractFloat,
        O <: Union{AbstractFloat, Nothing},
        R <: Union{AbstractArray, Nothing}} <: SemObs
    data::A
    obs_cov::B
    obs_mean::C
    n_man::D
    n_obs::O
    data_rowwise::R
end

function SemObsCommon(;
        specification = nothing,
        data = nothing,
        spec_colnames = nothing,
        obs_cov = nothing,
        data_colnames = nothing,
        meanstructure = false,
        rowwise = false,
        n_obs = nothing,
        kwargs...)

    # sort columns/rows
    if isnothing(spec_colnames) spec_colnames = get_colnames(specification) end

    data, obs_cov = reorder_observed(data, obs_cov, spec_colnames, data_colnames)

    # if no cov. matrix was given, compute one
    if isnothing(obs_cov) obs_cov = Statistics.cov(data) end
    isnothing(data) ? nothing : n_obs = convert(Float64, size(data, 1))
    n_man = Float64(size(obs_cov, 1))
    # if a meanstructure is needed, compute observed means
    meanstructure ?
        obs_mean = vcat(Statistics.mean(data, dims = 1)...) :
        obs_mean = nothing
    rowwise ?
        data_rowwise = [data[i, :] for i = 1:convert(Int64, n_obs)] :
        data_rowwise = nothing
    return SemObsCommon(data, obs_cov, obs_mean, n_man, n_obs, data_rowwise)
end

############################################################################
### Pretty Printing
############################################################################

function Base.show(io::IO, struct_inst::SemObsCommon)
    print_type_name(io, struct_inst)
    print_field_types(io, struct_inst)
end

############################################################################
### Recommended methods
############################################################################

n_obs(observed::SemObsCommon) = observed.n_obs
n_man(observed::SemObsCommon) = observed.n_man

############################################################################
### additional methods
############################################################################

get_data(observed::SemObsCommon) = observed.data
obs_cov(observed::SemObsCommon) = observed.obs_cov
obs_mean(observed::SemObsCommon) = observed.obs_mean
data_rowwise(observed::SemObsCommon) = observed.data_rowwise

############################################################################
### Additional functions
############################################################################

# specification colnames
function get_colnames(specification::ParameterTable)
    if !haskey(specification.variables, :sorted_vars) || (length(specification.variables[:sorted_vars]) == 0)
        colnames = specification.variables[:observed_vars]
    else
        is_obs = [var ∈ specification.variables[:observed_vars] for var in specification.variables[:sorted_vars]]
        colnames = specification.variables[:sorted_vars][is_obs]
    end
    return colnames
end

function get_colnames(specification::RAMMatrices)
    if isnothing(specification.colnames)
        @warn "Your RAMMatrices do not contain column names. Please make shure the order of variables in your data is correct!"
        return nothing
    else
        colnames = specification.colnames[specification.F_ind]
        return colnames
    end
end

function get_colnames(specification::Nothing)
    return nothing
end

# reorder data to match spec_colnames --------------------------------------------------------------
reorder_observed(data::Nothing, obs_cov, spec_colnames, data_colnames) = reorder_obs_cov(obs_cov, spec_colnames, data_colnames)
reorder_observed(data, obs_cov::Nothing, spec_colnames, data_colnames) = reorder_data(data, spec_colnames, data_colnames)
reorder_observed(data::Nothing, obs_cov, spec_colnames::Nothing, data_colnames) = nothing, obs_cov
reorder_observed(data, obs_cov::Nothing, spec_colnames::Nothing, data_colnames) = data, nothing

# too much or not enough data specified
reorder_observed(data, obs_cov, spec_colnames, data_colnames) =
    throw(ArgumentError("you specified both an observed dataset and an observed covariance matrix"))
reorder_observed(data::Nothing, obs_cov::Nothing, spec_colnames, data_colnames) =
    throw(ArgumentError("you specified neither an observed dataset nor an observed covariance matrix"))

# reorder data ------------------------------------------------------------------------------------------------------------
reorder_data(data::AbstractArray, spec_colnames, data_colnames::Nothing) =
    throw(ArgumentError("if your data format does not provide column names, please provide them via the `data_colnames = ...` argument."))

function reorder_data(data::AbstractArray, spec_colnames, data_colnames)
    if spec_colnames == data_colnames
        return data, nothing
    else
        new_position = [findall(x .== data_colnames)[1] for x in spec_colnames]
        data = Matrix(data[:, new_position])
        return data, nothing
    end
end

reorder_data(data::DataFrame, spec_colnames, data_colnames::Nothing) = Matrix(data[:, spec_colnames]), nothing
reorder_data(data::DataFrame, spec_colnames, data_colnames) =
    throw(ArgumentError("your data format has column names but you also provided column names via the `data_colnames = ...` argument."))

# reorder covariance matrices ---------------------------------------------------------------------------------------------
reorder_obs_cov(obs_cov::AbstractArray, spec_colnames, data_colnames::Nothing) =
    throw(ArgumentError("If an observed covariance is given, `data_colnames = ...` has to be specified."))

function reorder_obs_cov(obs_cov::AbstractArray, spec_colnames, data_colnames)
    new_position = [findall(x .== data_colnames)[1] for x in spec_colnames]
    indices = reshape([CartesianIndex(i, j) for j in new_position for i in new_position], size(obs_cov, 1), size(obs_cov, 1))
    obs_cov = obs_cov[indices]
    return nothing, obs_cov
end


text:
include("header.jl")

#=
# Streamlines
In this notebook, we will learn some basic ways to calculate and plot streamlines
=#

# ### Set up the module
using MAE103
#-
using Plots

#=
Streamlines are curves in space that are everywhere tangent to the local velocity field.
They indicate the lines of flow in a steady fluid flow. If the flow is unsteady,
the lines represent a snapshot of the flow, but those lines will tend to change
with time.

The basic equations for streamlines come from the fact that the streamline must
be locally parallel to the velocity. This means that the local slope
of a streamline, $\mathrm{d}y/\mathrm{d}x$, must be equal to the ratio of the
velocity components, $v/u$, if we use Cartesian components for the velocity,
$\boldsymbol{V} = u \boldsymbol{i} + v \boldsymbol{j}$. Sometimes streamlines may
not be describable as $y$ as a function of $x$. For example, if a streamline is vertical,
then it can't be written this way. In that case, it would be better to write it as $x(y)$.
Or better yet, we could write it *parametrically*: e.g., $x(s)$ and $y(s)$, where $s$
is a parameter that increases along the streamline curve. So the more general way
to write the streamline equation that covers all cases is

$$\dfrac{\mathrm{d}x}{u} =  \dfrac{\mathrm{d}y}{u}$$

From this, we can choose either $x$ or $y$ as our independent variable, or another
one if necessary. We will see a few ways in this notebook.
=#
#-
#=
# Example 1:
Suppose we are given the velocity field for a *stagnation flow*:
$u(x,y) = 10*x$, $v(x,y) = -10*y$
We seek the streamlines.
=#
#-
#=
We first write these as a
$$\dfrac{\mathrm{d}y}{\mathrm{d}x} = \dfrac{v}{u} = \dfrac{-10y}{10x} = -\dfrac{y}{x}$$

Then, we separate variables and integrate:
$$\dfrac{\mathrm{d}y}{y} = -\dfrac{\mathrm{d}x}{x}$$

This leads to $\ln y = c - \ln x$, where $c$ is an arbitrary constant. Raising these
to a power of $\mathrm{e}$, we get

$$y = k/x$$

where $k$ is an arbitrary constant. Different choices for $k$ correspond to different streamlines.
=#
#-
#=
Set up a function, which we'll call `ys`. We'll set a default value of `k`.
=#
ys(x;k=1) = k/x

#=
Now plot several for different choices of `k`:
=#
x = range(-3,3,length=101)
plot(xlim=(-3,3),ylim=(-3,3),ratio=1)
plot!(x,ys.(x;k=2),color=:black)
plot!(x,ys.(x;k=1),color=:black)
plot!(x,ys.(x;k=-1),color=:black)
plot!(x,ys.(x;k=-2),color=:black)
plot!(x,ys.(x;k=0),color=:black)

#=
The $x$ axis is a streamline ($k = 0$). The $y$ axis is also a streamline, but
we can't get it because we assumed that $y$ was a function of $x$. If we repeat
the process by assuming that $x$ is a function of $y$, we get

$$x = k/y$$
=#
xs(y;k=1) = k/y
y = range(-3,3,length=101)
plot!(xs.(y;k=0),y,color=:black)
#=
(Note that we swapped the order of arguments.)
=#

#=
# Example 2: A polar streamline plot
In this example, we wish to generate streamlines when the velocity is given
in polar coordinates: $\boldsymbol{V} = v_r \boldsymbol{e}_r + v_\theta \boldsymbol{e}_\theta$.
The velocity field is $v_r = 0$, $v_\theta = 10/r$. This velocity field has
only a circumferential component, which means that that flow only goes around in
circles; there is no radial flow.
=#
#-
#=
We still apply the fact that the streamline is tangent to the velocity. However,
now, we don't use $\mathrm{d}x$ and $\mathrm{d}y$, but rather $\mathrm{d}r$
and $r\mathrm{d}\theta$:

$$\dfrac{\mathrm{d}r}{r\mathrm{d}\theta} = \dfrac{v_r}{v_\theta} = 0$$
=#
#-
#=
In other words, $\mathrm{d}r/\mathrm{d}\theta = 0$ in this case, which simply
says that the streamlines are lines of constant $r$ --- they are circles!
That is, $r(\theta) = k$. Each choice of $k$ is a different streamline, and $k$
simply gives the radius of the streamline:
=#
rs(Θ;k=1) = k

#=
Let's plot these streamlines. For plotting purposes, it is easiest to
convert the polar coordinates into Cartesian:

$$x = r\cos\theta,\qquad y = r\sin\theta$$
=#
xs(Θ;k=1) = rs(Θ,k=k)*cos(Θ)
ys(Θ;k=1) = rs(Θ,k=k)*sin(Θ)

#-
Θ = range(0,2π,length=201)
plot(xlim=(-2,2),ylim=(-2,2),ratio=1)
plot!(xs.(Θ,k=1),ys.(Θ,k=1),color=:black)
plot!(xs.(Θ,k=1.5),ys.(Θ,k=1.5),color=:black)
plot!(xs.(Θ,k=0.5),ys.(Θ,k=0.5),color=:black)


text:
using Pkg
Pkg.activate(".")
Pkg.instantiate()
using FinEtoolsDeforNonlinear
using Profile

include("./examples/dynamics/transient/3-d/cantilever_dyn_examples.jl");

# Serial  execution
# cantilever_dyn_examples.neohookean_h8()
# Parallel execution
for NTHREADS in [1 2 4 8]
	cantilever_dyn_examples.neohookean_h8_thr(NTHREADS)
end

# exit()

# include("test/playground.jl")
# Pkg.test()

# include(".\\examples\\statics\\2-d\\tension_compression_examples.jl");
# tension_compression_examples.neohookeanad_q4()

# cd(".\\examples\\statics\\2-d")
# include(".\\bertoldi_compression_examples.jl");
# bertoldi_compression_examples.neohookeanad_q4()


text:
using DataFrames: DataFrame, sort, filter
using Dates: DateTime, Period, Day, now, format
using LegibleLambdas: @λ
using TryCatch: @try

export AtomicJob
export getstatus,
    getresult,
    description,
    ispending,
    isrunning,
    isexited,
    issucceeded,
    isfailed,
    isinterrupted,
    createdtime,
    starttime,
    stoptime,
    elapsed,
    outmsg,
    run!,
    interrupt!,
    initialize!,
    queue,
    query,
    isexecuted,
    ntimes

@enum JobStatus begin
    PENDING
    RUNNING
    SUCCEEDED
    FAILED
    INTERRUPTED
end

abstract type Job end
# Reference: https://github.com/cihga39871/JobSchedulers.jl/blob/aca52de/src/jobs.jl#L35-L69
mutable struct AtomicJob <: Job
    id::Int64
    def::Function
    desc::String
    user::String
    created_time::DateTime
    start_time::DateTime
    stop_time::DateTime
    max_time::Period
    status::JobStatus
    outmsg::String
    ref::Union{Task,Nothing}
    count::UInt64
    AtomicJob(def; desc = "", user = "", max_time = Day(1)) = new(
        generate_id(),
        def,
        desc,
        user,
        now(),
        DateTime(0),
        DateTime(0),
        max_time,
        PENDING,
        "",
        nothing,
        0,
    )
end
AtomicJob(cmd::Base.AbstractCmd; kwargs...) = AtomicJob(@λ(() -> run(cmd)); kwargs...)
AtomicJob(job::AtomicJob) =
    AtomicJob(job.def; desc = job.desc, user = job.user, max_time = job.max_time)

const JOB_REGISTRY = Job[]

isinitialized(job::AtomicJob) =
    job.start_time == job.stop_time == DateTime(0) &&
    job.status === PENDING &&
    job.ref === nothing

function run!(job::AtomicJob; attempts = 1, nap = 3)
    @assert isinteger(attempts) && attempts >= 1
    for _ in 1:attempts
        if !issucceeded(job)
            inner_run!(job)
            issucceeded(job) ? break : sleep(nap)
        end
    end
    return job
end
function inner_run!(job::AtomicJob)
    if isinitialized(job)
        job.ref = @async begin
            job.status = RUNNING
            job.start_time = now()
            if !isexecuted(job)
                push!(JOB_REGISTRY, job)
            end
            core_run!(job)
        end
        return job
    else
        job = initialize!(job)
        return inner_run!(job)
    end
end
function core_run!(job::AtomicJob)
    # See https://github.com/JuliaLang/julia/issues/21130#issuecomment-288423284
    @try begin
        global result = job.def()
    @catch e
        job.stop_time = now()
        @error "come across `$e` when running!"
        job.status = e isa InterruptException ? INTERRUPTED : FAILED
        return e
    @else
        job.stop_time = now()
        job.status = SUCCEEDED
        return result
    @finally
        job.count += 1
    end
end

function queue(; sortby = :created_time)
    @assert sortby in
            (:created_time, :user, :start_time, :stop_time, :elapsed, :status, :times)
    df = DataFrame(
        id = [job.id for job in JOB_REGISTRY],
        user = [job.user for job in JOB_REGISTRY],
        created_time = map(createdtime, JOB_REGISTRY),
        start_time = map(starttime, JOB_REGISTRY),
        stop_time = map(stoptime, JOB_REGISTRY),
        elapsed = map(elapsed, JOB_REGISTRY),
        status = map(getstatus, JOB_REGISTRY),
        times = map(ntimes, JOB_REGISTRY),
        desc = map(description, JOB_REGISTRY),
    )
    return sort(df, [:id, sortby])
end

query(id::Integer) = filter(row -> row.id == id, queue())
query(ids::AbstractVector{<:Integer}) = map(id -> query(id), ids)

isexecuted(job::Job) = job in JOB_REGISTRY

ntimes(id::Integer) = ntimes(first(filter(x -> x.id == id, JOB_REGISTRY)))
ntimes(job::Job) = Int(job.count)

getstatus(x::Job) = x.status

ispending(x::Job) = getstatus(x) === PENDING

isrunning(x::Job) = getstatus(x) === RUNNING

isexited(x::Job) = getstatus(x) in (SUCCEEDED, FAILED, INTERRUPTED)

issucceeded(x::Job) = getstatus(x) === SUCCEEDED

isfailed(x::Job) = getstatus(x) === FAILED

isinterrupted(x::Job) = getstatus(x) === INTERRUPTED

createdtime(x::Job) = x.created_time

starttime(x::Job) = ispending(x) ? nothing : x.start_time

stoptime(x::Job) = isexited(x) ? x.stop_time : nothing

function elapsed(x::Job)
    if ispending(x)
        return
    elseif isrunning(x)
        return now() - x.start_time
    else  # Exited
        return x.stop_time - x.start_time
    end
end

getresult(x::Job) = isexited(x) ? Some(fetch(x.ref)) : nothing

description(x::Job) = x.desc

outmsg(x::AtomicJob) = isexited(x) ? x.outmsg : nothing

# From https://github.com/cihga39871/JobSchedulers.jl/blob/aca52de/src/jobs.jl#L6-L10
function generate_id()
    time_value = (now().instant.periods.value - 63749462400000) << 16
    rand_value = rand(UInt16)
    return time_value + rand_value
end

function interrupt!(job::AtomicJob)
    if isexited(job)
        @info "the job $(job.id) has already exited!"
        return job
    elseif ispending(job)
        @info "the job $(job.id) has not started!"
        return job
    else
        schedule(job.ref, InterruptException(); error = true)
        return job
    end
end

function initialize!(job::AtomicJob)
    job.start_time = DateTime(0)
    job.stop_time = DateTime(0)
    job.status = PENDING
    job.outmsg = ""
    job.ref = nothing
    return job
end

Base.wait(x::Job) = wait(x.ref)

function Base.show(io::IO, job::AtomicJob)
    if get(io, :compact, false) || get(io, :typeinfo, nothing) == typeof(job)
        Base.show_default(IOContext(io, :limit => true), job)  # From https://github.com/mauro3/Parameters.jl/blob/ecbf8df/src/Parameters.jl#L556
    else
        println(io, summary(job))
        println(io, ' ', "id: ", job.id)
        if !isempty(job.desc)
            print(io, ' ', "description: ")
            show(io, job.desc)
            println(io)
        end
        print(io, ' ', "def: ")
        show(io, job.def)
        print(io, '\n', ' ', "status: ")
        printstyled(io, getstatus(job); bold = true)
        if !ispending(job)
            println(io, '\n', ' ', "from: ", format(starttime(job), "dd-u-YYYY HH:MM:SS.s"))
            print(io, ' ', "to: ")
            if isrunning(job)
                print(io, "still running...")
            else
                println(io, format(stoptime(job), "dd-u-YYYY HH:MM:SS.s"))
                print(io, ' ', "uses: ", elapsed(job))
            end
        end
    end
end


text:
export Register

"""
Creates register initializing with the |0...0⟩ state

- `N`     : Total number of qubits in circuit
- `state` : State vector for the 2^N Hilbert space

"""
mutable struct Register
    N::Int
    state::Array{ComplexF64}

    # Construct with state |0...0⟩
    function Register(N::Int)
        x = zeros(ComplexF64,2^N)
        x[1] = complex(1.0)
        new(N,x)
    end
end


text:
const zangwill = let res_init=zeros(3), jac_init=zeros(3,3), x_init=[100., -1., 2.5]

    function res(x, r)
        r[1] = x[1] - x[2] + x[3]
        r[2] = -x[1] + x[2] + x[3]
        r[3] = x[1] + x[2] - x[3]
        return r
    end

    function jac(x, jac)
        jac[:] .= 0.0
        for k = 1:3, i = 1:3
            jac[i, k] = 1.0
        end
        jac[1, 2] = -1.0
        jac[2, 1] = -1.0
        jac[3, 3] = -1.0
        return jac
    end

    f(;scale=1, verbose=false, print_steps=false) = testone("zangwill", res, jac,
                                                   res_init, jac_init, x_init;
                                                   scale=scale, verbose=verbose,
                                                   print_steps=print_steps)
end


text:
include("test_utilities.jl")

data, evts = loadtestdata("test_case_3a") #

data_r = reshape(data, (1, :))
data_e, times = Unfold.epoch(data = data_r, tbl = evts, τ = (0, 0.05), sfreq = 10) # cut the data into epochs

f = @formula 0 ~ 1 + conditionA + continuousA # 1
m_mul = fit(Unfold.UnfoldModel, Dict(Any=>(f,times)), evts, data_e)
##
@testset "Mass Univariate, all specified" begin

	# test simple case
	eff = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[0]),m_mul)
	@test size(eff,1) == 2 # we specified 2 levels @ 1 time point
	@test eff.conditionA ≈ [0.,1.] # we want to different levels
	@test eff.yhat ≈ [2.0,5.0] # these are the perfect predicted values

	# combination 2 levels /  6 values
	eff = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[-2,0,2]),m_mul)
	@test size(eff,1) == 6 # we want 6 values
	@test eff.conditionA ≈ [0.,0.,0.,1.,1.,1.]
	@test eff.continuousA ≈ [-2,0,2,-2,0,2.]
end
@testset "Mass Univariate, typified" begin
# testing typical value
	eff_man = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[mean(evts.continuousA)]),m_mul)
	eff_typ = Unfold.effects(Dict(:conditionA => [0,1]),m_mul)
	@test eff_man.yhat ≈ eff_typ.yhat
end

## Testing Splines
f_spl = @formula 0 ~ 1 + conditionA + spl(continuousA, 3) # 1
m_mul_spl = fit(UnfoldModel, f_spl, evts, data_e, times)

@testset "Mass Univariate, splines" begin

eff = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[0]),m_mul_spl)
@test size(eff,1) == 2 # we specified 2 levels @ 1 time point
@test eff.conditionA ≈ [0.,1.] # we want to different levels
@test eff.yhat ≈ [2.0,5.0] # these are the perfect predicted values

# combination 2 levels /  6 values
eff = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[-0.5,0,0.5]),m_mul_spl)
@test size(eff,1) == 6 # we want 6 values
@test eff.conditionA ≈ [0.,0.,0.,1.,1.,1.]
@test eff.continuousA ≈ [-0.5,0,0.5,-0.5,0,0.5]
@test eff.yhat ≈ [0,2,4,3,5,7]

# testing for safe predictions
eff = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[2]),m_mul_spl)
@test all(ismissing.(eff.yhat ))
end

## Timeexpansion
data, evts = loadtestdata("test_case_3a") #
f = @formula 0 ~ 1 + conditionA + continuousA # 1

@testset "Time Expansion, one event" begin

uf = fit(Unfold.UnfoldModel, Dict(Any=>(f,firbasis([0,0.1],10))), evts, data)
eff = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[0]),uf)
@test nrow(eff) == 4
@test eff.yhat ≈ [2., 2., 5., 5.]
@test eff.conditionA ≈ [0.,0.,1.,1.]
@test eff.continuousA ≈ [0.,0.,0.,0.]
end

data, evts = loadtestdata("test_case_4a") #
b1 = firbasis(τ = (0.0, 0.95), sfreq = 20, name = "basisA")
b2 = firbasis(τ = (0.0, 0.95), sfreq = 20, name = "basisB")
b3 = firbasis(τ = (0.0, 0.95), sfreq = 20, name = "basisC")
f = @formula 0 ~ 1 # 1
m_tul = fit(UnfoldModel, Dict("eventA"=>(f,b1),"eventB"=>(f,b2)), evts, data,eventcolumn="type")

@testset "Time Expansion, two events" begin

eff = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[0]),m_tul)
@test unique(eff.basisname)==["basisA","basisB"]
@test unique(eff.yhat) ≈ [2,3]
@test size(eff,1) == 2*2*20 # 2 basisfunctions, 2x conditionA, 1s a 20Hz

eff = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[-1,0,1]),m_tul)
@test size(eff,1) == 2*6*20

end
@testset "Time Expansion, three events" begin
    evts_3 = deepcopy(evts)
	evts_3.type[1:50] .= "eventC"
	evts_3.type[50:100] .= "eventD"
	m_tul_3 = fit(UnfoldModel, Dict("eventA"=>(f,b1),"eventB"=>(f,b2),"eventC"=>(f,b3)), evts_3, data,eventcolumn="type")

	eff = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[0]),m_tul_3)

	@test size(eff,1) == 3*2*20 # 2 basisfunctions, 2x conditionA, 1s a 20Hz

end

@testset "Time Expansion, two events different size + different formulas" begin

## Different sized events + different Formulas
data, evts = loadtestdata("test_case_4a") #
evts[!,:continuousA] = rand(MersenneTwister(42),nrow(evts))
b1 = firbasis(τ = (0.0, 0.95), sfreq = 20, name = "basisA")
b2 = firbasis(τ = (0.0, 0.5), sfreq = 20, name = "basisB")
f1 = @formula 0 ~ 1 # 1
f2 = @formula 0 ~ 1+continuousA # 1
m_tul = fit(UnfoldModel, Dict("eventA"=>(f1,b1),"eventB"=>(f2,b2)), evts, data,eventcolumn="type")
eff = Unfold.effects(Dict(:conditionA => [0,1],:continuousA =>[-1,0,1]),m_tul)
@test nrow(eff) == (length(Unfold.times(b1))-1+length(Unfold.times(b2))-1)*6
@test sum(eff.basisname .== "basisA") == 120
@test sum(eff.basisname .== "basisB") == 66

end

## Test two channels
data, evts = loadtestdata("test_case_3a") #

data_r = repeat(reshape(data, (1, :)),3,1)
data_r[2,:] = data_r[2,:] .* 2

data_e, times = Unfold.epoch(data = data_r, tbl = evts, τ = (0, 0.05), sfreq = 10) # cut the data into epochs

#
f = @formula 0 ~ 1 + conditionA + continuousA # 1
m_mul = fit(Unfold.UnfoldModel, Dict(Any=>(f,times)), evts, data_e)
m_tul = fit(Unfold.UnfoldModel, Dict(Any=>(f,firbasis([0,.05],10))), evts, data_r)
@testset "Two channels" begin

# test simple case
eff_m = Unfold.effects(Dict(:conditionA => [0,1,0,1],:continuousA =>[0]),m_mul)
eff_t = Unfold.effects(Dict(:conditionA => [0,1,0,1],:continuousA =>[0]),m_tul)

@test eff_m.yhat ≈ eff_t.yhat
@test length(unique(eff_m.channel)) == 3
@test eff_m[eff_m.channel .==1,:yhat] ≈ eff_m[eff_m.channel .==2,:yhat]./2
@test eff_m[eff_m.channel .==1,:yhat] ≈ [2,5,2,5.] # these are the perfect predicted values - note that we requested them twice

end

@testset "Timeexpansion, two events, typified" begin

data, evts = loadtestdata("test_case_4a") #
evts[!,:continuousA] = rand(MersenneTwister(42),nrow(evts))
evts[!,:continuousB] = rand(MersenneTwister(43),nrow(evts))
ixA = evts.type .== "eventA"
evts.continuousB[ixA] = evts.continuousB[ixA] .-mean(evts.continuousB[ixA]) .-5
evts.continuousB[.!ixA] = evts.continuousB[.!ixA] .-mean(evts.continuousB[.!ixA]) .+ 0.5
b1 = firbasis(τ = (0.0, 0.02), sfreq = 20, name = "basisA")
b2 = firbasis(τ = (1.0, 1.02), sfreq = 20, name = "basisB")
f1 = @formula 0 ~ 1+continuousA # 1
f2 = @formula 0 ~ 1+continuousB # 1
m_tul = fit(UnfoldModel, Dict("eventA"=>(f1,b1),"eventB"=>(f2,b2)), evts, data,eventcolumn="type")

m_tul.modelfit.estimate .= [0 -1 0 6]
eff = Unfold.effects(Dict(:continuousA => [0,1]),m_tul)
eff = Unfold.effects(Dict(:continuousA => [0,1],:continuousB =>[0.5]),m_tul)

@test eff.yhat[3] == eff.yhat[4]
@test eff.yhat[1] == 0.
@test eff.yhat[2] == -1.
@test eff.yhat[3] == 3
end

@testset "timeexpansion, Interactions, two events" begin
	data, evts = loadtestdata("test_case_4a") #
	evts[!,:continuousA] = rand(MersenneTwister(42),nrow(evts))
	evts[!,:continuousB] = ["m","x"][Int.(1 .+ round.(rand(MersenneTwister(43),nrow(evts))))]

	b1 = firbasis(τ = (0.0, 0.02), sfreq = 20, name = "basisA")
	b2 = firbasis(τ = (1.0, 1.02), sfreq = 20, name = "basisB")
	f1 = @formula 0 ~ 1+continuousA*continuousB # 1
	f2 = @formula 0 ~ 1+continuousB # 1
	m_tul = fit(UnfoldModel, Dict("eventA"=>(f1,b1),"eventB"=>(f2,b2)), evts, data,eventcolumn="type")

	m_tul.modelfit.estimate .= [0, -1, 0, 2., 0.,0.]'

	eff = Unfold.effects(Dict(:continuousA => [0,1]),m_tul)
	@test size(eff,1) == 4
	@test all(eff.basisname[1:2] .== "basisA")
	@test all(eff.basisname[4:end] .== "basisB")
	@test eff.yhat ≈ [0., mean(evts.continuousB[evts.type .== "eventA"].=="x") * coef(m_tul)[4] + 1*coef(m_tul)[2], 0., 0.]

	eff = Unfold.effects(Dict(:continuousB => ["m","x"]),m_tul)
	@test eff.yhat[1] == -eff.yhat[2]
	@test all(eff.yhat[3:4] .≈ 0.)

	eff = Unfold.effects(Dict(:continuousA => [0,1],:continuousB => ["m","x"]),m_tul)
	@test eff.yhat[3:4] == [-1,1]
	@test all(eff.yhat[1:2, 5:end] .== 0)
end

max_stars_repo_name: unfoldtoolbox/jlunfold

max_issues_repo_name: unfoldtoolbox/jlunfold

max_forks_repo_name: unfoldtoolbox/jlunfold


text:
spacing(::Periodic, n) = 1 / n
spacing(::NonPeriodic, n) = 1 / (n-1)

coordinate(top, n) = UnitRange(0, n-1) .* spacing(top, n)


text:
@with_kw mutable struct Scanner
    #Main
    B0::Real=1.5      # Main magnetic field [T]
    B1::Real=10e-6     # Max RF amplitude [T]
    Gmax::Real=60e-3  # Max Gradient [T/m]
    Smax::Real=500    # Max Slew-rate [mT/m/ms] or [T/m/s]
    #Sampling
    ADC_Δt::Real=2e-6 # ADC raster time
    seq_Δt::Real=1e-5 # Seq-block raster time
    GR_Δt::Real=1e-5  # GR raster time
    RF_Δt::Real=1e-6  # RF raster time
    #Secondary
    RF_ring_down_T::Real=20e-6  # RF ring down time [s]
    RF_dead_time_T::Real=100e-6 # RF dead tim [s]
    ADC_dead_time_T::Real=10e-6 # ADC dead time [s]
end

#Functions that check that Sequence satisfy hardware requirements:
#check_sys_req(seq::Sequence,sys::Scanner)


text:
# using Stars

@testset "gdalinfo" begin
    file = "data/MOD13A2_Henan_2015_2020_10km.tif"
    ga = st_read(file)
    @test gdalinfo(ga) == gdalinfo(file)
end


text:
using Plots

# Model
include_model("quadruped")
# model = free_time_model(model)

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)

# Horizon
T = 2

# Time step
h = 0.01

function initial_configuration(model::Quadruped, θ1, θ2, θ3)
    q1 = zeros(model.nq)
    q1[3] = pi / 2.0
    q1[4] = -θ1
    q1[5] = θ2

    q1[8] = -θ1
    q1[9] = θ2

    q1[2] = model.l_thigh1 * cos(q1[4]) + model.l_calf1 * cos(q1[5])

    q1[10] = -θ3
    q1[11] = acos((q1[2] - model.l_thigh2 * cos(q1[10])) / model.l_calf2)

    q1[6] = -θ3
    q1[7] = acos((q1[2] - model.l_thigh2 * cos(q1[6])) / model.l_calf2)

    return q1
end

θ1 = pi / 3.5
θ2 = pi / 3.5
θ3 = pi / 3.5

q1 = initial_configuration(model, θ1, θ2, θ3)
visualize!(vis, model, [q1])

# control
# ul <= u <= uu
_uu = Inf * ones(model.m)
_uu[model.idx_u] .= 33.5 * h
# _uu[end] = 2.0 * h
_ul = zeros(model.m)
_ul[model.idx_u] .= -33.5 * h
# _ul[end] = 0.75 * h
ul, uu = control_bounds(model, T, _ul, _uu)

xl, xu = state_bounds(model, T,
    x1 = [q1; q1],
    xT = [q1; q1])

# penalty on slack variable
obj_penalty = PenaltyObjective(1.0e4, model.m)

obj = MultiObjective([obj_penalty])

# Constraints
include_constraints(["stage", "contact"])

con_contact = contact_constraints(model, T)

con = multiple_constraints([con_contact])#, con_loop, con_pinned1, con_pinned2])

# Problem
prob = trajectory_optimization_problem(model,
               obj,
               T,
			   h = h,
               xl = xl,
               xu = xu,
               ul = ul,
               uu = uu,
               con = con)

# trajectory initialization
q_ref = linear_interpolation(q1, q1, T+1)
x0 = configuration_to_state(q_ref)
u0 = [1.0e-3 * rand(model.m) for t = 1:T-1] # random controls

# Pack trajectories into vector
z0 = pack(x0, u0, prob)

# Solve
@time z̄, info = solve(prob, copy(z0),
    nlp = :ipopt,
    tol = 1.0e-2, c_tol = 1.0e-2,
	max_iter = 2000,
    time_limit = 60 * 2, mapl = 5)

check_slack(z̄, prob)
x̄, ū = unpack(z̄, prob)
q̄ = state_to_configuration(x̄)

vis = Visualizer()
render(vis)
visualize!(vis, model, state_to_configuration(x̄), Δt = h)

u_stand = ū[1][1:model.m]
@save joinpath(@__DIR__, "quadruped_stand_100Hz.jld2") u_stand

max_issues_repo_name: thowell/DirectMotionPlanning


text:
## SUPERVISED

# This implementation of scitype for models and measures is highly
# experimental and not part of the public API

# only used in this file
const DefaultConvention = ScientificTypes.DefaultConvention
const STB = ScientificTypes.ScientificTypesBase

struct SupervisedScitype{input_scitype, target_scitype, prediction_type} end

STB.scitype(model::Deterministic, ::DefaultConvention) =
    SupervisedScitype{input_scitype(model),
                    target_scitype(model),
                    :deterministic}

STB.scitype(model::Probabilistic, ::DefaultConvention) =
    SupervisedScitype{input_scitype(model),
                    target_scitype(model),
                    :probabilistic}

STB.scitype(model::Interval, ::DefaultConvention) =
    SupervisedScitype{input_scitype(model),
                    target_scitype(model),
                    :interval}

function Base.getproperty(::SupervisedScitype{input_scitype, target_scitype, prediction_type},
                          field::Symbol) where {input_scitype, target_scitype, prediction_type}
    if field === :input_scitype
        return input_scitype
    elseif field === :target_scitype
        return target_scitype
    elseif field === :prediction_type
        return prediction_type
    else
        throw(ArgumentError("Unsupported property. "))
    end
end

# crashes julia:
# MLJBase.getproperty(S::Type{<:SupervisedScitype}, p::Symbol) =
#     getproperty(S(), p)

Base.propertynames(::SupervisedScitype) = (:input_scitype, :target_scitype, :prediction_type)

_as_named_tuple(s::SupervisedScitype) =
    NamedTuple{(:input_scitype, :target_scitype, :prediction_type)}((s.input_scitype, s.target_scitype, s.prediction_type))

function Base.show(io::IO, ::MIME"text/plain", S::SupervisedScitype)
    show(io, MIME("text/plain"), _as_named_tuple(S))
end

## UNSUPERVISED

struct UnsupervisedScitype{input_scitype, output_scitype} end

STB.scitype(model::Unsupervised, ::DefaultConvention) =
    UnsupervisedScitype{input_scitype(model),
                      MLJBase.output_scitype(model)}

function Base.getproperty(::UnsupervisedScitype{input_scitype, output_scitype},
                          field::Symbol) where {input_scitype, output_scitype}
    if field === :input_scitype
        return input_scitype
    elseif field === :output_scitype
        return output_scitype
    else
        throw(ArgumentError("Unsupported property. "))
    end
end

# crashes julia:
# MLJBase.getproperty(U::Type{<:UnsupervisedScitype}, p::Symbol) =
#     getproperty(U(), p)

Base.propertynames(::UnsupervisedScitype) = (:input_scitype, :output_scitype)

_as_named_tuple(s::UnsupervisedScitype) =
    NamedTuple{(:input_scitype, :output_scitype)}(
        (s.input_scitype, s.output_scitype))

function Base.show(io::IO, ::MIME"text/plain", S::UnsupervisedScitype)
    show(io, MIME("text/plain"), _as_named_tuple(S))
end


text:
# -------------------------------------------------------------------------------------
const _REGEXS = Dict()
function _set_regexs!()
    empty!(_REGEXS)

    # escape
    ESC_SEPS = _hex_escaped_seps()
    hex_plsep = ESC_SEPS[:PARAMS_LSEP]
    hex_prsep = ESC_SEPS[:PARAMS_RSEP]
    hex_elsep = ESC_SEPS[:ELEMT_SEP]
    hex_psep = ESC_SEPS[:PAIR_SEP]
    hex_extsep = ESC_SEPS[:EXT_SEP]

    # inval_chars
    inval_chars = [hex_plsep, hex_prsep, hex_elsep, hex_psep]
    val_char_rstr = "[^$(join(inval_chars))]"
    noextsep_char_rstr = "[^$(join([inval_chars; hex_extsep]))]"

    # ------------------------------------------------------------
    # EXT
    short_ext_rstr = string(
        "(?<extstr>",
            "$(hex_extsep)$(noextsep_char_rstr)*\$",
        ")",
    )
    _REGEXS[:SHORT_EXT_SPLIT_REGEX] = Regex(short_ext_rstr)

    long_ext_rstr = string(
        "$(hex_prsep)",
        "(?<extstr>",
            "(?:$(hex_extsep)$(noextsep_char_rstr)*)*\$",
        ")",
    )
    _REGEXS[:LONG_EXT_SPLIT_REGEX] = Regex(long_ext_rstr)

    # ------------------------------------------------------------
    # HEAD # IMPORTANT!!! it assumes the extension is being removed
    head_rstr = string(
        "^",
        "(?<headstr>",
            "(?:",
                "$(val_char_rstr)+",
                "(?:$(hex_elsep)|\$)",
            ")*",
        ")",
    )
    _REGEXS[:HEAD_SPLIT_REGEX] = Regex(head_rstr)

    # ------------------------------------------------------------
     # PARAMS # IMPORTANT!!! it assumes the extension is being removed
     params_rstr = string(
        "$(hex_plsep)",
            "(?<paramstr>",
                "(?:$(val_char_rstr)+$(hex_psep)$(val_char_rstr)+$(hex_elsep)?)+",
            ")",
        "$(hex_prsep)",
        "\$"
    )
    _REGEXS[:PARAMS_SPLIT_REGEX] = Regex(params_rstr)

end

# -------------------------------------------------------------------------------------
function _parse_regex(fname::String)
    _check__SEPS()

    _get(m, gk) = haskey(m, gk) ? m[gk] : ""
    _get(m::Nothing, gk) = ""

    # remove ext
    # try long
    ext_r = _REGEXS[:LONG_EXT_SPLIT_REGEX]
    m = match(ext_r, fname)
    extstr = _get(m, :extstr)
    if isempty(extstr)
        # try short
        ext_r = _REGEXS[:SHORT_EXT_SPLIT_REGEX]
        m = match(ext_r, fname)
        extstr = _get(m, :extstr)
    end
    replace_r = Regex("$(extstr)\$")
    fname = replace(fname, replace_r => "")

    # remove head
    head_r = _REGEXS[:HEAD_SPLIT_REGEX]
    m = match(head_r, fname)
    headstr = _get(m, :headstr)
    fname = replace(fname, head_r => "")

    # remove param body
    param_r = _REGEXS[:PARAMS_SPLIT_REGEX]
    m = match(param_r, fname)
    paramsstr = _get(m, :paramstr)
    fname = replace(fname, param_r => "")

    digest = fname
    return (;headstr, paramsstr, extstr, digest)
end


text:
using IcanProj
using NetCDF
using DataFrames
using PyPlot
using Statistics

# Global settings

path = "/data04/jmg/fsm_simulations/netcdf/fsmres_open/results_32"

# Load latent heat fluxes

variable = "latmo"

file_fine = joinpath(path, "$(variable)_1km.nc")

file_coarse = joinpath(path, "$(variable)_50km.nc")

df_links = link_results(file_fine, file_coarse)

latmo_coarse, latmo_aggregated = unify_results(file_fine, file_coarse, df_links, variable)

# Load snow water equivalents

variable = "swe"

file_fine = joinpath(path, "$(variable)_1km.nc")

file_coarse = joinpath(path, "$(variable)_50km.nc")

df_links = link_results(file_fine, file_coarse)

swe_coarse, swe_aggregated = unify_results(file_fine, file_coarse, df_links, variable)


text:
module GPLikelihoods

using Distributions
using AbstractGPs
using Random
using Functors
using StatsFuns: logistic, softmax

import Distributions

export BernoulliLikelihood,
    BinomialLikelihood,
    CategoricalLikelihood,
    ExponentialLikelihood,
    GaussianLikelihood,
    HeteroscedasticGaussianLikelihood,
    PoissonLikelihood,
    StudentsTLikelihood

# Likelihoods
include("likelihoods/bernoulli.jl")
include("likelihoods/binomial.jl")
include("likelihoods/categorical.jl")
include("likelihoods/exponential.jl")
include("likelihoods/gaussian.jl")
include("likelihoods/poisson.jl")
include("likelihoods/studentsT.jl")

end # module


text:
# destructive_add!
# ------------------------------------------------------------------------------

# destructive_add!{C}(ex::Number, c::Number, x::Number) = ex + c*x

#=
    Number
=#

JuMP.destructive_add!(ex::Number, c::C, x::ParameterRef) where C<:Number = PAE{C}(GAEv{C}(ex),GAEp{C}(zero(C), x => c))
JuMP.destructive_add!(ex::Number, x::ParameterRef, c::C) where C<:Number = JuMP.destructive_add!(ex, c, x)
JuMP.destructive_add!(ex::Number, c::C, x::PAE) where C<:Number = PAE{C}(JuMP.destructive_add!(ex, c, x.v), JuMP.destructive_add!(0.0, c, x.p))

#=
    VariableRef
=#

JuMP.destructive_add!(ex::JuMP.VariableRef, c::C, x::ParameterRef) where C<:Number = PAE{C}(GAEv{C}(zero(C), ex => one(C)), GAEp{C}(zero(C), x => c))
JuMP.destructive_add!(ex::JuMP.VariableRef, x::ParameterRef, c::C) where C<:Number = JuMP.destructive_add!(ex, c, x)

#=
    Parameter
=#

JuMP.destructive_add!(ex::ParameterRef, c::Number, x::Number) = c * x + ex

JuMP.destructive_add!(ex::ParameterRef, c::C, x::JuMP.VariableRef) where C<:Number = PAE{C}(GAEv{C}(zero(C), x => c), GAEp{C}(zero(C), ex => one(C)))
JuMP.destructive_add!(ex::ParameterRef, x::JuMP.VariableRef, c::Number) = JuMP.destructive_add!(ex, c, x)

JuMP.destructive_add!(ex::ParameterRef, c::C, x::ParameterRef) where {C<:Number} = PAE{C}(zero(GAEv{C}),  GAEp{C}(zero(C), ex => one(C), x => c))
JuMP.destructive_add!(ex::ParameterRef, x::ParameterRef, c::C) where {C<:Number} = JuMP.destructive_add!(ex, x, c)

#=
    GAEp
=#

JuMP.destructive_add!(aff::GAEp{C}, c::Number, x::Number) where {C} = PAE{C}(GAEv{C}(c*x), aff)

JuMP.destructive_add!(aff::GAEp{C}, x::Union{JuMP.VariableRef, GAEv{C}}, c::Number) where C = JuMP.destructive_add!(aff, c, x)
JuMP.destructive_add!(aff::GAEp{C}, c::Number, x::Union{JuMP.VariableRef, GAEv{C}}) where C = PAE{C}(GAEv{C}(zero(C), x => convert(C, c)), aff)

#=
    GAEv
=#

JuMP.destructive_add!(aff::GAEv{C}, x::Union{ParameterRef, GAEp{C}}, c::Number) where C = JuMP.destructive_add!(aff, c, x)
JuMP.destructive_add!(aff::GAEv{C}, c::Number, x::Union{ParameterRef, GAEp{C}}) where C = PAE{C}(aff, GAEp{C}(zero(C), x => convert(C, c)))

#=
    PAE
=#

JuMP.destructive_add!(aff::PAE, x::Union{JuMP.VariableRef, GAEv, ParameterRef, GAEp}, c::Number) = JuMP.destructive_add!(aff, c, x)
function JuMP.destructive_add!(aff::PAE, c::Number, x::Union{JuMP.VariableRef, GAEv})
    if !iszero(c)
        aff.v = JuMP.destructive_add!(aff.v, c, x)
    end
    aff
end
function JuMP.destructive_add!(aff::PAE, c::Number, x::Union{ParameterRef, GAEp})
    if !iszero(c)
        aff.p = JuMP.destructive_add!(aff.p, c, x)
    end
    aff
end
function JuMP.destructive_add!(aff::PAE, c::Number, x::Number)
    if !iszero(c)
        aff.v = JuMP.destructive_add!(aff.v, c, x)
    end
    aff
end

function JuMP.add_to_expression!(aff::PAE, other::Number)
    JuMP.add_to_expression!(aff.v, other)
end
function JuMP.add_to_expression!(aff::PAE, new_var::JuMP.VariableRef, new_coef)
    JuMP.add_to_expression!(aff.v, new_coef, new_var)
end
function JuMP.add_to_expression!(aff::PAE, new_coef, new_var::JuMP.VariableRef)
    JuMP.add_to_expression!(aff.v, new_coef, new_var)
end
function JuMP.add_to_expression!(aff::PAE, new_var::JuMP.VariableRef)
    JuMP.add_to_expression!(aff.v, new_var)
end
function JuMP.add_to_expression!(aff::PAE, new_param::ParameterRef, new_coef)
    JuMP.add_to_expression!(aff.p, new_coef, new_param)
end
function JuMP.add_to_expression!(aff::PAE, new_param::ParameterRef)
    JuMP.add_to_expression!(aff.p,new_param)
end
function JuMP.add_to_expression!(aff::PAE, new_coef, new_param::ParameterRef)
    JuMP.add_to_expression!(aff.p, new_coef, new_param)
end
function JuMP.add_to_expression!(lhs_aff::PAE, rhs_aff::PAE)
    JuMP.add_to_expression!(lhs_aff.p, rhs_aff.p)
    JuMP.add_to_expression!(lhs_aff.v, rhs_aff.v)
end


text:
using PythonCall, Test, Dates, Aqua

# The unbound_args test fails on methods with signature like foo(::Type{Tuple{Vararg{V}}}) where V
# Seems like a bug.
Aqua.test_all(PythonCall, unbound_args=false)

@testset "PythonCall.jl" begin
    @testset "abstract" begin
        include("abstract.jl")
    end
    @testset "concrete" begin
        include("concrete.jl")
    end
    @testset "convert" begin
        include("convert.jl")
    end
    @testset "jlwrap" begin
        include("jlwrap.jl")
    end
end