text:
include("bonds.jl")
include("itensors_timers.jl")

import LinearAlgebra
LinearAlgebra.BLAS.set_num_threads(1)

import ITensors
ITensors.disable_combine_contract()
#ITensors.enable_combine_contract()
using ITensors: Index, QN
using DelimitedFiles

K = 100 # number of repititions; scale up for more reliable testing.

# MPO Contractions

# TRIVIAL = no symmetry
N = length(mpodims_triv[1]) # number of different cases
mpo_triv_times = zeros(Float64, (K, N))
mpo_triv_times_gc = zeros(Float64, (K, N))
for i = 1:N
    Dmps = mpodims_triv[1][i]
    Dmpo = mpodims_triv[2][i]
    Dphys = mpodims_triv[3][i]

    Vmps = Index(Dmps)
    Vmpo = Index(Dmpo)
    Vphys = Index(Dphys)

    mpo_timer = ITensorsTimers.mpo_timer(; Vmps = Vmps, Vmpo = Vmpo, Vphys = Vphys)

    times, times_gc = mpo_timer(; outer = K)
    mpo_triv_times[:, i] = times
    mpo_triv_times_gc[:, i] = times_gc

    tavg = sum(times)/K

    println("MPO Contraction, trivial symmetry: case $i out of $N finished at average time $tavg")
end

writedlm("results/itensors_mpo_triv.txt", mpo_triv_times)
writedlm("results/itensors_mpo_triv_gc.txt", mpo_triv_times_gc)

# Z2 Symmetry
N = length(mpodims_z2[1]) # number of different cases
mpo_z2_times = zeros(Float64, (K, N))
mpo_z2_times_gc = zeros(Float64, (K, N))
for i = 1:N
    Dmps = distributeZ2(mpodims_z2[1][i])
    Dmpo = distributeZ2(mpodims_z2[2][i])
    Dphys = distributeZ2(mpodims_z2[3][i])

    Vmps = Index((QN(a,2)=>b for (a,b) in Dmps)...)
    Vmpo = Index((QN(a,2)=>b for (a,b) in Dmpo)...)
    Vphys = Index((QN(a,2)=>b for (a,b) in Dphys)...)

    mpo_timer = ITensorsTimers.mpo_timer(; Vmps = Vmps, Vmpo = Vmpo, Vphys = Vphys)

    times, times_gc = mpo_timer(; outer = K)
    mpo_z2_times[:, i] = times
    mpo_z2_times_gc[:, i] = times_gc

    tavg = sum(times)/K

    println("MPO Contraction, Z2 symmetry: case $i out of $N finished at average time $tavg")
end

writedlm("results/itensors_mpo_z2.txt", mpo_z2_times)
writedlm("results/itensors_mpo_z2_gc.txt", mpo_z2_times_gc)

# U1 Symmetry
N = length(mpodims_u1[1]) # number of different cases
mpo_u1_times = zeros(Float64, (K, N))
mpo_u1_times_gc = zeros(Float64, (K, N))
for i = 1:N
    Dmps = distributeU1(mpodims_u1[1][i])
    Dmpo = distributeU1(mpodims_u1[2][i])
    Dphys = distributeU1(mpodims_u1[3][i])

    Vmps = Index((QN(a)=>b for (a,b) in Dmps)...)
    Vmpo = Index((QN(a)=>b for (a,b) in Dmpo)...)
    Vphys = Index((QN(a)=>b for (a,b) in Dphys)...)

    mpo_timer = ITensorsTimers.mpo_timer(; Vmps = Vmps, Vmpo = Vmpo, Vphys = Vphys)

    times, times_gc = mpo_timer(; outer = K)
    mpo_u1_times[:, i] = times
    mpo_u1_times_gc[:, i] = times_gc

    tavg = sum(times)/K

    println("MPO Contraction, U1 symmetry: case $i out of $N finished at average time $tavg")
end

writedlm("results/itensors_mpo_u1.txt", mpo_u1_times)
writedlm("results/itensors_mpo_u1_gc.txt", mpo_u1_times_gc)

## PEPO Contractions

# TRIVIAL = no symmetry
N = length(pepodims_triv[1]) # number of different cases
pepo_triv_times = zeros(Float64, (K, N))
pepo_triv_times_gc = zeros(Float64, (K, N))
for i = 1:N
    Dpeps = pepodims_triv[1][i]
    Dpepo = pepodims_triv[2][i]
    Dphys = pepodims_triv[3][i]
    Denv = pepodims_triv[4][i]

    Vpeps = Index(Dpeps)
    Vpepo = Index(Dpepo)
    Vphys = Index(Dphys)
    Venv = Index(Denv)

    pepo_timer = ITensorsTimers.pepo_timer(; Vpeps = Vpeps, Vpepo = Vpepo, Vphys = Vphys, Venv = Venv)

    times, times_gc = pepo_timer(; outer = K)
    pepo_triv_times[:, i] = times
    pepo_triv_times_gc[:, i] = times_gc

    tavg = sum(times)/K

    println("PEPO Contraction, trivial symmetry: case $i out of $N finished at average time $tavg")
end

writedlm("results/itensors_pepo_triv.txt", pepo_triv_times)
writedlm("results/itensors_pepo_triv_gc.txt", pepo_triv_times_gc)

# Z2 Symmetry
N = length(pepodims_z2[1]) # number of different cases
pepo_z2_times = zeros(Float64, (K, N))
pepo_z2_times_gc = zeros(Float64, (K, N))
for i = 1:N
    Dpeps = distributeZ2(pepodims_z2[1][i])
    Dpepo = distributeZ2(pepodims_z2[2][i])
    Dphys = distributeZ2(pepodims_z2[3][i])
    Denv = distributeZ2(pepodims_z2[4][i])

    Vpeps = Index((QN(a,2)=>b for (a,b) in Dpeps)...)
    Vpepo = Index((QN(a,2)=>b for (a,b) in Dpepo)...)
    Vphys = Index((QN(a,2)=>b for (a,b) in Dphys)...)
    Venv = Index((QN(a,2)=>b for (a,b) in Denv)...)

    pepo_timer = ITensorsTimers.pepo_timer(; Vpeps = Vpeps, Vpepo = Vpepo, Vphys = Vphys, Venv = Venv)

    times, times_gc = pepo_timer(; outer = K)
    pepo_z2_times[:, i] = times
    pepo_z2_times_gc[:, i] = times_gc

    tavg = sum(times)/K

    println("PEPO Contraction, Z2 symmetry: case $i out of $N finished at average time $tavg")
end

writedlm("results/itensors_pepo_z2.txt", pepo_z2_times)
writedlm("results/itensors_pepo_z2_gc.txt", pepo_z2_times_gc)

# U1 Symmetry
N = length(pepodims_u1[1]) # number of different cases
pepo_u1_times = zeros(Float64, (K, N))
pepo_u1_times_gc = zeros(Float64, (K, N))
for i = 1:N
    Dpeps = distributeU1(pepodims_u1[1][i])
    Dpepo = distributeU1(pepodims_u1[2][i])
    Dphys = distributeU1(pepodims_u1[3][i])
    Denv = distributeU1(pepodims_u1[4][i])

    Vpeps = Index((QN(a)=>b for (a,b) in Dpeps)...)
    Vpepo = Index((QN(a)=>b for (a,b) in Dpepo)...)
    Vphys = Index((QN(a)=>b for (a,b) in Dphys)...)
    Venv = Index((QN(a)=>b for (a,b) in Denv)...)

    pepo_timer = ITensorsTimers.pepo_timer(; Vpeps = Vpeps, Vpepo = Vpepo, Vphys = Vphys, Venv = Venv)

    times, times_gc = pepo_timer(; outer = K)
    pepo_u1_times[:, i] = times
    pepo_u1_times_gc[:, i] = times_gc

    tavg = sum(times)/K

    println("PEPO Contraction, U1 symmetry: case $i out of $N finished at average time $tavg")
end

writedlm("results/itensors_pepo_u1.txt", pepo_u1_times)
writedlm("results/itensors_pepo_u1_gc.txt", pepo_u1_times_gc)

# ## MERA Contractions

# TRIVIAL = no symmetry
N = length(meradims_triv[1]) # number of different cases
mera_triv_times = zeros(Float64, (K, N))
mera_triv_times_gc = zeros(Float64, (K, N))
for i = 1:N
    Dmera = meradims_triv[1][i]

    Vmera = Index(Dmera)

    mera_timer = ITensorsTimers.mera_timer(; Vmera = Vmera)

    times, times_gc = mera_timer(; outer = K)

    mera_triv_times[:, i] = times
    mera_triv_times_gc[:, i] = times_gc

    tavg = sum(times)/K

    println("MERA Contraction, trivial symmetry: case $i out of $N finished at average time $tavg")
end

writedlm("results/itensors_mera_triv.txt", mera_triv_times)
writedlm("results/itensors_mera_triv_gc.txt", mera_triv_times_gc)

# Z2 Symmetry
N = length(meradims_z2[1]) # number of different cases
mera_z2_times = zeros(Float64, (K, N))
mera_z2_times_gc = zeros(Float64, (K, N))
for i = 1:N
    Dmera = distributeZ2(meradims_z2[1][i])

    Vmera = Index((QN(a,2)=>b for (a,b) in Dmera)...)

    mera_timer = ITensorsTimers.mera_timer(; Vmera = Vmera)

    times, times_gc = mera_timer(; outer = K)

    mera_z2_times[:, i] = times
    mera_z2_times_gc[:, i] = times_gc

    tavg = sum(times)/K

    println("MERA Contraction, Z2 symmetry: case $i out of $N finished at average time $tavg")
end

writedlm("results/itensors_mera_z2.txt", mera_z2_times)
writedlm("results/itensors_mera_z2_gc.txt", mera_z2_times_gc)

# U1 Symmetry
N = length(meradims_u1[1]) # number of different cases
mera_u1_times = zeros(Float64, (K, N))
mera_u1_times_gc = zeros(Float64, (K, N))
for i = 1:N
    Dmera = distributeU1(meradims_u1[1][i])

    Vmera = Index((QN(a)=>b for (a,b) in Dmera)...)

    mera_timer = ITensorsTimers.mera_timer(; Vmera = Vmera)

    times, times_gc = mera_timer(; outer = K)

    mera_u1_times[:, i] = times
    mera_u1_times_gc[:, i] = times_gc

    tavg = sum(times)/K

    println("MERA Contraction, U1 symmetry: case $i out of $N finished at average time $tavg")
end

writedlm("results/itensors_mera_u1.txt", mera_u1_times)
writedlm("results/itensors_mera_u1_gc.txt", mera_u1_times_gc)


text:
# Practical Machine learning
# Ensemble learning example
# Chapter 13

# MLBase transformers.
module MLBaseWrapper

importall Orchestra.Types
importall Orchestra.Util

import MLBase: Standardize, estimate, transform

export StandardScaler,
       fit!,
       transform!

# Standardizes each feature using (X - mean) / stddev.
# Will produce NaN if standard deviation is zero.
type StandardScaler <: Transformer
  model
  options

  function StandardScaler(options=Dict())
    default_options = {
      :center => true,
      :scale => true
    }
    new(nothing, nested_dict_merge(default_options, options))
  end
end

