text:
using Random

function create_modulation(modulation_dict)
  if isnothing(modulation_dict)
    return nothing
  end
  params_dict = get(modulation_dict, "params", Dict{String,Any}())
  modulation_name = modulation_dict["function"]
  modulation_params =  NamedTuple{Tuple(Symbol.(keys(params_dict)))}(values(params_dict))
  MocosSim.make_infection_modulation( modulation_name; modulation_params...)
end

function read_params(config, rng::AbstractRNG)
  population_path = config["population_path"] # <= JSON
  population_path::AbstractString # checks if it was indeed a string
  individuals_df = load(population_path)["individuals_df"]

  progression_params = MocosSim.make_progression_params()

  infection_modulation = get(config, "infection_modulation", nothing) |> create_modulation
  mild_detection_modulation = get(config, "mild_detection_modulation", nothing) |> create_modulation
  tracing_modulation = get(config, "tracing_modulation", nothing)

  constant_kernel_param = config["transmission_probabilities"]["constant"]  |> float
  household_kernel_param = config["transmission_probabilities"]["household"] |> float

  hospital_detections = get(config, "hospital_detections", true) |> Bool
  mild_detection_prob = config["mild_detection_prob"]  |> float
  mild_detection_delay = get(config, "mild_detection_delay", 2.0) |> float

  tracing_prob = config["contact_tracing"]["probability"]  |> float
  tracing_delay = config["contact_tracing"]["detection_delay"]  |> float
  testing_time = config["contact_tracing"]["testing_time"]  |> float
  quarantine_length = get(config, "quarantine_length", 14.0) |> float

  age_coupling_kernel_param = get(config["transmission_probabilities"], "age_coupling_param", nothing)
  age_coupling_data_path = get(config["transmission_probabilities"], "age_coupling_data_path", nothing)
  @assert isnothing(age_coupling_kernel_param) == isnothing(age_coupling_data_path)
  age_coupling_thresholds, age_coupling_weights, age_coupling_use_genders =
    isnothing(age_coupling_data_path) ? (nothing, nothing, false) : load(age_coupling_data_path, "age_thresholds", "contact_mat", "uses_genders")

  screening_params = if !haskey(config, "screening")
      nothing
    else
      screen = config["screening"]
      NamedTuple{Tuple(Symbol.(keys(screen)))}(values(screen))
      MocosSim.ScreeningParams(;NamedTuple{Tuple(Symbol.(keys(screen)))}(values(screen))...)
    end

  spreading = get(config, "spreading", nothing)
  spreading_alpha = isnothing(spreading) ? nothing : spreading["alpha"]
  spreading_x0 = isnothing(spreading) ? 1 : get(spreading, "x0", 1)
  spreading_truncation = isnothing(spreading) ? Inf : get(spreading, "truncation", Inf)

  phone_tracing = get(config, "phone_tracing", nothing)
  phone_tracing_usage = isnothing(phone_tracing) ? 0.0 : phone_tracing["usage"] |> float
  phone_tracing_testing_delay = isnothing(phone_tracing) ? 1.0 : phone_tracing["detection_delay"] |> float
  phone_tracing_usage_by_household = isnothing(phone_tracing) ? false : phone_tracing["usage_by_household"] |> Bool

  british_strain_multiplier = get(config["transmission_probabilities"], "british_strain_multiplier", 1.5) |> float
  delta_strain_multiplier = get(config["transmission_probabilities"], "delta_strain_multiplier", 1.5 * 1.5) |> float

  hospital_kernel_param = get(config["transmission_probabilities"], "hospital", 0.0) |> float
  healthcare_detection_prob, healthcare_detection_delay =  if !haskey(config, "healthare_detections")
    0.8, 1.0
  else
    float(config["healthare_detections"]["probability"]), float(config["healthare_detections"]["delay"])
  end

  MocosSim.make_params(
    rng;
    individuals_df = individuals_df,
    progression_params = progression_params,

    infection_modulation = infection_modulation,
    mild_detection_modulation = mild_detection_modulation,
    forward_tracing_modulation = tracing_modulation,
    backward_tracing_modulation = tracing_modulation,

    constant_kernel_param = constant_kernel_param,
    household_kernel_param = household_kernel_param,

    hospital_detections = hospital_detections,
    mild_detection_prob = mild_detection_prob,
    mild_detection_delay = mild_detection_delay,

    backward_tracing_prob = tracing_prob,
    backward_detection_delay = tracing_delay,

    forward_tracing_prob = tracing_prob,
    forward_detection_delay = tracing_delay,

    testing_time = testing_time,
    quarantine_length = quarantine_length,

    age_coupling_param = age_coupling_kernel_param,
    age_coupling_thresholds = age_coupling_thresholds,
    age_coupling_weights = age_coupling_weights,
    age_coupling_use_genders = age_coupling_use_genders,

    screening_params = screening_params,

    spreading_alpha=spreading_alpha,
    spreading_x0=spreading_x0,
    spreading_truncation=spreading_truncation,

    phone_tracing_usage = phone_tracing_usage,
    phone_detection_delay = phone_tracing_testing_delay,
    phone_tracing_usage_by_household = phone_tracing_usage_by_household,

    british_strain_multiplier=british_strain_multiplier,
    delta_strain_multiplier=delta_strain_multiplier,

    hospital_kernel_param = hospital_kernel_param,
    healthcare_detection_prob = healthcare_detection_prob,
    healthcare_detection_delay = healthcare_detection_delay,
  )
end


text:
"""
    `readrun(file)` reading running data.
"""
function readrun(file)
    data = readtable(file, header=false)
    inx = unique([find(x -> isna(x), data[1]);
                  find(x -> isna(x), data[2]);
                  find(x -> isna(x), data[3])])
    deleterows!(data, inx)
    date = DateTime(data[1])
    if length(find(x -> isna(x), data[4])) == 0
        ele = data[4]
    else
        ele = Vector{AbstractFloat}[]
    end
    if length(find(x -> isna(x), data[5])) == 0
        lap = data[5]
    else
        lap = Vector{Integer}[]
    end
    if length(find(x -> isna(x), data[6])) == 0
        hr = data[6]
    else
        hr = Vector{Integer}[]
    end
    run = Run(date, data[2], data[3], ele, lap, hr)
    return run
end

"""
    `readttwatch(file)` reading csv-data generated by ttwatch.

See: https://github.com/ryanbinns/ttwatch
"""
function readttwatch(file)
    data = readtable(file, header=false, skipstart=1)
    inx = unique([find(x -> isna(x), data[12]);
                  find(x -> isna(x), data[8]);
                  find(x -> isna(x), data[7])])
    deleterows!(data, inx)
    date = DateTime(data[12])
    if length(find(x -> isna(x), data[3])) == 0
        lap = data[3]
    else
        lap = Vector{Integer}[]
    end
    if length(find(x -> isna(x), data[9])) == 0
        ele = data[9]
    else
        ele = Vector{AbstractFloat}[]
    end
    if length(find(x -> isna(x), data[10])) == 0
        hr = data[10]
    else
        hr = Vector{Integer}[]
    end
    run = Run(lap, date, data[8], data[7], ele, hr)
    return run
end

ext: jl


text:
module Domains

# The domain is distributed over multiple processes

using Comm, FunRefs, Par
using Foldable, Functor, StencilFunctor
using Cells, Defs, Grids, Norms

import Base.show, Base.wait
import Norms.norm
import Cells.axpy, Cells.initial, Cells.error, Cells.rhs

export Domain
export show, wait
export axpy, norm, initial, error, rhs

immutable Domain
    t::Float64
    grids::Vector{FunRef{Grid}}
end

function show(io::IO, d::Domain)
    println(io, "domain: t=", d.t)
    for g in d.grids
        print(io, get_local(g)) # TODO
    end
end

function wait(d::Domain)
    fmap(wait, d.grids)
end

# Linear combination
function axpy(a::Float64, x::Domain, y::Domain)
    function gaxpy(x::FunRef{Grid}, y::FunRef{Grid})
        remote(y, R=Grid) do
            Grids.axpy(a, get_local(x), y[])
        end
    end
    Domain(a * x.t + y.t, fmap(R=FunRef{Grid}, gaxpy, x.grids, y.grids))
end

# Norm
function norm(d::Domain)
    function gnorm(g::FunRef{Grid})
        rpar(g, R=Norm) do; Grids.norm(g[]) end
    end
    freduce((n,g)->n+g[], Norm(), fmap(R=FunRef{Norm}, gnorm, d.grids))
end

# Initial condition
# (Also choose a domain decomposition)
function initial(t::Float64)
    is = 1:Defs.ncells_per_grid:Defs.ncells()
    function ginitial(i::Integer)
        remote(chooseproc(i,is), R=Grid) do; Grids.initial(t,i) end
    end
    Domain(t, FunRef{Grid}[ginitial(i) for i in is])
end
function chooseproc(i::Integer, r::Range)
    p = div(i - first(r), step(r)) + 1
    mod1(p, Comm.nprocs())
end

# Error
function error(d::Domain)
    function gerror(g::FunRef{Grid})
        remote(g, R=Grid) do; Grids.error(g[], d.t) end
    end
    Domain(d.t, fmap(R=FunRef{Grid}, gerror, d.grids))
end

# RHS
function rhs(d::Domain)
    function grhs(g::FunRef{Grid}, bm::FunRef{Cell}, bp::FunRef{Cell})
        remote(g, R=Grid) do
            Grids.rhs(g[], get_local(bm), get_local(bp))
        end
    end
    function gbnd(g::FunRef{Grid}, face::Bool)
        remote(g, R=Cell) do; Grids.getBoundary(g[], face) end
    end
    bm = par(R=Cell) do; Cells.boundary(d.t, Defs.xmin - 0.5 * Defs.dx()) end
    bp = par(R=Cell) do; Cells.boundary(d.t, Defs.xmax + 0.5 * Defs.dx()) end
    Domain(1.0,                 # dt/dt
           stencil_fmap(R=FunRef{Grid}, grhs, gbnd, d.grids, bm, bp))
end

end


text:
"""
		sparsify(x,nnzrate)

		replace matrices with at most `nnzrate` fraction of non-zeros with SparseMatrixCSC

```juliadoctest
julia> x = TreeNode((
				TreeNode((
					MatrixNode(randn(5,5)),
					MatrixNode(zeros(5,5))
						)),
				MatrixNode(zeros(5,5))
				))
julia> mapdata(i -> sparsify(i,0.05),x)

```
"""
sparsify(x,nnzrate) = x
sparsify(x::Matrix,nnzrate) = (mean(x .!= 0) <nnzrate) ? sparse(x) : x


text:
abstract type ScalarFunction <: AbstractFunction end
output_dim(fun::ScalarFunction) = 1

# Inplace reverts to scalar return
evaluate!(fun::ScalarFunction, y, x, u, p) = evaluate(fun, x, u, p)

