using Netpbm, ColorTypes, FixedPointNumbers, IndirectArrays, ImageCore, FileIO
using Test

@testset "IO" begin
    workdir = joinpath(tempdir(), "Images")
    isdir(workdir) && rm(workdir, recursive=true)
    mkdir(workdir)

    @testset "Gray pgm" begin
        af = rand(2, 3)
        for T in (N0f8, N4f12, N0f16,
                  Gray{N0f8}, Gray{N4f12}, Gray{N0f16})
            ac = convert(Array{T}, af)
            fn = File(format"PGMBinary", joinpath(workdir, "3by2.pgm"))
            Netpbm.save(fn, ac)
            b = Netpbm.load(fn)
            @test b == ac
        end
        a8 = convert(Array{N0f8}, af)
        for T in (Float32, Float64, Gray{Float32}, Gray{Float64})
            ac = convert(Array{T}, af)
            fn = File(format"PGMBinary", joinpath(workdir, "3by2.pgm"))
            Netpbm.save(fn, ac)
            b = Netpbm.load(fn)
            @test b == a8
        end

        # Issue #11
        @test_nowarn load("test_header_space.pgm")
    end

    @testset "Color ppm" begin
        af = rand(RGB{Float64}, 2, 3)
        for T in (RGB{N0f8}, RGB{N4f12}, RGB{N0f16})
            ac = convert(Array{T}, af)
            fn = File(format"PPMBinary", joinpath(workdir, "3by2.ppm"))
            Netpbm.save(fn, ac)
            b = Netpbm.load(fn)
            @test b == ac
        end
        a8 = convert(Array{RGB{N0f8}}, af)
        for T in (RGB{Float32}, RGB{Float64}, HSV{Float64})
            ac = convert(Array{T}, af)
            fn = File(format"PPMBinary", joinpath(workdir, "3by2.ppm"))
            Netpbm.save(fn, ac)
            b = Netpbm.load(fn)
            @test b == a8
        end
    end

    @testset "Colormap" begin
        datafloat = reshape(range(0.5, stop=1.5, length=6), 2, 3)
        dataint = map(x->round(UInt8, x), 254*(datafloat .- 0.5) .+ 1) # ranges 1 to 255
        # build our colormap
        b = RGB(0,0,1)
        w = RGB(1,1,1)
        r = RGB(1,0,0)
        cmaprgb = Array{RGB{N0f8}}(undef, 255)
        f = range(0, stop=1, length=128)
        cmaprgb[1:128] = [(1-x)*b + x*w for x in f]
        cmaprgb[129:end] = [(1-x)*w + x*r for x in f[2:end]]
        img = IndirectArray(dataint, cmaprgb)
        fn = File(format"PPMBinary", joinpath(workdir,"cmap.ppm"))
        Netpbm.save(fn, img)
        imgr = Netpbm.load(fn)
        @test imgr == img
        cmaprgb = Array{RGB}(undef, 255) # poorly-typed cmap, Images issue #336
        @test !isconcretetype(eltype(cmaprgb))
        cmaprgb[1:128] = RGB{N0f16}[(1-x)*b + x*w for x in f]
        cmaprgb[129:end] = RGB{N4f12}[(1-x)*w + x*r for x in f[2:end]]
        img = IndirectArray(dataint, cmaprgb)
        @test_throws ErrorException Netpbm.save(fn, img) # widens to unsupported type
        cmaprgb[129:end] = RGB{N0f8}[(1-x)*w + x*r for x in f[2:end]]
        img = IndirectArray(dataint, cmaprgb)
        Netpbm.save(fn, img)
        imgr = Netpbm.load(fn)
        @test imgr == img
    end

    # Images issue #256
    @testset "Clamping" begin
        A = rand(2,3)
        A[1,1] = -0.4
        fn = File(format"PGMBinary", joinpath(workdir, "2by3.pgm"))
        @test_throws InexactError Netpbm.save(fn, A)
        Netpbm.save(fn, A, mapf=clamp01nan)
        B = Netpbm.load(fn)
        A[1,1] = 0
        @test B == Gray{N0f8}.(A)
    end
end

nothing


export
    SyncLoopy,
    AsyncLoopy,
    CoherentLoopy

"""
    SyncLoopy(range_size = 10)

A window filter that uses a synchronous window with LoopyBP with the given range size.
"""
SyncLoopy(range_size = 10) = WindowFilter(SyncWindow(), LoopyBP(range_size))

"""
    AsyncLoopy(range_size = 10, T = Float64)

A window filter that uses an asynchronous window with LoopyBP with the given range size.
`T` represents the time type and must be the same as used in creation of the runtime.
"""
AsyncLoopy(range_size = 10, T = Float64)  = WindowFilter(AsyncWindow{T}(), LoopyBP(range_size))

"""
    CoherentLoopy(range_size = 10, T = Float64)

A window filter that uses a coherent window with LoopyBP with the given range size.
`T` represents the time type and must be the same as used in creation of the runtime.
"""
CoherentLoopy(range_size = 10, T = Float64)  = WindowFilter(CoherentWindow{T}(), LoopyBP(range_size))


function NewtonRaphsonSolver()
    return NewtonRaphsonSolver(())
end

function NewtonRaphsonSolver(arg0::jdouble)
    return NewtonRaphsonSolver((jdouble,), arg0)
end

function solve(obj::NewtonRaphsonSolver, arg0::jint, arg1::UnivariateDifferentiableFunction, arg2::jdouble, arg3::jdouble)
    return jcall(obj, "solve", jdouble, (jint, UnivariateDifferentiableFunction, jdouble, jdouble), arg0, arg1, arg2, arg3)
end


@testset "word vectors" begin
    wn = W2VNetwork(doc, 5, 100)
    wv = wn.wv
    a = cosine_similar_words(wn, "something", 2)
    println(a)

end


module WildBootTests
export BootTestResult, wildboottest, AuxWtType, PType, MAdjType, DistStatType,
       teststat, stattype, p, padj, reps, repsfeas, nbootclust, dof, dof_r, plotpoints, peak, CI, dist, statnumer, statvar, auxweights

using LinearAlgebra, Random, Distributions, SortingAlgorithms, LoopVectorization

include("StrBootTest.jl")
include("utilities.jl")
include("estimators.jl")
include("init.jl")
include("WRE.jl")
include("nonWRE.jl")
include("plot-CI.jl")
include("interface.jl")

# top-level computation routine for OLS/ARubin (and score BS on IV/2SLS); split off to reduce latency when just doing WRE
function boottestOLSARubin!(o::StrBootTest{T}) where T
  if !o.initialized
		Init!(o)
	elseif !o.null
		NoNullUpdate!(o)
		return
  end

  if o.Nw > 1  # if more than one weight group to save memory, make on every call to boottest(), not just once in Init!()
		Random.seed!(o.rng,o.seed)
		MakeWildWeights!(o, last(o.WeightGrp[1])-1, first=true)
	end

  MakeInterpolables!(o)  # make stuff that depends linearly on r, possibly by interpolating, for first weight group

  for w ∈ 1:o.Nw  # do group 1 first because it includes col 1, which is all that might need updating in constructing CI in WCU
		w > 1 && MakeWildWeights!(o, length(o.WeightGrp[w]), first=false)

		MakeNonWREStats!(o, w)

		!o.bootstrapt && UpdateBootstrapcDenom!(o, w)
  end

  o.BFeas = isnan(o.dist[1]) ? 0 : sum(.!(isnan.(o.dist) .| isinf.(o.dist))) - 1
  o.distCDR = zeros(T,0,0)
  nothing
end

# top-level computation routine for non-ARubin WRE; split off to reduce latency when just doing other tests
function boottestWRE!(o::StrBootTest{T}) where T
  if !o.initialized
		Init!(o)
	elseif !o.null
		NoNullUpdate!(o)
		return
  end

  if o.Nw > 1  # if more than one weight group to save memory, make on every call to boottest(), not just once in Init!()
		Random.seed!(o.rng,o.seed)
		MakeWildWeights!(o, last(o.WeightGrp[1])-1, first=true)
	end

  PrepWRE!(o)

  for w ∈ 1:o.Nw  # do group 1 first because it includes col 1, which is all that might need updating in constructing CI in WCU
		w > 1 && MakeWildWeights!(o, length(o.WeightGrp[w]), first=false)

		MakeWREStats!(o, w)

		!o.bootstrapt && UpdateBootstrapcDenom!(o, w)
  end

  o.BFeas = isnan(o.dist[1]) ? 0 : sum(.!(isnan.(o.dist) .| isinf.(o.dist))) - 1
  o.distCDR = zeros(T,0,0)
  nothing
end

# if not imposing null and we have returned to boottest!(), then dof=1 or 2; we're plotting or finding CI, and only test stat, not distribution, changes with r
function NoNullUpdate!(o::StrBootTest{T} where T)
  if o.WREnonARubin
		o.numer[:,1] = o.R * o.DGP.Rpar * o.β̂s[1] - o.r
  elseif o.ARubin
		EstimateARubin!(o.DGP, o, o.r)
		o.numer[:,1] = o.v_sd * @view o.DGP.β̈[o.kX₁+1:end,:]  # coefficients on excluded instruments in ARubin OLS
  else
		o.numer[:,1] = o.v_sd * (o.R * (o.ML ? o.β̈ : o.M.β̈) - o.r)  # Analytical Wald numerator; if imposing null then numer[:,1] already equals this. If not, then it's 0 before this
  end
  o.dist[1] = isone(o.dof) ? o.numer[1] / sqrt(o.statDenom[1]) : o.numer[:,1]'invsym(o.statDenom)*o.numer[:,1]
	nothing
end