function fit!(st::StandardScaler, instances::Matrix, labels::Vector)
  st_transform = estimate(Standardize, instances'; st.options...)
  st.model = {
    :standardize_transform => st_transform
  }
end

function transform!(st::StandardScaler, instances::Matrix)
  st_transform = st.model[:standardize_transform]
  transposed_instances = instances'
  return transform(st_transform, transposed_instances)'
end

end # module

ext: jl


text:
export demean!, demean, detrend!, detrend

@doc """
    demean!(S::SeisData[; chans=CC, irr=false])

Remove the mean from all channels `i` with `S.fs[i] > 0.0`. Specify `irr=true`
to also remove the mean from irregularly sampled channels (with S.fs[i] == 0.0).
Specifying a channel list with `chans=CC` restricts processing to channels CC.

    demean!(C::SeisChannel)

Remove the mean from data in `C`.

Ignores NaNs.
""" demean!
function demean!(S::GphysData;
  chans::ChanSpec=Int64[],
  irr::Bool=false)

  if chans == Int64[]
    chans = 1:S.n
  end

  for i in chans
    (irr==false && S.fs[i]<=0.0) && continue
    T = eltype(S.x[i])
    K = findall(isnan.(S.x[i]))
    if isempty(K)
      L = length(S.x[i])
      μ = T(sum(S.x[i]) / T(L))
      for j = 1:L
        S.x[i][j] -= μ
      end
    else
      J = findall(isnan.(S.x[i]) .== false)
      L = length(J)
      μ = T(sum(S.x[i][J])/T(L))
      for j in J
        S.x[i][j] -= μ
      end
    end
    proc_note!(S, i, string("demean!(S, chans=", i, ")"), "removed mean")
  end
  return nothing
end

function demean!(C::GphysChannel)
  T = eltype(C.x)
  K = findall(isnan.(C.x))
  if isempty(K)
    L = length(C.x)
    μ = T(sum(C.x) / T(L))
    for j = 1:L
      C.x[j] -= μ
    end
  else
    J = findall(isnan.(C.x) .== false)
    L = length(J)
    μ = T(sum(C.x[J])/T(L))
    for j in J
      C.x[j] -= μ
    end
  end
  proc_note!(C, "demean!(C)", "removed mean")
  return nothing
end

@doc (@doc demean!)
demean(S::GphysData,
  chans::ChanSpec=Int64[],
  irr::Bool=false) = (U = deepcopy(S); demean!(U, chans=chans, irr=irr); return U)
demean(C::GphysChannel) = (U = deepcopy(C); demean!(U); return U)

function dtr!(x::AbstractArray{T,1}, ti::Array{Int64,2}, fs::Float64, n::Int64) where T <: AbstractFloat
  L = length(x)

  # check for nans
  nf = false
  for i = 1:length(x)
    if isnan(x[i])
      nf = true
      break
    end
  end

  if nf
    t = tx_float(ti, fs)
    j = findall((isnan.(x)).==false)
    x1 = x[j]
    t1 = t[j]
    p = n == 1 ? linreg(t1, x1) : polyfit(t1, x1, n)
  else
    if n == 1 && size(ti,1) == 2 && fs > 0.0
      dt = 1.0/fs
      p = linreg(x, dt)
      v = zero(T)
      for i = 1:length(x)
        v = polyval(p, dt*i)
        x[i] -= v
      end
    else
      t = tx_float(ti, fs)
      p = n == 1 ? linreg(t, x) : polyfit(t, x, n)
      v = zero(T)
      for i = 1:length(x)
        v = polyval(p, t[i])
        x[i] -= v
      end
    end
  end
  return p
end

@doc """
    detrend!(S::SeisData[; chans=CC, n=1]))

Remove the linear trend from channels `CC`. Ignores NaNs.

To remove a higher-order polynomial fit than a linear trend, choose `n` >1.

    detrend!(C::SeisChanel[; n=1]))

Remove the linear trend from data in `C`. Ignores NaNs.

To remove a higher-order polynomial fit than a linear trend, choose n>1.

!!! warning

    detrend! does *not* check for data gaps; if this is problematic, call ungap!(S, m=true) first!
""" detrend!
function detrend!(S::GphysData;
  chans::ChanSpec=Int64[],
  n::Int64=1)

  if chans == Int64[]
    chans = 1:S.n
  end

  @inbounds for i in chans
    p = dtr!(S.x[i], S.t[i], S.fs[i], n)
    proc_note!(S, i, string("detrend!(S, chans=", i, ", n = ", n, ")"),
      string("detrended order ", n, " polynomial; polyfit result = ", p))
  end
  return nothing
end

function detrend!(C::GphysChannel; n::Int64=1)
  p = dtr!(C.x, C.t, C.fs, n)
  proc_note!(C, string("detrend!(C, n = ", n, ")"),
    string("detrended order ", n, " polynomial; polyfit result = ", p))
  return nothing
end

@doc (@doc detrend!)
detrend(S::GphysData;
  chans::ChanSpec=Int64[],
  n::Int64=1) = (U = deepcopy(S); detrend!(U, chans=chans, n=n); return U)
detrend(C::GphysChannel; n::Int64=1) = (U = deepcopy(C); detrend!(U, n=n); return U)


text:
s = """Veggies es bonus vobis, proinde vos postulo essum magis kohlrabi welsh onion daikon amaranth tatsoi tomatillo melon azuki bean garlic.

Gumbo beet greens corn soko endive gumbo gourd. Parsley shallot courgette tatsoi pea sprouts fava bean collard greens dandelion okra wakame tomato. Dandelion cucumber earthnut pea peanut soko zucchini.

Turnip greens yarrow ricebean rutabaga endive cauliflower sea lettuce kohlrabi amaranth water spinach avocado daikon napa cabbage asparagus winter purslane kale.
Celery potato scallion desert raisin horseradish spinach carrot soko. Lotus root water spinach fennel kombu maize bamboo shoot green bean swiss chard seakale pumpkin onion chickpea gram corn pea. Brussels sprout coriander water chestnut gourd swiss chard wakame kohlrabi beetroot carrot watercress.
Corn amaranth salsify bunya nuts nori azuki bean chickweed potato bell pepper artichoke.
"""

@testset "context" begin
    mess = F.context(s, 101)
    @test s[101] == 't'
    # println(mess)
    @test mess == "Context:\n\t...ikon amaranth tatsoi tomatillo melon azuki ... (near line 1)\n	                        ^---\n"

    mess = F.context(s, 211)
    @test s[211] == 't'
    # println(mess)
    @test mess == "Context:\n\t...ey shallot courgette tatsoi pea sprouts fav... (near line 2)\n	                        ^---\n"

    mess = F.context(s, 10)
    # println(mess)
    @test mess == "Context:\n\tVeggies es bonus vobis, proinde... (near line 1)\n	         ^---\n"

    mess = F.context(s, 880)
    # println(mess)
    @test mess == "Context:\n\t... potato bell pepper artichoke. (near line 5)\n	                        ^---\n"
end

@testset "show" begin
    ocbe = F.OCBlockError("foo", "bar")
    io = IOBuffer()
    Base.showerror(io, ocbe)
    r = String(take!(io))
    @test r == "foo\nbar"

    mcbe = F.MathBlockError("foo")
    @test mcbe.m == "foo"
end

@testset "ocbe" begin
    s = raw"""
        Foo $$ end.
        """
    @test_throws F.OCBlockError s |> fd2html_td
end

@testset "lxbe" begin
    s = raw"""
        Foo
        \newcommand{\hello}
        End
        """
    @test_throws F.LxDefError s |> fd2html_td

    s = raw"""
        Foo
        \newcommand{\hello} {goodbye}
        \hello
        End
        """ |> fd2html_td
    @test isapproxstr(s, "
Foo
 goodbye End
")

    s2 = raw"""
        Foo
        \newcommand{\hello}[ ]{goodbye}
        \hello
        End
        """ |> fd2html_td
    @test isapproxstr(s, s2)

    # should  fail XXX

    s2 = raw"""
        Foo
        \newcommand{\hello}b{goodbye}
        \hello
        End
        """
    @test_throws F.LxDefError s2 |> fd2html_td

    s2 = raw"""
        Foo
        \newcommand{\hello}b{goodbye #1}
        \hello{ccc}
        End
        """
    @test_throws F.LxDefError s2 |> fd2html_td

    s2 = raw"""
        Foo
        \newcommand{\hello}[bb]{goodbye #1}
        \hello{ccc}
        End
        """
    @test_throws F.LxDefError s2 |> fd2html_td

    # tolerated

    s2 = raw"""
        Foo
        \newcommand{\hello}  {goodbye}
        \hello
        End
        """ |> fd2html_td
    @test isapproxstr(s, s2)
end

@testset "Unfound" begin
    fs2();
    @test_throws F.FileNotFoundError F.resolve_rpath("./foo")
end

@testset "H-For" begin
    s = """
        {{for x in blah}}
        foo
        """
    @test_throws F.HTMLBlockError s |> F.convert_html
    s = """
        @def list = [1,2,3]
        ~~~
        {{for x in list}}
        foo
        ~~~
        """
    @test_throws F.HTMLBlockError s |> fd2html_td
    s = """
        @def list = [1,2,3]
        ~~~
        {{for x in list2}}
        foo
        {{end}}
        ~~~
        """
    @test_throws F.HTMLBlockError s |> fd2html_td
    s = """
        @def list = [1,2,3]
        ~~~
        {{for x in list}}
        {{if a}}
        foo
        {{end}}
        ~~~
        """
    @test_throws F.HTMLBlockError s |> fd2html_td
end

@testset "HToc" begin
    s = """
        {{toc 1}}
        """
    @test_throws F.HTMLFunctionError s |> F.convert_html
    s = """
        ~~~
        {{toc aa bb}}
        ~~~
        # Hello
        """
    @test_throws F.HTMLFunctionError s |> fd2html_td
end


text:
"""
    AbstractSampling

Abstract type for sparse sampling.

Encodes the "basis transformation" of a propagator from the truncated IR
basis coefficients `G_ir[l]` to time/frequency sampled on sparse points
`G(x[i])` together with its inverse, a least squares fit:

         ________________                   ___________________
        |                |    evaluate     |                   |
        |     Basis      |---------------->|     Value on      |
        |  coefficients  |<----------------|  sampling points  |
        |________________|      fit        |___________________|

"""
abstract type AbstractSampling{T,Tmat,F<:SVD} end

function cond(sampling::AbstractSampling)
    return first(sampling.matrix_svd.S) / last(sampling.matrix_svd.S)
end

function Base.show(io::IO, smpl::S) where {S<:AbstractSampling}
    println(io, S)
    print(io, "Sampling points: ")
    return println(io, smpl.sampling_points)
end

"""
    TauSampling <: AbstractSampling

Sparse sampling in imaginary time.

Allows the transformation between the IR basis and a set of sampling points
in (scaled/unscaled) imaginary time.
"""
struct TauSampling{T,Tmat,F<:SVD} <: AbstractSampling{T,Tmat,F}
    sampling_points::Vector{T}
    matrix::Matrix{Tmat}
    matrix_svd::F
end

"""
    TauSampling(basis[, sampling_points])

Construct a `TauSampling` object. If not given, the `sampling_points` are chosen as
the extrema of the highest-order basis function in imaginary time. This turns out
to be close to optimal with respect to conditioning for this size (within a few percent).
"""
function TauSampling(
    basis::AbstractBasis, sampling_points=default_tau_sampling_points(basis)
)
    matrix = eval_matrix(TauSampling, basis, sampling_points)
    sampling = TauSampling(sampling_points, matrix, svd(matrix))

    if iswellconditioned(basis) && cond(sampling) > 1e8
        @warn "Sampling matrix is poorly conditioned (cond = $(cond(sampling)))."
    end

    return sampling
end

const TauSampling64 = @static if VERSION < v"1.9-"
    TauSampling{Float64,Float64,SVD{Float64,Float64,Matrix{Float64}}}
else
    TauSampling{Float64,Float64,SVD{Float64,Float64,Matrix{Float64},Vector{Float64}}}
end

"""
    MatsubaraSampling <: AbstractSampling

Sparse sampling in Matsubara frequencies.

Allows the transformation between the IR basis and a set of sampling points
in (scaled/unscaled) imaginary frequencies.
"""
struct MatsubaraSampling{T,Tmat,F<:SVD} <: AbstractSampling{T,Tmat,F}
    sampling_points::Vector{T}
    matrix::Matrix{Tmat}
    matrix_svd::F
end

const MatsubaraSampling64 = @static if VERSION < v"1.9-"
    MatsubaraSampling{Int64,ComplexF64,SVD{ComplexF64,Float64,Matrix{ComplexF64}}}
else
    MatsubaraSampling{
        Int64,ComplexF64,SVD{ComplexF64,Float64,Matrix{ComplexF64},Vector{Float64}}
    }
end

"""
    MatsubaraSampling(basis[, sampling_points])

Construct a `MatsubaraSampling` object. If not given, the `sampling_points` are chosen as
the (discrete) extrema of the highest-order basis function in Matsubara. This turns out
to be close to optimal with respect to conditioning for this size (within a few percent).
"""
function MatsubaraSampling(
    basis::AbstractBasis, sampling_points=default_matsubara_sampling_points(basis)
)
    matrix = eval_matrix(MatsubaraSampling, basis, sampling_points)
    sampling = MatsubaraSampling(sampling_points, matrix, svd(matrix))

    if iswellconditioned(basis) && cond(sampling) > 1e8
        @warn "Sampling matrix is poorly conditioned (cond = $(cond(sampling)))."
    end

    return sampling
end

"""
    eval_matrix(T, basis, x)

Return evaluation matrix from coefficients to sampling points. `T <: AbstractSampling`.
"""
eval_matrix(::Type{TauSampling}, basis, x) = permutedims(basis.u(x))
eval_matrix(::Type{MatsubaraSampling}, basis, x) = permutedims(basis.uhat(x))

"""
    evaluate(sampling, al; dim=1)

Evaluate the basis coefficients `al` at the sparse sampling points.
"""
function evaluate(
    smpl::AbstractSampling{S,Tmat}, al::AbstractArray{T,N}; dim=1
) where {S,Tmat,T,N}
    if size(smpl.matrix, 2) ≠ size(al, dim)
        msg = "Number of columns (got $(size(smpl.matrix, 2))) has to match al's size in dim (got $(size(al, dim)))"
        throw(DimensionMismatch(msg))
    end
    bufsize = (size(al)[1:(dim - 1)]..., size(smpl.matrix, 1), size(al)[(dim + 1):end]...)
    buffer = Array{promote_type(Tmat, T),N}(undef, bufsize)
    return evaluate!(buffer, smpl, al; dim)
end

"""
    evaluate!(buffer::AbstractArray{T,N}, sampling, al; dim=1) where {T,N}

Like [`evaluate`](@ref), but write the result to `buffer`.
Please use dim = 1 or N to avoid allocating large temporary arrays internally.
"""
function evaluate!(buffer::AbstractArray, smpl::AbstractSampling, al; dim=1)
    bufsize = (size(al)[1:(dim - 1)]..., size(smpl.matrix, 1), size(al)[(dim + 1):end]...)
    if size(buffer) ≠ bufsize
        msg = "Buffer has the wrong size (got $(size(buffer)), expected $bufsize)"
        throw(DimensionMismatch(msg))
    end
    return matop_along_dim!(buffer, smpl.matrix, al, dim, mul!)
end

"""
    fit(sampling, al::AbstractArray{T,N}; dim=1)

Fit basis coefficients from the sparse sampling points
Please use dim = 1 or N to avoid allocating large temporary arrays internally.
"""
function fit(
    smpl::AbstractSampling{S,Tmat}, al::AbstractArray{T,N}; dim=1
) where {S,Tmat,T,N}
    if size(smpl.matrix, 1) ≠ size(al, dim)
        msg = "Number of rows (got $(size(smpl.matrix, 1))) has to match al's size in dim (got $(size(al, dim)))"
        throw(DimensionMismatch(msg))
    end
    bufsize = (size(al)[1:(dim - 1)]..., size(smpl.matrix, 2), size(al)[(dim + 1):N]...)
    buffer = Array{promote_type(Tmat, T),N}(undef, bufsize)
    return fit!(buffer, smpl, al; dim)
end

"""
    workarrlength(smpl::AbstractSampling, al; dim=1)

Return length of workarr for `fit!`.
"""
function workarrlength(smpl::AbstractSampling, al::AbstractArray; dim=1)
    return length(smpl.matrix_svd.S) * (length(al) ÷ size(al, dim))
end

"""
    fit!(buffer, sampling, al::Array{T,N}; dim=1)

Like [`fit`](@ref), but write the result to `buffer`.
Please use dim = 1 or N to avoid allocating large temporary arrays internally.
The length of `workarry` cannot be smaller than the returned value of `workarrlengthfit`.
"""
function fit!(
    buffer::Array{S,N}, smpl::AbstractSampling, al::Array{T,N};
    dim=1, workarr::Vector{S}=Vector{S}(undef, workarrlength(smpl, al; dim)),
) where {S,T,N}
    bufsize = (size(al)[1:(dim - 1)]..., size(smpl.matrix, 2), size(al)[(dim + 1):end]...)
    if size(buffer) ≠ bufsize
        msg = "Buffer has the wrong size (got $(size(buffer)), expected $bufsize)"
        throw(DimensionMismatch(msg))
    end
    length(workarr) ≥ workarrlength(smpl, al; dim) ||
        throw(ArgumentError("workarr too small"))
    return div_noalloc!(buffer, smpl.matrix_svd, al, workarr, dim)
end

"""
    movedim(arr::AbstractArray, src => dst)

Move `arr`'s dimension at `src` to `dst` while keeping the order of the remaining
dimensions unchanged.
"""
function movedim(arr::AbstractArray{T,N}, dims::Pair) where {T,N}
    src, dst = dims
    src == dst && return arr
    return permutedims(arr, getperm(N, dims))
end

function getperm(N, dims::Pair)
    src, dst = dims
    perm = collect(1:N)
    deleteat!(perm, src)
    insert!(perm, dst, src)
    return perm
end

"""
    matop_along_dim!(buffer, mat, arr::AbstractArray, dim::Integer, op)

Apply the operator `op` to the matrix `mat` and to the array `arr` along the dimension
`dim`, writing the result to `buffer`.
"""
function matop_along_dim!(buffer, mat, arr::AbstractArray{T,N}, dim, op) where {T,N}
    1 ≤ dim ≤ N || throw(DomainError(dim, "Dimension must be in [1, $N]"))

    if dim == 1
        matop!(buffer, mat, arr, op, 1)
    elseif dim != N
        # Move the target dim to the first position
        perm = getperm(N, dim => 1)
        arr_perm = permutedims(arr, perm)
        buffer_perm = permutedims(buffer, perm)
        matop!(buffer_perm, mat, arr_perm, op, 1)
        permutedims!(buffer, buffer_perm, getperm(N, 1 => dim))
    else
        # Apply the operator to the last dimension
        matop!(buffer, mat, arr, op, N)
    end
    return buffer
end

"""
    matop!(buffer, mat, arr::AbstractArray, op, dim)

Apply the operator `op` to the matrix `mat` and to the array `arr` along the first
dimension (dim=1) or the last dimension (dim=N).
"""
function matop!(
    buffer::AbstractArray{S,N}, mat, arr::AbstractArray{T,N}, op, dim
) where {S,T,N}
    if dim == 1
        flatarr = reshape(arr, (size(arr, 1), :))
        flatbuffer = reshape(buffer, (size(buffer, 1), :))
        op(flatbuffer, mat, flatarr)
    elseif dim == N
        flatarr = reshape(arr, (:, size(arr, N)))
        flatbuffer = reshape(buffer, (:, size(buffer, N)))
        op(flatbuffer, flatarr, transpose(mat))
    else
        throw(DomainError("Dimension must be 1 or $N"))
    end
    return buffer
end

function div_noalloc!(
    buffer::AbstractArray{S,N}, mat::SVD, arr::AbstractArray{T,N}, workarr, dim
) where {S,T,N}
    1 ≤ dim ≤ N || throw(DomainError(dim, "Dimension must be in [1, $N]"))

    if dim == 1
        flatarr = reshape(arr, (size(arr, 1), :))
        flatbuffer = reshape(buffer, (size(buffer, 1), :))
        ldiv_noalloc!(flatbuffer, mat, flatarr, workarr)
    elseif dim != N
        # Move the target dim to the first position
        arr_perm = movedim(arr, dim => 1)
        buffer_perm = movedim(buffer, dim => 1)
        flatarr = reshape(arr_perm, (size(arr_perm, 1), :))
        flatbuffer = reshape(buffer_perm, (size(buffer_perm, 1), :))
        ldiv_noalloc!(flatbuffer, mat, flatarr, workarr)
        buffer .= movedim(buffer_perm, 1 => dim)
    else
        flatarr = reshape(arr, (:, size(arr, N)))
        flatbuffer = reshape(buffer, (:, size(buffer, N)))
        rdiv_noalloc!(flatbuffer, flatarr, mat, workarr)
    end
    return buffer
end

function ldiv_noalloc!(Y::AbstractMatrix, A::SVD, B::AbstractMatrix, workarr)
    # Setup work space
    worksize = (size(A.U, 2), size(B, 2))
    worklength = prod(worksize)
    length(workarr) ≥ worklength ||
        throw(DimensionMismatch("size(workarr)=$(size(workarr)), min worksize=$worklength"))
    workarr_view = reshape(view(workarr, 1:worklength), worksize)

    mul!(workarr_view, A.U', B)
    workarr_view ./= A.S
    return mul!(Y, A.V, workarr_view)
end

function rdiv_noalloc!(Y::AbstractMatrix, A::AbstractMatrix, B::SVD, workarr)
    # Setup work space
    worksize = (size(A, 1), size(B.U, 2))
    worklength = prod(worksize)
    length(workarr) ≥ worklength ||
        throw(DimensionMismatch("size(workarr)=$(size(workarr)), min worksize=$worklength"))
    workarr_view = reshape(view(workarr, 1:worklength), worksize)

    # Note: conj creates a temporary matrix
    mul!(workarr_view, A, conj(B.U))
    workarr_view ./= reshape(B.S, 1, :)
    return mul!(Y, workarr_view, conj(B.Vt))
end


text:
module SwapLiterals

Base.Experimental.@optlevel 0

export @swapliterals

const FLOATS_USE_RATIONALIZE = Ref(false)

makedict(@nospecialize(pairs)) =
    foldl(Base.ImmutableDict, pairs; init=Base.ImmutableDict{Any,Any}())

makedict(@nospecialize(pairs::AbstractDict)) = pairs

const defaultswaps = makedict(Any[Float64 => "@big_str",
                                  Int     => :big,
                                  Int128  => :big])

"""
    floats_use_rationalize!(yesno::Bool=true)

If `true`, a `Float64` input is first converted to `Rational{Int}`
via `rationalize` before being further transformed.
"""
floats_use_rationalize!(yesno::Bool=true) = FLOATS_USE_RATIONALIZE[] = yesno

# swaps is a collection of pairs
function literalswapper(swaps)
    @nospecialize

    swaps = makedict(swaps)

    # Base.Callable might be overly restrictive for callables, this check should
    # probably be removed eventually
    all(kv -> kv[2] isa Union{String,Symbol,Nothing,Base.Callable}, swaps) ||
        throw(ArgumentError("unsupported type for swapper"))

    foreach(swaps) do kv
        kv[1] ∈ Any[Float32,Float64,Int,String,Char,
                    Base.BitUnsigned64_types...,Int128,UInt128,BigInt,
                    :braces, :tuple, :vect] ||
                        throw(ArgumentError("type $(kv[1]) cannot be replaced"))
    end

    function swapper(@nospecialize(ex::Union{Float32,Float64,Int,String,Char,
                                             Base.BitUnsigned64}), quoted=false)

        swap = get(swaps, typeof(ex), nothing)

        if swap === nothing
            requote(ex, quoted)
        elseif quoted
            ex
        elseif swap isa String
            Expr(:macrocall, Symbol(swap), nothing, string(ex))
        elseif ex isa AbstractFloat && FLOATS_USE_RATIONALIZE[]
            if swap == :big # big(1//2) doesn't return BigFloat
                swap = :BigFloat
            end
            :($swap(rationalize($ex)))
        elseif swap isa Symbol
            :($swap($ex))
        else
            swap(ex)
        end
    end

    function swapper(@nospecialize(ex::Expr), quoted=false)
        h = ex.head
        if h == :macrocall &&
            ex.args[1] isa GlobalRef &&
            ex.args[1].name ∈ (Symbol("@int128_str"),
                               Symbol("@uint128_str"),
                               Symbol("@big_str"))

            swap = get(swaps,
                       ex.args[1].name == Symbol("@big_str") ? BigInt :
                       ex.args[1].name == Symbol("@int128_str") ? Int128 : UInt128,
                       nothing)

            if swap === nothing
                requote(ex, quoted)
            elseif quoted
                ex
            else
                if swap == :big
                    swap = "@big_str"
                end
                if swap isa String
                    ex.args[1] = Symbol(swap)
                    ex
                elseif swap isa Symbol
                    :($swap($ex))
                else
                    swap(ex)
                end
            end
        elseif h ∈ (:braces, :tuple, :vect)
            ex = recswap(ex)
            swap = get(swaps, h, nothing)
            if swap === nothing
                requote(ex, quoted)
            elseif quoted
                ex
            elseif swap isa Symbol
                :($swap($ex))
            else
                swap(ex)
            end
        else
            ex = recswap(ex)
            if quoted
                Expr(:$, ex)
            else
                ex
            end
        end
    end

    function recswap(@nospecialize(ex))
        h = ex.head
        # copied from REPL.softscope
        if h in (:meta, :import, :using, :export, :module, :error, :incomplete, :thunk)
            ex
        elseif Meta.isexpr(ex, :$, 1)
            swapper(ex.args[1], true)
        else
            ex′ = Expr(h)
            map!(swapper, resize!(ex′.args, length(ex.args)), ex.args)
            ex′
        end
    end

    requote(@nospecialize(ex), quoted) = quoted ? Expr(:$, ex) : ex

    swapper(@nospecialize(ex), quoted=false) = requote(ex, quoted)

    swapper
end

# to save time loading SafeREPL
const default_literalswapper = literalswapper(defaultswaps)

## macro

transform_arg(mod, @nospecialize(x)) =
    if x isa QuoteNode
        x.value
    elseif x == :nothing
        nothing
    elseif x isa String
        x
    elseif x isa Symbol
        getfield(mod, x)
    elseif x isa Expr
        swap = mod.eval(x)
        ex -> Base.invokelatest(swap, ex)
    else
        throw(ArgumentError("invalid swapper type: $(typeof(x))"))
    end

macro swapliterals(swaps...)

    length(swaps) == 1 &&
        return literalswapper(defaultswaps)(esc(swaps[1]))

    if length(swaps) == 2 && swaps[1] isa Expr && swaps[1].head == :vect
        swaps = Any[swaps[1].args..., swaps[2]]
    end

    # either there are pairs/keyword arguments (handled first),
    # or positional arguments (handled second), but not both

    if swaps[1] isa Expr
        ex = esc(swaps[end])
        swaps = swaps[1:end-1]

        transform_src(x::Symbol) = getfield(Base, x)
        transform_src(x::QuoteNode) = x.value # for `:braces`, etc.

        if Meta.isexpr(swaps[1], :call, 3) # pairs
            swaps = map(swaps) do sw
                Meta.isexpr(sw, :call, 3) && sw.args[1] == :(=>) ||
                    throw(ArgumentError("invalid pair argument"))
                transform_src(sw.args[2]) => sw.args[3]
            end
        else
            swaps = map(swaps) do sw # keyword arguments
                Meta.isexpr(sw, :(=), 2) ||
                    throw(ArgumentError("invalid keyword argument"))
                transform_src(sw.args[1]) => sw.args[2]
            end
        end
    else
        for a in swaps[1:end-1]
            a isa Union{QuoteNode,String} || a == :nothing ||
                throw(ArgumentError("invalid argument: $a"))
        end

        ex = esc(swaps[end])

        if length(swaps) == 4
            swaps = Any[Float64=>swaps[1], Int=>swaps[2], Int128=>swaps[3]]
        elseif length(swaps) == 5
            swaps = Any[Float64=>swaps[1], Int=>swaps[2], Int128=>swaps[3], BigInt=>swaps[4]]
        else
            throw(ArgumentError("wrong number of arguments"))
        end
    end

    swaps = Any[k => transform_arg(__module__, v) for (k, v) in swaps]
    literalswapper(swaps)(ex)
end

end # module


text:
#------------- Basic Functionality ---------------#
"""
    save(job::DFJob)

Saves a DFJob, it's job file and all it's input files.
"""
function save(job::DFJob, local_dir=job.local_dir; kwargs...)
    local_dir = local_dir != "" ? local_dir : error("Please specify a valid local_dir!")
    if !ispath(local_dir)
        mkpath(local_dir)
        @info "$local_dir did not exist, it was created."
    end
    sanitize_magnetization!(job)
    sanitize_projections!(job)
    sanitizeflags!(job)
    job.local_dir = local_dir
    return writejobfiles(job; kwargs...)
end

"""
    submit(job::DFJob; server=job.server, server_dir=job.server_dir, rm_prev=true, kwargs...)

Saves the job locally, and then either runs it locally using `qsub` (when `job.server == "localhost"`) or sends it to the specified `job.server` in `job.server_dir`, and submits it using `qsub` on the server. Kwargs get passed through to `save(job; kwargs...)`. If `rm_prev == true` previous `job.tt` output files will be removed.
"""
function submit(job::DFJob; server=job.server, server_dir=job.server_dir, rm_prev=true, kwargs...)
    save(job; kwargs...)
    job.server = server
    job.metadata[:slurmid] = qsub(job; rm_prev=rm_prev)
end

"""
    abort(job::DFJob)

Will look for the job id inside it's metadata and try to remove it from the server queue.
If the lastrunning input happened to be a QE input, the correct abort file will be written.
If it's Wannier90 the job will be brutally removed from the slurm queue.
"""
function abort(job::DFJob)
    lastrunning = runninginput(job)
    if lastrunning == nothing
        error("Is this job running?")
    end
    if package(lastrunning) == QE
        length(inputs(job, QE)) > 1 &&
            @warn "It's absolutely impossible to guarantee a graceful abort of a multi job script with QE."

        abortpath = writeabortfile(job, lastrunning)
        while ispath(abortpath)
            continue
        end
        qdel(job)
    else
        qdel(job)
    end
end

"""
    setflow!(job::DFJob, should_runs...)

Sets whether or not calculations should be run. Calculations are specified using their indices.
"""
function setflow!(job::DFJob, should_runs...)
    for (name, run) in should_runs
        for input in inputs(job, name)
            input.run = run
        end
    end
    return job
end

"""
    setheaderword!(job::DFJob, old_new::Pair{String, String})

Replaces the specified word in the header with the new word.
"""
function setheaderword!(job::DFJob, old_new::Pair{String, String}; print=true)
    for (i, line) in enumerate(job.header)
        if occursin(first(old_new), line)
            job.header[i] = replace(line, old_new)
            s = """Old line:
            $line
            New line:
            $(job.header[i])
            """
            print && (@info s)
        end
    end
    return job
end
function progressreport(job::DFJob; kwargs...)
    dat = outputdata(job; kwargs...)
    plotdat = SymAnyDict(:fermi=>0.0)
    for (n, d) in dat
        i = input(job, n)
        if isbandscalc(i) && haskey(d, :bands)
            plotdat[:bands] = d[:bands]
        elseif isscfcalc(i) || isnscfcalc(i)
            haskey(d, :fermi) && (plotdat[:fermi] = d[:fermi])
            haskey(d, :accuracy) && (plotdat[:accuracy] = d[:accuracy])
        end
    end
    return plotdat
end

"""
Sets the server dir of the job.
"""
function setserverdir!(job, dir)
    job.server_dir = dir
    return job
end

"""
Sets the local dir of the job.
"""
function setlocaldir!(job, dir)
    if !isabspath(dir)
        dir = abspath(dir)
    end
    job.local_dir = dir
    for i in inputs(job)
        setdir!(i, dir)
    end
    return job
end

#-------------- Basic Interaction with DFInputs inside the DFJob ---------------#
setname!(job::DFJob, oldn, newn) = (input(job, oldn).name = newn)
Base.insert!(job::DFJob, index::Int, input::DFInput) = insert!(job.inputs, index, input)
Base.push!(job::DFJob, input::DFInput) = push!(job.inputs, input)
Base.pop!(job::DFJob) = pop!(job.inputs)

Base.append!(job::DFJob, args...) = append!(job.inputs, args...)

"""Access an input inside the job using it's name. E.g `job["scf"]`"""
function Base.getindex(job::DFJob, id::String)
    tmp = getfirst(x -> name(x)==id, inputs(job))
    if tmp != nothing
        return tmp
    else
        error("No Input with name $id")
    end
end

"""Searches through the inputs for the requested flag.
If a flag was found the input and value of the flag will be added to the returned Dict."""
function Base.getindex(job::DFJob, flg::Symbol)
    outdict = Dict()
    for i in inputs(job)
        tfl = flag(i, flg)
        if tfl != nothing
            outdict[name(i)] = tfl
        end
    end
    return outdict
end

"Set one flag in all the appropriate inputs. E.g `job[:ecutwfc] = 23.0`"
function Base.setindex!(job::DFJob, dat, key::Symbol)
    for input in inputs(job)
        input[key] = dat
    end
end

"Fuzzily search inputs in the job whose name contain the fuzzy."
searchinputs(job::DFJob, fuzzy::AbstractString) = inputs(job, fuzzy, true)

"Fuzzily search the first input in the job whose name contains the fuzzy."
searchinput(job::DFJob,  fuzzy::AbstractString) = input(job, fuzzy)

"Returns all inputs from a certain package."
searchinputs(job::DFJob, ::Type{P}) where {P <: Package} = filter(x->package(x) == P, inputs(job))

"""
    setflags!(job::DFJob, inputs::Vector{<:DFInput}, flags...; print=true)

Sets the flags in the names to the flags specified.
This only happens if the specified flags are valid for the names.
If necessary the correct control block will be added to the calculation (e.g. for QEInputs).

The values that are supplied will be checked whether they are valid.
"""
function setflags!(job::DFJob, inputs::Vector{<:DFInput}, flags...; print=true)
    found_keys = Symbol[]

    for calc in inputs
        t_, = setflags!(calc, flags..., print=print)
        push!(found_keys, t_...)
    end
    nfound = setdiff([k for (k, v) in flags], found_keys)
    if print && length(nfound) > 0
        f = length(nfound) == 1 ? "flag" : "flags"
        dfprintln("$f '$(join(":" .* String.(setdiff(flagkeys, found_keys)),", "))' were not found in the allowed input variables of the specified inputs!")
    end
    return job
end
setflags!(job::DFJob, flags...;kwargs...) =
    setflags!(job, inputs(job), flags...;kwargs...)
setflags!(job::DFJob, name::String, flags...; fuzzy=true, kwargs...) =
    setflags!(job, inputs(job, name, fuzzy), flags...; kwargs...)

""" data(job::DFJob, name::String, dataname::Symbol)

Looks through the calculation filenames and returns the data with the specified symbol.
"""
data(job::DFJob, name::String, dataname::Symbol) =
    data(input(job, name), dataname)

"""
    setdata!(job::DFJob, inputs::Vector{<:DFInput}, dataname::Symbol, data; option=nothing)

Looks through the calculation filenames and sets the data of the datablock with `data_block_name` to `new_block_data`.
if option is specified it will set the block option to it.
"""
function setdata!(job::DFJob, inputs::Vector{<:DFInput}, dataname::Symbol, data; kwargs...)
    setdata!.(inputs, dataname, data; kwargs...)
    return job
end
setdata!(job::DFJob, name::String, dataname::Symbol, data; fuzzy=true, kwargs...) =
    setdata!(job, inputs(job, name, fuzzy), dataname, data; kwargs...)

"""
    setdataoption!(job::DFJob, names::Vector{String}, dataname::Symbol, option::Symbol)

sets the option of specified data in the specified inputs.
"""
function setdataoption!(job::DFJob, names::Vector{String}, dataname::Symbol, option::Symbol; kwargs...)
    setdataoption!.(inputs(job, names), dataname, option; kwargs...)
    return job
end
setdataoption!(job::DFJob, n::String, name::Symbol, option::Symbol; kw...) =
    setdataoption!(job, [n], name, option; kw...)

"""
    setdataoption!(job::DFJob, name::Symbol, option::Symbol)

sets the option of specified data block in all calculations that have the block.
"""
setdataoption!(job::DFJob, n::Symbol, option::Symbol; kw...) =
    setdataoption!(job, name.(inputs(job)), n, option; kw...)

"""
    rmflags!(job::DFJob, inputs::Vector{<:DFInput}, flags...)

Looks through the input names and removes the specified flags.
"""
function rmflags!(job::DFJob, inputs::Vector{<:DFInput}, flags...; kwargs...)
    rmflags!.(inputs, flags...; kwargs...)
    return job
end
rmflags!(job::DFJob, name::String, flags...; fuzzy=true, kwargs...) =
    rmflags!(job, inputs(job, name, fuzzy), flags...; kwargs...)
rmflags!(job::DFJob, flags...; kwargs...) =
    rmflags!(job, inputs(job), flags...; kwargs...)

"Returns the executables attached to a given input."
execs(job::DFJob, name) =
    execs(input(job, name))

"""
    setexecflags!(job::DFJob, exec, flags...)

Goes through the calculations of the job and sets the exec flags to the specified ones.
"""
function setexecflags!(job::DFJob, exec, flags...)
	for i in job.inputs
	    setexecflags!(i, exec, flags...)
    end
end

"""
    rmexecflags!(job::DFJob, exec, flags...)

Goes through the calculations of the job and removes the specified `exec` flags.
"""
rmexecflags!(job::DFJob, exec, flags...) =
    rmexecflags!.(job.inputs, (exec, flags)...)

"Sets the directory of the specified executable."
setexecdir!(job::DFJob, exec, dir) =
    setexecdir!.(job.inputs, exec, dir)

"Finds the output files for each of the inputs of a job, and groups all found data into a dictionary."
function outputdata(job::DFJob, inputs::Vector{DFInput}; print=true, onlynew=false)
    datadict = Dict()
    stime = starttime(job)
    for input in inputs
        newout = hasnewout(input, stime)
        if onlynew && !newout
            continue
        end
        tout = outputdata(input; print=print, overwrite=newout)
        if !isempty(tout)
            datadict[name(input)] = tout
        end
    end
    datadict
end
outputdata(job::DFJob; kwargs...) = outputdata(job, inputs(job); kwargs...)
outputdata(job::DFJob, names::String...; kwargs...) =
    outputdata(job, inputs(job, names); kwargs...)
function outputdata(job::DFJob, n::String; fuzzy=true, kwargs...)
    dat = outputdata(job, inputs(job, n, fuzzy); kwargs...)
    if dat != nothing && haskey(dat, name(input(job, n)))
        return dat[name(input(job, n))]
    end
end

#------------ Specialized Interaction with DFInputs inside DFJob --------------#
"""
    setkpoints!(job::DFJob, name::String, k_points)

sets the data in the k point `DataBlock` inside the specified inputs.
"""
function setkpoints!(job::DFJob, name::String, k_points; print=true)
    for calc in inputs(job, n)
        setkpoints!(calc, k_points, print=print)
    end
    return job
end

"Reads throught the pseudo files and tries to figure out the correct cutoffs"
function setcutoffs!(job::DFJob)
    @assert job.server == "localhost" "Cutoffs can only be automatically set if the pseudo files live on the local machine."
    pseudofiles = map(x->x.name, filter(!isempty, [pseudo(at) for at in atoms(job)]))
    pseudodirs  = map(x->x.dir, filter(!isempty, [pseudo(at) for at in atoms(job)]))
    @assert !isempty(pseudofiles) "No atoms with pseudo files found."
    @assert !isempty(pseudodirs) "No valid pseudo directories found in the inputs."
    maxecutwfc = 0.0
    maxecutrho = 0.0
    for d in pseudodirs
        for f in pseudofiles
            pth = joinpath(d, f)
            if ispath(pth)
                ecutwfc, ecutrho = read_cutoffs_from_pseudofile(pth)
                if ecutwfc != nothing && ecutrho != nothing
                    maxecutwfc = ecutwfc > maxecutwfc ? ecutwfc : maxecutwfc
                    maxecutrho = ecutrho > maxecutrho ? ecutrho : maxecutrho
                end
            end
        end
    end
    setcutoffs!.(inputs(job), maxecutwfc, maxecutrho)
end

"""
    setwanenergies!(job::DFJob, nscf::DFInput{QE}, Emin::Real; Epad=5.0)

Will set the energy window limits correctly according to the projections specified in the
structure of the job and the specified Emin. The output of `nscf` will be used to determine the
DOS, and what the size of the frozen window needs to be to fit enough bands inside it,
depending on the projections.
"""
function setwanenergies!(job::DFJob, nscf::DFInput{QE}, Emin::Real; Epad=5.0)
    hasoutput_assert(nscf)
    iscalc_assert(nscf, "nscf")
    hasprojections_assert(structure(job))

    bands = readbands(nscf)
    wancalcs = searchinputs(job, Wannier90)
    @assert length(wancalcs) != 0 "Job ($(job.name)) has no Wannier90 calculations, nothing to do."
    nbnd = isnoncolincalc(nscf) ? 2 * nprojections(structure(job)) : nprojections(structure(job))
    @info "num_bands=$nbnd (inferred from provided projections)."
    winmin, frozmin, frozmax, winmax = wanenergyranges(Emin, nbnd, bands, Epad)
    map(x->setflags!(x, :dis_win_min => winmin, :dis_froz_min => frozmin, :dis_froz_max => frozmax, :dis_win_max => winmax, :num_wann => nbnd, :num_bands=>length(bands)), wancalcs)
    return job
end

"""
    setwanenergies!(job::DFJob, nscf::DFInput{QE}, projwfc::DFInput{QE}, threshold::Real; Epad=5.0)

Will set the energy window limits correctly according to the projections specified in the
structure of the job. The output of `projwfc` and the `threshold` will be used to determine
the minimum limit of the frozen energy window such that the interesting DOS of inside it exceeds
the threshold. `nscf` will be used to determine the DOS, and what the upper limit of the frozen window
needs to be to fit enough bands inside it, depending on the projections.
"""
function setwanenergies!(job::DFJob, nscf::DFInput{QE}, projwfc::DFInput{QE}, threshold::Real; Epad=5.0)
    hasoutput_assert(projwfc)
    hasexec_assert(projwfc, "projwfc.x")
    Emin = Emin_from_projwfc(job.structure, projwfc, threshold)
    setwanenergies!(job, nscf, Emin; Epad=Epad)
end

#--------------- Interacting with the Structure inside the DFJob ---------------#
"Returns the ith atom with id `atsym`."
atom(job::DFJob, atsym::Symbol, i=1) = filter(x -> x.name == atsym, atoms(job))[i]

"""
    atoms(job::DFJob)

Returns a list the atoms in the structure.
"""
atoms(job::DFJob) = atoms(job.structure)

"""
    setatoms!(job::DFJob, atoms::Dict{Symbol,<:Array{<:Point3,1}}, pseudo_setname=nothing, pseudospecifier=nothing, option=:angstrom)

Sets the data data with atomic positions to the new one. This is done for all calculations in the job that have that data.
If default pseudopotentials are defined, a set can be specified, together with a fuzzy that distinguishes between the possible multiple pseudo strings in the pseudo set.
These pseudospotentials are then set in all the calculations that need it.
All flags which specify the number of atoms inside the calculation also gets set to the correct value.
"""
function setatoms!(job::DFJob, atoms::Vector{<:AbstractAtom}; pseudoset=nothing, pseudospecifier="")
    job.structure.atoms = atoms
    pseudoset!=nothing && setpseudos!(job, pseudoset, pseudospecifier)
    return job
end

#automatically sets the cell parameters for the entire job, implement others
"""
    setcell!(job::DFJob, cell_::Mat3)

sets the cell parameters of the structure in the job.
"""
function setcell!(job::DFJob, cell_::Mat3)
    job.structure.cell = cell_
    return job
end

"sets the pseudopotentials to the specified one in the default pseudoset."
setpseudos!(job::DFJob, set::Symbol, specifier::String=""; kwargs...) =
    setpseudos!(job.structure, set, specifier; kwargs...)

"sets the pseudopotentials for the atom with name `atsym` to the specified one in the default pseudoset."
setpseudos!(job::DFJob, atsym::Symbol, set::Symbol, specifier::String=""; kwargs...) =
	setpseudos!(job.structure, atsym, set, specifier; kwargs...)

"sets the pseudopotentials to the specified one in the default pseudoset."
setpseudos!(job::DFJob, at_pseudos::Pair{Symbol, Pseudo}...; kwargs...) =
    setpseudos!(job.structure, at_pseudos...; kwargs...)

"Returns the projections inside the job for the specified `i`th atom in the job with id `atsym`."
projections(job::DFJob, atsym::Symbol, i=1) = projections(atom(job, atsym, i))

"Returns all the projections inside the job."
projections(job::DFJob) = projections(structure(job))

"""
sets the projections of the specified atoms inside the job structure.
"""
setprojections!(job::DFJob, projections...) =
    setprojections!(job.structure, projections...)

for hub_param in (:U, :J0, :α, :β)
	f = Symbol("set_Hubbard_$(hub_param)!")
	str = "$hub_param"
	@eval begin
		"""
			$($(f))(job::DFJob, ats_$($(str))s::Pair{Symbol, <:AbstractFloat}...; print=true)

		Set the Hubbard $($(str)) parameter for the specified atoms.

		Example:
			`$($(f))(job, :Ir => 2.1, :Ni => 1.0, :O => 0.0)`
		"""
		function $f(job::DFJob, $(hub_param)::Pair{Symbol, <:AbstractFloat}...; print=true)
			for (atsym, val) in $(hub_param)
				$f.(atoms(job, atsym), val; print=print)
			end
		end
		export $f
	end
end

"""
	set_Hubbard_J!(job::DFJob, ats_Js::Pair{Symbol, Vector{<:AbstractFloat}}...; print=true)

Set the Hubbard J parameter for the specified atom.

Example:
	`set_Hubbard_J(job, :Ir => [2.1], :Ni => [1.0])'
"""
function set_Hubbard_J!(job::DFJob, ats_Js::Pair{Symbol, <:Vector{<:AbstractFloat}}...; print=true)
	for (atsym, val) in ats_Js
		set_Hubbard_J!.(atoms(job, atsym), (val,); print=print)
	end
end

export set_Hubbard_J!

atoms(job::DFJob, atsym::Symbol) = atoms(job.structure, atsym)

"Rescales the unit cell."
scale_cell!(job::DFJob, s) =
	scale_cell!(job.structure, s)

set_magnetization!(job::DFJob, args...) =
	set_magnetization!(job.structure, args...)

Base.joinpath(job::DFJob, p) =
	joinpath(job.local_dir, p)

create_supercell(job::DFJob, args...) =
	create_supercell(structure(job), args...)

volume(job::DFJob) = volume(structure(job))

"""
	bandgap(job::DFJob, fermi=readfermi(getfirst(isscfcalc, inputs(job))))

Calculates the bandgap (possibly indirect) around the fermi level.
Uses the first found bands calculation, if there is none it uses the first found nscf calculation.
"""
function bandgap(job::DFJob, fermi=readfermi(getfirst(isscfcalc, inputs(job))))
	bandscalc = getfirst(isbandscalc, inputs(job))
	if bandscalc === nothing
		nscfcalc = getfirst(isnscfcalc, inputs(job))
		if nscfcalc === nothing
			error("No valid calculation found to calculate the bandgap.\nMake sure the job has either a valid bands or nscf calculation.")
		else
			bands = readbands(nscfcalc)
		end
	else
		bands = readbands(bandscalc)
	end
	return bandgap(bands, fermi)
end

# """
#     symmetry_operators(j::DFJob; maxsize=52, tolerance=$DEFAULT_TOLERANCE)
#     symmetry_operators(s::Structure; maxsize=52, tolerance=$DEFAULT_TOLERANCE)

# Finds and returns all the rotations and translations that are symmetry operators of the structure.
# """
# symmetry_operators(j::DFJob; kwargs...) = symmetry_operators(j.structure; kwargs...)

# """
#     international_symbol(j::DFJob; tolerance=$DEFAULT_TOLERANCE)
#     international_symbol(s::Structure; tolerance=$DEFAULT_TOLERANCE)

# Returns the international symbol of the space group of the structure.
# """
# international_symbol(j::DFJob; kwargs...) = international_symbol(j.structure; kwargs...)

# """
#     niggli_reduce(j::DFJob; tolerance=$DEFAULT_TOLERANCE)
#     niggli_reduce(s::Structure; tolerance=$DEFAULT_TOLERANCE)

# Returns the niggli reduced lattice cell.
# """
# niggli_reduce(j::DFJob; kwargs...) = niggli_reduce(j.structure; kwargs...)


text:
export Experiment

using .ArcadeLearningEnvironment
using Dates
using ReinforcementLearningCore
using .ReinforcementLearningEnvironments
using Flux
using BSON
using TensorBoardLogger
using Logging
using Statistics
using Random

function atari_env_factory(
    name,
    state_size,
    n_frames,
    max_episode_steps = 100_000;
    seed = nothing,
    repeat_action_probability = 0.25,
)
    AtariEnv(;
        name = string(name),
        grayscale_obs = true,
        noop_max = 30,
        frame_skip = 4,
        terminal_on_life_loss = false,
        repeat_action_probability = repeat_action_probability,
        max_num_frames_per_episode = n_frames * max_episode_steps,
        color_averaging = false,
        full_action_space = false,
        seed = seed,
    ) |>
    StateOverriddenEnv(
        ResizeImage(state_size...),  # this implementation is different from cv2.resize https://github.com/google/dopamine/blob/e7d780d7c80954b7c396d984325002d60557f7d1/dopamine/discrete_domains/atari_lib.py#L629
        StackFrames(state_size..., n_frames),
    ) |>
    StateCachedEnv |>
    RewardOverriddenEnv(r -> clamp(r, -1, 1))
end

function RLCore.Experiment(
    ::Val{:Dopamine},
    ::Val{:DQN},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
    seed = 123,
)

    @warn "Currently setting the `seed` will not guarantee the reproducibility. The instability seems to be caused by the `CrossCor` layer when calculating gradient."
    rng = MersenneTwister(seed)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "dopamine_DQN_atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = atari_env_factory(name, STATE_SIZE, N_FRAMES; seed = hash(seed + 1))
    N_ACTIONS = length(get_actions(env))
    init = glorot_uniform(rng)

    create_model() =
        Chain(
            x -> x ./ 255,
            CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
            CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
            CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
            x -> reshape(x, :, size(x)[end]),
            Dense(11 * 11 * 64, 512, relu; initW = init),
            Dense(512, N_ACTIONS; initW = init),
        ) |> gpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = DQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = create_model(),
                    optimizer = RMSProp(0.00025, 0.95),
                ),  # unlike TF/PyTorch RMSProp doesn't support center
                target_approximator = NeuralNetworkApproximator(model = create_model()),
                update_freq = 4,
                γ = 0.99f0,
                update_horizon = 1,
                batch_size = 32,
                stack_size = N_FRAMES,
                min_replay_history = 20_000,
                loss_func = huber_loss,
                target_update_freq = 8_000,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                ϵ_init = 1.0,
                ϵ_stable = 0.01,
                decay_steps = 250_000,
                kind = :linear,
                rng = rng,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1_000_000,
            state_type = Float32,
            state_size = STATE_SIZE,
        ),
    )

    evaluation_result = []
    EVALUATION_FREQ = 250_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    steps_per_episode = StepsPerEpisode()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        steps_per_episode,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env
            with_logger(lg) do
                @info "training" episode_length = steps_per_episode.steps[end] reward =
                    total_reward_per_episode.rewards[end] log_step_increment = 0
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            old_explorer = agent.policy.explorer
            agent.policy.explorer = EpsilonGreedyExplorer(0.001; rng = rng)  # set evaluation epsilon
            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
            s = @elapsed run(
                agent,
                atari_env_factory(
                    name,
                    STATE_SIZE,
                    N_FRAMES,
                    MAX_EPISODE_STEPS_EVAL;
                    seed = hash(seed + t),
                ),
                StopAfterStep(125_000; is_show_progress = false),
                h,
            )
            res = (
                avg_length = mean(h[2].steps[1:end-1]),
                avg_score = mean(h[1].rewards[1:end-1]),
            )
            push!(evaluation_result, res)
            agent.policy.explorer = old_explorer
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_length = res.avg_length avg_score =
                res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_length = res.avg_length avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent; is_save_trajectory = false)  # saving trajectory will take about 27G disk space each time
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_reward_per_episode time_per_step evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    N_TRAINING_STEPS = 50_000_000
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    description = """
    This experiment uses almost the same config in [dopamine](https://github.com/google/dopamine/blob/master/dopamine/agents/dqn/configs/dqn.gin). But do notice that there are some minor differences:

    - The RMSProp in Flux do not support center option (also the epsilon is not the same).
    - The image resize method used here is provided by ImageTransformers, which is not the same with the one in cv2.

    The testing environment is $name.
    Agent and statistic info will be saved to: `$(joinpath(save_dir, string(N_TRAINING_STEPS)))`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`

    To load the agent and statistic info:
    ```
    agent = RLCore.load("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", Agent)
    BSON.@load joinpath("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", "stats.bson") total_reward_per_episode time_per_step evaluation_result
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:Dopamine},
    ::Val{:Rainbow},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
    seed = 123,
)
    @warn "Currently setting the `seed` will not guarantee the reproducibility. The instability seems to be caused by the `CrossCor` layer when calculating gradient."
    rng = MersenneTwister(seed)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "Dopamine_Rainbow_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = atari_env_factory(name, STATE_SIZE, N_FRAMES; seed = hash(seed + 1))
    N_ACTIONS = length(get_actions(env))
    N_ATOMS = 51
    init = glorot_uniform(rng)

    create_model() =
        Chain(
            x -> x ./ 255,
            CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
            CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
            CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
            x -> reshape(x, :, size(x)[end]),
            Dense(11 * 11 * 64, 512, relu; initW = init),
            Dense(512, N_ATOMS * N_ACTIONS; initW = init),
        ) |> gpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = RainbowLearner(
                approximator = NeuralNetworkApproximator(
                    model = create_model(),
                    optimizer = ADAM(0.0000625),
                ),  # epsilon is not set here
                target_approximator = NeuralNetworkApproximator(model = create_model()),
                n_actions = N_ACTIONS,
                n_atoms = N_ATOMS,
                Vₘₐₓ = 10.0f0,
                Vₘᵢₙ = -10.0f0,
                update_freq = 4,
                γ = 0.99f0,
                update_horizon = 3,
                batch_size = 32,
                stack_size = N_FRAMES,
                min_replay_history = 20_000,
                loss_func = logitcrossentropy_unreduced,
                target_update_freq = 8_000,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                ϵ_init = 1.0,
                ϵ_stable = 0.01,
                decay_steps = 250_000,
                kind = :linear,
                rng = rng,
            ),
        ),
        trajectory = CircularCompactPSARTSATrajectory(
            capacity = 1_000_000,
            state_type = Float32,
            state_size = STATE_SIZE,
        ),
    )

    evaluation_result = []
    EVALUATION_FREQ = 250_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    steps_per_episode = StepsPerEpisode()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        steps_per_episode,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end] episode_length =
                    steps_per_episode.steps[end] log_step_increment = 0
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            old_explorer = agent.policy.explorer
            agent.policy.explorer = EpsilonGreedyExplorer(0.001; rng = rng)  # set evaluation epsilon
            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
            s = @elapsed run(
                agent,
                atari_env_factory(
                    name,
                    STATE_SIZE,
                    N_FRAMES,
                    MAX_EPISODE_STEPS_EVAL;
                    seed = hash(seed + t),
                ),
                StopAfterStep(125_000; is_show_progress = false),
                h,
            )
            res = (
                avg_length = mean(h[2].steps[1:end-1]),
                avg_score = mean(h[1].rewards[1:end-1]),
            )
            push!(evaluation_result, res)
            agent.policy.explorer = old_explorer
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_length = res.avg_length avg_score =
                res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_length = res.avg_length avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent; is_save_trajectory = false)  # saving trajectory will take about 27G disk space each time
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_reward_per_episode time_per_step evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    N_TRAINING_STEPS = 50_000_000
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    description = """
    This experiment uses alomost the same config in [dopamine](https://github.com/google/dopamine/blob/master/dopamine/agents/rainbow/configs/rainbow.gin). But do notice that there are some minor differences:

    - The epsilon in ADAM optimizer is not changed
    - The image resize method used here is provided by ImageTransformers, which is not the same with the one in cv2.

    The testing environment is $name.
    Agent and statistic info will be saved to: `$(joinpath(save_dir, string(N_TRAINING_STEPS)))`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`

    To load the agent and statistic info:
    ```
    agent = RLCore.load("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", Agent)
    BSON.@load joinpath("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", "stats.bson") total_reward_per_episode time_per_step evaluation_result
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:Dopamine},
    ::Val{:IQN},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
    seed = 123,
)
    @warn "Currently setting the `seed` will not guarantee the reproducibility. The instability seems to be caused by the `CrossCor` layer when calculating gradient."
    rng = MersenneTwister(seed)
    device_rng = CUDA.CURAND.RNG()
    Random.seed!(device_rng, hash(seed + 1))
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "Dopamine_IQN_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    MAX_STEPS_PER_EPISODE = 27_000

    env = atari_env_factory(name, STATE_SIZE, N_FRAMES; seed = hash(seed + 2))
    N_ACTIONS = length(get_actions(env))
    Nₑₘ = 64

    init = glorot_uniform(rng)

    create_model() =
        ImplicitQuantileNet(
            ψ = Chain(
                x -> x ./ 255,
                CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
                CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
                CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
                x -> reshape(x, :, size(x)[end]),
            ),
            ϕ = Dense(Nₑₘ, 11 * 11 * 64, relu; initW = init),
            header = Chain(
                Dense(11 * 11 * 64, 512, relu; initW = init),
                Dense(512, N_ACTIONS; initW = init),
            ),
        ) |> gpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = IQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = create_model(),
                    optimizer = ADAM(0.00005),  # epsilon is not set here
                ),
                target_approximator = NeuralNetworkApproximator(model = create_model()),
                κ = 1.0f0,
                N = 64,
                N′ = 64,
                Nₑₘ = Nₑₘ,
                K = 32,
                γ = 0.99f0,
                stack_size = 4,
                batch_size = 32,
                update_horizon = 3,
                min_replay_history = 20_000,
                update_freq = 4,
                target_update_freq = 8_000,
                default_priority = 1.0f2,
                rng = rng,
                device_rng = device_rng,
            ),
            explorer = EpsilonGreedyExplorer(
                ϵ_init = 1.0,
                ϵ_stable = 0.01,
                decay_steps = 250_000,
                kind = :linear,
                rng = rng,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1_000_000,
            state_type = Float32,
            state_size = STATE_SIZE,
        ),
    )

    evaluation_result = []
    EVALUATION_FREQ = 250_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    steps_per_episode = StepsPerEpisode()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        steps_per_episode,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end] episode_length =
                    steps_per_episode.steps[end] log_step_increment = 0
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            old_explorer = agent.policy.explorer
            agent.policy.explorer = EpsilonGreedyExplorer(0.001; rng = rng)  # set evaluation epsilon
            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
            s = @elapsed run(
                agent,
                atari_env_factory(
                    name,
                    STATE_SIZE,
                    N_FRAMES,
                    MAX_EPISODE_STEPS_EVAL;
                    seed = hash(seed + t),
                ),
                StopAfterStep(125_000; is_show_progress = false),
                h,
            )
            res = (
                avg_length = mean(h[2].steps[1:end-1]),
                avg_score = mean(h[1].rewards[1:end-1]),
            )
            push!(evaluation_result, res)
            agent.policy.explorer = old_explorer
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_length = res.avg_length avg_score =
                res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_length = res.avg_length avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent; is_save_trajectory = false)  # saving trajectory will take about 27G disk space each time
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_reward_per_episode time_per_step evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    N_TRAINING_STEPS = 50_000_000
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    description = """
    This experiment uses alomost the same config in [dopamine](https://github.com/google/dopamine/blob/master/dopamine/agents/implicit_quantile/configs/implicit_quantile.gin). But do notice that there are some minor differences:

    - The epsilon in ADAM optimizer is not changed
    - The image resize method used here is provided by ImageTransformers, which is not the same with the one in cv2.

    The testing environment is $name.
    Agent and statistic info will be saved to: `$(joinpath(save_dir, string(N_TRAINING_STEPS)))`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`

    To load the agent and statistic info:
    ```
    agent = RLCore.load("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", Agent)
    BSON.@load joinpath("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", "stats.bson") total_reward_per_episode time_per_step evaluation_result
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:rlpyt},
    ::Val{:A2C},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
    seed = 123,
)
    @warn "Currently setting the `seed` will not guarantee the reproducibility. The instability seems to be caused by the `CrossCor` layer when calculating gradient."
    rng = MersenneTwister(seed)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "rlpyt_A2C_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_ENV = 32
    UPDATE_FREQ = 5
    N_FRAMES = 4
    STATE_SIZE = (80, 104)
    env = MultiThreadEnv([
        atari_env_factory(
            name,
            STATE_SIZE,
            N_FRAMES;
            repeat_action_probability = 0,
            seed = hash(seed + i),
        ) for i in 1:N_ENV
    ])
    N_ACTIONS = length(get_actions(env[1]))

    init = orthogonal(rng)

    # share model
    model = Chain(
        x -> x ./ 255,
        CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 0, init = init),
        CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 1, init = init),
        x -> reshape(x, :, size(x)[end]),
        Dense(6912, 512, relu; initW = init),
    )

    agent = Agent(
        policy = RandomStartPolicy(
            num_rand_start = 1000,
            random_policy = RandomPolicy(get_actions(env); rng = rng),
            policy = QBasedPolicy(
                learner = A2CLearner(
                    approximator = ActorCritic(
                        actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
                        critic = Chain(model, Dense(512, 1; initW = init)),
                        optimizer = ADAM(3e-4),
                    ) |> gpu,
                    γ = 0.99f0,
                    max_grad_norm = 1.0f0,
                    actor_loss_weight = 1.0f0,
                    critic_loss_weight = 0.25f0,
                    entropy_loss_weight = 0.01f0,
                ),
                explorer = BatchExplorer(GumbelSoftmaxExplorer(; rng = rng)),
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (STATE_SIZE..., N_FRAMES, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )

    N_TRAINING_STEPS = 50_000_000 ÷ N_ENV
    EVALUATION_FREQ = N_TRAINING_STEPS ÷ 100
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    total_batch_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    batch_steps_per_episode = BatchStepsPerEpisode(N_ENV)
    evaluation_result = []

    hook = ComposedHook(
        total_batch_reward_per_episode,
        batch_steps_per_episode,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env
            learner = agent.policy.policy.learner
            with_logger(lg) do
                @info "training" loss = learner.loss actor_loss = learner.actor_loss critic_loss =
                    learner.critic_loss entropy_loss = learner.entropy_loss norm =
                    learner.norm log_step_increment = UPDATE_FREQ
            end
        end,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                rewards = [
                    total_batch_reward_per_episode.rewards[i][end]
                    for i in 1:length(env) if get_terminal(env[i])
                ]
                if length(rewards) > 0
                    @info "training" rewards = mean(rewards) log_step_increment = 0
                end
                steps = [
                    batch_steps_per_episode.steps[i][end]
                    for i in 1:length(env) if get_terminal(env[i])
                ]
                if length(steps) > 0
                    @info "training" steps = mean(steps) log_step_increment = 0
                end
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            h = TotalBatchRewardPerEpisode(N_ENV)
            s = @elapsed run(
                agent,
                MultiThreadEnv([
                    atari_env_factory(
                        name,
                        STATE_SIZE,
                        N_FRAMES,
                        MAX_EPISODE_STEPS_EVAL;
                        repeat_action_probability = 0,
                        seed = hash(seed + t + i),
                    ) for i in 1:N_ENV
                ]),
                StopAfterStep(27_000; is_show_progress = false),
                h,
            )
            res = (avg_score = mean(Iterators.flatten(h.rewards)),)
            push!(evaluation_result, res)
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_score = res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent;)
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_batch_reward_per_episode evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    description = """
    # Play Atari($name) with A2C

    The configurations use here should be aligned with [atari_ff_a2c.py](https://github.com/astooke/rlpyt/blob/master/rlpyt/experiments/configs/atari/pg/atari_ff_a2c.py) in rlpyt.

    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:rlpyt},
    ::Val{:PPO},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
    seed = 123,
)
    @warn "Currently setting the `seed` will not guarantee the reproducibility. The instability seems to be caused by the `CrossCor` layer when calculating gradient."
    rng = MersenneTwister(seed)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "rlpyt_PPO_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_ENV = 32
    UPDATE_FREQ = 64
    N_FRAMES = 4
    STATE_SIZE = (80, 104)
    env = MultiThreadEnv([
        atari_env_factory(
            name,
            STATE_SIZE,
            N_FRAMES;
            repeat_action_probability = 0,
            seed = seed + i,
        ) for i in 1:N_ENV
    ])
    N_ACTIONS = length(get_actions(env[1]))
    INIT_CLIP_RANGE = 0.1f0
    INIT_LEARNING_RATE = 1e-3

    init = orthogonal(rng)

    # share model
    model = Chain(
        x -> x ./ 255,
        CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 0, init = init),
        CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 1, init = init),
        x -> reshape(x, :, size(x)[end]),
        Dense(6912, 512, relu; initW = init),
    )

    agent = Agent(
        policy = RandomStartPolicy(
            num_rand_start = 1000,
            random_policy = RandomPolicy(get_actions(env); rng = rng),
            policy = QBasedPolicy(
                learner = PPOLearner(
                    approximator = ActorCritic(
                        actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
                        critic = Chain(model, Dense(512, 1; initW = init)),
                        optimizer = ADAM(INIT_LEARNING_RATE),  # decrease learning rate with a hook
                    ) |> gpu,
                    γ = 0.99f0,
                    λ = 0.98f0,
                    clip_range = INIT_CLIP_RANGE,  # decrease with a hook
                    max_grad_norm = 1.0f0,
                    n_microbatches = 4,
                    n_epochs = 4,
                    actor_loss_weight = 1.0f0,
                    critic_loss_weight = 0.5f0,
                    entropy_loss_weight = 0.01f0,
                    rng = rng,
                ),
                explorer = BatchExplorer(GumbelSoftmaxExplorer(; rng = rng)),
            ),
        ),
        trajectory = PPOTrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (STATE_SIZE..., N_FRAMES, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            action_log_prob_type = Float32,
            action_log_prob_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )

    N_TRAINING_STEPS = 50_000_000 ÷ N_ENV
    EVALUATION_FREQ = N_TRAINING_STEPS ÷ 100
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    total_batch_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    batch_steps_per_episode = BatchStepsPerEpisode(N_ENV)
    evaluation_result = []

    hook = ComposedHook(
        total_batch_reward_per_episode,
        batch_steps_per_episode,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env
            learner = agent.policy.policy.learner
            with_logger(lg) do
                @info "training" loss = mean(learner.loss) actor_loss =
                    mean(learner.actor_loss) critic_loss = mean(learner.critic_loss) entropy_loss =
                    mean(learner.entropy_loss) norm = mean(learner.norm) log_step_increment =
                    UPDATE_FREQ
            end
        end,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env
            decay = (N_TRAINING_STEPS - t) / N_TRAINING_STEPS
            agent.policy.policy.learner.approximator.optimizer.eta =
                INIT_LEARNING_RATE * decay
            agent.policy.policy.learner.clip_range = INIT_CLIP_RANGE * Float32(decay)
        end,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                rewards = [
                    total_batch_reward_per_episode.rewards[i][end]
                    for i in 1:length(env) if get_terminal(env[i])
                ]
                if length(rewards) > 0
                    @info "training" rewards = mean(rewards) log_step_increment = 0
                end
                steps = [
                    batch_steps_per_episode.steps[i][end]
                    for i in 1:length(env) if get_terminal(env[i])
                ]
                if length(steps) > 0
                    @info "training" steps = mean(steps) log_step_increment = 0
                end
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            # switch to GreedyExplorer?
            h = TotalBatchRewardPerEpisode(N_ENV)
            s = @elapsed run(
                agent,
                MultiThreadEnv([
                    atari_env_factory(
                        name,
                        STATE_SIZE,
                        N_FRAMES,
                        MAX_EPISODE_STEPS_EVAL;
                        repeat_action_probability = 0,
                        seed = seed + t + i,
                    ) for i in 1:4
                ]),
                StopAfterStep(27_000; is_show_progress = false),
                h,
            )
            res = (avg_score = mean(Iterators.flatten(h.rewards)),)
            push!(evaluation_result, res)
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_score = res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent;)
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_batch_reward_per_episode evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    description = """
    # Play Atari($name) with PPO
    The configurations use here should be aligned with [atari_ff_ppo.py](https://github.com/astooke/rlpyt/blob/master/rlpyt/experiments/configs/atari/pg/atari_ff_ppo.py) in rlpyt.
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    """

    Experiment(agent, env, stop_condition, hook, description)
end


text:
# # Two dimensional turbulence example
#
# In this example, we initialize a random velocity field and observe its viscous,
# turbulent decay in a two-dimensional domain. This example demonstrates:
#
#   * How to run a model with no buoyancy equation or tracers;
#   * How to create user-defined fields
#   * How to use differentiation functions

# ## Model setup

# For this example, we need `Plots` for plotting and `Statistics` for setting up
# a random initial condition with zero mean velocity.

using Plots, Statistics

# In addition to importing plotting and statistics packages, we import
# `Oceananigans`, its `AbstractOperations` and `Grids` submodules,
# and the `Face` and `Cell` types to will aid in the calculation
# and visualization of voriticty.

using Oceananigans, Oceananigans.AbstractOperations, Oceananigans.Grids

# `Face` and `Cell` represent "locations" on the staggered grid. We instantiate the
# model with a simple isotropic diffusivity.

model = IncompressibleModel(
        grid = RegularCartesianGrid(size=(128, 128, 1), extent=(2π, 2π, 2π)),
    buoyancy = nothing,
     tracers = nothing,
     closure = ConstantIsotropicDiffusivity(ν=1e-3, κ=1e-3)
)
nothing # hide

# ## Setting initial conditions

# Our initial condition randomizes `u` and `v`. We also ensure that both have
# zero mean for purely aesthetic reasons.

u₀ = rand(size(model.grid)...)
u₀ .-= mean(u₀)

set!(model, u=u₀, v=u₀)

# ## Calculating vorticity

# Next we create an object called an `Operation` that represents a vorticity calculation.
# We'll use this object to calculate vorticity on-line as the simulation progresses.

u, v, w = model.velocities
nothing # hide

# Create an object that represents the 'operation' required to compute vorticity.
vorticity_operation = ∂x(v) - ∂y(u)
nothing # hide

# The instance `vorticity_operation` is a binary subtraction between two derivative operations
# acting on `OffsetArrays` (the underyling representation of `u`, and `v`). In order to use
# `vorticity_operation` we create a field `ω` to store the result of the operation, and a
# `Computation` object for coordinate the computation of vorticity and storage in `ω`:

ω = Field(Face, Face, Cell, model.architecture, model.grid)

vorticity_computation = Computation(vorticity_operation, ω)
nothing # hide

# We ask for computation of vorticity by writing `compute!(vorticity_computation)`
# as shown below.

# ## Visualizing the simulation

# Finally, we run the model and animate the vorticity field.
simulation = Simulation(model, Δt=0.1, stop_iteration=0)

anim = @animate for i=1:100
    simulation.stop_iteration += 10
    run!(simulation)

    compute!(vorticity_computation)

    x, y = xnodes(ω)[:], ynodes(ω)[:]
    heatmap(x, y, interior(ω)[:, :, 1], xlabel="x", ylabel="y",
            color=:balance, clims=(-0.1, 0.1))
end

mp4(anim, "2d_turbulence_vorticity.mp4", fps = 15) # hide


text:
using RAISDKResultsWrapper
using Test

import RAI

ctx = RAI.Context(RAI.load_config()); config = (ctx, "nhd-test-1", "nhd-s");
config = (ctx, "nhd-test-1", "nhd-s");
r = RAI.exec(config..., """:a, (1;2;(3,"hi")); :b,"hi",:c,(range[2,15,1]); :x,true; :y,100 """)

@testset "RAISDKResultsWrapper.jl" begin
    cursor = output(r)

    tups = tuples(cursor)
    @test length(tups) == 19
    @test eltype(tups) == Tuple
    @test collect(tups) isa Vector{Tuple}
    @test length(collect(tups)) == 19

    show(tups)
    display(tups)
end


text:
using KUnet, HDF5, JLD, ArgParse, Compat

function initnet()
    args = parse_commandline()
    net = Op[]
    !isempty(args["dropout"]) && push!(net, Drop(args["dropout"][1]))
    for h in args["hidden"]
        append!(net, [Mmul(h), Bias(), Relu()])
        !isempty(args["dropout"]) && push!(net, Drop(args["dropout"][end]))
    end
    append!(net, [Mmul(args["nclass"]), Bias(), Logp(), LogpLoss()])
    for k in [fieldnames(Param)]
        haskey(args, string(k)) || continue
        v = args[string(k)]
        if isempty(v)
            continue
        elseif length(v)==1
            @eval setparam!($net; $k=$v[1])
        else
            @assert length(v)==length(net) "$k should have 1 or $(length(net)) elements"
            for i=1:length(v)
                @eval setparam!($net[$i]; $k=$v[$i])
            end
        end
    end
    savenet(args["netfile"], net)
    display(net)
end

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "netfile"
        help = "JLD file to save the net"
        required = true
        "--nclass", "-n"
        help = "Number of output classes"
        arg_type = Int
        required = true
        "--hidden"
        help = "One or more hidden layer sizes"
        nargs = '+'
        arg_type = Int
        "--dropout"
        help = "Dropout probability"
        arg_type = Float32
        nargs = '+'
        # default = [0.1f0, 0.5f0]
        "--adagrad"
        help = "If nonzero apply adagrad using arg as epsilon"
        arg_type = Float32
        nargs = '+'
        default = [1f-8]
        "--l1reg"
        help = "L1 regularization parameter"
        arg_type = Float32
        nargs = '+'
        "--l2reg"
        help = "L2 regularization parameter"
        arg_type = Float32
        nargs = '+'
        "--lr"
        help = "Learning rate"
        arg_type = Float32
        nargs = '+'
        default = [0.01f0]
        "--maxnorm"
        help = "If nonzero upper limit on weight matrix row norms"
        arg_type = Float32
        nargs = '+'
        "--momentum"
        help = "Momentum"
        arg_type = Float32
        nargs = '+'
        "--nesterov"
        help = "Apply nesterov's accelerated gradient"
        arg_type = Float32
        nargs = '+'
    end
    args = parse_args(s)
    for (arg,val) in args
        print("$arg:$val ")
    end
    println("")
    return args
end

initnet()


text:
"""
    Options

Options configurable by the user with recommended default values.
"""
@with_kw struct Options
    num_start_samples::Int = 5
    trace::Symbol = :compact
    sampling_plan_opt_gens::Int = 2_500
    rippa::Bool = true
    kerns = [ScatteredInterpolation.Gaussian,ScatteredInterpolation.InverseMultiquadratic,ScatteredInterpolation.InverseQuadratic]
    rbf_opt_gens::Int = 2_500
    constrained_seed_gens::Int = 2_500
    rbf_opt_pop::Int = 50
    rbf_opt_method::Symbol = :adaptive_de_rand_1_bin_radiuslimited
    rbf_dist_metric = Distances.Euclidean()
    variable_kernel_width::Bool = true
    variable_dim_scaling::Bool = true
    cond_max::Float64 = 5e12
    cond_check::Bool = false
    max_rbf_width::Float64 = 1.0
    min_rbf_width::Float64 = 0.0
    max_scale::Float64 = 1.0
    min_scale::Float64 = 0.0
    num_interpolants::Int = 10
    smooth = :single
    max_smooth::Float64 = 0.005
    smooth_user::Float64 = 0.0
    iterations::Int64 = 5
    num_infill_points::Int64 = 1
    parallel_surrogate::Bool = true
    infill_funcs::Array{Symbol,1} = [:std,:median,:wstdmed03,:median,:wstdmed06,:median,:wstdmed09,:median]
    infill_iterations::Int64 = 25_000
    create_final_surrogate::Bool = false

    @assert ((smooth == false) || (smooth == :variable) ||
    (smooth == :single) || (smooth == :single_user)) "Not supported smoothing option"

    @assert ((trace == :silent) || (trace == :compact) ||
    (trace == :verbose)) "Not supported trace option"
end

fast_opts() = Options(
    sampling_plan_opt_gens=100,
    kerns = [ScatteredInterpolation.Gaussian],
    rbf_opt_gens = 100,
    constrained_seed_gens = 0,
    rbf_opt_pop = 20,
    variable_kernel_width = false,
    num_interpolants = 5,
    smooth = :single_user,
    smooth_user = 1e-7,
    infill_iterations = 2_500
                        )

"""
    paper_bench_opts()

The options used for benchmarking in the cited paper.
Benchmark performed in release v0.4.0.
"""
paper_bench_opts() = Options(
    num_start_samples = 5,
    trace = :compact,
    sampling_plan_opt_gens = 1_000,
    rippa = true,
    kerns = [ScatteredInterpolation.Gaussian,ScatteredInterpolation.InverseMultiquadratic,ScatteredInterpolation.InverseQuadratic],
    rbf_opt_gens = 2_500,
    constrained_seed_gens = 2_500,
    rbf_opt_pop = 50,
    rbf_opt_method = :adaptive_de_rand_1_bin_radiuslimited,
    rbf_dist_metric = Distances.Euclidean(),
    variable_kernel_width = true,
    variable_dim_scaling = true,
    cond_max = 5e12,
    cond_check = false,
    max_rbf_width = 1.0,
    min_rbf_width = 0.0,
    max_scale = 1.0,
    min_scale = 0.0,
    num_interpolants = 10,
    smooth = :single,
    max_smooth = 0.005,
    smooth_user = 0.0,
    iterations = 95,
    num_infill_points = 1,
    parallel_surrogate = true,
    infill_funcs = [:std,:median,:wstdmed03,:median,:wstdmed06,:median,:wstdmed09,:median],
    infill_iterations = 25_000,
    create_final_surrogate = false
                        )

"""
    RBFHypers(width::S,kernelFunc,scaling::U,fitness::Float64)

Data structure to store results from the optimization of an RBF interpolation kernel
"""
struct RBFHypers{T,U}
    kernelFunc
    scaling::T
    smooth::U
end

"""
    SurrogateResult

Data structure to store the results of an optimization task.
"""
struct SurrogateResult
    lhc_samples::Array{Float64,2}
    lhc_plan::Array{Float64,2}
    sm_interpolant
    sm_interpolant_settings::Array{SurrogateModelOptim.RBFHypers{T,U},1} where U where T
    infill_samples::Array{Float64,2}
    infill_type::Array{Symbol,1}
    infill_plan::Array{Float64,2}
    infill_prediction::Array{Float64,1}
    options::Options
end

function Base.show(io::IO, res::SurrogateResult)
    println("Surrogate Model Optim Result")
    println(io, @sprintf("  Best fitness: \t%5.7g,  (worst %5.7g)",best_fitness(res),maximum([res.lhc_samples res.infill_samples])))
    println(io, "  Best candidate: \t",best_candidate(res))
    println(io, "  Function calls: \t",f_calls(res))
    println(io, "  Iterations: \t\t",res.options.iterations)
    println(io, "  LHC sampling points: \t",size(res.lhc_plan,2))
    println(io, "  Infill criteria: \t",unique(res.infill_type))
    if res.sm_interpolant_settings[1].scaling != false
        sc = mean([opt.scaling./opt.scaling[1] for opt in res.sm_interpolant_settings])
        println(io, "  Mean axis scaling: \t",sc," (relative to first dimension)")
    else
        println(io, "  Mean axis scaling: \t","scaling is fixed")
    end
    println(io, "  Smooth: \t\t",res.options.smooth, " (set to :single if noise is expected)")
    println(io, "  Returned surrogate contains all samples: ",res.options.create_final_surrogate)
end


text:
# this package
using JSOSolverTemplate

using Logging
using NLPModels
using Test

include("unc.jl")