# Gradients
function jacobian!(sig::FunctionSignature, diff::DiffMethod, fun::ScalarFunction, J, y, z)
    @assert length(y) == 1
    gradient!(sig, diff, fun, J', z)
end

gradient!(::FunctionSignature, diff::DiffMethod, fun::ScalarFunction, grad, z) =
    gradient!(diff, fun, grad, z)
gradient!(diff::UserDefined, fun::ScalarFunction, grad, z) =
    gradient!(fun, grad, z)
gradient!(fun::AbstractFunction, grad, z::AbstractKnotPoint) =
    gradient!(fun, grad, state(z), control(z), getparams(z))
gradient!(fun::AbstractFunction, grad, x, u, p) =
    gradient!(fun, grad, x, u)
gradient!(fun::AbstractFunction, grad, x, u) =
    throw(NotImplementedError("Gradient not implemented for scalar function $(typeof(fun))"))

# Hessian
function ∇jacobian!(sig::FunctionSignature, diff::DiffMethod, fun::ScalarFunction, H, b, y, z)
    @assert b[1] ≈ 1.0
    @assert length(y) == 1
    hessian!(sig, diff, fun, H, z)
end

hessian!(::FunctionSignature, diff::DiffMethod, fun::ScalarFunction, hess, z) =
    hessian!(diff, fun, hess, z)
hessian!(diff::UserDefined, fun::ScalarFunction, hess, z) = hessian!(fun, hess, z)
hessian!(fun::ScalarFunction, hess, z::AbstractKnotPoint) =
    hessian!(fun, hess, state(z), control(z), getparams(z))
hessian!(fun::ScalarFunction, hess, x, u, p) =
    hessian!(fun, hess, x, u)
hessian!(fun::ScalarFunction, hess, x, u) =
    throw(NotImplementedError("Hessian not implemented for scalar function $(typeof(fun))"))


text:
MP.differentiate(::V, ::V) where {V <: Variable} = 1
MP.differentiate(v::Variable, ::Variable) = 0

@pure inmonomial(V) = false
@pure inmonomial(V, Vars...) = (V === first(Vars)) || inmonomial(V, Base.tail(Vars)...)

function MP.differentiate(m::Monomial{Vars}, v::V) where {Vars, V <: Variable}
    if inmonomial(v, Vars...)
        _diff(m, v)
    else
        0
    end
end

_diff(m::Monomial, v::Variable) = _diff(m, exponents(m), v)

function find_variable_index(var::Variable, vars::Tuple, i=1)
    if var === first(vars)
        return i
    else
        return find_variable_index(var, Base.tail(vars), i + 1)
    end
end

find_variable_index(v::Variable, ::Tuple{}, i) = error("Could not find variable $v")

function _diff(m::Monomial{Vars},
               exponents::NTuple{N, Integer},
               v::Variable) where {Vars, N}
    vi = find_variable_index(v, Vars)
    # vi = findfirst(var -> name(var) == name(v), Vars)
    exponents[vi] * Monomial{Vars, N}(ntuple(i -> begin
        if i == vi
            exponents[i] - 1
        else
            exponents[i]
        end
    end, Val{N}()))
end


text:
@testset "Math operators" begin

# Test built-in math and comparison operators
@test resolve(@julog(1 == 1), Clause[])[1] == true
@test resolve(@julog(1 == 2), Clause[])[1] == false
@test resolve(@julog(1 != 2), Clause[])[1] == true
@test resolve(@julog(1 < 2), Clause[])[1] == true
@test resolve(@julog(3 > 2), Clause[])[1] == true
@test resolve(@julog(3 <= 4), Clause[])[1] == true
@test resolve(@julog(5 >= 4), Clause[])[1] == true

@test resolve(@julog(1 + 1 == 2), Clause[])[1] == true
@test resolve(@julog(5 - 6 == -1), Clause[])[1] == true
@test resolve(@julog(6 * 9 != 42), Clause[])[1] == true
@test resolve(@julog(8 * 8 > 7 * 9), Clause[])[1] == true
@test resolve(@julog(5 / 2 == 2.5), Clause[])[1] == true
@test resolve(@julog(mod(5, 2) == 1), Clause[])[1] == true

# Define addition using built-in operators
clauses = @julog [
    add(X, Y, Z) <<= is(Z, X + Y),
    add(X, Y, Z) <<= is(X, Z - Y),
    add(X, Y, Z) <<= is(Y, Z - X)
]

# Is 1 + 1 = 2?
@test resolve(@julog(add(1, 1, 2)), clauses)[1] == true
# What is 2 + 3?
@test @varsub({Z => 5}) in resolve(@julog(add(2, 3, Z)), clauses)[2]
# What is 5 - 2?
@test @varsub({X => 2}) in resolve(@julog(add(X, 3, 5)), clauses)[2]
# What is 5 - 3?
@test @varsub({Y => 3}) in resolve(@julog(add(2, Y, 5)), clauses)[2]
# Using is/2 doesn't allow us ask for all the ways to add to 5
@test resolve(@julog(add(X, Y, 5)), clauses)[1] == false

# Test normal evaluation of operators
query = @julog [is(X, 1), is(Y, 2), (X < Y)]
@test resolve(query, Clause[])[1] == true
# Test deferred evaluation of operators
query = @julog [(X < Y), is(X, 1), is(Y, 2)]
@test resolve(query, Clause[]; defer_eval=true)[1] == true

end


text:
using Revise
using Plots
using RadiativeTransfer
using RadiativeTransfer.Absorption
using RadiativeTransfer.PhaseFunction
using RadiativeTransfer.RTM
using Distributions
using BenchmarkTools
using Test
using CUDA
using StatsPlots

device!(3)

FT = Float32
"Generate aerosol optical properties"

# Wavelength (just one for now)
λ = FT(0.770)       # Incident wavelength
depol = FT(0.0)
# Truncation
Ltrunc = 20             # Truncation
truncation_type   = PhaseFunction.δBGE{Float32}(Ltrunc, 2.0)

# polarization_type
polarization_type = Stokes_IQU{FT}()

# Quadrature points for RTM
# Nquad, qp_μ, wt_μ = rt_set_streams(RTM.RadauQuad(), Ltrunc, FT(60.0), FT[0.0, 15.0, 30., 45., 60.])
# Nquad, qp_μ, wt_μ = rt_set_streams(RTM.GaussQuadFullSphere(), Ltrunc, FT(60.0), FT[0.0, 15.0, 30., 45., 60.])

# Aerosol particle distribution and properties
μ            = [1.3]    # [0.3,2.0]       # Log mean radius
σ            = [2.0]    # [2.0,1.8]       # Log stddev of radius
r_max        = [30.0]   # [30.0,30.0]     # Maximum radius
nquad_radius = [2500]   # [2500,2500]     # Number of quadrature points for integrating of size dist.
nᵣ           = [1.3]    # [1.3, 1.66]     # Real part of refractive index
nᵢ           = [0.00000001]  # [0.001,0.0003]  # Imag part of refractive index

# Aerosol vertical distribution profiles
p₀          = FT[90000.]  # [50000., 20000.] # Pressure peak [Pa]
σp          = FT[5000.]   # [5000., 2000.]   # Pressure peak width [Pa]

size_distribution = [LogNormal(log(μ[1]), log(σ[1]))] # [LogNormal(log(μ[1]), log(σ[1])), LogNormal(log(μ[2]), log(σ[2]))]

# Create the aerosols (needs to be generalized through loops):
aero1 = make_univariate_aerosol(size_distribution[1], r_max[1], nquad_radius[1], nᵣ[1], nᵢ[1])
# aero2 = make_univariate_aerosol(size_distribution[2], r_max[2], nquad_radius[2], nᵣ[2], nᵢ[2])

# Define some details, run aerosol optics
model_NAI2_aero1 = make_mie_model(NAI2(), aero1, λ, polarization_type, truncation_type)
aerosol_optics_NAI2_aero1 = compute_aerosol_optical_properties(model_NAI2_aero1, FT);

# Truncate:
aerosol_optics_trunc_aero1 = PhaseFunction.truncate_phase(truncation_type, aerosol_optics_NAI2_aero1; reportFit=true)

# Define some details, run aerosol optics
# model_NAI2_aero2 = make_mie_model(NAI2(), aero2, λ, polarization_type, truncation_type)
# aerosol_optics_NAI2_aero2 = compute_aerosol_optical_properties(model_NAI2_aero2);
# Truncate:
# aerosol_optics_trunc_aero2 = PhaseFunction.truncate_phase(truncation_type, aerosol_optics_NAI2_aero2)

# Rayleigh Greek
GreekRayleigh = PhaseFunction.get_greek_rayleigh(depol)

# In[ ]:

vza = FT[60., 45., 30., 15., 0., 15., 30., 45., 60.]
vaz = FT[180., 180., 180., 180., 0., 0., 0., 0., 0.]
sza = FT(60.)

obs_geom = ObsGeometry(FT(1000.0), sza, vza, vaz)

Nquad, qp_μ, wt_μ = rt_set_streams(RTM.RadauQuad(), Ltrunc, obs_geom);
# Nquad, qp_μ, wt_μ = rt_set_streams(RTM.GaussQuadFullSphere(), Ltrunc, sza, vza);
@show Nquad

# In[ ]:

" Atmospheric Profiles, basics, needs to be refactore entirely"
file = "/net/fluo/data1/ftp/XYZT_ESE156/Data/MERRA300.prod.assim.inst6_3d_ana_Nv.20150613.hdf.nc4"
# file = "MERRA300.prod.assim.inst6_3d_ana_Nv.20150613.hdf.nc4"
timeIndex = 2 # There is 00, 06, 12 and 18 in UTC, i.e. 6 hourly data stacked together

# What latitude do we want?
myLat = 34.1377;
myLon = -118.1253;

# Read profile (and generate dry/wet VCDs per layer)
profile_caltech_hr = RTM.read_atmos_profile(file, myLat, myLon, timeIndex);
profile_caltech = RTM.reduce_profile(20, profile_caltech_hr);
# Compute layer optical thickness for Rayleigh (surface pressure in hPa)
τRayl =  RTM.getRayleighLayerOptProp(profile_caltech.psurf / 100, λ, depol, profile_caltech.vcd_dry);
ϖRayl = ones(FT, length(τRayl));

# Compute Naer aerosol optical thickness profiles
τAer_1 = RTM.getAerosolLayerOptProp(1.0, p₀[1], σp[1], profile_caltech.p_levels)
# τAer_2 = RTM.getAerosolLayerOptProp(0.3, p₀[2], σp[2], profile_caltech.p_levels)

# Can be done with arbitrary length later:
τAer = FT(0.2) * τAer_1; # [τAer_1 τAer_2]
@show sum(τAer)# , sum(τAer_2)
ϖAer = FT[aerosol_optics_NAI2_aero1.ω̃]; # [aerosol_optics_NAI2_aero1.ω̃ aerosol_optics_NAI2_aero2.ω̃];
fᵗ   = FT[aerosol_optics_trunc_aero1.fᵗ]; # [aerosol_optics_trunc_aero1.fᵗ aerosol_optics_trunc_aero2.fᵗ];

aerosol_optics = [aerosol_optics_trunc_aero1] # [aerosol_optics_trunc_aero1 aerosol_optics_trunc_aero2]
# Aer𝐙⁺⁺ = [aero1_Z⁺⁺] # [aero1_Z⁺⁺, aero2_Z⁺⁺];
# Aer𝐙⁻⁺ = [aero1_Z⁻⁺] # [aero1_Z⁻⁺, aero2_Z⁻⁺];

maxM = 3

#################################################################
# PLOTTING RUNTIME VS. N_SPEC (CPU vs GPU)
#################################################################

elapsed_absorption_time = []
elapsed_RT_time_cpu = []
elapsed_RT_time_gpu_32 = []

grid_sizes = [100, 500, 1000, 2500, 5000, 7500, 10000]

for grid_size in grid_sizes

    @show grid_size

    grid = range(1e7 / 780, 1e7 / 755, length=grid_size)
    τ_abs = zeros(FT, length(grid), length(profile_caltech.p));

    hitran_data = read_hitran(artifact("O2"), iso=1)
    model = make_hitran_model(hitran_data, Voigt(), wing_cutoff=100, CEF=HumlicekWeidemann32SDErrorFunction(), architecture=Absorption.GPU(), vmr=0.21)

    compute_absorption_profile!(τ_abs, model, grid, profile_caltech);
    # push!(elapsed_absorption_time, @elapsed compute_absorption_profile!(grid, τ_abs, profile_caltech))

    push!(elapsed_RT_time_gpu_32, @elapsed R_GPU, T_GPU = RTM.rt_run(polarization_type, obs_geom, τRayl, ϖRayl, τAer, ϖAer, fᵗ, qp_μ, wt_μ, maxM, aerosol_optics, GreekRayleigh, τ_abs, RadiativeTransfer.Architectures.GPU()))

    # push!(elapsed_RT_time_cpu, @elapsed RTM.run_RTM(polarization_type, sza, vza, vaz, τRayl, ϖRayl, τAer, ϖAer, fᵗ, qp_μ, wt_μ, maxM, aerosol_optics, GreekRayleigh, τ_abs, RadiativeTransfer.Architectures.CPU()))

end

plot(grid_sizes, elapsed_absorption_time)
plot(grid_sizes, elapsed_RT_time_gpu)
plot(grid_sizes, elapsed_RT_time_cpu)

groupedbar([elapsed_RT_time_gpu_32 elapsed_RT_time_gpu_64], bar_position = :dodge, bar_width=0.7, xticks=(1:12, grid_sizes), legend=:topleft, label=["32" "64"], ylabel="Time (s)", xlabel="N_Spec", title="Runtime with Float32 vs Float64")

# elapsed_absorption_time = [0.282189253, 0.283802194, 0.307112032, 0.3477742, 0.38201981, 0.372605842, 0.468313961]
# elapsed_RT_time_gpu = [1.368995249, 4.598695304, 30.890917859, 30.341784914, 71.295454828, 107.297501462, 143.935854327]
elapsed_RT_time_gpu_64 =
elapsed_RT_time_gpu_32 = [1.208653814, 3.607311384 , 29.202259484, 25.388859926, 40.688672364, 74.440302864, 97.265656638]
# elapsed_RT_time_cpu = [14.577631018, 72.8522157, 138.689995059, 349.936594999, 683.309987181, 1013.949079029, 1645.13642589]

#################################################################
# PLOTTING R/T VS. SZA (GPU)
#################################################################

# szas = collect(-90.0:1.0:90.0)
# Rs = zeros(9,3,10000,length(szas))
# Ts = zeros(9,3,10000,length(szas))

# grid = range(1e7 / 780, 1e7 / 755, length=5000)
# τ_abs = zeros(length(grid), length(profile_caltech.p))
# compute_absorption_profile!(grid, τ_abs, profile_caltech)

# for sza_idx in 1:length(szas)

#     @show sza_idx

#     model_NAI2_aero1 = make_mie_model(NAI2(), aero1, λ, polarization_type, truncation_type)
#     aerosol_optics_NAI2_aero1 = compute_aerosol_optical_properties(model_NAI2_aero1);

#     # Truncate:
#     aerosol_optics_trunc_aero1 = PhaseFunction.truncate_phase(truncation_type, aerosol_optics_NAI2_aero1; reportFit=true)

#     fᵗ   = [aerosol_optics_trunc_aero1.fᵗ];
#     aerosol_optics = [aerosol_optics_trunc_aero1]

#     Nquad, qp_μ, wt_μ = rt_set_streams(RTM.RadauQuad(), Ltrunc, szas[sza_idx], vza);
#     R, T = RTM.run_RTM(polarization_type, szas[sza_idx], vza, vaz, τRayl, ϖRayl, τAer, ϖAer, fᵗ, qp_μ, wt_μ, maxM, aerosol_optics, GreekRayleigh, τ_abs, RadiativeTransfer.Architectures.GPU())

#     R[:,:,:,sza_idx] = R
#     T[:,:,:,sza_idx] = T
# end

#################################################################
# PLOTTING R/T VS. LTRUNC (GPU)
#################################################################

Ltruncs = collect(2:30)
Rs = zeros((9,3,10000,length(Ltruncs)))
Ts = zeros((9,3,10000,length(Ltruncs)))

grid = range(1e7 / 780, 1e7 / 755, length=10000)
τ_abs = zeros(length(grid), length(profile_caltech.p))
compute_absorption_profile!(grid, τ_abs, profile_caltech)

for ltrunc_idx in 1:length(Ltruncs)

    @show ltrunc_idx

    Ltrunc = Ltruncs[ltrunc_idx]
    truncation_type   = PhaseFunction.δBGE{Float32}(Ltrunc, 2.0)

    model_NAI2_aero1 = make_mie_model(NAI2(), aero1, λ, polarization_type, truncation_type)
    aerosol_optics_NAI2_aero1 = compute_aerosol_optical_properties(model_NAI2_aero1);

    # Truncate:
    aerosol_optics_trunc_aero1 = PhaseFunction.truncate_phase(truncation_type, aerosol_optics_NAI2_aero1; reportFit=true)

    fᵗ   = [aerosol_optics_trunc_aero1.fᵗ];
    aerosol_optics = [aerosol_optics_trunc_aero1]

    Nquad, qp_μ, wt_μ = rt_set_streams(RTM.RadauQuad(), Ltrunc, sza, vza);
    R, T = RTM.run_RTM(polarization_type, sza, vza, vaz, τRayl, ϖRayl, τAer, ϖAer, fᵗ, qp_μ, wt_μ, maxM, aerosol_optics, GreekRayleigh, τ_abs, RadiativeTransfer.Architectures.GPU())

    Rs[:,:,:,ltrunc_idx] = R
    Ts[:,:,:,ltrunc_idx] = T
end

anim = @animate for i = 1:24
    plot(1:nSpec, Rs[1,1,:,i], ylims=(-0.1, 0.2), title="Reflectance, Ltrunc=$(Ltruncs[i])")
end
gif(anim, "reflectance.gif", fps = 5)


text:
module BoundedStreams

export BoundedInputStream, BoundedOutputStream

const CLOSE = typemin(Int) # special value for optional `close` argument

# abstract base type for Input and Output
abstract type BoundedStream{T} <: IO end

"""
    BoundedInputStream(source::IO, nbytes::Integer; offset=0, close=nbytes)

Provide the `IO` interface for reading the source stream `source`. Restrict the number of
bytes to to `nbytes`. A `BoundedInputStream` can be considered as a view on the wrapped
stream, which starts at the position of that at the time of creation, and extends by the
given size. Trying to read beyond this area throws `EOFError`.

The optional integer argument `offset` shifts the starting point off the
current position of the source stream.

The optional argument `close` determines the position of the source stream after
this stream is closed. The special value `BoundedStreams.CLOSE` closes
the source stream in this case.
"""
struct BoundedInputStream{T} <: BoundedStream{T}
    source::T
    offset::Int
    length::Int
    close::Int
    function BoundedInputStream(io::T, nb::Integer;
                                offset::Integer=0, close::Integer=nb) where T

        new{T}(io, initposition!(io, offset), nb, close)
    end
end

"""
    BoundedOutputStream(source::IO, nbytes::Integer; offset=0, close=nbytes)

Provide the `IO` interface for writing the source stream `source`. Restrict the number of
bytes to to `nbytes`. A `BoundedOutputStream` can be considered as a view on the wrapped
stream, which starts at the position of that at the time of creation, and extends by the
given size. Trying to write beyond this area throws `EOFError`.

The optional integer argument `offset` shifts the starting point off the
current position of the source stream.

The optional argument `close` determines the position of the source stream after
this stream is closed. The special value `BoundedStreams.CLOSE` closes
the source stream in this case.
"""
struct BoundedOutputStream{T} <: BoundedStream{T}
    source::T
    offset::Int
    length::Int
    close::Int
    function BoundedOutputStream(io::T, nb::Integer;
                                offset::Integer=0, close::Integer=nb) where T

        new{T}(io, initposition!(io, offset), nb, close)
    end
end

# set position of stream, if required. prefer skip over seek if possible. return position.
function initposition!(io::IO, offset::Integer)
    offset = Int(offset)
    if offset == 0
        position(io)
    elseif offset > 0
        skip(io, offset)
        position(io)
    elseif offset < 0
        pos = position(io) + offset
        seek(io, pos)
        pos
    end
end

function Base.bytesavailable(io::BoundedStream)
    a = io. length - position(io)
    0 <= a <= io.length ? a : 0
end
Base.position(io::BoundedStream) = position(io.source) - io.offset

function Base.eof(io::BoundedStream)
    bytesavailable(io) <= 0 || eof(io.source)
end

function Base.close(io::BoundedStream)
    source = io.source
    if io.close == CLOSE
        close(source)
    else
        p = position(source)
        q = io.offset + io.close
        if p < q
            skip(source, q - p)
        elseif p > q
            seek(source, q)
        end
    end
    nothing
end

function Base.skip(io::BoundedStream, nb::Integer)
    stream = io.source
    nbmax = bytesavailable(io)
    0 <= nb <= nbmax || throw(ArgumentError("cannot skip outside BoundedStream"))
    skip(stream, nb)
    io
end

function Base.seek(io::BoundedStream, nb::Integer)
    0 <= nb <= io.length || throw(ArgumentError("cannot seek outside BoundedStream"))
    seek(io.source, io.offset + nb)
    io
end
Base.seekend(io::BoundedStream) = seek(io, io.length)

Base.isreadable(io::BoundedStream) = io isa BoundedInputStream
Base.iswritable(io::BoundedStream) = io isa BoundedOutputStream
Base.mark(io::BoundedStream) = mark(io.source)
Base.unmark(io::BoundedStream) = unmark(io.source)
Base.reset(io::BoundedStream) = reset(io.source)
Base.ismarked(io::BoundedStream) = ismarked(io.source)

function Base.readbytes!(io::BoundedInputStream, buf::Vector{UInt8}, nb::Integer)
    readbytes!(io.source, buf, min(nb, bytesavailable(io)))
end

function Base.read(io::BoundedInputStream, T::Type{UInt8})
    nb = bytesavailable(io)
    nb >= sizeof(T) || throw(EOFError())
    read(io.source, T)
end

function Base.unsafe_write(io::BoundedOutputStream, p::Ptr{UInt8}, n::UInt)
    nb = bytesavailable(io)
    nb >= n || throw(EOFError())
    Base.unsafe_write(io.source, p, n)
end

function Base.write(io::BoundedOutputStream, x::UInt8)
    nb = bytesavailable(io)
    nb >= 1 || throw(EOFError())
    write(io.source, x)
end

end # module


text:
############################################################################################
### connect to NLopt.jl as backend
############################################################################################

# wrapper to define the objective
function sem_wrap_nlopt(par, G, model::AbstractSem)
    need_gradient = length(G) != 0
    if need_gradient
        return objective_gradient!(G, model, par)
    else
        return objective!(model, par)
    end
end

mutable struct NLoptResult
    result
    problem
end

optimizer(res::NLoptResult) = res.problem.algorithm
n_iterations(res::NLoptResult) = res.problem.numevals
convergence(res::NLoptResult) = res.result[3]

# construct SemFit from fitted NLopt object
function SemFit_NLopt(optimization_result, model::AbstractSem, start_val, opt)
    return SemFit(
        optimization_result[1],
        optimization_result[2],
        start_val,
        model,
        NLoptResult(optimization_result, opt)
    )
end

# sem_fit method
function sem_fit(model::Sem{O, I, L, D}; start_val = start_val, kwargs...) where {O, I, L, D <: SemOptimizerNLopt}

    # starting values
    if !isa(start_val, Vector)
        start_val = start_val(model; kwargs...)
    end

    # construct the NLopt problem
    opt = construct_NLopt_problem(
        model.optimizer.algorithm,
        model.optimizer.options,
        length(start_val))
    set_NLopt_constraints!(opt, model.optimizer)
    opt.min_objective = (par, G) -> sem_wrap_nlopt(par, G, model)

    if !isnothing(model.optimizer.local_algorithm)
        opt_local = construct_NLopt_problem(
            model.optimizer.local_algorithm,
            model.optimizer.local_options,
            length(start_val))
        opt.local_optimizer = opt_local
    end

    # fit
    result = NLopt.optimize(opt, start_val)

    return SemFit_NLopt(result, model, start_val, opt)
end

function sem_fit(model::SemEnsemble{N, T , V, D, S}; start_val = start_val, kwargs...) where {N, T, V, D <: SemOptimizerNLopt, S}

    # starting values
    if !isa(start_val, Vector)
        start_val = start_val(model; kwargs...)
    end

    # construct the NLopt problem
    opt = construct_NLopt_problem(
        model.optimizer.algorithm,
        model.optimizer.options,
        length(start_val))
    set_NLopt_constraints!(opt, model.optimizer)
    opt.min_objective = (par, G) -> sem_wrap_nlopt(par, G, model)

    if !isnothing(model.optimizer.local_algorithm)
        opt_local = construct_NLopt_problem(
            model.optimizer.local_algorithm,
            model.optimizer.local_options,
            length(start_val))
        opt.local_optimizer = opt_local
    end

    # fit
    result = NLopt.optimize(opt, start_val)

    return SemFit_NLopt(result, model, start_val, opt)
end

############################################################################################
### additional functions
############################################################################################

function construct_NLopt_problem(algorithm, options, npar)
    opt = Opt(algorithm, npar)

    for key in keys(options)
        setproperty!(opt, key, options[key])
    end

    return opt

end

function set_NLopt_constraints!(opt, optimizer::SemOptimizerNLopt)
    for con in optimizer.inequality_constraints
        inequality_constraint!(opt::Opt, con.f, con.tol)
    end
    for con in optimizer.equality_constraints
        equality_constraint!(opt::Opt, con.f, con.tol)
    end
end

############################################################################################
# pretty printing
############################################################################################

function Base.show(io::IO, result::NLoptResult)
    print(io, "Optimizer status: $(result.result[3]) \n")
    print(io, "Minimum:          $(round(result.result[1]; digits = 2)) \n")
    print(io, "Algorithm:        $(result.problem.algorithm) \n")
    print(io, "No. evaluations:  $(result.problem.numevals) \n")
end


text:
@doc raw"""
    Bbox2

An object `b` of the type [`Bbox2`](@ref) is a bounding box in the
two-dimensional Euclidean plane ``Ε^2``.
"""
Bbox2

@doc raw"""
    Bbox2()

Introduces an *empty* bounding box with lower left corner point at ``(∞, ∞)``
and with upper right corner point at ``(-∞, -∞)``, ``∞`` being [`Inf`](@ref).
"""
Bbox2()

"""
    Bbox2(xmin::Real, ymin::Real, xmax::Real, ymax::Real)

Introduces a bounding box `b` with lower left corner at `(xmin, ymin)` and with
upper right corner at `(xmax, ymax)`.
"""
Bbox2(xmin::Real, ymin::Real, xmax::Real, ymax::Real) =
    Bbox2(convert(Float64, xmin), convert(Float64, ymin),
          convert(Float64, xmax), convert(Float64, ymax))

"""
    dimension(b::Bbox2)

Returns 2.
"""
dimension(b::Bbox2)

"""
    min(b::Bbox2, i::Integer)

Returns [`xmin(b)`](@ref) if `i == 0` or [`ymin(b)`](@ref) if `i == 1`.

!!! info "Precondition"

    `i == 0` or `i == 1`.
"""
min(b::Bbox2, i::Integer)

"""
    max(b::Bbox2, i::Integer)

Returns [`xmax(b)`](@ref) if `i == 0` or [`ymax(b)`](@ref) if `i == 1`.

!!! info "Precondition"

    `i == 0` or `i == 1`.
"""
max(b::Bbox2, i::Integer)

"""
    +(b::Bbox2, c::Bbox2)

Returns a bounding box of `b` and `c`.
"""
+(b::Bbox2, c::Bbox2)

"""
    dilate(b::Bbox2, dist::Integer)

Dilates the bounding box by a specified number of ULP.
"""
dilate(b::Bbox2, dist::Integer)


text:
module Matheur

using JuMP
using DataFrames: DataFrame
using ..JOH
export Matheuristic, evaluate, Model_DF, get_DF_row

abstract type Matheuristic end

function evaluate(model::Model, matheuristic::T; id=1) where {T<:Matheuristic}
	matheuristic.executor(model, matheuristic, id=id)
end

include("utils.jl")
include("status.jl")
include("SSIT.jl")

end

max_stars_repo_name: Dibillilia/JuliaOptHeuristics

max_issues_repo_name: Dibillilia/JuliaOptHeuristics

max_forks_repo_name: Dibillilia/JuliaOptHeuristics


text:
## parameter sweep for comparison with optimization trajectories
using OptiMode
using HDF5
using Zygote: @ignore
using Base.Filesystem

# task_id = Base.parse(Int, ENV["SLURM_ARRAY_TASK_ID"])

function sum_Δng²_FHSH(ωs, p; Δx=6.0,Δy=4.0,Nx=128,Ny=128, fpath=nothing, grp_name=nothing,
	pidx=nothing,verbose=true)
	grid = Grid(Δx,Δy,Nx,Ny)
	rwg_pe(x) = ridge_wg_partial_etch(x[1],x[2],x[3],x[4],0.5,MgO_LiNbO₃,SiO₂,Δx,Δy) # partially etched ridge waveguide with dispersive materials, x[3] is partial etch fraction of top layer, x[3]*x[2] is etch depth, remaining top layer thickness = x[2]*(1-x[3]).
	nω = length(ωs)
    ngs_FHSH = solve_n(
				vcat(ωs, 2*ωs ),
				rwg_pe(p),
				grid,
			)[2]
    ngs_FH = ngs_FHSH[1:nω]
	ngs_SH = ngs_FHSH[nω+1:2*nω]
	Δng² = abs2.(ngs_SH .- ngs_FH)
	@ignore begin
		# print results for job log
		if verbose
			println("")
			println("pidx: $pidx")
			println("p: $p")
			println("\tngs_FH: $ngs_FH")
			println("\tngs_SH: $ngs_SH")
			println("\tsum(Δng²): $(sum(Δng²))")
			println("")
		end
		# save results to a data file (if provided)
		if !isnothing(fpath)
			grp_name = isnothing(grp_name) ? "/" : grp_name*"/"
			h5open(fpath,"r+") do fid
				fid[grp_name*"ngs_FH"][:,pidx.I...]		=	ngs_FH
				fid[grp_name*"ngs_SH"][:,pidx.I...]		=	ngs_SH
				fid[grp_name*"sumΔng²"][pidx.I...]	=	sum(Δng²)
			end
		end
	end
    return sum(Δng²)
end

"""
	create_data_file(fname; data_dir, groups, datasets)

Create a data file with filename `fname` in directory `data_dir` and return the path.

Currently this always creates HDF5 files, for which optional inputs `groups` and `datasets`
specify names of group and dataset objects to create in the new HDF5 file.
"""
function create_data_file(fname; data_dir=joinpath(homedir(),"data"), groups=nothing,
	datasets=nothing)
	fpath 	= joinpath(data_dir, fname)
	fid 	= h5open(fpath, "cw")			# "cw"== "read-write, create if not existing"
	if !isnothing(groups)
		group_refs = [create_group(fid, gr_name) for gr_name in groups]
		ds_refs = [create_dataset(fid, ds_name) for ds_name in datasets]
	end
	close(fid)
	return fpath
end

function sweep_params(p_upper,p_lower,nps,ωs; fname=nothing, data_dir=joinpath(homedir(),"data"), grp_name=nothing)
	ranges = [range(p_lower[i],p_upper[i],length=nps[i]) for i=1:length(nps)]
	ps = collect(Iterators.product(ranges...)) # collect(Iterators.product(reverse(ranges)...))

	nω = length(ωs)
	fpath 	= joinpath(data_dir, fname)
	h5open(fpath, "cw") do fid			# "cw"== "read-write, create if not existing"
		ds_names = ["ngs_FH", "ngs_SH", "sumΔng²"]
		ds_sizes = [(nω,size(ps)...),(nω,size(ps)...),size(ps) ]
		ds_refs = [create_dataset(fid, nm, rand(sz...)) for (nm,sz) in zip(ds_names,ds_sizes)]
	end

	for pidx in CartesianIndices(ps)
		pp = ps[pidx]
		try
			sum_Δng²_FHSH(ωs,pp; pidx, fpath, grp_name)
		catch
			println("")
			println("pidx: $pidx")
			println("p: $pp")
			println("\tngs_FH: Error")
			println("\tngs_SH: Error")
			println("\tsum(Δng²): Error")
			println("")
		end
	end
end

λs = reverse(1.4:0.01:1.6)
ωs = 1 ./ λs
p_lower 	= 	[	0.4,	0.3, 	0., 	0.,		]
p_upper 	= 	[	2., 	2., 	1., 	π/4.,	]
nps			=	[	20, 	20, 	5, 		3,		]
# nps			=	[	2, 	2, 	2, 		2,		]
sweep_params(p_upper,p_lower,nps,ωs;fname="sweep_LN_test3.h5")


text:
# gkl.jl

mutable struct GKLFactorization{T, S<:Real}
    k::Int # current Krylov dimension
    U::OrthonormalBasis{T} # basis of length k
    V::OrthonormalBasis{T} # basis of length k
    αs::Vector{S}
    βs::Vector{S}
    r::T
end

Base.length(F::GKLFactorization) = F.k
Base.sizehint!(F::GKLFactorization, n) = begin
    sizehint!(F.U, n)
    sizehint!(F.V, n)
    sizehint!(F.αs, n)
    sizehint!(F.βs, n)
    return F
end
Base.eltype(F::GKLFactorization) = eltype(typeof(F))
Base.eltype(::Type{<:GKLFactorization{<:Any,S}}) where {S} = S

function basis(F::GKLFactorization, which::Symbol)
    length(F.U) == F.k || error("Not keeping vectors during GKL bidiagonalization")
    which == :U || which == :V || error("invalid flag for specifying basis")
    return which == :U ? F.U : F.V
end
rayleighquotient(F::GKLFactorization) = Bidiagonal(view(F.αs,1:F.k),
                                        view(F.βs,1:(F.k-1)),:L)
residual(F::GKLFactorization) = F.r
@inbounds normres(F::GKLFactorization) = F.βs[F.k]
rayleighextension(F::GKLFactorization) = SimpleBasisVector(F.k, F.k)

# GKL iteration for constructing the orthonormal basis of a Krylov subspace.
struct GKLIterator{F,T,O<:Orthogonalizer} <: KrylovIterator{F,T}
    operator::F
    u₀::T
    orth::O
    keepvecs::Bool
    function GKLIterator{F,T,O}(operator::F, u₀::T, orth::O, keepvecs::Bool) where
        {F,T,O<:Orthogonalizer}

        if !keepvecs && isa(orth, Reorthogonalizer)
            error("Cannot use reorthogonalization without keeping all Krylov vectors")
        end
        new{F,T,O}(operator, u₀, orth, keepvecs)
    end
end
GKLIterator(operator::F, u₀::T, orth::O = KrylovDefaults.orth, keepvecs::Bool = true) where
                {F,T,O<:Orthogonalizer} = GKLIterator{F,T,O}(operator, u₀, orth, keepvecs)
GKLIterator(A::AbstractMatrix, u₀::AbstractVector, orth::O = KrylovDefaults.orth,
                keepvecs::Bool = true) where {O<:Orthogonalizer} =
    GKLIterator((x,flag)->flag ? A'*x : A*x, u₀, orth, keepvecs)

Base.IteratorSize(::Type{<:GKLIterator}) = Base.SizeUnknown()
Base.IteratorEltype(::Type{<:GKLIterator}) = Base.EltypeUnknown()

function Base.iterate(iter::GKLIterator)
    state = initialize(iter)
    return state, state
end
function Base.iterate(iter::GKLIterator, state::GKLFactorization)
    nr = normres(state)
    if nr < eps(typeof(nr))
        return nothing
    else
        state = expand!(iter, deepcopy(state))
        return state, state
    end
end

function initialize(iter::GKLIterator; verbosity::Int = 0)
    # initialize without using eltype
    u₀ = iter.u₀
    β₀ = norm(u₀)
    iszero(β₀) && throw(ArgumentError("initial vector should not have norm zero"))
    v₀ = iter.operator(u₀, true) # apply adjoint operator, might change eltype
    α = norm(v₀)/β₀
    Av₀ = iter.operator(v₀, false) # apply operator
    α² = dot(u₀, Av₀)/β₀^2
    α² ≈ α*α || throw(ArgumentError("operator and its adjoint are not compatible"))
    T = typeof(α²)
    # this line determines the type that we will henceforth use
    u = (one(T)/β₀)*u₀ # u = mul!(similar(u₀, T), u₀, 1/β₀)
    if typeof(v₀) == typeof(u)
        v = rmul!(v₀, 1/(α*β₀))
    else
        v = mul!(similar(u), v₀, 1/(α*β₀))
    end
    if typeof(Av₀) == typeof(u)
        r = rmul!(Av₀, 1/(α*β₀))
    else
        r = mul!(similar(u), Av₀, 1/(α*β₀))
    end
    r = axpy!(-α, u, r)
    β = norm(r)

    U = OrthonormalBasis([u])
    V = OrthonormalBasis([v])
    S = real(T)
    αs = S[α]
    βs = S[β]
    if verbosity > 0
        @info "GKL iteration step 1: normres = $β"
    end

    return GKLFactorization(1, U, V, αs, βs, r)
end
function initialize!(iter::GKLIterator, state::GKLFactorization; verbosity::Int = 0)
    V = state.V
    while length(U) > 1
        pop!(U)
    end
    V = empty!(state.V)
    αs = empty!(state.αs)
    βs = empty!(state.βs)

    u = mul!(V[1], iter.u₀, 1/norm(iter.u₀))
    v = iter.operator(u, true)
    α = norm(v)
    rmul!(v, 1/α)
    r = iter.operator(v, false) # apply operator
    r = axpy!(-α, u, r)
    β = norm(r)

    state.k = 1
    push!(V, v)
    push!(αs, α)
    push!(βs, β)
    state.r = r
    if verbosity > 0
        @info "GKL iteration step 1: normres = $β"
    end

    return state
end
function expand!(iter::GKLIterator, state::GKLFactorization; verbosity::Int = 0)
    βold = normres(state)
    U = state.U
    V = state.V
    r = state.r
    U = push!(U, rmul!(r, 1/βold))
    v, r, α, β = gklrecurrence(iter.operator, U, V, βold, iter.orth)

    push!(V, v)
    push!(state.αs, α)
    push!(state.βs, β)

    #!iter.keepvecs && popfirst!(state.V) # remove oldest V if not keepvecs

    state.k += 1
    state.r = r
    if verbosity > 0
        @info "GKL iteration step $(state.k): normres = $β"
    end

    return state
end
function shrink!(state::GKLFactorization, k)
    length(state) == length(state.V) ||
        error("we cannot shrink GKLFactorization without keeping vectors")
    length(state) <= k && return state
    U = state.U
    V = state.V
    while length(V) > k+1
        pop!(U)
        pop!(V)
    end
    pop!(V)
    r = pop!(U)
    resize!(state.αs, k)
    resize!(state.βs, k)
    state.k = k
    state.r = rmul!(r, normres(state))
    return state
end

# Golub-Kahan-Lanczos recurrence relation
function gklrecurrence(operator, U::OrthonormalBasis, V::OrthonormalBasis, β,
                        orth::Union{ClassicalGramSchmidt,ModifiedGramSchmidt})
    u = U[end]
    v = operator(u, true)
    v = axpy!(-β, V[end], v)
    α = norm(v)
    rmul!(v, inv(α))

    r = operator(v, false)
    r = axpy!(-α, u, r)
    β = norm(r)
    return v, r, α, β
end
function gklrecurrence(operator, U::OrthonormalBasis, V::OrthonormalBasis, β,
                        orth::ClassicalGramSchmidt2)
    u = U[end]
    v = operator(u, true)
    v = axpy!(-β, V[end], v) # not necessary if we definitely reorthogonalize next step and previous step
    # v, = orthogonalize!(v, V, ClassicalGramSchmidt())
    α = norm(v)
    rmul!(v, inv(α))

    r = operator(v, false)
    r = axpy!(-α, u, r)
    r, = orthogonalize!(r, U, ClassicalGramSchmidt())
    β = norm(r)
    return v, r, α, β
end
function gklrecurrence(operator, U::OrthonormalBasis, V::OrthonormalBasis, β,
                        orth::ModifiedGramSchmidt2)
    u = U[end]
    v = operator(u, true)
    v = axpy!(-β, V[end], v)
    # for q in V # not necessary if we definitely reorthogonalize next step and previous step
    #     v, = orthogonalize!(v, q, ModifiedGramSchmidt())
    # end
    α = norm(v)
    rmul!(v, inv(α))

    r = operator(v, false)
    r = axpy!(-α, u, r)
    for q in U
        r, = orthogonalize!(r, q, ModifiedGramSchmidt())
    end
    β = norm(r)
    return v, r, α, β
end
function gklrecurrence(operator, U::OrthonormalBasis, V::OrthonormalBasis, β,
                        orth::ClassicalGramSchmidtIR)
    u = U[end]
    v = operator(u, true)
    v = axpy!(-β, V[end], v)
    α = norm(v)
    nold = sqrt(abs2(α) + abs2(β))
    while α < orth.η * nold
        nold = α
        v, = orthogonalize!(v, V, ClassicalGramSchmidt())
        α = norm(v)
    end
    rmul!(v, inv(α))

    r = operator(v, false)
    r = axpy!(-α, u, r)
    β = norm(r)
    nold = sqrt(abs2(α) + abs2(β))
    while eps(one(β)) < β < orth.η * nold
        nold = β
        r, = orthogonalize!(r, U, ClassicalGramSchmidt())
        β = norm(r)
    end

    return v, r, α, β
end
function gklrecurrence(operator, U::OrthonormalBasis, V::OrthonormalBasis, β,
                        orth::ModifiedGramSchmidtIR)
    u = U[end]
    v = operator(u, true)
    v = axpy!(-β, V[end], v)
    α = norm(v)
    nold = sqrt(abs2(α) + abs2(β))
    while eps(one(α)) < α < orth.η * nold
        nold = α
        for q in V
            v, = orthogonalize!(v, q, ModifiedGramSchmidt())
        end
        α = norm(v)
    end
    rmul!(v, inv(α))

    r = operator(v, false)
    r = axpy!(-α, u, r)
    β = norm(r)
    nold = sqrt(abs2(α) + abs2(β))
    while eps(one(β)) < β < orth.η * nold
        nold = β
        for q in U
            r, = orthogonalize!(r, q, ModifiedGramSchmidt())
        end
        β = norm(r)
    end

    return v, r, α, β
end


text:
using Pkg;
Pkg.activate(joinpath(@__DIR__, "..")); Pkg.instantiate()
Pkg.activate(); Pkg.instantiate()

pushfirst!(LOAD_PATH, joinpath(@__DIR__, ".."))

using Documenter, Photon

makedocs(modules=[Photon],
         sitename = "Photon",
         pages = ["Home" => "index.md",
                  "API" => [
                        "Core"    => "core.md",
                        "Layers"  => "layers.md",
                        "Losses"  => "losses.md",
                        "Metrics" => "metrics.md",
                        "Data"    => "data.md"
                  ],
                  "Community" => "community.md"],
         format = Documenter.HTML(
                  prettyurls = get(ENV, "CI", nothing) == "true"
         ))

deploydocs(
        repo = "github.com/neurallayer/Photon.jl.git",
        target = "build")


text:
using CUDA
using GemmKernels
using LinearAlgebra

CUDA.CUBLAS.cublasSetMathMode(CUBLAS.handle(), CUBLAS.CUBLAS_TENSOR_OP_MATH)

@test_if "blas" @testset "BLAS API" begin
    @testset "WMMA GEMM $(A_type)*$(B_type)+$(CD_type)=$(CD_type) ($( !transpose_a ? 'N' : 'T' )$( !transpose_b ? 'N' : 'T' ))" for transpose_a = [false, true],
        transpose_b = [false, true],
        (A_type, B_type, CD_type, min_dimension) in [(Float16, Float16, Float16, 256), (Float16, Float16, Float32, 128)]

        @testset "(M = $M, N = $N, K = $K)" for M in min_dimension .* [1, 2],
            N in min_dimension .* [1, 2],
            K in min_dimension .* [1, 2]

            alpha = rand(A_type)
            beta  = rand(CD_type)

            a_h = rand(A_type, (M, K)) / sqrt(A_type(K))
            b_h = rand(B_type, (K, N)) / sqrt(B_type(K))
            c_h = rand(CD_type, (M, N))

            # Transpose input if necessary
            a_h = transpose_a ? transpose(a_h) : a_h
            b_h = transpose_b ? transpose(b_h) : b_h

            a   = CuArray(a_h)
            b   = CuArray(b_h)

            c_gemmkernels = CuArray(c_h)
            GemmKernels.BLAS.gemmEx!(!transpose_a ? 'N' : 'T', !transpose_b ? 'N' : 'T', alpha, a, b, beta, c_gemmkernels)

            c_cublas = CuArray(c_h)
            CUDA.CUBLAS.gemmEx!(!transpose_a ? 'N' : 'T', !transpose_b ? 'N' : 'T', alpha, a, b, beta, c_cublas)

            @test all(isapprox.(Array(c_gemmkernels), Array(c_cublas); rtol=sqrt(eps(A_type))));
        end
    end

    @testset "WMMA GEMM (A = diagonal, B = $( !transpose_b ? 'N' : 'T' ))" for transpose_b = [false, true]
        @testset "(M = $M, N = $N, K = $K)" for M in [128, 256],
            N in [128, 256],
            K in [M]

            transpose_a = false

            alpha = rand(Float32)
            beta = rand(Float32)

            a_h = rand(Float16, M);
            b_h = rand(Float16, (K, N)) / sqrt(Float16(K))
            c_h = rand(Float32, (M, N))

            # Transpose input if necessary
            a_h = transpose_a ? transpose(a_h) : a_h
            b_h = transpose_b ? transpose(b_h) : b_h

            a   = Diagonal(CuArray(a_h))
            b   = CuArray(b_h)

            c_gemmkernels = CuArray(c_h)
            GemmKernels.BLAS.gemmEx!(!transpose_a ? 'N' : 'T', !transpose_b ? 'N' : 'T', alpha, a, b, beta, c_gemmkernels)

            c_cublas = CuArray(c_h)
            CUDA.CUBLAS.gemmEx!(!transpose_a ? 'N' : 'T', !transpose_b ? 'N' : 'T', alpha, CuArray(Diagonal(a_h)), b, beta, c_cublas)

            @test all(isapprox.(Array(c_gemmkernels), Array(c_cublas); rtol=sqrt(eps(Float16))));
        end
    end
end


text:
function lift(a::GenRes)
  return a.data
end

function ^(a::ResElem, f::fmpz)
  f==0 && return one(parent(a))
  f==1 && return a
  if f<0
    f=-f
    a = inv(a)
  end
  if f<(1<<30)
    return a^Int(f)
  end
  b = a^(div(f, 2))
  b = b^2
  if isodd(f)
    b *= a
  end
  return b
end

function zero!(z::fq_nmod)
  ccall((:fq_nmod_zero, :libflint), Void,
          (Ptr{fq_nmod}, Ptr{FqNmodFiniteField}), &z, &parent(z))
end

function add!(z::fq_nmod, x::fq_nmod, y::fq_nmod)
  ccall((:fq_nmod_add, :libflint), Void,
          (Ptr{fq_nmod}, Ptr{fq_nmod}, Ptr{fq_nmod},
           Ptr{FqNmodFiniteField}), &z, &x, &y, &parent(z))
end

function set!(z::fq_nmod, x::fq_nmod)
  ccall((:fq_nmod_set, :libflint), Void,
          (Ptr{fq_nmod}, Ptr{fq_nmod}, Ptr{FqNmodFiniteField}),
          &z, &x, &parent(z))
end


text:
using jInv.Mesh;
using ForwardHelmholtz
using ForwardHelmholtz.PointSourceADR

plotting = true;

if plotting
	using PyPlot;
	# close("all");
end

include("getModels.jl");

NeumannAtFirstDim = true;
Sommerfeld = true;
gamma_val = 0.00001*2*pi;

# model = "const";
# model = "linear";
# model = "linear squared slowness";
# model = "Gaussian"
# model = "Wave-guide"
model = "Wedge"
domain = [0.0,1.0,0.0,1.0];
	n_cells = [512,512];
src = div(n_cells,2)+1;
src[2] = 1;

if model=="const"
	NeumannAtFirstDim = false;
	Sommerfeld = true;
	f = 40.0;
	src = div(n_cells,2)+1;
	m_coarse = getConstModel(n_cells+1);
	pad_cells = [20;20];
elseif model=="linear"
	f = 40.0
	m_coarse = getLinearModel(1.6,3.5,n_cells+1);m_coarse = 1./m_coarse.^2
	pad_cells = [20;20];
elseif model=="linear squared slowness"
	f = 70.0
	m_coarse = getInvLinearModel(1.6,3.5,n_cells+1);
	pad_cells = [20;20];
elseif model=="Gaussian"
	f = 40.0
	pad_cells = [50;50];
	Minv = getRegularMesh(domain,n_cells);
	m_coarse = getSmoothGaussianMedium(n_cells+1,Minv.h)[1];
elseif model=="Wave-guide"
	pad_cells = [64;64];
	f = 20.0;
	m_coarse = waveGuideModel(n_cells+1)[1];
elseif model=="Wedge"
	f = 20.0;
	pad_cells = [36;36];
	n_cells = [512,512];  (m_coarse,Minv) = getWedgeModel(n_cells+1);domain = Minv.domain;
else
	error("unknown model");
end
Minv = getRegularMesh(domain,n_cells);

# m_coarse = getLinearModel(1.6,3.5,n_cells+1);
# m_coarse = 1./m_coarse.^2
# m_coarse = getLinearModel(1./((1.6)^2),1./(3.5)^2,n_cells+1);
# m_coarse = getSmoothGaussianMedium(n_cells+1,Minv.h)[1];
#

w = 2*pi*f

figure()
imshow(m_coarse'); colorbar();
title(string("The ",model," model"))
# error("Here")

n_tup = tuple((n_cells+1)...);
q_coarse = zeros(Complex,n_tup)
q_coarse[src[1],src[2]] = 1/(Minv.h[2]^2);

# src2 = div(3*n_cells,4)+1;
# q_coarse[src2[1],src2[2]] = 1/(Minv.h[2]^2);
# src = src2;

println("omega*h:");
println(w*Minv.h*sqrt(maximum(m_coarse)));

# m = readdlm("SEGmodel2Dsalt.dat"); m = m'; m = m*1e-3;
maxOmega = getMaximalFrequency(m_coarse,Minv);
ABLamp = maxOmega;

### GENERATE FINE REFERENCE ##############################################################################
factor = 4;
n_cells_fine = factor*n_cells;
pad_cells_fine = factor*pad_cells;

Minv_fine = getRegularMesh(domain,n_cells_fine);
src_fine = div(n_cells_fine,2)+1;
if src[2]==1
	src_fine[2] = 1;
end
n_tup = tuple((n_cells_fine+1)...);
q = zeros(Complex,n_tup)
q[src_fine[1],src_fine[2]] = 1/(Minv_fine.h[2]^2);

# src2 = div(3*n_cells_fine,4)+1;
# q[src2[1],src2[2]] = 1/(Minv_fine.h[2]^2);

m=0;
if model=="const"
	m = getConstModel(n_cells_fine+1);
elseif model=="linear"
	m = getLinearModel(1.6,3.5,n_cells_fine+1);m = 1./m.^2
elseif model=="linear squared slowness"
	m = getInvLinearModel(1.6,3.5,n_cells_fine+1);
elseif model=="Gaussian"
	m = getSmoothGaussianMedium(n_cells_fine+1,Minv_fine.h)[1];
elseif model=="Wave-guide"
	m = waveGuideModel(n_cells_fine+1)[1];
elseif model=="Wedge"
	m = getWedgeModel(n_cells_fine+1)[1];
end
#
#

H, = GetHelmholtzOperator(Minv_fine, m, w, gamma_val*ones(Float32,size(m)),NeumannAtFirstDim,pad_cells_fine+1,2.0*ABLamp,Sommerfeld);
xh_fine = (H\q[:]);
xh_fine = reshape(xh_fine,n_tup);

n_tup = tuple(n_cells+1...);
xh_ref = zeros(Complex128,n_tup);
xh_ref = xh_fine[1:factor:end,1:factor:end];
Iref = xh_ref
figure()
imshow(real(xh_ref)');colorbar();
### FINE REFERENCE DONE ##############################################################################

m = m_coarse;

gamma = gamma_val*ones(Float32,size(m)) + getABL(Minv,NeumannAtFirstDim,pad_cells+1,ABLamp);
(ADR_long,ADR_short,T) = getHelmholtzADR(true,Minv, m, w, gamma,NeumannAtFirstDim,src,Sommerfeld);
(ADR_longUp,ADR_short,T) = getHelmholtzADR(false,Minv, m, w, gamma,NeumannAtFirstDim,src,Sommerfeld);

figure();
imshow(real(exp(1im*w*reshape(T,n_tup)))');

H, = GetHelmholtzOperator(Minv, m, w, gamma_val*ones(Float32,size(m)),NeumannAtFirstDim,pad_cells+1,ABLamp,Sommerfeld);
S =  GetHelmholtzShiftOP(m, w,0.2);

q = q_coarse;
xh = (H\q[:])
q = q[:].*exp(1im*w*T);
xl = (ADR_long\q[:]).*exp(-1im*w*T);
xs = ((ADR_short )\q[:]).*exp(-1im*w*T);
us = (H+S)\q[:];
alu = (ADR_longUp\q[:]);
xlu = alu.*exp(-1im*w*T);

I1 = reshape(xh,n_tup);
I2 = reshape(xl,n_tup);
I3 = reshape(xlu,n_tup);
I4 = reshape(xs,n_tup);
I5 = reshape(us,n_tup);

b = maximum(real([I1[:];I2[:];I3[:]]));
a = minimum(real([I1[:];I2[:];I3[:]]));
a_err = 1.0;
N = abs(Iref);
N = N + 0.1*mean(N);
# N[:]=1.0;

E1 = abs(Iref-I1)./N;
E2 = abs(Iref-I2)./N;
E3 = abs(Iref-I3)./N;
E4 = abs(Iref-I4)./N;
# E1[src[1],src[2]] = 0.0;
# E2[src[1],src[2]] = 0.0;
# E3[src[1],src[2]] = 0.0;
# E4[src[1],src[2]] = 0.0;

if plotting
figure();
subplot(1,4,1);
imshow(real(I2'));
subplot(1,4,2);
imshow(reshape(T,n_tup)');
subplot(1,4,3);
imshow(reshape(real(alu),n_tup)');
subplot(1,4,4);
imshow(reshape(real(exp(-1im*w*T)),n_tup)');

figure();
subplot(1,3,1);
imshow(real(I1')); title("Full Helmholtz solution")
subplot(1,3,2);
imshow(reshape(real(us),n_tup)'); title("Shifted Helmholtz solution")
subplot(1,3,3);
imshow(reshape(real(xs),n_tup)'); title("First order upwind ADR solution")

# figure()
# subplot(2,4,1);
# imshow(real(I1'),clim = [a,b]); colorbar();title("Standard discretization")
# subplot(2,4,2);
# imshow(real(I2'),clim = [a,b]);colorbar();title("ADR discretization - central")
# subplot(2,4,3);
# imshow(real(I3'),clim = [a,b]);colorbar();title("ADR discretization - 2nd upwind")
# subplot(2,4,4);
# imshow(real(I4'),clim = [a,b]);colorbar();title("ADR discretization - 1st upwind")

# subplot(2,4,5);
# imshow(E1',clim = [0.0,a_err]);colorbar();title("Relative error - standard")
# subplot(2,4,6);
# imshow(E2',clim = [0.0,a_err]);colorbar();title("Relative error - ADR central")
# subplot(2,4,7);
# imshow(E3',clim = [0.0,a_err]);colorbar();title("Relative error - ADR upwind")
# subplot(2,4,8);
# imshow(E4',clim = [0.0,a_err]);colorbar();

figure()
subplot(2,3,1);
imshow(real(I1'),clim = [a,b]); colorbar();title("Standard discretization");
subplot(2,3,2);
imshow(real(I2'),clim = [a,b]);colorbar();title("ADR discretization - central")
subplot(2,3,3);
imshow(real(I3'),clim = [a,b]);colorbar();title("ADR discretization - 2nd order upwind")

subplot(2,3,4);
imshow(E1',clim = [0.0,a_err]);colorbar();title("Relative error - standard")
subplot(2,3,5);
imshow(E2',clim = [0.0,a_err]);colorbar();title("Relative error - ADR central")
subplot(2,3,6);
imshow(E3',clim = [0.0,a_err]);colorbar();title("Relative error - ADR upwind")

# figure()
# subplot(2,3,1);
# imshow(reshape(T,n[1],n[2])');colorbar();
# subplot(2,3,2);
# imshow(reshape(real(exp(1im*w*T)),n[1],n[2])'); colorbar();
# subplot(2,3,3);
# imshow(reshape(G1,n[1],n[2])');colorbar();
# subplot(2,3,4);
# imshow(reshape(G2,n[1],n[2])');colorbar();
# subplot(2,3,5);
# imshow(reshape(G1.^2 + G2.^2,n[1],n[2])');colorbar();
# subplot(2,3,6);
# imshow(reshape(LT,n[1],n[2])');colorbar();

end

error("ET")


text:
module Day17

export day17

function parseinput(::Val{3}, filename)
    n = length(readline(filename))
    space = zeros(Int,n+16,n+16,n+16)
    for (i,line) in enumerate(readlines(filename))
        for (j,ch) in enumerate(line)
            if ch == '.'
                space[7+i,7+j,8] = 0
            else
                space[7+i,7+j,8] = 1
            end
        end
    end
    space
end

function parseinput(::Val{4}, filename)
    n = length(readline(filename))
    space = zeros(Int,n+16,n+16,n+16, n+16)
    for (i,line) in enumerate(readlines(filename))
        for (j,ch) in enumerate(line)
            if ch == '.'
                space[7+i,7+j,8,8] = 0
            else
                space[7+i,7+j,8,8] = 1
            end
        end
    end
    space
end

function count_neighbours(::Val{3}, space, x,y,z)
    count = - space[x,y,z]
    for i in -1:1, j in -1:1, k in -1:1
        count += space[x+i, y+j, z+k]
    end
    count
end

function count_neighbours(::Val{4}, space, x,y,z,w)
    count = - space[x,y,z,w]
    for i in -1:1, j in -1:1, k in -1:1, s in -1:1
        count += space[x+i, y+j, z+k, w+s]
    end
    count
end

function simulate(::Val{3}, space)
    new_space = copy(space)
    for x in 2:size(space,1)-1, y in 2:size(space,2)-1, z in 2:size(space,3)-1
        n = count_neighbours(Val(3),space,x,y,z)
        if space[x,y,z] == 1 && n != 2 && n != 3
            new_space[x,y,z] = 0
        elseif space[x,y,z] == 0 && n == 3
            new_space[x,y,z] = 1
        end
    end
    new_space
end

function simulate(::Val{4}, space)
    new_space = copy(space)
    for x in 2:size(space,1)-1, y in 2:size(space,2)-1, z in 2:size(space,3)-1, w in 2:size(space,4)-1
        n = count_neighbours(Val(4),space,x,y,z,w)
        if space[x,y,z,w] == 1 && n != 2 && n != 3
            new_space[x,y,z,w] = 0
        elseif space[x,y,z,w] == 0 && n == 3
            new_space[x,y,z,w] = 1
        end
    end
    new_space
end

function day17(filename, dim)
    space = parseinput(Val(dim), filename)
    for cycle in 1:6
        space = simulate(Val(dim), space)
    end
    sum(space)
end

end


text:
module CreateTableTweets

import Migration: create_sequence, create_table, column_id, column,
                  column_id_sequence, add_index, drop_table, drop_sequence,
                  constraint, nextval

function up()
  create_sequence(:tweets, :id)
  create_table(:tweets) do
    [
      column_id(constraint = constraint(:tweets, :id), nextval = nextval(:tweets, :id))
      column(:tweet_id, :string, limit = 24)
      column(:text, :string, limit = 512)
      column(:user_name, :string, limit = 32)
      column(:screen_name, :string, limit = 32)
      column(:avatar_url, :string, limit = 512)
      column(:favorite_count, :integer)
      column(:retweet_count, :integer)
      column(:polarity, :integer)
      column(:subjective, :boolean)
    ]
  end
  column_id_sequence(:tweets, :id)

  add_index(:tweets, :tweet_id, unique = true)
  add_index(:tweets, :text)
  add_index(:tweets, :polarity)
  add_index(:tweets, :subjective)
end

function down()
  drop_table(:tweets)
end

end

max_stars_repo_name: essenciary/genie-demo-tweet-stats

max_issues_repo_name: essenciary/genie-demo-tweet-stats

max_forks_repo_name: essenciary/genie-demo-tweet-stats


text:
export PBSManager, SGEManager, QRSHManager, addprocs_pbs, addprocs_sge, addprocs_qrsh

struct PBSManager <: ClusterManager
    np::Integer
    queue
    wd
end

struct SGEManager <: ClusterManager
    np::Integer
    queue
    wd
end

struct QRSHManager <: ClusterManager
    np::Integer
    queue
    wd
end

function launch(manager::Union{PBSManager, SGEManager, QRSHManager},
        params::Dict, instances_arr::Array, c::Condition)
    try
        dir = params[:dir]
        exename = params[:exename]
        exeflags = params[:exeflags]
        wd = manager.wd
        isPBS = isa(manager, PBSManager)

        np = manager.np
        queue = manager.queue

        jobname = `julia-$(getpid())`

        if isa(manager, QRSHManager)
          cmd = `cd $dir '&&' $exename $exeflags $(worker_arg())`
          qrsh_cmd = `qrsh $queue -V -N $jobname -wd $wd -now n "$cmd"`

          stream_proc = [open(qrsh_cmd) for i in 1:np]

          for i in 1:np
              config = WorkerConfig()
              config.io, io_proc = stream_proc[i]
              config.userdata = Dict{Symbol, Any}(:task => i,
                                                  :process => io_proc)
              push!(instances_arr, config)
              notify(c)
          end

        else  # PBS & SGE
            cmd = `cd $dir '&&' $exename $exeflags $(worker_arg())`
            qsub_cmd = pipeline(`echo $(Base.shell_escape(cmd))` , (isPBS ?
                    `qsub -N $jobname -wd $wd -j oe -k o -t 1-$np $queue` :
                    `qsub -N $jobname -wd $wd -terse -j y -R y -t 1-$np -V $queue`))

            out = open(qsub_cmd)
            if !success(out)
              throw(error()) # qsub already gives a message
            end

            id = chomp(split(readline(out),'.')[1])
            if endswith(id, "[]")
                id = id[1:end-2]
            end

            filenames(i) = "$wd/julia-$(getpid()).o$id-$i","$wd/julia-$(getpid())-$i.o$id","$wd/julia-$(getpid()).o$id.$i"

            println("Job $id in queue.")
            for i=1:np
                # wait for each output stream file to get created
                fnames = filenames(i)
                j = 0
                while (j=findfirst(x->isfile(x),fnames))==nothing
                    sleep(1.0)
                end
                fname = fnames[j]

                # Hack to get Base to get the host:port, the Julia process has already started.
                cmd = `tail -f $fname`

                config = WorkerConfig()

                config.io = open(detach(cmd))

                config.userdata = Dict{Symbol, Any}(:job=>id, :task=>i, :iofile=>fname)
                push!(instances_arr, config)
                notify(c)
            end
            println("Running.")
        end

    catch e
        println("Error launching workers")
        println(e)
    end
end

function manage(manager::Union{PBSManager, SGEManager, QRSHManager},
        id::Int64, config::WorkerConfig, op::Symbol)
end

function kill(manager::Union{PBSManager, SGEManager, QRSHManager}, id::Int64, config::WorkerConfig)
    remotecall(exit,id)
    close(config.io)

    if isa(manager, QRSHManager)
      kill(config.userdata[:process],15)
      return
    end

    if isfile(config.userdata[:iofile])
        rm(config.userdata[:iofile])
    end
end

addprocs_pbs(np::Integer; qsub_flags=``, wd=ENV["HOME"]) = addprocs(PBSManager(np, qsub_flags, wd))

addprocs_sge(np::Integer; qsub_flags=``, wd=ENV["HOME"]) = addprocs(SGEManager(np, qsub_flags, wd))

addprocs_qrsh(np::Integer; qsub_flags=``, wd=ENV["HOME"]) = addprocs(QRSHManager(np, qsub_flags, wd))


text:
#=

This file wraps `compute_aerosol_optical_properties` so that autodiff users and non-autodiff
users can call the same function with just a keyword argument change.

=#

""" Function used by auto-differentiation to convert a jacobian result to an AerosolOptics type"""
function convert_jacobian_result_to_aerosol_optics(result)

    value = result.value
    derivs = result.derivs[1]

    greek_length = Int64((length(value) - 2)/6)

    α = value[1:greek_length]
    β = value[greek_length + 1 : 2greek_length]
    γ = value[2greek_length + 1 : 3greek_length]
    δ = value[3greek_length + 1 : 4greek_length]
    ϵ = value[4greek_length + 1 : 5greek_length]
    ζ = value[5greek_length + 1 : 6greek_length]

    greek_coefs = GreekCoefs(α, β, γ, δ, ϵ, ζ)

    ω̃ = value[end-1]
    k = value[end]

    return AerosolOptics(greek_coefs=greek_coefs, ω̃=ω̃, k=k, derivs=derivs, fᵗ=eltype(ω̃)(1))
end

"""
    $(FUNCTIONNAME)(model::MieModel{FDT})

Reference: Suniti Sanghavi 2014, https://doi.org/10.1016/j.jqsrt.2013.12.015

This function enables user to specify whether to perform auto-differentiation (using either computation type)
Input: MieModel, holding all computation and aerosol properties & autodiff flag (whether to perform
auto-differentiation)
Output: AerosolOptics, holding all Greek coefficients and Cross-Sectional information
"""
function compute_aerosol_optical_properties(model::MieModel ; autodiff=false)

    # This function takes in the "x-vector" along with the input model so that ForwardDiff will work
    function compute_aerosol_optical_properties_autodiff(x ; model::MieModel = model)

        if length(x) !== 4
            @error "Must receive four aerosol parameters for auto-differentiation (μ, σ, nᵣ, nᵢ)" x
        end

        # Make sure that 𝐱 and model match
        @assert (model.aerosol.size_distribution.μ == log(x[1]))
        @assert (model.aerosol.size_distribution.σ == log(x[2]))
        @assert (model.aerosol.nᵣ == x[3])
        @assert (model.aerosol.nᵢ == x[4])

        # Unpack the model and aerosol
        @unpack computation_type, aerosol, λ, polarization_type, truncation_type, r_max, nquad_radius, wigner_A, wigner_B = model
        @unpack size_distribution, nᵣ, nᵢ = aerosol

        aerosol_x = Aerosol(LogNormal(log(x[1].value), log(x[2].value)), x[3].value, x[4].value)
        model_x = MieModel(computation_type, aerosol_x, λ, polarization_type, truncation_type, r_max, nquad_radius, wigner_A, wigner_B)

        aerosol_optics = compute_aerosol_optical_properties(model_x, ForwardDiff.Dual);

        return [aerosol_optics.greek_coefs.α
                aerosol_optics.greek_coefs.β
                aerosol_optics.greek_coefs.γ
                aerosol_optics.greek_coefs.δ
                aerosol_optics.greek_coefs.ϵ
                aerosol_optics.greek_coefs.ζ
                aerosol_optics.ω̃
                aerosol_optics.k]
    end

    if (autodiff)

        x = [exp(model.aerosol.size_distribution.μ),
            exp(model.aerosol.size_distribution.σ),
            model.aerosol.nᵣ,
            model.aerosol.nᵢ]

        # Get length of greek coefs
        r, wᵣ = gauleg(model.nquad_radius, 0.0, model.r_max ; norm=true)
        N_max = get_n_max(2 * π * model.r_max/ model.λ)
        greek_length = 2 * N_max - 1

        result = DiffResults.JacobianResult(zeros(6 * greek_length + 2), x)
        ForwardDiff.jacobian!(result, compute_aerosol_optical_properties_autodiff, x);
        return convert_jacobian_result_to_aerosol_optics(result);

    else
        return compute_aerosol_optical_properties(model)
    end

end


text:
using Printf
using LinearAlgebra

using DifferentialEquations
using Plots

ENV["GKSwstype"] = "100"

""" Returns a discrete 1D derivative operator for cell center to cell (f)aces. """
function Dᶠ(N, Δ)
    D = zeros(N, N+1)
    for k in 1:N
        D[k, k]   = -1.0
        D[k, k+1] =  1.0
    end
    D = 1/Δ * D
    return D
end

""" Returns a discrete 1D derivative operator for cell faces to cell (c)enters. """
function Dᶜ(N, Δ)
    D = zeros(N+1, N)
    for k in 2:N
        D[k, k-1] = -1.0
        D[k, k]   =  1.0
    end
    D = 1/Δ * D
    return D
end

function ℑᶠ(N, Δ)
    ℑ = zeros(N, N+1)
    for k in 1:N
        ℑ[k, k]   = 0.5
        ℑ[k, k+1] = 0.5
    end
    return ℑ
end

function ℑᶜ(N, Δ)
    ℑ = zeros(N+1, N)
    for k in 2:N
        ℑ[k, k-1] = 0.5
        ℑ[k, k]   = 0.5
    end
    return ℑ
end

function convective_adjustment(T, p, t)
    N, κ, Qₛ = p.N, p.κ, p.Qₛ
    Dzᶠ, Dzᶜ = p.Dzᶠ, p.Dzᶜ

    ∂T∂z = Dzᶜ * T

    wT = @. min(0, κ * ∂T∂z)
    wT[N+1] += Qₛ

    return Dzᶠ * wT
end

function surface_flux(T, p, t)
    N, Qₛ, Dzᶠ = p.N, p.Qₛ, p.Dzᶠ

    wT = zeros(N+1)
    wT[N+1] += Qₛ

    return Dzᶠ * wT
end

N = 32
L = 100
Δ = L/N
z = range(Δ/2, L-Δ/2, length=N)

Dzᶠ = Dᶠ(N, Δ)
Dzᶜ = Dᶜ(N, Δ)
ℑzᶠ = ℑᶠ(N, Δ)
ℑzᶜ = ℑᶜ(N, Δ)

ρ₀ = 1027
cₚ = 4000
Q  = -100 / (ρ₀ * cₚ)

T₀ = [19 + zₖ/L for zₖ in z]

day = 86400
tspan = (0.0, 8day)
tsave = range(tspan...; length=101)

params = (N=N, κ=1000, Qₛ=Q, Dzᶠ=Dzᶠ, Dzᶜ=Dᶜ(N, Δ))

prob = ODEProblem(convective_adjustment, T₀, tspan, params)
@time sol = solve(prob, Rodas5(), reltol=1e-5, saveat=tsave, progress=true)

anim = @animate for n in 1:length(sol)
    @info "Frame $n/$(length(sol))"
    time_str = @sprintf("%.2f days", sol.t[n] / day)

    plot(sol[n], z, linewidth=2, xlim=(19, 20), ylim=(0, 100),
         xlabel="Temperature (°C)", ylabel="Depth z (meters)",
         title="Convective adjustment: $time_str", legend=:bottomright, show=false)
end

mp4(anim, "convective_adjustment_stiff.mp4", fps=15)

function convective_adjustment!(integrator)
    T = integrator.u
    ∂T∂z = Dzᶜ * T
    ℑ∂T∂z = ℑzᶠ * ∂T∂z

    κ = zeros(N)
    for j in 1:N
        κ[j] = ℑ∂T∂z[j] < 0 ? K : 0
    end

    ld = [-Δt/Δz^2 * κ[j]   for j in 2:N]
    ud = [-Δt/Δz^2 * κ[j+1] for j in 1:N-1]

    d = zeros(N)
    for j in 1:N-1
        d[j] = 1 + Δt/Δz^2 * (κ[j] + κ[j+1])
    end
    d[N] = 1 + Δt/Δz^2 * κ[N]

    𝓛 = Tridiagonal(ld, d, ud)
    T .= 𝓛 \ T

    return nothing
end

Δz = Δ
Δt = 3600.0
K  = 1000
time_steps = Int(8day / Δt)

prob = ODEProblem(surface_flux, T₀, tspan, params)
integrator = init(prob, Tsit5(), adaptive=false, dt=Δt, saveat=0:Δt:Δt*time_steps)

@time begin
    for _ in 1:time_steps
        step!(integrator)
        convective_adjustment!(integrator)
    end
end

anim = @animate for n in 1:time_steps
    @info "frame $n/$time_steps"
    time_str = @sprintf("%.2f days", integrator.sol.t[n] / day)
    plot(integrator.sol[n], z, linewidth=2, xlim=(19, 20), ylim=(0, 100),
         xlabel="Temperature (°C)", ylabel="Depth z (meters)",
         title="Convective adjustment: $time_str", legend=:bottomright, show=false)
end

mp4(anim, "convective_adjustment_implicit.mp4", fps=30)


text:
## Particle filter update ##
export pf_update!

"""
    pf_update!(state::ParticleFilterState, new_args::Tuple,
               argdiffs::Tuple, observations::ChoiceMap)

Perform a particle filter update, where the model arguments are adjusted and
new observations are conditioned upon. New latent choices are sampled from
the model's default proposal.
"""
function pf_update!(state::ParticleFilterView, new_args::Tuple,
                    argdiffs::Tuple, observations::ChoiceMap)
    n_particles = length(state.traces)
    for i=1:n_particles
        state.new_traces[i], increment, _, discard =
            update(state.traces[i], new_args, argdiffs, observations)
        if !isempty(discard)
            error("Choices were updated or deleted: $discard")
        end
        state.log_weights[i] += increment
    end
    update_refs!(state)
    return state
end

"""
    pf_update!(state::ParticleFilterState, translator; translator_args...)

Perform a particle filter update using an arbitrary trace `translator`, which
takes in each previous trace and returns a new trace and incremental log.
importance weight. `translator_args` are additional keyword arguments which
are passed to the `translator` when it is called.
"""
function pf_update!(state::ParticleFilterView, translator; translator_args...)
    n_particles = length(state.traces)
    for i=1:n_particles
        state.new_traces[i], log_weight =
            translator(state.traces[i]; translator_args...)
        state.log_weights[i] += log_weight
    end
    update_refs!(state)
    return state
end

"""
    pf_update!(state::ParticleFilterState, new_args::Tuple,
               argdiffs::Tuple, observations::ChoiceMap,
               proposal::GenerativeFunction, proposal_args::Tuple,
               [transform::TraceTransformDSLProgram])

Perform a particle filter update, where the model arguments are adjusted and
new observations are conditioned upon. New latent choices are sampled from a
custom `proposal` distribution in conjuction with the model's default proposal.
For each particle:
* `proposal` is evaluated with arguments `(t_old, proposal_args...)`,
   where `t_old` is the old model trace, and produces its own trace `t_prop`.
* The old model trace is replaced by a new model trace `t_new`, constructed
  by merging the choices in `t_old` and `t_prop`, and sampling any remaining
  choices from the model's default proposal.

The choicemap of `t_new` satisfies the following conditions:
1. `observations` is a subset of `get_choices(t_new)`;
2. `get_choices(t_old)` is a subset of `get_choices(t_new)`;
3. `get_choices(t_prop)` is a subset of `get_choices(t_new)`.
where one choicemap `a` is a "subset" of another choicemap `b`, when all keys
that occur in `a` also occur in `b`, and the values at those addresses are
equal. It is an error if no trace `t_new` satisfying the above conditions
exists in the support of the model (with the new arguments).

If a deterministic trace `transform` is also provided, `t_prop` is transformed
before its choices are merged with `t_old`.
"""
function pf_update!(state::ParticleFilterView, new_args::Tuple,
                    argdiffs::Tuple, observations::ChoiceMap,
                    proposal::GenerativeFunction, proposal_args::Tuple)
    n_particles = length(state.traces)
    for i=1:n_particles
        (prop_choices, prop_weight, _) =
            propose(proposal, (state.traces[i], proposal_args...))
        constraints = merge(observations, prop_choices)
        state.new_traces[i], up_weight, _, discard =
            update(state.traces[i], new_args, argdiffs, constraints)
        if !isempty(discard)
            error("Choices were updated or deleted: $discard")
        end
        state.log_weights[i] += up_weight - prop_weight
    end
    update_refs!(state)
    return state
end

function pf_update!(state::ParticleFilterView, new_args::Tuple,
                    argdiffs::Tuple, observations::ChoiceMap,
                    proposal::GenerativeFunction, proposal_args::Tuple,
                    transform::TraceTransformDSLProgram)
    n_particles = length(state.traces)
    translator = GenParticleFilters.ExtendingTraceTranslator(
        p_new_args=new_args, p_argdiffs=argdiffs, new_observations=observations,
        q_forward=proposal, q_forward_args=proposal_args, transform=transform)
    return pf_update!(state, translator)
end

"""
    pf_update!(state::ParticleFilterState, new_args::Tuple,
               argdiffs::Tuple, observations::ChoiceMap,
               fwd_proposal::GenerativeFunction, fwd_args::Tuple,
               bwd_proposal::GenerativeFunction, bwd_args::Tuple,
               [transform::TraceTransformDSLProgram, check::Bool=false])

Perform a particle filter update, with a custom forward and backward kernel.
New latent choices are sampled from `fwd_proposal`, and any discarded choices
are evaluated under `bwd_proposal`. For each particle:
* `fwd_proposal` is evaluated with arguments `(t_old, fwd_args...)`,
   where `t_old` is the old model trace, and produces its own trace `t_fwd`.
* The old model trace is replaced by a new model trace `t_new`, constructed
  by merging the choices in `t_old` and `t_fwd`, sampling any remaining
  choices from the model's default proposal, and discarding `disc_choices`,
  the choices in `t_old` inconsistent with those in `t_fwd`.
* The probability of `disc_choices` is assessed under `bwd_proposal` with the
  arguments `(t_new, bwd_args)`, giving a backward weight. This weight is used
  as a correction within the importance weight update, ensuring that the
  particle filter remains a valid approximation of the posterior.

The choicemap of `t_new` and `disc_choices` satisfy the following conditions:
1. `observations` is a subset of `get_choices(t_new)`;
2. `get_choices(t_old) ∖ disc_choices` is a subset of `get_choices(t_new)`;
3. `get_choices(t_fwd)` is a subset of `get_choices(t_new)`.
4. `disc_choices` is  within the support of `bwd_proposal`

For valid posterior inference conditioned on prior observations, note that
`fwd_proposal` should not cause any of those observations to be discarded,
(i.e., `disc_choices` should not contain any `observations` given in previous
calls to `pf_update!`).

If a trace `transform` is also provided, then more general forward and backward
kernels can be used: `t_new` (the model's new trace) and `t_bwd` (the trace
for the backward kernel) are constructed as a function of `t_old`, `t_fwd`,
and any new `observations`. The `check` argument can also be set to `true`
to check for correctness.

Similar functionality is provided by [`move_reweight`](@ref), except that
`pf_update!` also allows model arguments to be updated.
"""
function pf_update!(state::ParticleFilterView, new_args::Tuple,
                    argdiffs::Tuple, observations::ChoiceMap,
                    fwd_proposal::GenerativeFunction, fwd_args::Tuple,
                    bwd_proposal::GenerativeFunction, bwd_args::Tuple)
    n_particles = length(state.traces)
    for i=1:n_particles
        (fwd_choices, fwd_weight, _) =
            propose(fwd_proposal, (state.traces[i], fwd_args...))
        constraints = merge(observations, fwd_choices)
        state.new_traces[i], up_weight, _, discard =
            update(state.traces[i], new_args, argdiffs, constraints)
        bwd_weight, _ = isempty(discard) ? (0.0, nothing) :
            assess(bwd_proposal, (state.new_traces[i], bwd_args...), discard)
        state.log_weights[i] += up_weight - fwd_weight + bwd_weight
    end
    update_refs!(state)
    return state
end

function pf_update!(state::ParticleFilterView, new_args::Tuple,
                    argdiffs::Tuple, observations::ChoiceMap,
                    fwd_proposal::GenerativeFunction, fwd_args::Tuple,
                    bwd_proposal::GenerativeFunction, bwd_args::Tuple,
                    transform::TraceTransformDSLProgram, check::Bool=false)
    translator = GenParticleFilters.UpdatingTraceTranslator(
        p_new_args=new_args, p_argdiffs=argdiffs, new_observations=observations,
        q_forward=fwd_proposal, q_forward_args=fwd_args,
        q_backward=bwd_proposal, q_backward_args=bwd_args, transform=transform)
    return pf_update!(state, translator; check=check)
end


text:
export MultiSpanProcessor,
    SimpleSpanProcessor,
    InMemorySpanExporter,
    ConsoleSpanExporter

abstract type AbstractSpanProcessor end

struct MultiSpanProcessor <: AbstractSpanProcessor
    span_processors::Vector
    is_spawn::Bool
    MultiSpanProcessor(ps...;is_spawn=true) = new([ps...], is_spawn)
end

for f in (:on_start, :on_end, :shutdown!)
    @eval function $f(sp::MultiSpanProcessor, args...)
        @sync for p in sp.span_processors
            if sp.is_spawn
                Threads.@spawn $f(p, args...)
            else
                @async $f(p, args...)
            end
        end
    end
end

function force_flush!(sp::MultiSpanProcessor, timeout_millis=30_000)
    res = fill(false, length(sp.span_processors))
    @sync for (i,p) in enumerate(sp.span_processors)
        if sp.is_spawn
            Threads.@spawn res[i] = force_flush!(p, timeout_millis)
        else
            @async res[i] = force_flush!(p, timeout_millis)
        end
    end
    res
end

abstract type AbstractSpanExporter end

@enum SpanExportResult begin
    SUCCESS
    FAILURE
end

Base.@kwdef struct InMemorySpanExporter <: AbstractSpanExporter
    finished_spans::Vector = []
    is_shut_down::Ref{Bool} = Ref(false)
end

Base.empty!(se::InMemorySpanExporter) = empty!(se.finished_spans)

shut_down!(se::InMemorySpanExporter) = se.is_shut_down[] = true

function export!(se::InMemorySpanExporter, sps::Vector)
    if se.is_shut_down[]
        FAILURE
    else
        append!(se.finished_spans, sps)
        SUCCESS
    end
end

Base.@kwdef struct ConsoleSpanExporter <: AbstractSpanExporter
    io::IO = stdout
end

function export!(se::ConsoleSpanExporter, sps)
    for s in sps
        write(se.io, s)
    end
    flush(se.io)
    SUCCESS
end

struct SimpleSpanProcessor <: AbstractSpanProcessor
    span_exporter
end

on_start(ssp::SimpleSpanProcessor, span) = nothing

function on_end(ssp::SimpleSpanProcessor, span::Span)
    if span.span_context.trace_flag.sampled
        export!(ssp.span_exporter, span)
    end
end

shut_down!(ssp::SimpleSpanProcessor) = shut_down!(ssp.span_exporter)

force_flush!(ssp::SimpleSpanProcessor, args...) = true

struct SpanWrapper <: API.AbstractSpan
    span::Span
    span_processor::MultiSpanProcessor
    instrumentation_info::InstrumentationInfo
    resource::Resource
    function SpanWrapper(
        ;span_processor,
        instrumentation_info,
        resource,
        kw...
    )
        sp = Span(;kw...)
        sw = new(sp, span_processor, instrumentation_info, resource)
        on_start(span_processor, sp)
        sw
    end
end

for f in (:record_exception!, :end!, :update_name!, :set_status!, :add_event!, :is_recording, :span_context)
    @eval API.$f(s::SpanWrapper, args...;kw...) = $f(s.span, args...;kw...)
end

Base.setindex!(s::SpanWrapper, args...) = setindex!(s.span, args...)


text:
using FlowOverCircle
using Test

@testset "FlowOverCircle" begin
    include("data.jl")
end