# compute bootstrap-c denominator from all bootstrap numerators
function UpdateBootstrapcDenom!(o::StrBootTest{T} where T, w::Integer)
  if isone(w)
		tmp = o.numer[:,1]
		o.statDenom = o.numer * o.numer' - tmp * tmp'
		o.numersum = rowsum(o.numer) - tmp
  else
		statDenom .+= o.numer * o.numer'
		numersum .+= rowsum(o.numer)
  end
  if w == o.Nw  # last weight group?
		o.statDenom .= (o.statDenom - o.numersum * o.numersum' / o.B) / o.B
		if o.sqrt
			o.dist .= o.numer ./ sqrtNaN.(o.statDenom)
		else
			negcolquadform!(o.dist, -invsym(o.statDenom), o.numer)  # to reduce latency by minimizing @tturbo instances, work with negative of colquadform in order to fuse code with colquadformminus!
		end
  end
	nothing
end

# include("precompile_WildBootTests.jl")  # source: https://timholy.github.io/SnoopCompile.jl/stable/snoopi_deep_parcel/#SnoopCompile.write
# _precompile_()

end

using StableRNGs, Random, BenchmarkTools
N=1_000_00; G=40; k=2; l=4
Random.seed!(1231)
β=rand(); γ=rand(k); Π=rand(l)
W = rand(N,l)
u₂ = randn(N)
y₂ = W * Π + u₂
u₁ = u₂ + randn(N)
Z = rand(N,k)
y₁ = y₂ * β + Z * γ + u₁
ID = floor.(Int8, collect(0:N-1) / (N/G))
R = [zeros(1,k) 1]; r = [0]
WildBootTests.wildboottest(Float32, R,.46; R1 = [1 zeros(1,k)], r1=[0], resp=y₁, predexog=Z, predendog=y₂, inst=W, clustid=ID, reps=999, issorted=true, rng=StableRNG(1231), LIML=true)

N=1_000_000; G=40; k=12; l=40
Random.seed!(1231)
β=rand(); γ=rand(k); Π=rand(l)
W = rand(N,l)
u₂ = randn(N)
y₂ = W * Π + u₂
u₁ = u₂ + randn(N)
Z = rand(N,k)
y₁ = y₂ * β + Z * γ + u₁
ID = floor.(Int8, collect(0:N-1) / (N/G))
R = [zeros(1,k) 1]; r = [0]
@btime WildBootTests.wildboottest(Float32, R,.4; resp=y₁, predexog=Z, predendog=y₂, inst=W, clustid=ID, reps=9999, issorted=true, getCI=false);

@btime WildBootTests.wildboottest(Float32, R,.4; resp=y₁, predexog=Z, predendog=y₂, inst=W, clustid=ID, reps=9999, issorted=true, getCI=true);


##
using Revise
using GLMakie
using Pkg
Pkg.activate(".")
using ProteinStructureReconstruction
using HDF5
using Images
import Statistics: median
import GeometryBasics: Point2f0
##
filepath = "./Data/Fed_X63_Z3_SIM.h5"
h5file = h5open(filepath, "r")
##
# TODO: don't load everything in memory at once -  just draw overlay first, then
##
x_range, y_range = 2000:2300, 1600:1900;
img_data = h5file["Data"][x_range, y_range, :, :] |> x -> reinterpret(N0f16, x);


# Compared to cudnnAddTensor!(copy(a),b), cudnnOpTensor is ~50% faster on
# (14,14,256,32)+(1,1,256,1), ~50% slower on (1,1,100,100)+(1,1,100,1) Unlike cudnnAddTensor
# it supports all broadcasting shapes up to ndims=5 as described in the documentation.

"""
    cudnnOpTensor(x1, x2; op, compType, nanOpt, alpha1, alpha2)
    cudnnOpTensor(x1, x2, d::cudnnOpTensorDescriptor; alpha1, alpha2)
    cudnnOpTensor!(y, x1, x2; op, compType, nanOpt, alpha1, alpha2, beta)
    cudnnOpTensor!(y, x1, x2, d::cudnnOpTensorDescriptor; alpha1, alpha2, beta)

Return the result of the specified broadcasting operation applied to `x1` and `x2`.
Optionally `y` holds the result and `d` specifies the operation. Each dimension of the input
tensor `x1` must match the corresponding dimension of the destination tensor `y`, and each
dimension of the input tensor `x2` must match the corresponding dimension of the destination
tensor `y` or must be equal to 1. Keyword arguments:

* `alpha1=1, alpha2=1, beta=0` are used for scaling, i.e. `y .= beta*y .+ op.(alpha1*x1, alpha2*x2)`

Keyword arguments used when `cudnnOpTensorDescriptor` is not specified:

* `op = CUDNN_OP_TENSOR_ADD`, ADD can be replaced with MUL, MIN, MAX, SQRT, NOT; SQRT and NOT performed only on x1; NOT computes 1-x1
* `compType = (eltype(x1) <: Float64 ? Float64 : Float32)`: Computation datatype (see cudnn docs for available options)
* `nanOpt = CUDNN_NOT_PROPAGATE_NAN`: NAN propagation policy. The other option is `CUDNN_PROPAGATE_NAN`.
"""
cudnnOpTensor, cudnnOpTensor!

# Public methods:
cudnnOpTensor(x1,x2; o...)    = cudnnOpTensorWithDefaults(x1,x2; o...)
cudnnOpTensor!(y,x1,x2; o...) = cudnnOpTensorWithDefaults(x1,x2; y, o...)
cudnnOpTensor(x1,x2,d::cudnnOpTensorDescriptor; o...)    = cudnnOpTensorWithDefaults(x1,x2; opTensorDesc=d, o...)
cudnnOpTensor!(y,x1,x2,d::cudnnOpTensorDescriptor; o...) = cudnnOpTensorWithDefaults(x1,x2; y, opTensorDesc=d, o...)

# Private method:
function cudnnOpTensorWithDefaults(
    x1, x2;
    y = similar(x1),
    op::cudnnOpTensorOp_t = CUDNN_OP_TENSOR_ADD,
    compType::DataType = (eltype(x1) <: Float64 ? Float64 : Float32),
    nanOpt::cudnnNanPropagation_t = CUDNN_NOT_PROPAGATE_NAN,
    opTensorDesc::cudnnOpTensorDescriptor = cudnnOpTensorDescriptor(op, cudnnDataType(compType), nanOpt),
    alpha1::Real = 1,
    alpha2::Real = 1,
    beta::Real = 0,
    x1Desc::cudnnTensorDescriptor = cudnnTensorDescriptor(x1),
    x2Desc::cudnnTensorDescriptor = cudnnTensorDescriptor(x2),
    yDesc::cudnnTensorDescriptor = cudnnTensorDescriptor(y)
)
    @assert ndims(x1) <= 5
    @assert size(y) == size(x1)
    @assert all(size(x2,i) == size(x1,i) || size(x2,i) == 1 for i in 1:ndims(x2))
    T = eltype(x1)
    alpha1, alpha2, beta = scalingParameter(T,alpha1), scalingParameter(T,alpha2), scalingParameter(T,beta)
    cudnnOpTensorAD(x1, x2; opTensorDesc, alpha1, x1Desc, alpha2, x2Desc, beta, yDesc, y)
end

# AD method: This method aids gradient definition, please do not remove!
function cudnnOpTensorAD(x1, x2; opTensorDesc, alpha1, x1Desc, alpha2, x2Desc, beta, yDesc, y)
    cudnnOpTensor(handle(), opTensorDesc, alpha1, x1Desc, x1, alpha2, x2Desc, x2, beta, yDesc, y)
    return y
end

# Deprecated:
function cudnnOpTensor(op::cudnnOpTensorOp_t,
                       A::DenseCuArray{T,N}, B::DenseCuArray{T,N}, C::DenseCuArray{T,N};
                       alpha1=true, alpha2=true, beta=false) where {T,N}
    @warn "cudnnOpTensor(op,A,B,C) is deprecated, please use one of the methods in `@doc cudnnOpTensor`." maxlog=1
    cudnnOpTensorWithDefaults(A, B; y=C, op, alpha1, alpha2, beta)
end


function H5ReadPhysicalParameter2D(
  h5obj::PyCall.PyObject;
  phptype::H5PhPType,
  xkey::String)

  h5geo = pyimport("h5geopy._h5geo")

  nullValue = h5obj.getNullValue()

  if phptype == VELOCITY
    php = transpose(h5obj.getTrace(0, typemax(Int), 0, typemax(Int), "km/s"))
    # Slowness squared [s^2/km^2]
    php = (1f0 ./ php).^2
  elseif phptype == DENSITY
    php = transpose(h5obj.getTrace(0, typemax(Int), 0, typemax(Int), "g/cm^3"))
  elseif phptype == THETA || phptype == PHI
    php = transpose(h5obj.getTrace(0, typemax(Int), 0, typemax(Int), "rad"))
  else
    php = transpose(h5obj.getTrace(0, typemax(Int), 0, typemax(Int)))
  end

  if isempty(php)
    @error "Unable to read PhysicalParameter. Probably data units incorrect/missing\n"
    return
  end

  if ndims(php) != 2
    @error "PhysicalParameter must be two dimensional array\n"
    return
  end

  # JUDI falls down if NaN are passed
  meanValue = mean(filter(!isnan, php))
  php[isnan.(php)] .= meanValue

  x = h5obj.getTraceHeader(xkey, 0, typemax(Int), h5obj.getLengthUnits(), "m")
  if length(x) < 2
    @error "$(xkey) can't have less than 2 points\n"
    return
  end

  # in case PhysicalParameter is not sorted
  x_ind, x = h5geo.sortv(x)
  php = php[x_ind .+ 1, :]   # C++ returned indexes starts from 0

  # Set up PhysicalParameter structure
  n = size(php)   # (x,y,z) or (x,z)
  d = (x[2]-x[1], abs(h5obj.getSampRate("m")))
  o = (x[1], h5obj.getSRD("m") * (-1))

  return JUDI.PhysicalParameter(php, n, d, o), 0   # the second parameter is the orientation
end

function H5ReadPhysicalParameter3D(
  h5obj::PyCall.PyObject;
  phptype::H5PhPType,
  xkey::String,
  ykey::String)

  h5geo = pyimport("h5geopy._h5geo")

  keylist = ["INLINE", "XLINE"]
  minlist = [-Inf, -Inf]
  maxlist = [Inf, Inf]
  if !h5obj.hasPKeySort(keylist[1])
    @info "Updating trace header limits and adding pkey sort: $pkey..."
    h5obj.updateTraceHeaderLimits(Int(1e7))
    h5obj.addPKeySort(keylist[1])
  end

  if phptype == VELOCITY
    php, il_xl, ind = h5obj.getSortedData(keylist, minlist, maxlist, 0, typemax(Int), "km/s")
    php = (1f0 ./ php).^2
  elseif phptype == DENSITY
    php, il_xl, ind = h5obj.getSortedData(keylist, minlist, maxlist, 0, typemax(Int), "g/cm^3")
  elseif phptype == THETA || phptype == PHI
    php, il_xl, ind = h5obj.getSortedData(keylist, minlist, maxlist, 0, typemax(Int), "rad")
  else
    php, il_xl, ind = h5obj.getSortedData(keylist, minlist, maxlist, 0, typemax(Int))
  end

  if isempty(php)
    @error "Unable to read PhysicalParameter. Probably data units incorrect/missing\n"
    return
  end

  # JUDI falls down if NaN are passed
  meanValue = mean(filter(!isnan, php))
  php[isnan.(php)] .= meanValue

  keylist = [xkey, ykey]
  xy = h5obj.getXYTraceHeaders(keylist, ind, "m", true)
  if isempty(xy)
    @error "Unable to read $(xkey) and $(ykey) trace headers. Probably length units incorrect/missing\n"
    return
  end

  _, uil, uil_from_size = h5geo.sort_unique(il_xl[:,1])
  _, uxl, _ = h5geo.sort_unique(il_xl[:,2])
  nsamp = h5obj.getNSamp()

  nil = length(uil)
  nxl = length(uxl)

  # if the number of XL of each IL is not constant return false
  if !all(uil_from_size[:,2] .== uil_from_size[1,2])
    @error "Selected traces can't be represented as cube\n"
    return
  end

  # p1/p2 - second point on the first IL/XL
  origin = Vector{Float32}(undef, 2)
  origin[1] = xy[1,1]
  origin[2] = xy[1,2]

  p1 = Vector{Float32}(undef, 2)
  p2 = Vector{Float32}(undef, 2)
  if nxl > 1
    p1[1] = xy[2,1]
    p1[2] = xy[2,2]
  else
    p1 = origin
  end

  if nil > 1
    p2[1] = xy[nxl+1,1]
    p2[2] = xy[nxl+1,2]
  else
    p2 = origin
  end

  orientation = 0
  if nxl > 1
    orientation = atan((p1[2]-origin[2])/(p1[1]-origin[1]))
  elseif nil > 1
    orientation = atan((p2[2]-origin[2])/(p2[1]-origin[1]))
  end

  p1_rot = rotate_2xN_array(p1-origin, -orientation) + origin
  p2_rot = rotate_2xN_array(p2-origin, -orientation) + origin
  dx = 0
  dy = 0
  if nxl > 1
    dx = p1_rot[1]-origin[1]
    if dx < 0
      dx = abs(dx)
      origin[1] = origin[1]-dx
      p1_rot[1] = p1_rot[1]+dx
    end
    dy = p2_rot[2]-origin[2]
    if dy < 0
      dy = abs(dy)
      origin[2] = origin[2]-dy
      p2_rot[2] = p2_rot[2]+dy
    end
  elseif nil > 1
    dx = p2_rot[1]-origin[1]
    if dx < 0
      dx = abs(dx)
      origin[1] = origin[1]-dx
      p2_rot[1] = p2_rot[1]+dx
    end
    dy = p1_rot[2]-origin[2]
    if dy < 0
      dy = abs(dy)
      origin[2] = origin[2]-dy
      p1_rot[2] = p1_rot[2]+dy
    end
  end

  php = reshape(php, nsamp, nxl, nil)
  php = permutedims(php, [2,3,1])

  # Set up PhysicalParameter structure
  n = size(php)   # (x,y,z)
  d = (dx, dy, abs(h5obj.getSampRate("m")))
  o = (origin[1], origin[2], h5obj.getSRD("m") * (-1))

  return JUDI.PhysicalParameter(php, n, d, o), orientation
end

function H5ReadPhysicalParameter(
  h5obj::PyCall.PyObject;
  phptype::H5PhPType,
  xkey::String,
  ykey::String)

  h5geo = pyimport("h5geopy._h5geo")

  if h5obj.getSurveyType() == h5geo.SurveyType.TWO_D
    return H5ReadPhysicalParameter2D(h5obj, phptype=phptype, xkey=xkey)
  elseif h5obj.getSurveyType() == h5geo.SurveyType.THREE_D
    return H5ReadPhysicalParameter3D(h5obj, phptype=phptype, xkey=xkey, ykey=ykey)
  else
    @error "H5Object is neither TWO_D nor THREE_D\n"
    return
  end
end

function H5WritePhysicalParameter(;
  cntName::String,
  objName::String,
  cntCreationType::PyCall.PyObject,
  objCreationType::PyCall.PyObject,
  php::JUDI.PhysicalParameter)

  if isnothing(cntName) || length(cntName) < 1
    @error "Container name is empty\n"
    return
  end

  if isnothing(objName) || length(objName) < 1
    @error "Object name is empty\n"
    return
  end

  if isnothing(php) || length(php.data) < 1
    @error "PhysicalParameter is empty\n"
    return
  end

  colada = pyimport("colada")
  h5geo = pyimport("h5geopy._h5geo")

  cnt = h5geo.createSeisContainerByName(cntName, cntCreationType)
  if isnothing(cnt)
    @error "Unable to create SeisContainer: $cntName\n"
    return
  end

  authName = colada.Util().CRSAuthName()
  authCode = string(colada.Util().CRSCode())

  p = h5geo.SeisParam()
  p.spatialReference = "$authName:$authCode"
  p.lengthUnits = colada.Util().lengthUnits()
  p.temporalUnits = colada.Util().timeUnits()
  p.domain = h5geo.Domain.TVD
  p.dataType = h5geo.SeisDataType.STACK

  dims = length(size(php.data))
  TRACE = Matrix{Float32}
  x0 = 0.0; y0 = 0.0; z0 = 0.0
  dx = 0.0; dy = 0.0; dz = 0.0
  nx = 1; ny = 1; nz = 1
  if dims == 2
    p.surveyType = h5geo.SurveyType.TWO_D
    x0 = php.o[1]; z0 = -php.o[2]
    dx = php.d[1]; dz = php.d[2]
    nx = php.n[1]; nz = php.n[2]
    nTrc = nx
    nSamp = nz
    TRACE = permutedims(php.data, (2,1))
  elseif dims == 3
    p.surveyType = h5geo.SurveyType.THREE_D
    x0 = php.o[1]; y0 = php.o[2]; z0 = -php.o[3]
    dx = php.d[1]; dy = php.d[2]; dz = php.d[3]
    nx = php.n[1]; ny = php.n[2]; nz = php.n[3]
    nTrc = nx*ny
    nSamp = nz
    TRACE = reshape(permutedims(php.data, (3,2,1)), (nz, nx*ny))
  else
    @error "PhysicalParameter dimensions is neither 2D or 3D\n"
    return
  end

  p.nTrc = nTrc
  p.nSamp = nSamp

  if nTrc < 1000
    p.trcChunk = nTrc
  else
    p.trcChunk = 1000
  end

  seis = cnt.createSeis(objName, p, objCreationType)
  if isnothing(seis)
    @error "Unable to create Seis: $objName\n"
    return
  end

  # don't forget to set samp rate
  seis.setSampRate(length(php.d) == 2 ? -php.d[2] : -php.d[3], "m")

  val = seis.generateSTKGeometry(x0, dx, nx, y0, dy, ny, z0)
  if isnothing(val)
    @error "Unable to create generate STK geometry for Seis: $objName\n"
    return
  end

  val = seis.writeTrace(TRACE)
  if isnothing(val)
    @error "Unable to write traces: $objName\n"
    return
  end

  seis.getH5File().flush()
end


@testset "send_retx.jl" begin
    @testset "Retx SYN twice at the right times, then ack" begin
        isn = rand(UInt32(0):typemax(UInt32))
        rto = rand(UInt16(30):UInt16(10000))
        window_size = rand(UInt16(50000):UInt16(63000))

        sender = TCPSender(fixed_isn=WrappingInt32(isn), retx_timeout = rto)
        fill_window!(sender)

        expect_seg(sender, no_flag=true, syn=true, payload_size=0, seqno=WrappingInt32(isn))
        tick!(sender, rto - 1)
        expect_no_seg(sender)
        tick!(sender, 1)
        expect_seg(sender, no_flag=true, syn=true, payload_size=0, seqno=WrappingInt32(isn))
        expect_state(sender, JLSponge.SYN_SENT)
        @test sender.bytes_in_flight == 1

        tick!(sender, 2rto - 1)
        expect_no_seg(sender)
        tick!(sender, 1)
        expect_seg(sender, no_flag=true, syn=true, payload_size=0, seqno=WrappingInt32(isn))
        expect_state(sender, JLSponge.SYN_SENT)
        @test sender.bytes_in_flight == 1
        ack_received_test(sender, WrappingInt32(isn+1))
        @test sender.bytes_in_flight == 0
        tick!(sender, 2rto - 1)
        @test sender.ticks == 0
    end

    @testset "Retx SYN until too many retransmissions" begin
        MAX_RETX_ATTEMPTS = 10
        isn = rand(UInt32(0):typemax(UInt32))
        rto = rand(UInt16(30):UInt16(10000))
        window_size = rand(UInt16(50000):UInt16(63000))

        sender = TCPSender(fixed_isn=WrappingInt32(isn), retx_timeout = rto)
        fill_window!(sender)

        expect_seg(sender, no_flag=true, syn=true, payload_size=0, seqno=WrappingInt32(isn))
        expect_no_seg(sender)
        expect_state(sender, JLSponge.SYN_SENT)

        for i in 0:MAX_RETX_ATTEMPTS-1
            tick!(sender, (Int(rto) << i) - 1)
            @test consecutive_retransmissions(sender) <= MAX_RETX_ATTEMPTS
            tick!(sender, 1)
            expect_seg(sender, no_flag=true, syn=true, payload_size=0, seqno=WrappingInt32(isn))
            expect_no_seg(sender)
            expect_state(sender, JLSponge.SYN_SENT)
            @test sender.bytes_in_flight == 1
        end
        tick!(sender, (Int(rto) << MAX_RETX_ATTEMPTS) - 1)
        @test sender.consecutive_retransmissions <= MAX_RETX_ATTEMPTS
        tick!(sender, 1)
        @test sender.consecutive_retransmissions > MAX_RETX_ATTEMPTS
    end

    @testset "Send some data, the retx and succeed, then retx till limit" begin
        MAX_RETX_ATTEMPTS = 10
        isn = rand(UInt32(0):typemax(UInt32))
        rto = rand(UInt16(30):UInt16(10000))
        window_size = rand(UInt16(50000):UInt16(63000))

        sender = TCPSender(fixed_isn=WrappingInt32(isn), retx_timeout = rto)
        fill_window!(sender)

        expect_seg(sender, no_flag=true, syn=true, payload_size=0, seqno=WrappingInt32(isn))
        expect_no_seg(sender)
        ack_received_test(sender, WrappingInt32(isn+1))
        expect_state(sender, JLSponge.SYN_ACKED)
        write_bytes!(sender, "abcd")
        expect_seg(sender, payload_size=4)
        expect_no_seg(sender)
        ack_received_test(sender, WrappingInt32(isn+5))
        @test sender.bytes_in_flight == 0
        write_bytes!(sender, "efgh")
        expect_seg(sender, payload_size=4)
        expect_no_seg(sender)
        @test sender.consecutive_retransmissions <= MAX_RETX_ATTEMPTS
        ack_received_test(sender, WrappingInt32(isn+9))
        @test sender.bytes_in_flight == 0

        write_bytes!(sender, "efgh")
        expect_seg(sender, payload_size=4, seqno=WrappingInt32(isn+9))
        expect_no_seg(sender)

        for i in 0:MAX_RETX_ATTEMPTS-1
            tick!(sender, (Int(rto) << i) - 1)
            @test sender.consecutive_retransmissions <= MAX_RETX_ATTEMPTS
            expect_no_seg(sender)
            tick!(sender, 1)
            expect_seg(sender, payload_size=4, seqno=WrappingInt32(isn + 9))
            expect_no_seg(sender)
            expect_state(sender, JLSponge.SYN_ACKED)
            @test sender.bytes_in_flight == 4
        end

        tick!(sender, (Int(rto) << MAX_RETX_ATTEMPTS) - 1)
        @test sender.consecutive_retransmissions <= MAX_RETX_ATTEMPTS
        tick!(sender, 1)
        @test sender.consecutive_retransmissions > MAX_RETX_ATTEMPTS
    end
end


using Random
using Images

include("../../configs.jl")

mutable struct Image_Reader
    augment::Bool
    pre_scales::Array{Float64}
    function Image_Reader(augment::Bool=false)
        crop_ratios = [0.6, 0.8, 1.0]
        return new(augment, crop_ratios)
    end
end

# r is the reader object
function read_img(dir; len=nothing, boxes=nothing, r=nothing)

    img = convert(Array{Int64}, rawview(channelview(load(dir))))
    new_boxes = boxes === nothing ? nothing : deepcopy(boxes)
    c, h, w = size(img)

    if len === nothing
        len = max(w, h) # makes image square without changing its size, for evaluation
        div = Int(ceil(len / 64))
        len = 64 * div
    end

    if r !== nothing && r.augment
        img = convert(Array{Float32}, img) ./ 255
        roi = _find_roi(r, img)
        # horizontal flip
        if rand() < 0.5 img, new_boxes = flip_horizontal(img, new_boxes) end
        # color distortion
        img = distort_color(img)
        # cropping the picture with a scale
        img, new_boxes = crop_image(img, new_boxes, roi)
        # resizing to <len>
        img, new_boxes = resize_square_img(img, new_boxes, len, (roi[3] - roi[1]))
        img .*= 255
    else
        img = convert(Array{Float32}, img) ./ 255
        img = squaritize_img(img, len)
        # img, new_boxes = resize_square_img(img, new_boxes, len, maxlen)
        img .*= 255
    end

    if new_boxes !== nothing
        new_boxes[new_boxes .< 0] .= -1
        bbox_indices = getindex.(findall(new_boxes[1,:] .>= 0)) # images having bboxes
        new_boxes = new_boxes[:,bbox_indices]
    end

    img .-= avg_img
    return img, new_boxes
end

function _find_roi(r::Image_Reader, img)
    c, h, w = size(img)
    maxlen = max(h, w); minlen = min(h, w)
    scale = r.pre_scales[rand(1:length(r.pre_scales))]
    new_short = Int(ceil(minlen * scale))
    x1 = 1; y1 = 1
    if scale < 1
        x1 = rand(1:max(1, w-new_short))
        y1 = rand(1:max(1, h-new_short))
    end
    x2 = (new_short-1) + x1
    y2 = (new_short-1) + y1
    #region of interest is determined
    return [x1 y1 x2 y2]
end

function crop_image(img, new_boxes, roi)
    img = img[:, roi[2]:roi[4], roi[1]:roi[3]]
    box_size = size(new_boxes)

    for person in 1:box_size[2]
        # setting landmarks outside of the region to -1
        for lm in 5:2:box_size[1]-2
            # if not in the cropped area
            if (new_boxes[lm, person] > roi[3]) ||
                (new_boxes[lm, person] < roi[1]) ||
                (new_boxes[lm + 1, person] > roi[4]) ||
                (new_boxes[lm + 1, person] < roi[2])
                new_boxes[lm, person] = -1
                new_boxes[lm + 1, person] = -1
            # if in the cropped area
            else
                new_boxes[lm, person] -= roi[1]
                new_boxes[lm + 1, person] -= roi[2]
            end
        end
        if !all(>=(0), new_boxes[5:14,person])
            new_boxes[15, person] = -1.0
        end
        # setting bounding box coordinates

        # completely off the cropped area
        if (new_boxes[3,person] <= roi[1]) || (new_boxes[1,person] >= roi[3])
            new_boxes[1:4, person] .= -1
        elseif (new_boxes[4,person] <= roi[2]) || (new_boxes[2,person] >= roi[4])
            new_boxes[1:4, person] .= -1
        else
            # partial fits
            if new_boxes[3,person] > roi[3]
                new_boxes[3,person] = roi[3]
            end
            if new_boxes[4,person] > roi[4]
                new_boxes[4,person] = roi[4]
            end

            if 0 <= new_boxes[1,person] && new_boxes[1,person] < roi[1]
                new_boxes[1,person] = 0
            else
                new_boxes[1,person] -= roi[1]
            end
            new_boxes[3,person] -= roi[1]

            if 0 <= new_boxes[2,person] && new_boxes[2,person] < roi[2]
                new_boxes[2,person] = 0
            else
                new_boxes[2,person] -= roi[2]
            end
            new_boxes[4,person] -= roi[2]
        end
    end

    return img, new_boxes
end

function resize_square_img(img, new_boxes, new_len, old_len)
    img = imresize(img, (3, new_len, new_len))
    ratio = new_len / old_len
    if new_boxes !== nothing
        new_boxes = ratio .* new_boxes
    end
    return img, new_boxes
end

function squaritize_img(img, len)
    c, h, w = size(img)
    full_img = zeros(3, len, len)
    full_img[:,1:h,1:w] .= img
    return full_img
end

function flip_horizontal(img, new_boxes)
    img = reverse(img, dims=3)
    new_boxes[1:2:13,:] .= size(img, 3) .- new_boxes[1:2:13,:]
    temp = new_boxes[3,:]
    new_boxes[3,:] .= new_boxes[1,:]
    new_boxes[1,:] .= temp
    new_boxes[new_boxes .> size(img, 3)] .= -1
    return img, new_boxes
end

function distort_color(img)
    # brightness distortion
    if rand() < 0.5 img = _convert(img, beta=rand(-0.125:0.125)) end
    # contrast distortion
    if rand() < 0.5 img = _convert(img, alpha=rand(0.5:1.5)) end

    img = channelview(colorview(HSV, float.(img)))
    # saturation distortion
    if rand() < 0.5 img[2,:,:] = _convert(img[2,:,:], alpha=rand(0.5:1.5)) end
    # hue distortion
    if rand() < 0.5 img[1,:,:] = _convert(img[1,:,:], beta=rand(-(18/256):(18/256))) end

    img = channelview(colorview(RGB, colorview(HSV, img)))
    return img
end

function _convert(image; alpha=1, beta=0)
    image = image .* alpha .+ beta
    image[image .< 0] .= 0
    image[image .> 1] .= 1
    return image
end


using BenchmarkTools, AMDGPU

DAT = Float32

if DAT==Float64
    sc = 1
    DAT_Int = Int32
elseif DAT==Float32
    sc = 2
    DAT_Int = Int32
end

@inbounds function memcopy_triad_pow_int!(A, B, C, s, pow_int, ni)
    ix = (workgroupIdx().x - 1) * workgroupDim().x + workitemIdx().x
    iy = (workgroupIdx().y - 1) * workgroupDim().y + workitemIdx().y
    for ii=1:ni
        A[ix,iy] = B[ix,iy] + s*C[ix,iy]^pow_int
    end
    return
end

@inbounds function memcopy_triad_pow_float!(A, B, C, s, pow_float, ni)
    ix = (workgroupIdx().x - 1) * workgroupDim().x + workitemIdx().x
    iy = (workgroupIdx().y - 1) * workgroupDim().y + workitemIdx().y
    for ii=1:ni
        A[ix,iy] = B[ix,iy] + s*C[ix,iy]^pow_float
    end
    return
end

function run_bench()
    fact      = 24
    ni        = 100
    nx, ny    = sc*fact*1024, fact*1024
    threads   = (256, 2)
    grid      = (nx, ny)
    A         = ROCArray(zeros(DAT, nx, ny))
    B         = ROCArray( rand(DAT, nx, ny))
    C         = ROCArray( ones(DAT, nx, ny))
    pow_int   = DAT_Int(3)
    pow_float = DAT(3.75)
    s         = rand(DAT)
    lam       = rand(DAT)
    _dx, _dy  = DAT(1.0), DAT(1.0)
    dt        = DAT(1.0/10.0/4.1)
    println("nx, ny = $(nx), $(ny)")
    # run test 2
    println("test $(DAT)^$(DAT_Int)")
    t_it = @belapsed begin wait( @roc groupsize=$threads gridsize=$grid memcopy_triad_pow_int!($A, $B, $C, $s, $pow_int, $ni) ) end
    T_tot = 3*1/1e9*nx*ny*sizeof(DAT)/t_it
    Flops = 3*1/1e9*nx*ny*ni/t_it
    println("T_tot triad2D pow_int   = $(round(T_tot,sigdigits=7)) GB/s")
    println("Flops triad2D pow_int   = $(round(Flops,sigdigits=7)) GFLOP/s")
    # run test 3
    println("test $(DAT)^$(DAT)")
    t_it = @belapsed begin wait( @roc groupsize=$threads gridsize=$grid memcopy_triad_pow_float!($A, $B, $C, $s, $pow_float, $ni) ) end
    T_tot = 3*1/1e9*nx*ny*sizeof(DAT)/t_it
    Flops = 3*1/1e9*nx*ny*ni/t_it
    println("T_tot triad2D pow_float = $(round(T_tot,sigdigits=7)) GB/s")
    println("Flops triad2D pow_float = $(round(Flops,sigdigits=7)) GFLOP/s")
    return
end

run_bench()


# This file is a part of SimilaritySearch.jl

using Random
export BeamSearch

"""
    BeamSearch(bsize::Integer=16, beam=KnnResult(bsize))

BeamSearch is an iteratively improving local search algorithm that explores the graph using blocks of `bsize` elements and neighborhoods at the time.
Multithreading applications must have copies of this object due to shared cache objects.

- `bsize`: The size of the beam.
"""
@with_kw mutable struct BeamSearch <: LocalSearchAlgorithm
    bsize::Int32 = 8  # size of the search beam
end

Base.copy(bsearch::BeamSearch;
        bsize=bsearch.bsize
    ) = BeamSearch(; bsize)

function Base.copy!(dst::BeamSearch, src::BeamSearch)
    dst.bsize = src.bsize
end

Base.string(s::BeamSearch) = """{BeamSearch: bsize=$(s.bsize)}"""

const GlobalBeamKnnResult = [KnnResult(32)]  # see __init__ function
@inline function getbeam(bs::BeamSearch)
    @inbounds beam = GlobalBeamKnnResult[Threads.threadid()]
    empty!(beam, bs.bsize)
    beam
end

function __init__beamsearch()
    for i in 2:Threads.nthreads()
        push!(GlobalBeamKnnResult, KnnResult(32))
    end
end

### local search algorithm

function beamsearch_queue(index::SearchGraph, q, res::KnnResult, objID, vstate)
    @inbounds if !visited(vstate, objID)
        visit!(vstate, objID)
        d = evaluate(index.dist, q, index[objID])
        push!(res, objID, d)
    end
end

function beamsearch_init(bs::BeamSearch, index::SearchGraph, q, res::KnnResult, hints, vstate)
    for objID in hints
        beamsearch_queue(index, q, res, objID, vstate)
    end

    if length(res) == 0
        _range = 1:length(index)
        for i in 1:bs.bsize
           objID = rand(_range)
           beamsearch_queue(index, q, res, objID, vstate)
       end
    end
end

function beamsearch_inner(index::SearchGraph, q, res::KnnResult, beam::KnnResult, vstate)
    while length(beam) > 0
        prev_id, prev_dist = popfirst!(beam)
        @inbounds for childID in keys(index.links[prev_id])
            if !visited(vstate, childID)
                visit!(vstate, childID)
                d = evaluate(index.dist, q, index[childID])
                push!(res, childID, d) && push!(beam, childID, d)
                #d <= 0.9 * farthest(res).dist && push!(beam, childID, d)
            end
        end
    end
end

"""
    search(bs::BeamSearch, index::SearchGraph, q, res::KnnResult, hints, vstate)

Tries to reach the set of nearest neighbors specified in `res` for `q`.
- `bs`: the parameters of `BeamSearch`
- `index`: the local search index
- `q`: the query
- `res`: The result object, it stores the results and also specifies the kind of query
- `hints`: Starting points for searching, randomly selected when it is an empty collection
- `vstate`: A dictionary like object to store the visiting state of vertices
"""
function search(bs::BeamSearch, index::SearchGraph, q, res::KnnResult, hints, vstate)
    beamsearch_init(bs, index, q, res, hints, vstate)
    beam = getbeam(bs)
    push!(beam, first(res))
    beamsearch_inner(index, q, res, beam, vstate)
    res
end


struct ZygoteAD <: ADBackend end
ADBackend(::Val{:zygote}) = ZygoteAD
function setadbackend(::Val{:zygote})
    ADBACKEND[] = :zygote
end

function gradient_logp(
    backend::ZygoteAD,
    θ::AbstractVector{<:Real},
    vi::VarInfo,
    model::Model,
    sampler::AbstractSampler = SampleFromPrior(),
)
    T = typeof(getlogp(vi))

    # Specify objective function.
    function f(θ)
        new_vi = VarInfo(vi, sampler, θ)
        return getlogp(runmodel!(model, new_vi, sampler))
    end

    # Compute forward and reverse passes.
    l::T, ȳ = Zygote.pullback(f, θ)
    ∂l∂θ::typeof(θ) = ȳ(1)[1]

    return l, ∂l∂θ
end

Zygote.@nograd DynamicPPL.updategid!


@testset "Inverter Components" begin
    converter = AverageConverter(690.0, 2750000.0) #S_rated goes in Watts
    @test converter isa PowerSystems.DynamicComponent
    dc_source = FixedDCSource(600.0) #Not in the original data, guessed.
    @test dc_source isa PowerSystems.DynamicComponent
    filter = LCLFilter(0.08, 0.003, 0.074, 0.2, 0.01)
    @test filter isa PowerSystems.DynamicComponent
    pll = KauraPLL(500.0, 0.084, 4.69)
    @test pll isa PowerSystems.DynamicComponent
    virtual_H = VirtualInertia(2.0, 400.0, 20.0, 2 * pi * 50.0)
    @test virtual_H isa PowerSystems.DeviceParameter
    Q_control = ReactivePowerDroop(0.2, 1000.0)
    @test Q_control isa PowerSystems.DeviceParameter
    outer_control = OuterControl(virtual_H, Q_control)
    @test outer_control isa PowerSystems.DynamicComponent
    vsc = CurrentControl(0.59, 736.0, 0.0, 0.0, 0.2, 1.27, 14.3, 0.0, 50.0, 0.2)
    @test vsc isa PowerSystems.DynamicComponent
    BESS_source = ZeroOrderBESS(
        (sqrt(8) / sqrt(3)) * 690.0,
        (sqrt(3) / sqrt(8)) * 2750000.0,
        370.0,
        0.001,
        4.63,
        3200.0,
        0.6,
        4.0,
        0.39,
        10.34,
        1.08,
    )
    @test BESS_source isa PowerSystems.DynamicComponent
end

@testset "Dynamic Inverter" begin
    sys = create_system_with_dynamic_inverter()
    inverters = collect(get_components(DynamicInverter, sys))
    @test length(inverters) == 1
    test_inverter = inverters[1]
    @test test_inverter isa PowerSystems.Component
end


using Random
using MultiAgentPathFinding
using MultiAgentAllocationTransit
using BenchmarkTools
using Logging

global_logger(SimpleLogger(stderr, Logging.Warn))

## Hard-code a bunch of depots and many more sites
const DEPOT1 = LatLonCoords((lat = 37.762892, lon = -122.472193))
const DEPOT2 = LatLonCoords((lat = 37.751751, lon = -122.410654))
const DEPOT3 = LatLonCoords((lat = 37.718779, lon = -122.462401))

const SITE1 = LatLonCoords((lat = 37.789290, lon = -122.426797)) # Top right
const SITE2 = LatLonCoords((lat = 37.739611, lon = -122.492203)) # Way left
const SITE3 = LatLonCoords((lat = 37.780593, lon = -122.434555)) # Under SITE 1
const SITE4 = LatLonCoords((lat = 37.739011, lon = -122.430059)) # Near the bottom
const SITE5 = LatLonCoords((lat = 37.749018, lon = -122.462966)) # 2nd most left
const SITE1DUMMY = LatLonCoords((lat = 37.788042, lon = -122.429579))

const TRANSIT_CAP_RANGE = (1, 1)
const ECBS_WEIGHT = 1.05

stop_coords_file = "data/sfmta/stop_to_coords.json"
trips_file = "data/sfmta/trips.json"
drone_params_file = "data/drone_params.toml"

depots = [DEPOT1, DEPOT2, DEPOT3]
sites = [SITE1, SITE2, SITE3, SITE4, SITE5, SITE1DUMMY]

agent_tasks = [AgentTask((origin = 1, site = 1, dest = 2)),
               AgentTask((origin = 2, site = 6, dest = 1)),
               AgentTask((origin = 1, site = 3, dest = 2)),
               AgentTask((origin = 2, site = 4, dest = 3)),
               AgentTask((origin = 3, site = 5, dest = 1))
               ]

rng = MersenneTwister(1234)

# Load drone params and setup ENV!!
drone_params = parse_drone_params(drone_params_file)

# Generate OTG and load TG and drone params
otg = OffTransitGraph(depots = depots, sites = sites)
tg = load_transit_graph_latlong(stop_coords_file, trips_file, TRANSIT_CAP_RANGE, rng)

tg, stop_idx_to_trips = transit_graph_preprocessing(tg, MultiAgentAllocationTransit.distance_lat_lon_euclidean, drone_params)

aug_trips_fws_dists = augmented_trip_meta_graph_fws_dists(tg, MultiAgentAllocationTransit.distance_lat_lon_euclidean,
                                                      length(depots), length(sites),
                                                      vcat(depots, sites),
                                                      drone_params)

state_graph, depot_sites_to_vtx, trip_to_vtx_range = setup_state_graph(tg, otg)

agent_states = [AgentState(task=agt_task) for agt_task in agent_tasks]

env = MAPFTransitEnv(off_transit_graph = otg, transit_graph = tg, state_graph = state_graph,
                     agent_states = agent_states, depot_sites_to_vtx = depot_sites_to_vtx, trip_to_vtx_range = trip_to_vtx_range,
                     stop_idx_to_trips = stop_idx_to_trips, aug_trips_fws_dists = aug_trips_fws_dists,
                     plan_ref_times = zeros(length(agent_tasks)),
                     drone_params = drone_params, dist_fn = MultiAgentAllocationTransit.distance_lat_lon_euclidean,
                     curr_site_points = [0 for _ = 1:length(agent_tasks)])

# Define ECBS Solver; Run search
initial_states = Vector{MAPFTransitVertexState}(undef, length(agent_tasks))
for i = 1:length(agent_tasks)
    initial_states[i] = env.state_graph.vertices[env.depot_sites_to_vtx[string("d-", agent_tasks[i].origin)]]
end

solver = ECBSSolver{MAPFTransitVertexState,MAPFTransitAction,Float64,Makespan,MAPFTransitConflict,MAPFTransitConstraints,MAPFTransitEnv}(env = env, weight = ECBS_WEIGHT)

solution = search!(solver, initial_states)
#
# # NEED TO RE_INIT SOLVER!!!!
agent_states = [AgentState(task=agt_task) for agt_task in agent_tasks]

env = MAPFTransitEnv(off_transit_graph = otg, transit_graph = tg, state_graph = state_graph,
                     agent_states = agent_states, depot_sites_to_vtx = depot_sites_to_vtx, trip_to_vtx_range = trip_to_vtx_range,
                     stop_idx_to_trips = stop_idx_to_trips, aug_trips_fws_dists = aug_trips_fws_dists,
                     plan_ref_times = zeros(length(agent_tasks)), valid_transit_options = zeros(length(agent_tasks)),
                     drone_params = drone_params, dist_fn = MultiAgentAllocationTransit.distance_lat_lon_euclidean,
                     curr_site_points = [0 for _ = 1:length(agent_tasks)])

solver = ECBSSolver{MAPFTransitVertexState,MAPFTransitAction,Float64,Makespan,MAPFTransitConflict,MAPFTransitConstraints,MAPFTransitEnv}(env = env, weight = ECBS_WEIGHT)
@time solution = search!(solver, initial_states)


"""
    luminance_spec(λ,r)`

Luminance spectrum of type `RSpec`
"""
function luminance_spec(λ,l)
    return LSpec(λ,l)
end

"""
`reflectance_spec(λ,r)`

Reflectance spectrum of type `RSpec`
"""
function reflectance_spec(λ,r)
    return RSpec(λ,r)
end

"""
    transmittance_spec(λ,t)

Tansmittance spectrum at unit thickness (x = 1.0).
"""
function transmittance_spec(λ,t)
    return TSpec(λ,t,1.0)
end

"""
    transmittance_spec(λ, t, x)

Transmittance spectrum with thickness `x` given relative to unit thickness, according to Bougouer’s Law/Lambert’s Law.

See: "Color Measurement" by David L. MacAdam, chapter 3.2: "Bouguer’s Law"
"""
function transmittance_spec(λ::Vector{Real},t::Vector{Real},x::Real)
    return TSpec(λ, t, x)
end

"""
    transmittance_spec(λ,t,d,d0)

Transmittance spectrum with thickness term `d`. Unit thickness is `d0`, according to Bouguer’s Law/Lambert’s Law.

See: "Color Measurement" by David L. MacAdam, chapter 3.2: "Bougouer’s Law"
"""
function transmittance_spec(λ::Vector{Real},t::Vector{Real},d::Real,d0::Real)
    transmittance_spec(λ,t,(d/d0))
end

"""
`reinterpolate(spectrum,λmin,Δλ,λmax,style=:natural)`
"""
function reinterpolate(spectrum,λmin,Δλ,λmax,istyle=:natural,estyle=:linear)
    if typeof(spectrum) == RSpec
        spl = cubicspline(spectrum.λ,spectrum.r,istyle)
        epsl = extrap(spl,)
        extrapolate(spl,xrange,extr::Symbol)
        l,s = interp(spl,(λmin:Δλ:λmax))
        return RSpec(l,s)
    elseif typeof(spectrum) == LSpec
        spl = cubicspline(spectrum.λ,spectrum.l,istyle)
        l,s = interp(spl,(λmin:Δλ,λmax))
        return LSpec(l,s)
    end
end

"""
    adaptated_color(spectrum)
"""
function adapted_color(illuminant1,λrange)
    lum=normalize_spec(blackbody_illuminant(T))
    wp=blackbody_whitepoint(SPECENV, T)
end


using Jtac

chain = @stack MNKGame{5, 5, 4} stack_input=true begin
  Conv(256, relu)
  Dense(128, relu)
end

model = NeuralModel(MNKGame{5, 5, 4}, chain) |> to_gpu |> Async
player = MCTSPlayer(model, power = 25)

loss = Loss(value = 2.5, regularization = 1e-4)

mcts_opponents = [MCTSPlayer(power = p) for p in [10, 50, 100]]#, 200, 500, 1000]]
self_opponents = [MCTSPlayer(model, power = p, name = "current$p") for p in [10, 50]]

with_contest( train_self!
            , player
            , loss = loss
            , playings = 100
            , batchsize = 100
            , iterations = 10
            , branching = 0.1
            , epochs = 100
            , opponents = [mcts_opponents; self_opponents]
            , interval = 5
            , length = 2000 )


"""
A generic, framework agnostic interface for neural networks.
"""
module Network

export AbstractNetwork, OptimiserSpec, CyclicNesterov, Adam

using ..AlphaZero

using Base: @kwdef
using Statistics: mean

"""
    AbstractNetwork

Abstract base type for a neural network.

# Constructor

Any subtype `Network` must implement `Base.copy` along with
the following constructor:

    Network(game_spec, hyperparams)

where the expected type of `hyperparams` is given by
[`HyperParams(Network)`](@ref HyperParams).
"""
abstract type AbstractNetwork end

#####
##### Interface
#####

"""
    HyperParams(::Type{<:AbstractNetwork})

Return the hyperparameter type associated with a given network type.
"""
function HyperParams end

"""
    hyperparams(::AbstractNetwork)

Return the hyperparameters of a network.
"""
function hyperparams end

"""
    game_spec(::AbstractNetwork)

Return the game specification that was passed to the network's constructor.
"""
function game_spec end

"""
    to_gpu(::AbstractNetwork)

Return a copy of the given network that has been transferred to the GPU
if one is available. Otherwise, return the given network untouched.
"""
function to_gpu end

"""
    to_cpu(::AbstractNetwork)

Return a copy of the given network that has been transferred to the CPU
or return the given network untouched if it is already on CPU.
"""
function to_cpu end

"""
    on_gpu(::AbstractNetwork) :: Bool

Test whether or not a network is located on GPU.
"""
function on_gpu end

"""
    set_test_mode!(mode=true)

Put a network in test mode or in training mode.
This is relevant for networks featuring layers such as
batch normalization layers.
"""
function set_test_mode! end

"""
    convert_input(::AbstractNetwork, input)

Convert an array (or number) to the right format so that it can be used
as an input by a given network.
"""
function convert_input end

function convert_input_tuple(nn::AbstractNetwork, input::Tuple)
  return map(input) do arr
    convert_input(nn, arr)
  end
end

"""
    convert_output(::AbstractNetwork, output)

Convert an array (or number) produced by a neural network
to a standard CPU array (or number) type.
"""
function convert_output end

function convert_output_tuple(nn::AbstractNetwork, output::Tuple)
  return map(output) do arr
    convert_output(nn, arr)
  end
end

"""
    forward(::AbstractNetwork, states)

Compute the forward pass of a network on a batch of inputs.

Expect a `Float32` tensor `states` whose batch dimension is the last one.

Return a `(P, V)` triple where:

  - `P` is a matrix of size `(num_actions, batch_size)`. It is allowed
    to put weight on invalid actions (see [`evaluate`](@ref)).
  - `V` is a row vector of size `(1, batch_size)`
"""
function forward end

"""
    regularized_params(::AbstractNetwork)

Return the collection of regularized parameters of a network.
This usually excludes neuron's biases.
"""
function regularized_params end

"""
    params(::AbstractNetwork)

Return the collection of trainable parameters of a network.
"""
function params end

"""
    gc(::AbstractNetwork)

Perform full garbage collection and empty the GPU memory pool.
"""
function gc end

# Optimizers and training

"""
    OptimiserSpec

Abstract type for an optimiser specification.
"""
abstract type OptimiserSpec end

"""
    CyclicNesterov(; lr_base, lr_high, lr_low, momentum_low, momentum_high)

SGD optimiser with a cyclic learning rate and cyclic Nesterov momentum.

  - During an epoch, the learning rate goes from `lr_low`
    to `lr_high` and then back to `lr_low`.
  - The momentum evolves in the opposite way, from high values
    to low values and then back to high values.
"""
@kwdef struct CyclicNesterov <: OptimiserSpec
  lr_base :: Float32
  lr_high :: Float32
  lr_low  :: Float32
  momentum_low :: Float32
  momentum_high :: Float32
end

"""
    Adam(;lr)

Adam optimiser.
"""
@kwdef struct Adam <: OptimiserSpec
  lr :: Float32
end

"""
    train!(callback, ::AbstractNetwork, opt::OptimiserSpec, loss, batches, n)

Update a given network to fit some data.
  - [`opt`](@ref OptimiserSpec) specifies which optimiser to use.
  - `loss` is a function that maps a batch of samples to a tracked real.
  - `data` is an iterator over minibatches.
  - `n` is the number of minibatches. If `length` is defined on `data`,
     we must have `length(data) == n`. However, not all finite
     iterators implement `length` and thus this argument is needed.
  - `callback(i, loss)` is called at each step with the batch number `i`
     and the loss on last batch.
"""
function train! end

#####
##### Derived functions
#####

"""
    num_parameters(::AbstractNetwork)

Return the total number of parameters of a network.
"""
function num_parameters(nn::AbstractNetwork)
  return sum(length(p) for p in params(nn))
end

"""
    num_regularized_parameters(::AbstractNetwork)

Return the total number of regularized parameters of a network.
"""
function num_regularized_parameters(nn::AbstractNetwork)
  return sum(length(p) for p in regularized_params(nn))
end

"""
    mean_weight(::AbstractNetwork)

Return the mean absolute value of the regularized parameters of a network.
"""
function mean_weight(nn::AbstractNetwork)
  sw = sum(sum(abs.(p)) for p in regularized_params(nn))
  sw = convert_output(nn, sw)
  return sw / num_regularized_parameters(nn)
end

"""
    forward_normalized(network::AbstractNetwork, states, actions_mask)

Evaluate a batch of vectorized states. This function is a wrapper
on [`forward`](@ref) that puts a zero weight on invalid actions.

# Arguments

  - `states` is a tensor whose last dimension has size `bach_size`
  - `actions_mask` is a binary matrix of size `(num_actions, batch_size)`

# Returned value

Return a `(P, V, Pinv)` triple where:

  - `P` is a matrix of size `(num_actions, batch_size)`.
  - `V` is a row vector of size `(1, batch_size)`.
  - `Pinv` is a row vector of size `(1, batch_size)`
     that indicates the total probability weight put by the network
     on invalid actions for each sample.

All tensors manipulated by this function have elements of type `Float32`.
"""
function forward_normalized(nn::AbstractNetwork, state, actions_mask)
  p, v = forward(nn, state)
  p = p .* actions_mask
  sp = sum(p, dims=1)
  p = p ./ (sp .+ eps(eltype(p)))
  p_invalid = 1 .- sp
  return (p, v, p_invalid)
end

to_singletons(x) = reshape(x, size(x)..., 1)
from_singletons(x) = reshape(x, size(x)[1:end-1])

"""
    evaluate(::AbstractNetwork, state)

    (nn::AbstractNetwork)(state) = evaluate(nn, state)

Evaluate the neural network as an MCTS oracle on a single state.

Note, however, that evaluating state positions once at a time is slow and so you
may want to use a `BatchedOracle` along with an inference server that uses
[`evaluate_batch`](@ref).
"""
function evaluate(nn::AbstractNetwork, state)
  gspec = game_spec(nn)
  actions_mask = GI.actions_mask(GI.init(gspec, state))
  x = GI.vectorize_state(gspec, state)
  a = Float32.(actions_mask)
  xnet, anet = to_singletons.(convert_input_tuple(nn, (x, a)))
  net_output = forward_normalized(nn, xnet, anet)
  p, v, _ = from_singletons.(convert_output_tuple(nn, net_output))
  return (p[actions_mask], v[1])
end

(nn::AbstractNetwork)(state) = evaluate(nn, state)

"""
    evaluate_batch(::AbstractNetwork, batch)

Evaluate the neural network as an MCTS oracle on a batch of states at once.

Take a list of states as input and return a list of `(P, V)` pairs as defined in the
MCTS oracle interface.
"""
function evaluate_batch(nn::AbstractNetwork, batch)
  gspec = game_spec(nn)
  X = Util.superpose((GI.vectorize_state(gspec, b) for b in batch))
  A = Util.superpose((GI.actions_mask(GI.init(gspec, b)) for b in batch))
  Xnet, Anet = convert_input_tuple(nn, (X, Float32.(A)))
  P, V, _ = convert_output_tuple(nn, forward_normalized(nn, Xnet, Anet))
  return [(P[A[:,i],i], V[1,i]) for i in eachindex(batch)]
end

"""
    copy(::AbstractNetwork; on_gpu, test_mode)

A copy function that also handles CPU/GPU transfers and
test/train mode switches.
"""
function copy(network::AbstractNetwork; on_gpu, test_mode)
  network = Base.copy(network)
  network = on_gpu ? to_gpu(network) : to_cpu(network)
  set_test_mode!(network, test_mode)
  return network
end

end


inpt = """NOT dq -> dr
kg OR kf -> kh
ep OR eo -> eq
3176 -> b
NOT gs -> gt
dd OR do -> dp
eg AND ei -> ej
y AND ae -> ag
jx AND jz -> ka
lf RSHIFT 2 -> lg
z AND aa -> ac
dy AND ej -> el
bj OR bi -> bk
kk RSHIFT 3 -> km
NOT cn -> co
gn AND gp -> gq
cq AND cs -> ct
eo LSHIFT 15 -> es
lg OR lm -> ln
dy OR ej -> ek
NOT di -> dj
1 AND fi -> fj
kf LSHIFT 15 -> kj
NOT jy -> jz
NOT ft -> fu
fs AND fu -> fv
NOT hr -> hs
ck OR cl -> cm
jp RSHIFT 5 -> js
iv OR jb -> jc
is OR it -> iu
ld OR le -> lf
NOT fc -> fd
NOT dm -> dn
bn OR by -> bz
aj AND al -> am
cd LSHIFT 15 -> ch
jp AND ka -> kc
ci OR ct -> cu
gv AND gx -> gy
de AND dk -> dm
x RSHIFT 5 -> aa
et RSHIFT 2 -> eu
x RSHIFT 1 -> aq
ia OR ig -> ih
bk LSHIFT 1 -> ce
y OR ae -> af
NOT ca -> cb
e AND f -> h
ia AND ig -> ii
ck AND cl -> cn
NOT jh -> ji
z OR aa -> ab
1 AND en -> eo
ib AND ic -> ie
NOT eh -> ei
iy AND ja -> jb
NOT bb -> bc
ha OR gz -> hb
1 AND cx -> cy
NOT ax -> ay
ev OR ew -> ex
bn RSHIFT 2 -> bo
er OR es -> et
eu OR fa -> fb
jp OR ka -> kb
ea AND eb -> ed
k AND m -> n
et RSHIFT 3 -> ev
et RSHIFT 5 -> ew
hz RSHIFT 1 -> is
ki OR kj -> kk
NOT h -> i
lv LSHIFT 15 -> lz
as RSHIFT 1 -> bl
hu LSHIFT 15 -> hy
iw AND ix -> iz
lf RSHIFT 1 -> ly
fp OR fv -> fw
1 AND am -> an
ap LSHIFT 1 -> bj
u LSHIFT 1 -> ao
b RSHIFT 5 -> f
jq AND jw -> jy
iu RSHIFT 3 -> iw
ih AND ij -> ik
NOT iz -> ja
de OR dk -> dl
iu OR jf -> jg
as AND bd -> bf
b RSHIFT 3 -> e
jq OR jw -> jx
iv AND jb -> jd
cg OR ch -> ci
iu AND jf -> jh
lx -> a
1 AND cc -> cd
ly OR lz -> ma
NOT el -> em
1 AND bh -> bi
fb AND fd -> fe
lf OR lq -> lr
bn RSHIFT 3 -> bp
bn AND by -> ca
af AND ah -> ai
cf LSHIFT 1 -> cz
dw OR dx -> dy
gj AND gu -> gw
jg AND ji -> jj
jr OR js -> jt
bl OR bm -> bn
gj RSHIFT 2 -> gk
cj OR cp -> cq
gj OR gu -> gv
b OR n -> o
o AND q -> r
bi LSHIFT 15 -> bm
dy RSHIFT 1 -> er
cu AND cw -> cx
iw OR ix -> iy
hc OR hd -> he
0 -> c
db OR dc -> dd
kk RSHIFT 2 -> kl
eq LSHIFT 1 -> fk
dz OR ef -> eg
NOT ed -> ee
lw OR lv -> lx
fw AND fy -> fz
dz AND ef -> eh
jp RSHIFT 3 -> jr
lg AND lm -> lo
ci RSHIFT 2 -> cj
be AND bg -> bh
lc LSHIFT 1 -> lw
hm AND ho -> hp
jr AND js -> ju
1 AND io -> ip
cm AND co -> cp
ib OR ic -> id
NOT bf -> bg
fo RSHIFT 5 -> fr
ip LSHIFT 15 -> it
jt AND jv -> jw
jc AND je -> jf
du OR dt -> dv
NOT fx -> fy
aw AND ay -> az
ge LSHIFT 15 -> gi
NOT ak -> al
fm OR fn -> fo
ff AND fh -> fi
ci RSHIFT 5 -> cl
cz OR cy -> da
NOT ey -> ez
NOT ju -> jv
NOT ls -> lt
kk AND kv -> kx
NOT ii -> ij
kl AND kr -> kt
jk LSHIFT 15 -> jo
e OR f -> g
NOT bs -> bt
hi AND hk -> hl
hz OR ik -> il
ek AND em -> en
ao OR an -> ap
dv LSHIFT 1 -> ep
an LSHIFT 15 -> ar
fo RSHIFT 1 -> gh
NOT im -> in
kk RSHIFT 1 -> ld
hw LSHIFT 1 -> iq
ec AND ee -> ef
hb LSHIFT 1 -> hv
kb AND kd -> ke
x AND ai -> ak
dd AND do -> dq
aq OR ar -> as
iq OR ip -> ir
dl AND dn -> do
iu RSHIFT 5 -> ix
as OR bd -> be
NOT go -> gp
fk OR fj -> fl
jm LSHIFT 1 -> kg
NOT cv -> cw
dp AND dr -> ds
dt LSHIFT 15 -> dx
et RSHIFT 1 -> fm
dy RSHIFT 3 -> ea
fp AND fv -> fx
NOT p -> q
dd RSHIFT 2 -> de
eu AND fa -> fc
ba AND bc -> bd
dh AND dj -> dk
lr AND lt -> lu
he RSHIFT 1 -> hx
ex AND ez -> fa
df OR dg -> dh
fj LSHIFT 15 -> fn
NOT kx -> ky
gk OR gq -> gr
dy RSHIFT 2 -> dz
gh OR gi -> gj
lj AND ll -> lm
x OR ai -> aj
bz AND cb -> cc
1 AND lu -> lv
as RSHIFT 3 -> au
ce OR cd -> cf
il AND in -> io
dd RSHIFT 1 -> dw
NOT lo -> lp
c LSHIFT 1 -> t
dd RSHIFT 3 -> df
dd RSHIFT 5 -> dg
lh AND li -> lk
lf RSHIFT 5 -> li
dy RSHIFT 5 -> eb
NOT kt -> ku
at OR az -> ba
x RSHIFT 3 -> z
NOT lk -> ll
lb OR la -> lc
1 AND r -> s
lh OR li -> lj
ln AND lp -> lq
kk RSHIFT 5 -> kn
ea OR eb -> ec
ci AND ct -> cv
b RSHIFT 2 -> d
jp RSHIFT 1 -> ki
NOT cr -> cs
NOT jd -> je
jp RSHIFT 2 -> jq
jn OR jo -> jp
lf RSHIFT 3 -> lh
1 AND ds -> dt
lf AND lq -> ls
la LSHIFT 15 -> le
NOT fg -> fh
at AND az -> bb
au AND av -> ax
kw AND ky -> kz
v OR w -> x
kk OR kv -> kw
ks AND ku -> kv
kh LSHIFT 1 -> lb
1 AND kz -> la
NOT kc -> kd
x RSHIFT 2 -> y
et OR fe -> ff
et AND fe -> fg
NOT ac -> ad
jl OR jk -> jm
1 AND jj -> jk
bn RSHIFT 1 -> cg
NOT kp -> kq
ci RSHIFT 3 -> ck
ev AND ew -> ey
1 AND ke -> kf
cj AND cp -> cr
ir LSHIFT 1 -> jl
NOT gw -> gx
as RSHIFT 2 -> at
iu RSHIFT 1 -> jn
cy LSHIFT 15 -> dc
hg OR hh -> hi
ci RSHIFT 1 -> db
au OR av -> aw
km AND kn -> kp
gj RSHIFT 1 -> hc
iu RSHIFT 2 -> iv
ab AND ad -> ae
da LSHIFT 1 -> du
NOT bw -> bx
km OR kn -> ko
ko AND kq -> kr
bv AND bx -> by
kl OR kr -> ks
1 AND ht -> hu
df AND dg -> di
NOT ag -> ah
d OR j -> k
d AND j -> l
b AND n -> p
gf OR ge -> gg
gg LSHIFT 1 -> ha
bn RSHIFT 5 -> bq
bo OR bu -> bv
1 AND gy -> gz
s LSHIFT 15 -> w
NOT ie -> if
as RSHIFT 5 -> av
bo AND bu -> bw
hz AND ik -> im
bp AND bq -> bs
b RSHIFT 1 -> v
NOT l -> m
bp OR bq -> br
g AND i -> j
br AND bt -> bu
t OR s -> u
hz RSHIFT 5 -> ic
gk AND gq -> gs
fl LSHIFT 1 -> gf
he RSHIFT 3 -> hg
gz LSHIFT 15 -> hd
hf OR hl -> hm
1 AND gd -> ge
fo OR fz -> ga
id AND if -> ig
fo AND fz -> gb
gr AND gt -> gu
he OR hp -> hq
fq AND fr -> ft
ga AND gc -> gd
fo RSHIFT 2 -> fp
gl OR gm -> gn
hg AND hh -> hj
NOT hn -> ho
gl AND gm -> go
he RSHIFT 5 -> hh
NOT gb -> gc
hq AND hs -> ht
hz RSHIFT 3 -> ib
hz RSHIFT 2 -> ia
fq OR fr -> fs
hx OR hy -> hz
he AND hp -> hr
gj RSHIFT 5 -> gm
hf AND hl -> hn
hv OR hu -> hw
NOT hj -> hk
gj RSHIFT 3 -> gl
fo RSHIFT 3 -> fq
he RSHIFT 2 -> hf"""

function buildCode(sentence)
	tokens = split(sentence)
	expression = uppercase(tokens[end])
	expression *= "xxx = "
	for x in 1:length(tokens)-2
		tokens[x] = uppercase(tokens[x])
		if tokens[x] == "AND"
			tokens[x] = "&"
		elseif tokens[x] == "OR"
			tokens[x] = "|"
		elseif tokens[x] == "NOT"
			tokens[x] = "~"
		elseif tokens[x] == "LSHIFT"
			tokens[x] = "<<"
		elseif tokens[x] == "RSHIFT"
			tokens[x] = ">>"
		elseif !(typeof(parse(tokens[x])) <: Number)
			tokens[x] = tokens[x] * "xxx" # to protect variable names from clashes with builtins
		end
	end
	expression *= join(tokens[1:end-2], " ")
	println(expression)
	return parse(expression)
end

grid = Expr[]
for sen in split(inpt, "\n")
	sent = buildCode(sen)
	push!(grid, sent)
end

function isDefined(expr)
	if typeof(expr) == Symbol
		# println(expr)
		if !(expr in names(Main) || expr in names(Base))
			return false
		end
	elseif typeof(expr) == Expr
		for a in expr.args
			if !isDefined(a)
				return false
			end
		end
	# else
		# println(typeof(expr), expr)
	end
	return true
end

while length(grid) > 0
	for (i,e) in enumerate(grid)
		if isDefined(grid[i].args[2])
			p = splice!(grid, i)
			println(p)
			eval(p)
			break
		end
	end
end
println(Axxx)


type CNS # central nervous system
    # sympathetic, parasympathetic activations
    np::Vector{Float64}
    ns::Vector{Float64}
    # heart rate and maximum ventriclular elastance
    H::Vector{Float64}
    Emaxlv::Vector{Float64}
    Emaxrv::Vector{Float64}
    # lower/upper body compartmental resistances, compliances, unstressed vols.
    R2L::Vector{Float64}
    R3L::Vector{Float64}
    R2U::Vector{Float64}
    R3U::Vector{Float64}
    C4L::Vector{Float64}
    C5L::Vector{Float64}
    C4U::Vector{Float64}
    C5U::Vector{Float64}
    V4L::Vector{Float64}
    V5L::Vector{Float64}
    V4U::Vector{Float64}
    V5U::Vector{Float64}
    # feedback model weights (alpha:ns, beta:np, gamma:bias)
    alphaH::Float64
    betaH::Float64
    gammaH::Float64
    alphalv::Float64
    gammalv::Float64
    alpharv::Float64
    gammarv::Float64
    alphaR2L::Float64
    gammaR2L::Float64
    alphaR3L::Float64
    gammaR3L::Float64
    alphaR2U::Float64
    gammaR2U::Float64
    alphaR3U::Float64
    gammaR3U::Float64
    alphaC4L::Float64
    gammaC4L::Float64
    alphaC5L::Float64
    gammaC5L::Float64
    alphaC4U::Float64
    gammaC4U::Float64
    alphaC5U::Float64
    gammaC5U::Float64
    alphaV4L::Float64
    gammaV4L::Float64
    alphaV5L::Float64
    gammaV5L::Float64
    alphaV4U::Float64
    gammaV4U::Float64
    alphaV5U::Float64
    gammaV5U::Float64
    # time constants
    tauH::Float64
    tauEmax::Float64
    tauC::Float64
    tauR::Float64
    tauV::Float64
    # activation function parameters
    Ptarget::Float64
    Paverage::Vector{Float64}
    baselinefirerate::Float64
    a::Float64
    b::Float64
    c::Float64
    # indices for baroreflex pressure averaging
    avgindexstart::Int32
    avgindexend::Int32

    function CNS()
        this = new()
        this.ns = Vector{Float64}[];
        this.np = Vector{Float64}[];
        this.H = Vector{Float64}[];
        this.Emaxlv = Vector{Float64}[];
        this.Emaxrv = Vector{Float64}[];
        this.R2L = Vector{Float64}[];
        this.R3L = Vector{Float64}[];
        this.R2U = Vector{Float64}[];
        this.R3U = Vector{Float64}[];
        this.C4L = Vector{Float64}[];
        this.C5L = Vector{Float64}[];
        this.C4U = Vector{Float64}[];
        this.C5U = Vector{Float64}[];
        this.V4L = Vector{Float64}[];
        this.V5L = Vector{Float64}[];
        this.V4U = Vector{Float64}[];
        this.V5U = Vector{Float64}[];
        this.tauH = 2;
        this.tauEmax = 2;
        this.tauC = 20;
        this.tauR = 6;
        this.tauV = 20;
        this.Ptarget = 94.6*mmHgToPa;
        this.Paverage = Vector{Float64}[];
        this.avgindexstart = 1;
        this.avgindexend = 0;
        this.baselinefirerate = 0.25;
        this.a = 1;
        this.b = -log(this.baselinefirerate);
        this.c = 6;
        return this
    end
end


@testset "FIR" begin
    println("VITUIX FIR TESTS")

    Hs = load_vituix_FIRs((joinpath(test_dir, "FIR", wav) for wav in ["T.wav", "M.wav", "W.wav"])...; Δt=-371.519ms, names=[:FIR_T, :FIR_M, :FIR_W])

    test_set = [
        (src="T/polar_noXO", target="T/polar", freqrange=(1kHz, 20kHz), T=Dict(:H => Hs[:, :FIR_T]), max_mag=0.6, mae_mag=0.15, max_phi=180.0, mae_phi=2.5),
        (src="T/polar_noXO_r20", target="T/polar_r20", freqrange=(1kHz, 20kHz), T=Dict(:H => Hs[:, :FIR_T]), max_mag=0.6, mae_mag=0.15, max_phi=180.0, mae_phi=3.7),
        (src="T/polar_noXO", target="T/polar_r10_t15", freqrange=(1kHz, 20kHz), T=Dict([:H => Hs[:, :FIR_T], :r => 10.0, :t => 15.0]), max_mag=4.7, mae_mag=0.5, max_phi=180.0, mae_phi=11.0),

        (src="M/polar_noXO", target="M/polar", freqrange=(0kHz, 20kHz), T=Dict(:H => Hs[:, :FIR_M]), max_mag=1.9, mae_mag=0.06, max_phi=2, mae_phi=1.0),
        (src="M/polar_noXO_r170", target="M/polar_r170", freqrange=(0kHz, 20kHz), T=Dict(:H => Hs[:, :FIR_M]), max_mag=1.9, mae_mag=0.06, max_phi=180, mae_phi=2.2),
        (src="M/polar_noXO_y-118_r170", target="M/polar_y-118_r170", freqrange=(0kHz, 20kHz), T=Dict(:H => Hs[:, :FIR_M]), max_mag=1.9, mae_mag=0.06, max_phi=2, mae_phi=1.0),
        (src="M/polar_noXO", target="M/polar_x95_y-118_z150_r12_t-15", freqrange=(0kHz, 20kHz), T=Dict([:H => Hs[:, :FIR_M], :x => 0.095, :y => -0.118, :z => 0.150, :r => 12.0, :t => -15.0]), max_mag=6.0, mae_mag=0.6, max_phi=180.0, mae_phi=13.0),

        (src="W/polar_noXO", target="W/polar", freqrange=(0kHz, 2kHz), T=Dict(:H => Hs[:, :FIR_W]), max_mag=0.28, mae_mag=0.04, max_phi=180.0, mae_phi=1.2),
        (src="W/polar_noXO_y-348_r14", target="W/polar_y-348_r14", freqrange=(0kHz, 2kHz), T=Dict(:H => Hs[:, :FIR_W]), max_mag=0.25, mae_mag=0.03, max_phi=180.0, mae_phi=1.4),
        (src="W/polar_noXO", target="W/polar_x-150_y-348_z-35_r-90_t3", freqrange=(0kHz, 2kHz), T=Dict([:H => Hs[:, :FIR_W], :x => -0.15, :y => -0.348, :z => -0.035, :r => -90.0, :t => 3.0]), max_mag=4.5, mae_mag=0.3, max_phi=180.0, mae_phi=1.5),
    ]

    for matchFIR in [true, false]
        println("match FIR: $matchFIR")
        for e in test_set
            println("  \"$(e.src)\" -> \"$(e.target)\"")
            src, target = joinpath(test_dir, e.src), joinpath(test_dir, e.target)
            print("\tA:")
            @test test_driver_transform(src, target, e.freqrange, compare_magdb, e.max_mag, e.mae_mag, matchFIR; e.T...)
            print("\tΦ:")
            @test test_driver_transform(src, target, e.freqrange, compare_phasedeg, e.max_phi, e.mae_phi, matchFIR; e.T...)
        end
    end
end


module Optimization

using Evolutionary

import BlackBoxOptim

import ..QuerySynthesisStrategies:
    check_epsilon,
    data_boundaries

include("optimization_base.jl")
include("BlackBoxOptimization.jl")
include("EvolutionaryOptimization.jl")
include("ParticleSwarmOptimization.jl")

export
    QuerySynthesisOptimizer,
    BlackBoxOptimization,
    EvolutionaryOptimization,
    ParticleSwarmOptimization,

    query_synthesis_optimize
end


module DiffEqMonteCarlo

using DiffEqBase, RecursiveArrayTools, StaticArrays, Distributed, Statistics
import DiffEqBase: MonteCarloSummary

include("solve.jl")
include("analysis.jl")

export MonteThreads, MonteDistributed, MonteSplitThreads, MonteSerial

export get_timestep, get_timepoint, apply_timestep, apply_timepoint,
       componentwise_vectors_timestep, componentwise_vectors_timepoint

export componentwise_mean, componentwise_meanvar

export timestep_mean, timestep_median, timestep_quantile, timestep_meanvar,
       timestep_meancov, timestep_meancor, timestep_weighted_meancov

export timeseries_steps_mean, timeseries_steps_median, timeseries_steps_quantile,
       timeseries_steps_meanvar, timeseries_steps_meancov,
       timeseries_steps_meancor, timeseries_steps_weighted_meancov

export timepoint_mean, timepoint_median, timepoint_quantile,
       timepoint_meanvar, timepoint_meancov,
       timepoint_meancor, timepoint_weighted_meancov

export timeseries_point_mean, timeseries_point_median, timeseries_point_quantile,
       timeseries_point_meanvar, timeseries_point_meancov,
       timeseries_point_meancor, timeseries_point_weighted_meancov

end # module


module JVODE

using LinearAlgebra

using Parameters: @unpack
using DiffEqBase: ODEProblem, @..

using Reexport: @reexport
@reexport using DiffEqBase
export Adams, BDF

###
### Types
###
abstract type JVNLSolver <: DiffEqBase.AbstractNLSolverAlgorithm
end
struct JVFunctional <: JVNLSolver
end
struct JVNewton <: JVNLSolver
end

abstract type JVODEAlgorithm <: DiffEqBase.AbstractODEAlgorithm
end
struct Adams{NLAlg} <: JVODEAlgorithm
    nlsolve::NLAlg
end
Adams() = Adams(JVFunctional())
struct BDF{NLAlg} <: JVODEAlgorithm
    nlsolve::NLAlg
end
BDF() = BDF(JVNewton())

# TODO: Optional Inputs and Outputs (in `cvode.h`)

# Basic JVODE constants
const ADAMS_Q_MAX = 12            # max value of q for lmm == Adams
const BDF_Q_MAX   =  5            # max value of q for lmm == BDF
const Q_MAX       =  ADAMS_Q_MAX  # max value of q for either lmm
const L_MAX       =  (Q_MAX+1)    # max value of L for either lmm
const NUM_TESTS   =  5            # number of error test quantities

mutable struct JVOptions{Rtol,Atol}
    reltol::Rtol
    abstol::Atol
end

mutable struct JVIntegrator{Alg,uType,tType,uEltype,solType,Rtol,Atol,F,P} <: DiffEqBase.AbstractODEIntegrator{Alg,true,uType,tType}
    sol::solType
    opts::JVOptions{Atol,Rtol}
    zn::NTuple{L_MAX,uType} # Nordsieck history array

    # vectors with length `length(u0)`
    ewt::uType # error weight vector
    u::uType
    acor::uType
    tempv::uType
    ftemp::uType

    # step data
    q::Int                        # current order
    qprime::Int                   # order to be used on the next step {q-1, q, q+1}
    qwait::Int                    # number of steps to wait before order change
    #L::Int                        # L = q+1
    dt::tType                     # current step size
    dtprime::tType                # next step size
    eta::tType                    # eta = dtprime / dt
    dtscale::tType                # the step size information in `zn`
    t::tType                      # current time
    tau::NTuple{L_MAX,tType}      # tuple of previous `q+1` successful step sizes
    tq::NTuple{NUM_TESTS,tType}   # tuple of test quantities
    coeff::NTuple{L_MAX,uEltype}  # coefficients of l(x)
    rl2::uEltype                  # 1/l[2]
    gamma::uEltype                # gamma = h * rl2
    gammap::uEltype               # `gamma` at the last setup call
    gamrat::uEltype               # gamma/gammap
    crate::uEltype                # estimated corrector convergence rate
    acnrm::uEltype                # | acor | wrms
    mnewt::Int                    # Newton iteration counter

    # Limits
    qmax::Int                     # q <= qmax
    mxstep::Int                   # maximum number of internal steps for one user call
    maxcor::Int                   # maximum number of `nlsolve`
    mxhnil::Int                   # maximum number of warning messages issued to the
                                  # user that `t + h == t` for the next internal step

    dtmin::tType                  # |h| >= hmin
    dtmax_inv::tType              # |h| <= 1/hmax_inv
    etamax::tType                 # eta <= etamax

    # counters
    nst::Int                      # number of internal steps taken
    nfe::Int                      # number of f calls
    ncfn::Int                     # number of corrector convergence failures
    netf::Int                     # number of error test failures
    nni::Int                      # number of Newton iterations performed
    nsetups::Int                  # number of setup calls
    nhnil::Int                    # number of messages issued to the user that
                                  # `t + h == t` for the next iternal step
    lrw::Int                      # number of real words in CVODE work vectors
    liw::Int                      # no. of integer words in CVODE work vectors

    # saved vales
    qu::Int                       # last successful q value used
    nstlp::Int                    # step number of last setup call
    dtu::tType                    # last successful h value used
    saved_tq5::tType              # saved value of tq[5]
    jcur::Bool                    # Is the Jacobian info used by
                                  # linear solver current?
    tolsf::Float64                # tolerance scale factor
    setupNonNull::Bool            # Does setup do something?

    # Arrays for Optional Input and Optional Output

    #long int *cv_iopt::Int  # long int optional input, output */
    #real     *cv_ropt::Int  # real optional input, output     */
    function JVIntegrator(prob::ODEProblem, ::solType, ::Alg, ::Rtol, ::Atol) where {solType,Alg,Rtol,Atol}
        @unpack f, u0, tspan, p = prob
        obj = new{Alg,typeof(u0),eltype(tspan),eltype(u0),solType,Rtol,Atol,typeof(f),typeof(p)}()
        return obj
    end
end

###
### Default Constants
###
const HMIN_DEFAULT     = 0.0
const HMAX_INV_DEFAULT = 0.0
const MXHNIL_DEFAULT   = 10
const MXSTEP_DEFAULT   = 5000

###
### Routine-Specific Constants
###
macro define(name, val)
    esc(:(const $name = $val))
end
begin

    # CVodeDky */

    @define FUZZ_FACTOR 100.0

    # CVHin */

    @define HLB_FACTOR 100.0
    @define HUB_FACTOR 0.1
    @define H_BIAS     0.5
    @define MAX_ITERS  4

    # CVSet */

    @define CORTES 0.1

    # CVStep return values */

    @define SUCCESS_STEP      0
    @define REP_ERR_FAIL     -1
    @define REP_CONV_FAIL    -2
    @define SETUP_FAILED     -3
    @define SOLVE_FAILED     -4

    # CVStep control constants */

    @define PREDICT_AGAIN    -5
    @define DO_ERROR_TEST     1

    # CVStep */

    @define THRESH 1.5
    @define ETAMX1 10000.0
    @define ETAMX2 10.0
    @define ETAMX3 10.0
    @define ETAMXF 0.2
    @define ETAMIN 0.1
    @define ETACF  0.25
    @define ADDON  0.000001
    @define BIAS1  6.0
    @define BIAS2  6.0
    @define BIAS3  10.0
    @define ONEPSM 1.000001

    @define SMALL_NST    10   # nst > SMALL_NST => use ETAMX3          */
    @define MXNCF        10   # max no. of convergence failures during */
    # one step try                           */
    @define MXNEF         7   # max no. of error test failures during  */
    # one step try                           */
    @define MXNEF1        3   # max no. of error test failures before  */
    # forcing a reduction of order           */
    @define SMALL_NEF     2   # if an error failure occurs and         */
    # SMALL_NEF <= nef <= MXNEF1, then       */
    # reset eta =  MIN(eta, ETAMXF)          */
    @define LONG_WAIT    10   # number of steps to wait before         */
    # considering an order change when       */
    # q==1 and MXNEF1 error test failures    */
    # have occurred                          */

    # CVnls return values */

    @define SOLVED            0
    @define CONV_FAIL        -1
    @define SETUP_FAIL_UNREC -2
    @define SOLVE_FAIL_UNREC -3

    # CVnls input flags */

    @define FIRST_CALL      0
    @define PREV_CONV_FAIL -1
    @define PREV_ERR_FAIL  -2

    # CVnls other constants */

    @define FUNC_MAXCOR 3  # maximum no. of corrector iterations   */
    # for iter == FUNCTIONAL                */
    @define NEWT_MAXCOR 3  # maximum no. of corrector iterations   */
    # for iter == NEWTON                    */

    @define CRDOWN 0.3 # constant used in the estimation of the   */
    # convergence rate (crate) of the          */
    # iterates for the nonlinear equation      */
    @define DGMAX  0.3 # iter == NEWTON, |gamma/gammap-1| > DGMAX */
    # => call lsetup                           */

    @define RDIV        2  # declare divergence if ratio del/delp > RDIV  */
    @define MSBP       20  # max no. of steps between lsetup calls        */

    @define TRY_AGAIN  99  # control constant for CVnlsNewton - should be */
    # distinct from CVnls return values            */
end

###
### CVODE Implementation
###
function DiffEqBase.__init(prob::ODEProblem, alg::JVODEAlgorithm; reltol=1e-3, abstol=1e-6)
    # analogous to `CVodeMalloc`
    @unpack f, u0, tspan, p = prob
    # copy input paramters
    uType = typeof(u0)
    uEltype = eltype(u0)
    tType = eltype(tspan)
    timeseries = uType[]
    ts = tType[]

    destats = DiffEqBase.DEStats(0)
    # TODO: dense output
    dense = true
    #ks = Vector{uType}(undef, 0)
    #id = InterpolationData(f,timeseries,ts,ks,dense,cache)

    sol = DiffEqBase.build_solution(prob,alg,ts,timeseries,
                                    dense=dense, #k=ks, interp=id, TODO
                                    calculate_error = false, destats=destats)
    integrator = JVIntegrator(prob, sol, alg, reltol, abstol)

    integrator.sol = sol
    integrator.t = prob.tspan |> first
    integrator.opts = JVOptions(reltol, abstol)
    integrator.dtmin = HMIN_DEFAULT
    integrator.dtmax_inv = HMAX_INV_DEFAULT
    integrator.mxhnil = MXHNIL_DEFAULT
    integrator.mxstep = MXSTEP_DEFAULT
    integrator.maxcor = alg.nlsolve isa JVNewton ? NEWT_MAXCOR : FUNC_MAXCOR

    maxorder = alg isa Adams ? ADAMS_Q_MAX : BDF_Q_MAX

    # allocate the vectors
    integrator.u = copy(u0)
    integrator.zn = ntuple(i->i<=maxorder+1 ? similar(integrator.u) : similar(integrator.u, 0), Val(L_MAX))
    integrator.ewt = similar(integrator.u)
    integrator.acor = similar(integrator.u)
    integrator.tempv = similar(integrator.u)
    integrator.ftemp = similar(integrator.u)

    # set the `ewt` vector
    setewt!(integrator, integrator.u)

    # set step paramters
    integrator.q = 1
    integrator.qwait = integrator.q + 1
    integrator.etamax = maxorder

    # TODO: set the linear solver

    # initialize `zn[1]` in the history array
    copyto!(integrator.zn[1], u0)

    # initialize all counters
    integrator.nst = integrator.nfe = integrator.ncfn = integrator.netf = integrator.nni = integrator.nsetups = integrator.nhnil = integrator.lrw = integrator.liw = 0

    # initialize misc
    integrator.qu = 0
    integrator.dtu = zero(integrator.dtu)
    integrator.tolsf = 1
    return integrator
end

DiffEqBase.has_reinit(integrator::JVIntegrator) = true
function DiffEqBase.reinit!(integrator::JVIntegrator, u0=integrator.uprev;
                            alg=integrator.alg, tspan=integrator.tspan,
                            reltol=integrator.reltol, abstol=integrator.abstol,
                            p=integrator.p) # TODO: opts
    # TODO: `CVReInit`
end

function DiffEqBase.step!(integrator::JVIntegrator, dt=nothing)#, stop_at_tdt=false)
    @unpack f, p, tspan = integrator.sol.prob
    @unpack zn, t, dtmax_inv, dtmin, nst, mxstep, mxhnil = integrator
    tout = dt === nothing ? tspan[end] : t + dt
    if nst === 0 # first step
        f(zn[2], zn[1], p, t)
        integrator.nfe = 1
        integrator.dt = zero(integrator.dt)
        # TODO: user initdt
        if iszero(integrator.dt)
            dt_ok = initdt!(integrator, tout)
            dt_ok || error("tout=$tout too close to t0=$t to start integration")
        end
        #rh = abs(integrator.dt) * dtmax_inv
        #rh > one(rh) && (integrator.dt /= rh)
        abs(integrator.dt) < dtmin && (integrator.dt *= dtmin / abs(integrator.dt))
        integrator.dtscale = integrator.dt
        lmul!(integrator.dt, zn[2])
    end
    # If not the first call, check if tout already reached
    if dt !== nothing && nst > 0 && (integrator.t - tout)*integrator.dt >= zero(integrator.dt)
        # integrator.t = tout
        # interpolate???
        return integrator
    end

    nstloc = 0
    while true
        nextdt = integrator.dt
        nextq  = integrator.q

        # reset and check ewt
        nst > 0 && setewt!(integrator, zn[1])

        # check for too many steps
        if nstloc >= mxstep
            @warn "At t=$t, mxstep=$mxstep steps taken on this call before reaching tout=$tout"
            copyto!(integrator.u, integrator.zn[1])
            break
        end

        # TODO: check for too much accuracy requested

        # check for `dt` below roundoff
        t = integrator.t
        if t + integrator.dt == t
            integrator.nhnil += 1
            integrator.nhnil <= mxhnil && @warn "internal t=$t and step size dt=$(integrator.dt) are such that t + dt = t on the next step. The solver will continue anyway."
            integrator.nhnil == mxhnil && @warn "The above warning has been issued $mxhnil times and will not be issued again for this problem."
        end

        kflag = jvstep(integrator)
        t = integrator.t
        if kflag !== SUCCESS_STEP
            handlefailure(integrator, kflag)
            copyto!(integrator.u, zn[1])
            break
        end

        nstloc += 1

        if dt === nothing
            copyto!(integrator.u, zn[1])
            nextq  = integrator.qprime
            nextdt = integrator.dtprime
        end

        if (t-tout)*integrator.dt >= zero(integrator.t)
            copyto!(integrator.u, zn[1])
            nextq  = integrator.qprime
            nextdt = integrator.dtprime
            break
        end
    end
    return integrator
end

function jvstep(integrator::JVIntegrator)
    # TODO
    return SUCCESS_STEP
end

function initdt!(integrator::JVIntegrator, tout)::Bool
    @unpack zn, t = integrator
    # test for tout too close to t
    tdiff = tout - t
    iszero(tdiff) && return false
    tdir = sign(tdiff)
    tdist = abs(tdiff)
    tround = eps() * max(abs(t), abs(tout))
    tdist < 2tround && return false

    # Set lower and upper bounds on h0, and take geometric mean. Exit with this
    # value if the bounds cross each other.
    hlb = HLB_FACTOR * tround
    hub = let integrator=integrator, tdist=tdist
        temp1 = integrator.tempv
        temp2 = integrator.acor
        map!(abs, temp1, zn[1])
        map!(abs, temp2, zn[2])
        @.. temp1 = temp2 / muladd(HUB_FACTOR, temp1, integrator.opts.abstol)
        hub_inv = norm(temp1, Inf)
        hub = HUB_FACTOR * tdist
        hub*hub_inv > one(hub) && (hub = inv(hub_inv))
        hub
    end
    hg = sqrt(hlb*hub)
    if hub < hlb
        tdir == -1 && (hg = -hg)
        integrator.dt = hg
        return true
    end

    # Loop up to MAX_ITERS times to find h0.
    # Stop if new and previous values differ by a factor < 2.
    # Stop if hnew/hg > 2 after one iteration, as this probably means
    # that the ydd value is bad because of cancellation error.
    count = 0
    yddnorm = (integrator, hg) -> begin
        # This routine computes an estimate of the second derivative of y
        # using a difference quotient, and returns its WRMS norm.
        @unpack zn, tempv, ewt, u, t = integrator
        @unpack f, p = integrator.sol.prob
        @.. u = muladd(hg, zn[2], zn[1])
        f(tempv, u, p, t+hg)
        integrator.nfe += 1
        @. tempv = tempv - zn[2]
        lmul!(inv(hg), tempv)
        yddnrm = wrmsnorm(tempv, ewt)
    end
    local hnew
    while true
        hgs = hg*tdir
        yddnrm = yddnorm(integrator, hgs)
        hnew = (yddnrm*hub*hub > 2) ? sqrt(2/yddnrm) : sqrt(hg*hub)
        count += 1
        count >= MAX_ITERS && break
        hrat = hnew/hg
        (hrat > 0.5) && (hrat < 2) && break
        if (count >= 2) && (hrat > 2)
            hnew = hg
            break
        end
        hg = hnew
    end

    # Apply bounds, bias factor, and attach sign */

    h0 = H_BIAS*hnew
    h0 < hlb && (h0 = hlb)
    h0 > hub && (h0 = hub)
    tdir == -1 && (h0 = -h0)
    integrator.dt = h0
    return true
end

setewt!(integrator::JVIntegrator, u) = (@.. integrator.ewt = inv(integrator.opts.reltol * abs(u) + integrator.opts.abstol); return)
wrmsnorm(x, w) = sum(((x,w),)->abs2(x*w), zip(x, w))/length(x) |> sqrt

"""
    jvode(f, u0, tspan; rtol=1e-3, atol=1e-6)

DVODE: Variable-coefficient Ordinary Differential Equation solver,
with fixed-leading-coefficient implementation.
This version is in double precision.

DVODE solves the initial value problem for stiff or nonstiff
systems of first order ODEs,
    dy/dt = f(t,y) ,  or, in component form,
    dy(i)/dt = f(i) = f(i,t,y(1),y(2),...,y(NEQ)) (i = 1,...,NEQ).
DVODE is a package based on the EPISODE and EPISODEB packages, and
on the ODEPACK user interface standard, with minor modifications.
----------------------------------------------------------------------
Authors:
              Peter N. Brown and Alan C. Hindmarsh
              Center for Applied Scientific Computing, L-561
              Lawrence Livermore National Laboratory
              Livermore, CA 94551
and
              George D. Byrne
              Illinois Institute of Technology
              Chicago, IL 60616
----------------------------------------------------------------------
References:

1. P. N. Brown, G. D. Byrne, and A. C. Hindmarsh, "VODE: A Variable
   Coefficient ODE Solver," SIAM J. Sci. Stat. Comput., 10 (1989),
   pp. 1038-1051.  Also, LLNL Report UCRL-98412, June 1988.
2. G. D. Byrne and A. C. Hindmarsh, "A Polyalgorithm for the
   Numerical Solution of Ordinary Differential Equations,"
   ACM Trans. Math. Software, 1 (1975), pp. 71-96.
3. A. C. Hindmarsh and G. D. Byrne, "EPISODE: An Effective Package
   for the Integration of Systems of Ordinary Differential
   Equations," LLNL Report UCID-30112, Rev. 1, April 1977.
4. G. D. Byrne and A. C. Hindmarsh, "EPISODEB: An Experimental
   Package for the Integration of Systems of Ordinary Differential
   Equations with Banded Jacobians," LLNL Report UCID-30132, April
   1976.
5. A. C. Hindmarsh, "ODEPACK, a Systematized Collection of ODE
   Solvers," in Scientific Computing, R. S. Stepleman et al., eds.,
   North-Holland, Amsterdam, 1983, pp. 55-64.
6. K. R. Jackson and R. Sacks-Davis, "An Alternative Implementation
   of Variable Step-Size Multistep Formulas for Stiff ODEs," ACM
   Trans. Math. Software, 6 (1980), pp. 295-318.
"""
function jvode(f, u0, tspan; rtol=1e-3, atol=1e-6)
    maxorder = (12, 5)
    maxstep0 = 500
    maxwarns = 10
end

end # module


"""
REINFORCE
---------

Williams, 1992

This implementation follows the pseudo code given in
Sutton Barto, "RL an Intro", Chapter 13 Monte Carlo Policy Gradient.

The policy parameters θ as well as the feature vector ϕ are
column vectors. Here we use a states x actions matrix for simpler access.
Also, the feature vector ϕ is a one-hot vector in our case. Therefore we
do not need to keep a vector but access the entries by the
state/action pair [s_t, a_t]

See learner/REINFORCElearner.jl for the core implementation.

Note
----
get_all_puzzles(), get_all_episodes_in_puzzle(), get_data_by_puzzle_episode()
are in rl_data_toolsforalgos.jl
"""
function getLL_REINFORCE(data::Array{Int,2}, alpha::Float64,
                        gamma::Float64, temp::Float64; do_debug_print = false)
    data_stats, nr_states, nr_actions, nr_samples = get_data_stats_and_dimensions(data)
    if data_stats.nr_subjects != 1
        error("Invalid argument: getLL_REINFORCE supports only one subject")
    end
    if nr_actions!=2; error("REINFORCE supports two actions only") end
    all_puzzles = get_all_puzzles(data)

    LL_unbiased = 0.
    action_counter = 0
    probs_unbiased = zeros(Float64, nr_samples)
    Δpolicy_record = zeros(Float64, nr_samples)

    PolicyParametersImbalance_record = zeros(Float64, nr_samples)
    PolicyImbalance_record = zeros(Float64, nr_samples)
    PolicyParametersImbalanceAllstates_record = Array{Array{Float64,1}, 1}(undef, nr_samples)
    [PolicyParametersImbalanceAllstates_record[i] = zeros(Float64, nr_states) for i in 1:nr_samples]
    PolicyParametersDiffChosenVsUnchosen_record = zeros(Float64, nr_samples)
    PolicyDiffChosenVsUnchosen_record = zeros(Float64, nr_samples)

    modulatingfactor = 1.
    learner = []
    for i_puzzle in all_puzzles
        # Puzzle = independent block of the experiment.
        # The code was made for general purposes, in this experiment we
        # have only ONE puzzle.
        learner = REINFORCE(nr_states = nr_states,
                        nr_actions = nr_actions,
                        α = alpha,
                        γ = gamma,
                        do_debug_print = do_debug_print)

        all_epi_in_puzzle = get_all_episodes_in_puzzle(data, i_puzzle)
        for i_episode in all_epi_in_puzzle
            episode_data = get_data_by_puzzle_episode(data, i_puzzle, i_episode)
            path_length = size(episode_data)[1]  # = T
            if do_debug_print
                println("i_puzzle:$i_puzzle | i_episode:$i_episode | path_length:$path_length")
            end
            reward = 0. # we have a single reward at goal state. If the episode was
            # not completed, this code still correctly works (updates are 0)
            for t in 1:path_length
                action_counter += 1
                state = episode_data[t, Int(c_state)]
                action = episode_data[t, Int(c_action)]
                reward = episode_data[t, Int(c_reward)]
                next_state = episode_data[t, Int(c_nextState)]
                done = reward>0 ? true : false

                p_actions = getSoftmaxProbs([learner.θ[state,1],learner.θ[state,2]], temp, [0., 0.])
                LL_unbiased += log(p_actions[action])
                probs_unbiased[action_counter] = copy(p_actions[action])

                PolicyParametersImbalance_record[action_counter] = abs(learner.θ[state, 1] - learner.θ[state, 2])
                PolicyImbalance_record[action_counter] = abs(p_actions[1] - p_actions[2])
                PolicyParametersImbalanceAllstates_record[action_counter] = abs.(learner.θ[:, 1] - learner.θ[:, 2])
                PolicyParametersDiffChosenVsUnchosen_record[action_counter] = learner.θ[state, action] - learner.θ[state, mod(action,2)+1]
                PolicyDiffChosenVsUnchosen_record[action_counter] = p_actions[action] - p_actions[mod(action,2)+1]

                θ_before = deepcopy(learner.θ)

                update!(learner, state, action, reward, p_actions[action], modulatingfactor)
                Δpolicy_record[action_counter] = sum(abs.(learner.θ .- θ_before))

                if do_debug_print
                    println("t:$t | s:$state | a:$action | R:$reward | s':$next_state | p(s,a):$(round(p_action,digits=2))")
                end
            end # path_length (in single episode)
        end # episode
    end # puzzzle
  #---- Without bias
  LL = LL_unbiased
  trial_by_trial_p_sa = (probs_unbiased)

  algo = "REINFORCE"
  param_v = [alpha, gamma, temp]
  param_str = "alpha=$(round(alpha,digits=4)), gamma=$(round(gamma,digits=4)), temp=$(round(temp,digits=4))"

  results_dict = Dict{String,Any}("final_Theta" => learner.θ,
                                "DeltaPolicy" => Δpolicy_record,
                                "PolicyParametersImbalance" => PolicyParametersImbalance_record,
                                "PolicyImbalance" => PolicyImbalance_record,
                                "PolicyParametersImbalanceAllstates" => PolicyParametersImbalanceAllstates_record,
                                "PolicyParametersDiffChosenVsUnchosen" => PolicyParametersDiffChosenVsUnchosen_record,
                                "PolicyDiffChosenVsUnchosen" => PolicyDiffChosenVsUnchosen_record)

  signals = Step_by_step_signals(
    data_stats.all_subject_ids[1],
    [], # no RPE
    [], # no SPE
    trial_by_trial_p_sa,
    LL,
    data_stats,
    algo,
    param_v,
    param_str,
    99, #biased_action_idx, dummy value
    results_dict)
  return LL, signals

end


import UncertainData:
    AbstractUncertainIndexValueDataset,
    ConstrainedIndexValueResampling,
    SequentialResampling,
    SequentialInterpolatedResampling,
    constrain,
    resample

function causality(source::AbstractUncertainIndexValueDataset,
    target::AbstractUncertainIndexValueDataset,
    test::CausalityTest,
    resampling::ConstrainedIndexValueResampling{N_VARIABLES, N_DATASETS}) where {N_VARIABLES, N_DATASETS}
    @warn "Performing causality test on uncertain index-value datasets without explicitly providing sequential sampling constraints. Results may be nonsensical, because the values are not sampled according to increasing indices."

    N_DATASETS != 2 ? error("Too many constraints ($N_DATASETS) provided. This `causality` method only takes `source` and `target` as data inputs. Thus, you should provide an `ConstrainedIndexValueResampling` containing two (`index_constraints::CONSTRAINT_TYPES, value_constraints::CONSTRAINT_TYPES)`-tuples. For example, `ConstrainedIndexValueResampling((TruncateStd(0.5), TruncateQuantiles(0.3, 0.7)), (TruncateStd(0.7), TruncateStd(1.1)))` would work.") : nothing

    # Take the (index_constraints, values_constraints)-tuple for the source and target
    # and represent as separate `IndexValueResampling` instances.
    cs_src = ConstrainedIndexValueResampling(resampling[1])
    cs_targ = ConstrainedIndexValueResampling(resampling[2])

    constrained_source = constrain(source, cs_src)
    constrained_target = constrain(target, cs_targ)

    [causality(resample(constrained_source.values), resample(constrained_target.values), test) for i = 1:resampling.n]
end

function causality(source::AbstractUncertainIndexValueDataset,
    target::AbstractUncertainIndexValueDataset,
    test::CausalityTest,
    constrained_resampling::ConstrainedIndexValueResampling{N_VARIABLES, N_DATASETS},
    sequential_resampling::SequentialInterpolatedResampling{S, G}) where {N_VARIABLES, N_DATASETS, S, G}

    N_DATASETS != 2 ? error("Too many constraints ($N_DATASETS) provided. This `causality` method only takes `source` and `target` as data inputs. Thus, you should provide an `ConstrainedIndexValueResampling` containing two (`index_constraints::CONSTRAINT_TYPES, value_constraints::CONSTRAINT_TYPES)`-tuples. For example, `ConstrainedIndexValueResampling((TruncateStd(0.5), TruncateQuantiles(0.3, 0.7)), (TruncateStd(0.7), TruncateStd(1.1)))` would work.") : nothing

    # Take the (index_constraints, values_constraints)-tuple for the source and target
    # and represent as separate `IndexValueResampling` instances.
    cs_src = ConstrainedIndexValueResampling(constrained_resampling[1])
    cs_targ = ConstrainedIndexValueResampling(constrained_resampling[2])

    constrained_source = constrain(source, cs_src)
    constrained_target = constrain(target, cs_targ)

    [causality(resample(constrained_source, sequential_resampling)[2],
               resample(constrained_target, sequential_resampling)[2],
               test) for i = 1:constrained_resampling.n]
end

function causality(source::AbstractUncertainIndexValueDataset,
    target::AbstractUncertainIndexValueDataset,
    test::CausalityTest,
    constrained_resampling::ConstrainedIndexValueResampling{N_VARIABLES, N_DATASETS},
    sequential_resampling::SequentialResampling{S}) where {N_VARIABLES, N_DATASETS, S}

    N_DATASETS != 2 ? error("Too many constraints ($N_DATASETS) provided. This `causality` method only takes `source` and `target` as data inputs. Thus, you should provide an `ConstrainedIndexValueResampling` containing two (`index_constraints::CONSTRAINT_TYPES, value_constraints::CONSTRAINT_TYPES)`-tuples. For example, `ConstrainedIndexValueResampling((TruncateStd(0.5), TruncateQuantiles(0.3, 0.7)), (TruncateStd(0.7), TruncateStd(1.1)))` would work.") : nothing

    # Take the (index_constraints, values_constraints)-tuple for the source and target
    # and represent as separate `IndexValueResampling` instances.
    cs_src = ConstrainedIndexValueResampling(constrained_resampling[1])
    cs_targ = ConstrainedIndexValueResampling(constrained_resampling[2])

    constrained_source = constrain(source, cs_src)
    constrained_target = constrain(target, cs_targ)

    [causality(resample(constrained_source, sequential_resampling)[2],
               resample(constrained_target, sequential_resampling)[2],
               test) for i = 1:constrained_resampling.n]
end


using OperatorFlux
using Test
using FFTW
using Flux
using LinearAlgebra
using ChainRulesCore
using ChainRulesTestUtils

@testset "ChainRulesCore Extensions" begin
    # FFTW.r2r
    M = 4
    trafo = ChebyshevTransform(modes = (4,))
    c = rand(3, M, 5)
    test_rrule(
        FFTW.r2r,
        c,
        FFTW.REDFT00,
        2:2,
        check_thunked_output_tangent = false,
    )

    # truncate modes
    M = 4
    trafo = FourierTransform(modes = (2, 2))
    c = rand(3, M, M, 5)
    test_rrule(
        OperatorFlux.truncate_modes,
        trafo,
        c,
        check_thunked_output_tangent = false,
    )

    trafo = ChebyshevTransform(modes = (2, 2))
    test_rrule(
        OperatorFlux.truncate_modes,
        trafo,
        c,
        check_thunked_output_tangent = false,
    )

    for mode in [(5,), (5, 6), (5, 7, 9)]
        z = rand(3, mode..., 5)
        trafos = (LegendreTransform(modes = mode),)
        for trafo in trafos
            test_rrule(
                OperatorFlux.truncate_modes,
                trafo,
                z,
                check_thunked_output_tangent = false,
            )
        end
    end

    # pad modes
    c = rand(3, 2, 8, 5)
    trafo = FourierTransform(modes = (2, 4))
    test_rrule(
        OperatorFlux.pad_modes,
        trafo,
        c,
        (4, 8),
        check_thunked_output_tangent = false,
    )

    trafo = ChebyshevTransform(modes = (2, 4))
    c = rand(3, 2, 4, 5)
    test_rrule(
        OperatorFlux.pad_modes,
        trafo,
        c,
        (2, 4),
        check_thunked_output_tangent = false,
    )

    for mode in [(5,), (5, 6), (5, 7, 9)]
        z = rand(3, mode..., 5)
        trafos = (LegendreTransform(modes = mode),)
        for trafo in trafos
            test_rrule(
                OperatorFlux.pad_modes,
                trafo,
                z,
                mode,
                check_thunked_output_tangent = false,
            )
        end
    end
end