import QuantumOptics.timeevolution
import QuantumOptics.stochastic
using QuantumOptics: Basis, Operator, CompositeBasis, Ket, StateVector
using QuantumOptics.timeevolution: DecayRates

#############################################################################################
# Wrap QuantumOptics solvers in order to implement our form of basis check
#############################################################################################

# function timeevolution.master_dynamic(tspan, rho0::T, f::Function;
#             rates::DecayRates=nothing, fout::Union{Function,Nothing}=nothing, kwargs...
#         ) where {B<:Basis,T<:Operator{B,B}}
#     check_bases(f(0., 0.).basis_l, rho0.basis_l)
#     timeevolution.master_dynamic(tspan, rho0, f; rates=rates, fout=fout, kwargs...)
# end

# function timeevolution.mcwf_dynamic(tspan, psi0::T, f::Function;
#     seed=rand(UInt), rates::DecayRates=nothing,
#     fout=nothing, display_beforeevent=false, display_afterevent=false,
#     kwargs...) where {T<:Ket}
#     check_bases(f(0., 0.).basis_l, psi0.basis)
#     timeevolution.mcwf_dynamic(tspan, psi0, f; seed=seed, rates=rates, fout=fout,
#             display_beforeevent=display_beforeevent, display_afterevent=display_afterevent,
#             kwargs...
#         )
# end

# function stochastic.schroedinger_dynamic(tspan, psi0::T, fdeterm::Function, fstoch::Function;
#             fout::Union{Function,Nothing}=nothing, noise_processes::Int=0,
#             normalize_state::Bool=false, kwargs...
#         ) where T<:Ket
#     check_bases(fdeterm(0.0, 0.0), fstoch)
#     stochastic.schroedinger_dynamic(tspan, psi0, fdeterm, fstoch;
#             fout=fout, noise_processes=noise_processes, normalize_state=normalize_state,
#             kwargs...
#         )
# end

# function stochastic.master_dynamic(tspan, rho0::T, fdeterm::Function, fstoch::Function;
#             rates::DecayRates=nothing, fout::Union{Function,Nothing}=nothing,
#             noise_processes::Int=0, kwargs...
#         ) where {B<:Basis,T<:Operator{B,B}}
#         check_bases()
#         stochastic.master_dynamic(tspan, rho0, fdeterm, fstoch; rates=rates, fout=fout,
#                 noise_processes=noise_processes, kwargs...
#             )
# end

# function timeevolution.schroedinger_dynamic(tspan, psi0::T, f::Function;
#             fout::Union{Function,Nothing}=nothing, kwargs...
#         ) where T<:StateVector
#     check_bases(f(0., 0.).basis_l, psi0.basis)
#     timeevolution.schroedinger_dynamic(tspan, psi0, f; fout=fout, kwargs...)
# end

# when solving for unitary evolution, it is convenient to switch back and forth between an
# initial state that is pure and one that is mixed
function timeevolution.schroedinger_dynamic(tspan, rho0::T, f::Function;
            fout::Union{Function,Nothing}=nothing, kwargs...
        ) where {B<:Basis,T<:Operator{B,B}}
    check_bases(f(0., 0.).basis_l, rho0.basis_l)
    Jvec = []
    timeevolution.master_dynamic(tspan, rho0, (t,rho) -> (f(t, rho), Jvec, Jvec);
        fout=fout, kwargs...)
end

check_bases(b1::CompositeBasis, b2::CompositeBasis) = b1 == b2 || throw(IncompatibleBases())

mutable struct IncompatibleBases <: Exception end


import Base: getindex, setindex!, copy

mutable struct Board
    data::Array{Int}
    score::Int
    level::Int
    lines_to_goal::Int
    tile::Tile
    location::Vector{Int}
    orientation::Int
    nexttiles::Vector{Tile}
    holdtile::Tile
    allowhold::Bool
end
Board(i=0) = Board(ones(Int, 20, 10)*i, 0, 1+i, 5,
            rand(Tiles)(), [4,1], 0,
            [rand(Tiles)() for i in 1:3], rand(Tiles)(), true)

copy(b::Board) = Board(copy(b.data), b.score, b.level, b.lines_to_goal,
                    b.tile, b.location, b.orientation,
                    copy(b.nexttiles), b.holdtile, b.allowhold)

function getindex(b::Board)
    dy,dx = size(rotatedtile(b)) .-1
    x,y = b.location
    return b.data[y:y+dy, x:x+dx]
end

function setindex!(b::Board, s::AbstractArray)
    dy,dx = size(rotatedtile(b)) .-1
    x,y = b.location
    b.data[y:y+dy, x:x+dx] = s
end

rotatedtile(b::Board) = rotr90(data(b.tile), b.orientation)

function nexttile!(b::Board)
    push!(b.nexttiles, rand(Tiles)())
    b.tile = popfirst!(b.nexttiles)
end

function add_tile!(b::Board)
    nexttile!(b)
    b.location = start_location(b.tile)
    b.orientation = 0
    if all(b[] .== 0)
        oldboard = copy(b)
        b[] += data(b.tile)
        update_board!(oldboard, b)
        return true
    end
    false
end

function delete_lines!(b::Board)
    oldboard = copy(b)
    nr_lines = 0
    for i in 1:20
        if all(b.data[i, :] .!= 0)
            b.data[2:i, :] = b.data[1:i-1, :]
            b.data[1,:] .= 0
            nr_lines += 1
        end
    end
    b.lines_to_goal -= nr_lines
    b.score += [0 1 3 5 8][nr_lines+1] * b.level * 100
    if b.lines_to_goal ≤ 0
        b.level += 1
        b.lines_to_goal += b.level*5
    end
    update_board!(oldboard, b)
end


@everywhere using BioSequences
@everywhere using DataStructures: DefaultDict

@everywhere function twin(km)
    DNAKmer(reverse_complement(DNASequence(km)))
end

function contig_to_string(c)
  return "$(c[1])" * join(["$(x[end])" for x in c[2:end]])
end

@everywhere function get_contig(::Type{DNAKmer{k}}, kmers, km) where {k}
  contig_fw = get_contig_forward(DNAKmer{k}, kmers, km)
  contig_bw = get_contig_forward(DNAKmer{k}, kmers, twin(km))

  if km in neighbors(contig_fw[end])
    return contig_fw
  else
    return [[twin(x) for x in reverse(contig_bw)[1:end-1]] ; contig_fw]
  end
end

@everywhere function get_contig_forward(::Type{DNAKmer{k}}, kmers, km) where {k}
    c_fw = DNAKmer[km]
    while true
        # check if forward exists, and only 1:
        fw_neighbors = [kmer for kmer in neighbors(c_fw[end]) if canonical(kmer) in kmers]
        if length(fw_neighbors) != 1
          break
        end
        candidate = fw_neighbors[1]
        if candidate == km || candidate == twin(km) || candidate == twin(c_fw[end])
          break
        end
        bw_neighbors = [kmer for kmer in neighbors(twin(candidate)) if canonical(kmer) in kmers]
        if length(bw_neighbors) != 1
          break
        end
        push!(c_fw, candidate)
    end
    return c_fw
end

@everywhere function all_contigs(::Type{DNAKmer{k}}, kmers) where {k}
  done = Set{DNAKmer}()
  contigs = []
  for kmer in kmers
    if kmer in done
      continue
    end
    contig = get_contig(DNAKmer{k}, kmers, kmer)
      for c_kmer in contig
      push!(done, canonical(c_kmer))
    end
    push!(contigs, contig)
  end
  return contigs
end

### classifies each kmer with his colors (present alleles)
@everywhere function find_all_kmer_colors(::Type{DNAKmer{k}}, fastafile) where {k}
  # record = FASTASeqRecord{DNASequence}()
  errors = 0
  record = FASTA.Record()
  kmer_class = DefaultDict{DNAKmer, Vector{Int16}}(() -> Int16[])
  allele_ids = Int16[]
  allele_idx::Int16 = 1
  n_kmers = Int[]
  seen = Set{DNAKmer}()
  # reader = open(FASTA.Reader{DNASequence}, fastafile)
  # reader = FASTA.Reader{DNASequence}, fastafile)
  reader = open(FASTA.Reader, fastafile)
  local seq_id, pos
  while !eof(reader)
    try
      read!(reader, record)
      seq_id = FASTA.identifier(record)
      seen = Set{DNAKmer}()
      for (pos, kmer) in each(DNAKmer{k}, FASTA.sequence(record))
        can_kmer = canonical(kmer)
        if !in(can_kmer, seen)
          push!(kmer_class[can_kmer], allele_idx)
          push!(seen, can_kmer)
        end
      end
      push!(n_kmers, length(seen)) # number of unique kmers for this allele;
      # find the separator; will assume that if I see a "_", that's it, otherwise try "-";
      separator = in('_', seq_id) ? "_" : "-"
      # update idx; the counter idx is incremental (1,2, ...) because we need the array sorted.
      # But this is not always in the allele ordering, so we have to save the original id to restore it later;
      allele_id = parse(Int16,split(FASTA.identifier(record), separator)[end])
      push!(allele_ids, allele_id)
    catch e
      @error("Error parsing file $fastafile, around record $seq_id, most likely some unkown characters are present; this allele will be skipped. \nThis might make the results on this locus unreliable; please fix this FASTA file and rerun.")
      errors += 1
      if errors > 2
        @error("Too many unrecoverable errors on this FASTA file; skipping this locus ... ")
        exit(-1) # TODO: Could I try to skip just this locus?
      end
    end
    allele_idx += 1
  end
  close(reader)
  return kmer_class, allele_ids, n_kmers
end

@everywhere function build_db_graph(::Type{DNAKmer{k}}, fastafile, coverage_p=1) where {k}
  # get kmers and colors (alleles)
  kmer_class, allele_ids, n_kmers = find_all_kmer_colors(DNAKmer{k}, fastafile)
  locus = splitext(basename(fastafile))[1]

  # Solve a coverage ILP to find the kmers:
  selected_kmers = DNAKmer[]
  actual_coverages = []
  # Coverage per allele:
  coverages = [ Int(round(n_k * coverage_p)) for n_k in n_kmers]
  # If coverage < 1, solve the ilp, otherwise just select all keys (all kmers)
  selected_kmers, actual_coverages = coverage_p < 1 ? Base.invokelatest(kmer_coverage_ilp, locus, kmer_class, allele_ids, coverages) : (keys(kmer_class), coverages)
  # filter more kmers with de Bruijn graph contigs
  filtered_kmer_class, kmer_weights = filter_kmers_with_db_graph(DNAKmer{k}, kmer_class, selected_kmers)
  # debug log:
  sel_kmers, contig_kmers = length(selected_kmers), length(filtered_kmer_class)
  # println("$locus\t$coverage_p\tAll kmers\t$sel_kmers")
  # println("$locus\t$coverage_p\tdb Graph\t$contig_kmers")

  return (filtered_kmer_class, allele_ids, kmer_weights, actual_coverages)

end

@everywhere function filter_kmers_with_db_graph(::Type{DNAKmer{k}}, kmer_class, selected_kmers) where {k}
  # select 1 kmer per contig and set weight as length of contig; save kmers in filtered_kmer_class
  contig_list = all_contigs(DNAKmer{k}, selected_kmers)
  kmer_weights = Dict{UInt64, Int}() # Kmer is converted to UInt64, because DNAKmer apparently does not support write(), needed for the parallel pmap() call;
  filtered_kmer_class = Dict{UInt64, Vector{Int16}}()
  for contig in contig_list
      kmer = canonical(contig[1])
      kmer_uint = convert(UInt64,kmer)
      kmer_weights[kmer_uint] = length(contig)
      filtered_kmer_class[kmer_uint] = kmer_class[kmer]
  end
  return filtered_kmer_class, kmer_weights
end


function BodyCenterPointing(arg0::Frame, arg1::Ellipsoid)
    return BodyCenterPointing((Frame, Ellipsoid), arg0, arg1)
end

function get_attitude(obj::GroundPointing, arg0::FieldPVCoordinatesProvider, arg1::FieldAbsoluteDate, arg2::Frame)
    return jcall(obj, "getAttitude", FieldAttitude, (FieldPVCoordinatesProvider, FieldAbsoluteDate, Frame), arg0, arg1, arg2)
end

function get_attitude(obj::GroundPointing, arg0::PVCoordinatesProvider, arg1::AbsoluteDate, arg2::Frame)
    return jcall(obj, "getAttitude", Attitude, (PVCoordinatesProvider, AbsoluteDate, Frame), arg0, arg1, arg2)
end

function get_body_frame(obj::GroundPointing)
    return jcall(obj, "getBodyFrame", Frame, ())
end

function get_target_pv(obj::BodyCenterPointing, arg0::FieldPVCoordinatesProvider, arg1::FieldAbsoluteDate, arg2::Frame)
    return jcall(obj, "getTargetPV", TimeStampedFieldPVCoordinates, (FieldPVCoordinatesProvider, FieldAbsoluteDate, Frame), arg0, arg1, arg2)
end

function get_target_pv(obj::BodyCenterPointing, arg0::PVCoordinatesProvider, arg1::AbsoluteDate, arg2::Frame)
    return jcall(obj, "getTargetPV", TimeStampedPVCoordinates, (PVCoordinatesProvider, AbsoluteDate, Frame), arg0, arg1, arg2)
end


# push!(LOAD_PATH, "/home/ian/Grad/Research/Disks/DiskJockey/")
push!(LOAD_PATH, "/home/ian/Grad/Research/Disks/DiskJockey/movie/")
# push!(LOAD_PATH, "/n/home07/iczekala/DiskJockey/")
# push!(LOAD_PATH, "/pool/scout0/DiskJockey/")

using consonance

using ArgParse

s = ArgParseSettings()
@add_arg_table s begin
    "--run_index", "-r"
    help = "Output run index"
    default = 0
    arg_type = Int
end

parsed_args = parse_args(ARGS, s)
run_index = parsed_args["run_index"]

imgdir = moviedir * @sprintf("g%02d/", run_index)
scratchdir = moviedir * "scratch/"

import YAML
config = YAML.load(open("config.yaml"))

# make the directory where RADMC will make all its files exist
if !ispath(imgdir)
    println("Creating ", imgdir)
    mkdir(imgdir)
end

if !ispath(scratchdir)
    mkdir(scratchdir)
end

cd(scratchdir)
println("Now in ", pwd())

files = ["lines.inp", "molecule_co.inp", "wavelength_micron.inp", "radmc3d.inp"]
for file in files
    src = moviedir * file
    dst = scratchdir * file
    cp(src, dst, remove_destination=true)
end

using DiskJockey.constants
using DiskJockey.image
using DiskJockey.model

species = "12CO"

function make_image(pars, id::Int)

    write_model(pars, "", grid, species)

    run(pipeline(`radmc3d image incl $(pars.incl) posang $(pars.PA) npix $npix loadlambda`, DevNull))

    src = "image.out"
    dst = imgdir * @sprintf("gimage%04d.out", id)

    cp(src, dst, remove_destination=true)

end

grd = config["grid"]

const global npix = config["npix"] # number of pixels
const global grid = Grid(grd["nr"], grd["ntheta"], grd["r_in"], grd["r_out"], true)
write_grid("", grid)
write_lambda(lams, "")

start = run_index * nframes_per_proc + 1
ids = Int[i for i=start:(start + nframes_per_proc)]

for i in ids
    if i <= nframes
        make_image(pars[i], i)
    else
        break
    end
end


# Julia wrapper for header: /usr/include/scip/cons_bivariate.h
# Automatically generated using Clang.jl wrap_c

function SCIPincludeConshdlrBivariate(scip)
    ccall((:SCIPincludeConshdlrBivariate, libscip), SCIP_RETCODE, (Ptr{SCIP_},), scip)
end

function SCIPcreateConsBivariate(scip, cons, name, f, convextype, z, zcoef, lhs, rhs, initial, separate, enforce, check, propagate, _local, modifiable, dynamic, removable, stickingatnode)
    ccall((:SCIPcreateConsBivariate, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{Ptr{SCIP_CONS}}, Cstring, Ptr{SCIP_EXPRTREE}, SCIP_BIVAR_CONVEXITY, Ptr{SCIP_VAR}, Cdouble, Cdouble, Cdouble, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32), scip, cons, name, f, convextype, z, zcoef, lhs, rhs, initial, separate, enforce, check, propagate, _local, modifiable, dynamic, removable, stickingatnode)
end

function SCIPcreateConsBasicBivariate(scip, cons, name, f, convextype, z, zcoef, lhs, rhs)
    ccall((:SCIPcreateConsBasicBivariate, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{Ptr{SCIP_CONS}}, Cstring, Ptr{SCIP_EXPRTREE}, SCIP_BIVAR_CONVEXITY, Ptr{SCIP_VAR}, Cdouble, Cdouble, Cdouble), scip, cons, name, f, convextype, z, zcoef, lhs, rhs)
end

function SCIPgetLinearVarBivariate(scip, cons)
    ccall((:SCIPgetLinearVarBivariate, libscip), Ptr{SCIP_VAR}, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetLinearCoefBivariate(scip, cons)
    ccall((:SCIPgetLinearCoefBivariate, libscip), Cdouble, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetExprtreeBivariate(scip, cons)
    ccall((:SCIPgetExprtreeBivariate, libscip), Ptr{SCIP_EXPRTREE}, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetLhsBivariate(scip, cons)
    ccall((:SCIPgetLhsBivariate, libscip), Cdouble, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end

function SCIPgetRhsBivariate(scip, cons)
    ccall((:SCIPgetRhsBivariate, libscip), Cdouble, (Ptr{SCIP_}, Ptr{SCIP_CONS}), scip, cons)
end


#=
/* Switches for the triangulator.                                            */

-p Triangulates a Planar Straight Line Graph (.poly file).
-r Refines a previously generated mesh.
-q Quality mesh generation with no angles smaller than 20 degrees. An alternate minimum angle may be specified after the `q'.
-a Imposes a maximum triangle area constraint. A fixed area constraint (that applies to every triangle) may be specified after the `a', or varying area constraints may be read from a .poly file or .area file.
-u Imposes a user-defined constraint on triangle size.
-A Assigns a regional attribute to each triangle that identifies what segment-bounded region it belongs to.
-c Encloses the convex hull with segments.
-D Conforming Delaunay: use this switch if you want all triangles in the mesh to be Delaunay, and not just constrained Delaunay; or if you want to ensure that all Voronoi vertices lie within the triangulation.
-j Jettisons vertices that are not part of the final triangulation from the output .node file (including duplicate input vertices and vertices ``eaten'' by holes).
-e Outputs (to an .edge file) a list of edges of the triangulation.
-v Outputs the Voronoi diagram associated with the triangulation. Does not attempt to detect degeneracies, so some Voronoi vertices may be duplicated.
-n Outputs (to a .neigh file) a list of triangles neighboring each triangle.
-g Outputs the mesh to an Object File Format (.off) file, suitable for viewing with the Geometry Center's Geomview package.
-B Suppresses boundary markers in the output .node, .poly, and .edge output files.
-P Suppresses the output .poly file. Saves disk space, but you lose the ability to maintain constraining segments on later refinements of the mesh.
-N Suppresses the output .node file.
-E Suppresses the output .ele file.
-I Suppresses mesh iteration numbers.
-O Suppresses holes: ignores the holes in the .poly file.
-X Suppresses exact arithmetic.
-z Numbers all items starting from zero (rather than one). Note that this switch is normally overrided by the value used to number the first vertex of the input .node or .poly file. However, this switch is useful when calling Triangle from another program.
-o2 Generates second-order subparametric elements with six nodes each.
-Y Prohibits the insertion of Steiner points on the mesh boundary. If specified twice (-YY), it prohibits the insertion of Steiner points on any segment, including internal segments.
-S Specifies the maximum number of added Steiner points.
-i Uses the incremental algorithm for Delaunay triangulation, rather than the divide-and-conquer algorithm.
-F Uses Steven Fortune's sweepline algorithm for Delaunay triangulation, rather than the divide-and-conquer algorithm.
-l Uses only vertical cuts in the divide-and-conquer algorithm. By default, Triangle uses alternating vertical and horizontal cuts, which usually improve the speed except with vertex sets that are small or short and wide. This switch is primarily of theoretical interest.
-s Specifies that segments should be forced into the triangulation by recursively splitting them at their midpoints, rather than by generating a constrained Delaunay triangulation. Segment splitting is true to Ruppert's original algorithm, but can create needlessly small triangles. This switch is primarily of theoretical interest.
-C Check the consistency of the final mesh. Uses exact arithmetic for checking, even if the -X switch is used. Useful if you suspect Triangle is buggy.
-Q Quiet: Suppresses all explanation of what Triangle is doing, unless an error occurs.
-V Verbose: Gives detailed information about what Triangle is doing. Add more `V's for increasing amount of detail. `-V' gives information on algorithmic progress and detailed statistics.
-h Help: Displays complete instructions.

/*  If the size of the object file is important to you, you may wish to      */
/*  generate a reduced version of triangle.o.  The REDUCED symbol gets rid   */
/*  of all features that are primarily of research interest.  Specifically,  */
/*  the -DREDUCED switch eliminates Triangle's -i, -F, -s, and -C switches.  */
/*  The CDT_ONLY symbol gets rid of all meshing algorithms above and beyond  */
/*  constrained Delaunay triangulation.  Specifically, the -DCDT_ONLY switch */
/*  eliminates Triangle's -r, -q, -a, -u, -D, -Y, -S, and -s switches.       */
=#

mutable struct TriangulateOptions
    pslg::Bool #p
    regionattrib::Bool # A
    convex::Bool # c
    jettison::Bool # j
    firstnumberiszero::Bool # z
    edgesout::Bool # e
    voronoi::Bool # v
    neighbors::Bool # n
    nobound::Bool # B
    nopolywritten::Bool # P
    nonodewritten::Bool # N
    noelewritten::Bool # E
    noiterationnum::Bool # I
    noholes::Bool # O
    noexactaritmetic::Bool # X
    order::Bool # o
    orderHow::Int64 # 1...2...3
    dwyer::Bool # l
    quiet::Bool # Q
    verbose::Bool # V
    TriangulateOptions() = new(false,
    false, false, false, false, false, false, false,
    # No_xyz_ selector(s)
    true, true, false, false, true, false, false,
    # order
    false, 0,
    false,
    # Quiet Verbose
    true, false)
end

function getTriangulateStringOptions(self::TriangulateOptions)
    output_stri = ""

    if self.pslg
        output_stri = output_stri * "p"
    end

    if self.regionattrib
        output_stri = output_stri * "A"
    end

    if self.convex
        output_stri = output_stri * "c"
    end

    if self.jettison
        output_stri = output_stri * "j"
    end

    if self.firstnumberiszero
        output_stri = output_stri * "z"
    end

    if self.edgesout
        output_stri = output_stri * "e"
    end

    if self.voronoi
        output_stri = output_stri * "v"
    end

    if self.neighbors
        output_stri = output_stri * "n"
    end

    if self.nobound
        output_stri = output_stri * "B"
    end

    if self.nopolywritten
        output_stri = output_stri * "P"
    end

    if self.nonodewritten
        output_stri = output_stri * "N"
    end

    if self.noelewritten
        output_stri = output_stri * "E"
    end

    if self.noiterationnum
        output_stri = output_stri * "I"
    end

    if self.noholes
        output_stri = output_stri * "O"
    end

    if self.noexactaritmetic
        output_stri = output_stri * "X"
    end

    if self.order && self.orderHow > 0
        output_stri = output_stri * "o" * string(self.orderHow)
    end

    if self.dwyer
        output_stri = output_stri * "l"
    end

    if self.quiet
        output_stri = output_stri * "Q"
    end

    if self.verbose
        output_stri = output_stri * "V"
    end

    return output_stri
end


const DUMMY_PRECISION = 1e-12

Base.isfinite(a::AbstractVec) = all(isfinite, a)
Base.isinf(a::AbstractVec) = any(isinf, a)
Base.isnan(a::AbstractVec) = any(isnan, a)

function StaticArrays.similar_type(::Type{V}, ::Type{F}, size::Size{N}) where {V<:AbstractVec, F<:AbstractFloat, N}
    if size == Size(V) && eltype(V) == F
        return V
    else # convert back to SArray
        return SArray{Tuple{N...},F,length(size),prod(size)}
    end
end

Base.abs(a::AbstractVec) = error("abs(v::AbstractVec) has been removed. Use norm(v) to get the norm; abs.(v) to get the element-wise absolute value.")


using ClimateMachine.Writers

function ivdc_default_diagnostics_init()
end
function ivdc_default_diagnostics_fini()
end
function ivdc_default_diagnostics_collect()
    mpicomm = Diagnostics.Settings.mpicomm
    dg1 = Diagnostics.Settings.dg.modeldata.diagsconf.myivdcdg
    dg2 = Diagnostics.Settings.dg.modeldata.diagsconf.myivdcdg2
    Q = Diagnostics.Settings.dg.modeldata.diagsconf.Q
end

function setup_single_column_default_diagnostics(::IVDCConfigType, interval, out_pref; writer=NetCDFWriter(),interpol=nothing)

	return DiagnosticsGroup(
				"IVDCdefault",
				ivdc_default_diagnostics_init,
				ivdc_default_diagnostics_fini,
				ivdc_default_diagnostics_collect,
				interval,
				out_pref,
				writer,
				interpol,
				)

end


module CCBladeOpenMDAOExample

include("derivatives.jl")
include("openmdao.jl")

end # module


"""
    Maybe{X}

Type shortcut for "`X` or nothing" or "nullable `X`" in javaspeak. Name
got inspired by our functional friends.
"""
const Maybe{X} = Union{Nothing,X}

"""
Part of a measurement unit definition that corresponds to the SBML definition
of `Unit`. For example, the unit "per square megahour", Mh^(-2), is written as:

    UnitPart("second",  # base SI unit, this says we are measuring time
             -2,        # exponent, says "per square"
             6,         # log-10 scale of the unit, says "mega"
             1/3600)    # second-to-hour multiplier

Compound units (such as "volt-amperes" and "dozens of yards per ounce") are
built from multiple `UnitPart`s; see the definition of field `units` in
[`Model`](@ref).
"""
struct UnitPart
    kind::String
    exponent::Int
    scale::Int
    multiplier::Float64
    UnitPart(k, e, s, m) = new(k, e, s, m)
end

"""
Abstract type for all kinds of gene product associations
"""
abstract type GeneProductAssociation end

"""
Gene product reference in the association expression
"""
struct GPARef <: GeneProductAssociation
    gene_product::String
end

"""
Boolean binary "and" in the association expression
"""
struct GPAAnd <: GeneProductAssociation
    terms::Vector{GeneProductAssociation}
end

"""
Boolean binary "or" in the association expression
"""
struct GPAOr <: GeneProductAssociation
    terms::Vector{GeneProductAssociation}
end

"""
Reaction with stoichiometry that assigns reactants and products their relative
consumption/production rates (accessible in field `stoichiometry`), lower/upper
bounds (in tuples `lb` and `ub`, with unit names), and objective coefficient
(`oc`). Also may contains `notes` and `annotation`.
"""
struct Reaction
    stoichiometry::Dict{String,Float64}
    lb::Tuple{Float64,String}
    ub::Tuple{Float64,String}
    oc::Float64
    gene_product_association::Maybe{GeneProductAssociation}
    notes::Maybe{String}
    annotation::Maybe{String}
    Reaction(s, l, u, o, as, n = nothing, an = nothing) = new(s, l, u, o, as, n, an)
end

"""
Species metadata -- contains a human-readable `name`, a `compartment`
identifier, `formula`, `charge`, and additional `notes` and `annotation`.
"""
struct Species
    name::String
    compartment::String
    formula::Maybe{String}
    charge::Maybe{Int}
    notes::Maybe{String}
    annotation::Maybe{String}
    Species(na, co, f, ch, no = nothing, a = nothing) = new(na, co, f, ch, no, a)
end

"""
Gene product metadata.
"""
struct GeneProduct
    name::Maybe{String}
    label::Maybe{String}
    notes::Maybe{String}
    annotation::Maybe{String}
    GeneProduct(na, l, no = nothing, a = nothing) = new(na, l, no, a)
end

"""
Structure that collects the model-related data. Contains `parameters`, `units`,
`compartments`, `species` and `reactions` and `gene_products`, and additional
`notes` and `annotation` (also present internally in some of the data fields).
The contained dictionaries are indexed by identifiers of the corresponding
objects.
"""
struct Model
    parameters::Dict{String,Float64}
    units::Dict{String,Vector{UnitPart}}
    compartments::Vector{String}
    species::Dict{String,Species}
    reactions::Dict{String,Reaction}
    gene_products::Dict{String,GeneProduct}
    notes::Maybe{String}
    annotation::Maybe{String}
    Model(p, u, c, s, r, g, n = nothing, a = nothing) = new(p, u, c, s, r, g, n, a)
end


module synthtest
using JSON

include("../src/Ebic.jl")
include("metrics.jl")

using .Ebic: run_ebic
using .metrics: eval_metrics

const OUTPUT_DIR = "results/Ebic.jl"

function test_dataset(dataset_path)
    out_path = joinpath(OUTPUT_DIR, splitpath(dataset_path)[end])
    println(out_path)
    isdir(out_path) || mkpath(out_path)

    for (root, _, files) in walkdir(dataset_path)
        isempty(files) && continue

        println("""
        ####################################
        Starting test case: '$(basename(root))'
        ####################################""")

        input_paths = Vector()
        biclusters_paths = Vector()
        for file in files
            if occursin("hidden", file)
                push!(biclusters_paths, joinpath(root, file))
            else
                push!(input_paths, joinpath(root, file))
            end
        end

        if length(input_paths) != length(biclusters_paths)
            throw(ErrorException("Something is wrong in '$root'."))
        end

        test_case_results = Vector()
        for (input_path, ground_truth_path) in zip(input_paths, biclusters_paths)
            println("##################################")
            println("Testing: '$(basename(input_path))'")
            println("Ground truth: '$(basename(ground_truth_path))'")

            ground_truth::Vector = JSON.parsefile(ground_truth_path)

            result = run_ebic(
                input_path,
                verbose = true,
                max_iterations = 20_000,
                max_biclusters = length(ground_truth),
                overlap_threshold = 0.75,
                negative_trends = true,
                approx_trends_ratio = 0.85f0,
                best_bclrs_stats = false,
            )

            result["input_data"] = input_path
            result["ground_truth"] = ground_truth_path

            relevance, recovery, ce =
                eval_metrics(result["biclusters"], input_path, ground_truth)

            result["relevance"] = relevance
            result["recovery"] = recovery
            result["ce"] = ce

            println("Prelic relevance   : $(relevance)")
            println("Prelic recovery    : $(recovery)")
            println("Clustering error   : $(ce)")
            println("Last iter tabu hits: $(result["last_iter_tabu_hits"])")

            push!(test_case_results, result)
        end

        open(joinpath(out_path, "$(basename(root))_res.json"), "w") do f
            JSON.print(f, test_case_results)
        end
    end
end

test_unibic() = test_dataset("data/unibic/")
test_recbic_sup() = test_dataset("data/recbic_sup/")
test_recbic_main() = test_dataset("data/recbic_maintext/")

function main()
    test_unibic()
    test_recbic_main()
    test_recbic_sup()
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end
end


export
    KnitroProblem,
    createProblem, freeProblem,
    initializeProblem,
    solveProblem,
    restartProblem,
    setFuncCallback, setGradCallback, setHessCallback,
    setCallbacks, setMIPCallback,
    loadOptionsFile,
    loadTunerFile,
    setOption, getOption,
    applicationReturnStatus

mutable struct KnitroProblem
    # For KNITRO
    env::Ptr{Nothing} # pointer to KTR_context
    eval_status::Int32 # scalar input used only for reverse comms
    status::Int32  # Final status
    mip::Bool # whether it is a Mixed Integer Problem

    # For MathProgBase
    x::Vector{Float64}  # Starting and final solution
    lambda::Vector{Float64}
    g::Vector{Float64}  # Final constraint values
    obj_val::Vector{Float64}  # (length 1) Final objective

    # Callbacks
    eval_f::Function
    eval_g::Function
    eval_grad_f::Function
    eval_jac_g::Function
    eval_h::Function
    eval_hv::Function
    eval_mip_node::Function

    function KnitroProblem()
        kp = new(newcontext(),
                    0,
                    100, # Code for :Uninitialized
                    false)
        finalizer(freeProblem, kp)
        kp
    end
end

createProblem() = KnitroProblem()

function freeProblem(kp::KnitroProblem)
    kp.env == C_NULL && return
    return_code = @ktr_ccall(free, Int32, (Ptr{Nothing},), [kp.env])
    if return_code != 0
        error("KNITRO: Error freeing memory")
    end
    kp.env = C_NULL
end

function initializeKP(kp, x0, lambda0, g; mip = false)
    kp.status = 101 # code for :Initialized
    kp.mip = mip
    kp.x = x0
    kp.lambda = lambda0
    kp.g = g
    kp.obj_val = zeros(Float64, 1)
end

function initializeProblem(kp, objGoal, objType, x_l, x_u, c_Type, g_lb,
                            g_ub, jac_var, jac_con, hess_row, hess_col;
                            initial_x = C_NULL, initial_lambda = C_NULL)
    initializeKP(kp, (initial_x != C_NULL) ? initial_x :
                        zeros(Float64, length(x_l)),
                        (initial_lambda != C_NULL) ? initial_lambda :
                        zeros(Float64, length(x_l) + length(g_lb)),
                        zeros(Float64, length(g_lb)))
    init_problem(kp, objGoal, objType, x_l, x_u, c_Type, g_lb, g_ub,
                    jac_var, jac_con, hess_row, hess_col; initial_x=kp.x,
                    initial_lambda=kp.lambda)
end

function initializeProblem(kp, objGoal, objType, x_l, x_u, c_Type, g_lb,
                            g_ub, jac_var, jac_con; initial_x = C_NULL,
                            initial_lambda = C_NULL)
    hessopt = zeros(Int32 , 1)
    getOption(kp, "hessopt", hessopt)
    @assert hessopt[1] != KTR_HESSOPT_EXACT
    # KNITRO documentation:
    # If user option hessopt is not set to KTR_HESSOPT_EXACT, then Hessian
    # nonzeros will not be used. In this case, set nnzH=0, and pass NULL
    # pointers for hessIndexRows and hessIndexCols.
    initializeKP(kp, (initial_x != C_NULL) ? initial_x :
                        zeros(Float64, length(x_l)),
                        (initial_lambda != C_NULL) ? initial_lambda :
                        zeros(Float64, length(x_l) + length(g_lb)),
                        zeros(Float64, length(g_lb)))
    init_problem(kp, objGoal, objType, x_l, x_u, c_Type, g_lb, g_ub,
                    jac_var, jac_con; initial_x=kp.x,initial_lambda=kp.lambda)
end

# Initialization for MIP
function initializeProblem(kp, objGoal, objType, objFnType,
                            x_Type, x_l, x_u, c_Type, c_FnType, g_lb,
                            g_ub, jac_var, jac_con, hess_row, hess_col;
                            initial_x = C_NULL, initial_lambda = C_NULL)
    initializeKP(kp, (initial_x != C_NULL) ? initial_x :
                        zeros(Float64, length(x_l)),
                        (initial_lambda != C_NULL) ? initial_lambda :
                        zeros(Float64, length(x_l) + length(g_lb)),
                        zeros(Float64, length(g_lb)),
                        mip=true)
    mip_init_problem(kp, objGoal, objType, objFnType, x_Type, x_l, x_u,
                        c_Type, c_FnType, g_lb, g_ub, jac_var, jac_con,
                        hess_row, hess_col; initial_x=kp.x,
                        initial_lambda=kp.lambda)
end

function initializeProblem(kp, objGoal, objType, objFnType,
                            x_Type, x_l, x_u, c_Type, c_FnType, g_lb,
                            g_ub, jac_var, jac_con; initial_x = C_NULL,
                            initial_lambda = C_NULL)
    hessopt = Array(Int32, 1)
    getOption(kp, "hessopt", hessopt)
    @assert hessopt[1] != KTR_HESSOPT_EXACT
    # KNITRO documentation:
    # If user option hessopt is not set to KTR_HESSOPT_EXACT, then Hessian
    # nonzeros will not be used. In this case, set nnzH=0, and pass NULL
    # pointers for hessIndexRows and hessIndexCols.
    initializeKP(kp, (initial_x != C_NULL) ? initial_x :
                        zeros(Float64, length(x_l)),
                        (initial_lambda != C_NULL) ? initial_lambda :
                        zeros(Float64, length(x_l) + length(g_lb)),
                        zeros(Float64, length(g_lb)),
                        mip=true)
    mip_init_problem(kp, objGoal, objType, objFnType, x_Type, x_l, x_u,
                        c_Type, c_FnType, g_lb, g_ub, jac_var, jac_con;
                        initial_x=kp.x, initial_lambda=kp.lambda)
end

function solveProblem(kp::KnitroProblem)
    # Ensure that Knitro will not segfault when using Julia callbacks.
    setOption(kp, "par_numthreads", 1)
    setOption(kp, "par_msnumthreads", 1)
    if kp.mip
        kp.status = mip_solve_problem(kp, kp.x, kp.lambda, kp.eval_status,
                                        kp.obj_val)
    else
        kp.status = solve_problem(kp, kp.x, kp.lambda, kp.eval_status,
                                    kp.obj_val)
    end
end

function solveProblem(kp::KnitroProblem,
                        cons::Vector{Float64},
                        objGrad::Vector{Float64},
                        jac::Vector{Float64},
                        hess::Vector{Float64},
                        hessVector::Vector{Float64})
    # Ensure that Knitro will not segfault when using Julia callbacks.
    setOption(kp, "par_numthreads", 1)
    setOption(kp, "par_msnumthreads", 1)
    if kp.mip
        kp.status = mip_solve_problem(kp, kp.x, kp.lambda, kp.eval_status,
                                        kp.obj_val, cons, objGrad, jac, hess,
                                        hessVector)
    else
        kp.status = solve_problem(kp, kp.x, kp.lambda, kp.eval_status,
                                    kp.obj_val, cons, objGrad, jac, hess,
                                    hessVector)
    end
end

function restartProblem(kp, x0, lambda0)
    kp.status = 101 # code for :Initialized
    kp.eval_status = 0
    restart_problem(kp, x0, lambda0)
end

function restartProblem(kp)
    kp.status = 101 # code for :Initialized
    kp.eval_status = 0
    restart_problem(kp)
end

# -----
# Callback Wrappers
# -----
function eval_fc_wrapper(evalRequestCode::Cint,
                            n::Cint,
                            m::Cint,
                            nnzJ::Cint,
                            nnzH::Cint,
                            x_::Ptr{Cdouble},
                            lambda_::Ptr{Cdouble},
                            obj_::Ptr{Cdouble},
                            c_::Ptr{Cdouble},
                            g_::Ptr{Cdouble},
                            J_::Ptr{Cdouble},
                            H_::Ptr{Cdouble},
                            HV_::Ptr{Cdouble},
                            userParams_::Ptr{Nothing})
    if evalRequestCode != KTR_RC_EVALFC
        return KTR_RC_CALLBACK_ERR
    end
    kp = unsafe_pointer_to_objref(userParams_)::KnitroProblem
    x = unsafe_wrap(Array,x_,n)

    # calculate the new objective function value
    unsafe_store!(obj_, kp.eval_f(x))
    # calculate the new constraint values
    kp.eval_g(x,unsafe_wrap(Array,c_,m))

    Int32(0)
end

function eval_ga_wrapper(evalRequestCode::Cint,
                            n::Cint,
                            m::Cint,
                            nnzJ::Cint,
                            nnzH::Cint,
                            x_::Ptr{Cdouble},
                            lambda_::Ptr{Cdouble},
                            obj_::Ptr{Cdouble},
                            c_::Ptr{Cdouble},
                            g_::Ptr{Cdouble},
                            J_::Ptr{Cdouble},
                            H_::Ptr{Cdouble},
                            HV_::Ptr{Cdouble},
                            userParams_::Ptr{Nothing})
    if evalRequestCode != KTR_RC_EVALGA
        return KTR_RC_CALLBACK_ERR
    end
    kp = unsafe_pointer_to_objref(userParams_)::KnitroProblem
    x = unsafe_wrap(Array,x_,n)

    # evaluate the gradient
    kp.eval_grad_f(x,unsafe_wrap(Array,g_,n))
    # evaluate the jacobian
    if m > 0
        kp.eval_jac_g(x,unsafe_wrap(Array,J_,nnzJ))
    end

    Int32(0)
end

function eval_hess_wrapper(evalRequestCode::Cint,
                            n::Cint,
                            m::Cint,
                            nnzJ::Cint,
                            nnzH::Cint,
                            x_::Ptr{Cdouble},
                            lambda_::Ptr{Cdouble},
                            obj_::Ptr{Cdouble},
                            c_::Ptr{Cdouble},
                            g_::Ptr{Cdouble},
                            J_::Ptr{Cdouble},
                            H_::Ptr{Cdouble},
                            HV_::Ptr{Cdouble},
                            userParams_::Ptr{Nothing})
    kp = unsafe_pointer_to_objref(userParams_)::KnitroProblem
    x = unsafe_wrap(Array,x_, n)
    lambda = unsafe_wrap(Array,lambda_, m+n)

    if evalRequestCode == KTR_RC_EVALH
        kp.eval_h(x, lambda, 1.0, unsafe_wrap(Array,H_, nnzH))
    elseif evalRequestCode == KTR_RC_EVALH_NO_F
        kp.eval_h(x, lambda, 0.0, unsafe_wrap(Array,H_, nnzH))
    elseif evalRequestCode == KTR_RC_EVALHV
        kp.eval_hv(x, lambda, 1.0, unsafe_wrap(Array,HV_, n))
    elseif evalRequestCode == KTR_RC_EVALHV_NO_F
        kp.eval_hv(x, lambda, 0.0, unsafe_wrap(Array,HV_, n))
    else
        return KTR_RC_CALLBACK_ERR
    end
    Int32(0)
end

function eval_mip_node_wrapper(evalRequestCode::Cint,
                            n::Cint,
                            m::Cint,
                            nnzJ::Cint,
                            nnzH::Cint,
                            x_::Ptr{Cdouble},
                            lambda_::Ptr{Cdouble},
                            obj_::Ptr{Cdouble},
                            c_::Ptr{Cdouble},
                            g_::Ptr{Cdouble},
                            J_::Ptr{Cdouble},
                            H_::Ptr{Cdouble},
                            HV_::Ptr{Cdouble},
                            userParams_::Ptr{Nothing})
    kp = unsafe_pointer_to_objref(userParams_)::KnitroProblem
    obj = unsafe_load(obj_)
    kp.eval_mip_node(kp,obj)
    Int32(0)
end

function setFuncCallback(kp::KnitroProblem,
                            eval_f::Function,
                            eval_g::Function)
    kp.eval_f = eval_f
    kp.eval_g = eval_g
    set_func_callback(kp,eval_fc_wrapper)
end

function setGradCallback(kp::KnitroProblem,
                            eval_grad_f::Function,
                            eval_jac_g::Function)
    kp.eval_grad_f = eval_grad_f
    kp.eval_jac_g = eval_jac_g
    set_grad_callback(kp,eval_ga_wrapper)
end

function setHessCallback(kp::KnitroProblem,
                            eval_h::Function,
                            eval_hv::Function)
    kp.eval_h = eval_h
    kp.eval_hv = eval_hv
    set_hess_callback(kp,eval_hess_wrapper)
end

function setCallbacks(kp::KnitroProblem,
                        eval_f::Function,
                        eval_g::Function,
                        eval_grad_f::Function,
                        eval_jac_g::Function,
                        eval_h::Function,
                        eval_hv::Function)
    setFuncCallback(kp, eval_f, eval_g)
    setGradCallback(kp, eval_grad_f, eval_jac_g)
    setHessCallback(kp, eval_h, eval_hv)
end

function setMIPCallback(kp::KnitroProblem, eval_mip_node::Function)
    kp.eval_mip_node = eval_mip_node
    set_mip_node_callback(kp,eval_mip_node_wrapper)
end

# Getters and Setters for Parameters/Options
loadOptionsFile(kp, filename) = load_param_file(kp, filename)
loadTunerFile(kp, filename) = load_tuner_file(kp, filename)
setOption(args...) = set_param(args...)
getOption(args...) = get_param(args...)

function applicationReturnStatus(kp::KnitroProblem)
    if kp.status == 100
        # chosen not to clash with any of the KTR_RC_* codes
        return :Uninitialized
    elseif kp.status == 101
        # chosen not to clash with any of the KTR_RC_* codes
        return :Initialized
    elseif kp.status == 0
        return :Optimal
    elseif 1 <= kp.status <= 11
        return :ReverseComms
    elseif -199 <= kp.status <= -100
        return :FeasibleApproximate
    elseif -299 <= kp.status <= -200
        return :Infeasible
    elseif kp.status == -300
        return :Unbounded
    elseif -499 <= kp.status <= -400
        return :UserLimit
    elseif -599 <= kp.status <= -500
        return :KnitroError
    else
        return :Undefined
    end
end


__precompile__()

module SCIP

if isfile(joinpath(dirname(@__FILE__),"..","deps","deps.jl"))
    include("../deps/deps.jl")
else
    error("SCIP.jl not properly installed. Please run Pkg.build(\"SCIP\")")
end

importall MathProgBase.SolverInterface

include("types.jl")
include("wrapper.jl")
include("mpb_interface.jl")
include("params.jl")

function CSIPversion()
    "$(_majorVersion()).$(_minorVersion()).$(_patchVersion())"
end

function __init__()
    version = VersionNumber(CSIPversion())
    if !(version.major == 0 && version.minor == 3 && version.patch == 4)
        depsdir = realpath(joinpath(dirname(@__FILE__),"..","deps"))
        error("Current CSIP version installed is $(version), but we require version 0.3.4. On Linux, delete the contents of the `$depsdir` directory except for `build.jl`, then rerun Pkg.build(\"SCIP\").")
    end
end

end


function InitCoords(configuration::Configuration,balancingMPI::BalancingMPI)
    n_total=sum(configuration.moles_length)
    coords=zeros(n_total,3)
    index=1
    for comm_mole_i in balancingMPI.mole_index_comms
        for mole_i in comm_mole_i
            molecule=configuration.moles[mole_i]
            mole_length=size(molecule.coords)[1]
            coords[index:mole_length+index-1,1:3]=molecule.coords
            index=index+mole_length
        end
    end
    @info "Initialize Coordinates Size $(n_total)"
    return coords
end

function InitMasses(configuration::Configuration,balancingMPI::BalancingMPI)
    n_total=sum(configuration.moles_length)
    masses=zeros(n_total)
    index=1
    for comm_mole_i in balancingMPI.mole_index_comms
        for mole_i in comm_mole_i
            molecule=configuration.moles[mole_i]
            mole_length=size(molecule.coords)[1]
            masses[index:mole_length+index-1]=molecule.mass
            index=index+mole_length
        end
    end
    @info "Initialize Masses Size $(n_total)"
    return masses
end

function InitTypes(configuration::Configuration,balancingMPI::BalancingMPI)
    n_total=sum(configuration.moles_length)
    types=zeros(Int64,n_total)
    index=1
    for comm_mole_i in balancingMPI.mole_index_comms
        for mole_i in comm_mole_i
            molecule=configuration.moles[mole_i]
            mole_length=size(molecule.coords)[1]
            types[index:mole_length+index-1]=molecule.type
            index=index+mole_length
        end
    end
    @info "Initialize Types Size $(n_total)"
    return types
end

function InitBonds(configuration::Configuration,balancingMPI::BalancingMPI)
    n_total=0
    for molei in configuration.moles
        n_total+=size(molei.bonds)[1]
    end
    bonds=zeros(Int64,n_total,3)
    index=1
    for comm_mole_i in balancingMPI.mole_index_comms
        total_atom_index=0
        for mole_i in comm_mole_i
            molecule=configuration.moles[mole_i]
            bond_length=size(molecule.bonds)[1]
            mole_length=size(molecule.coords)[1]
            if bond_length!=0
                bonds[index:bond_length+index-1,1]=molecule.bonds[:,1]
                bonds[index:bond_length+index-1,2:3]=molecule.bonds[:,2:3].+total_atom_index
            end
            total_atom_index+=mole_length
            index=index+bond_length
        end
    end
    @info "Initialize Bonds Size $(n_total)"
    return bonds
end

function InitVels(configuration::Configuration,balancingMPI::BalancingMPI,velocity::ZeroVelocity)
    n_total=sum(configuration.moles_length)
    vels=zeros(n_total,3)
    @info "Initialize Velocities Size $(n_total)"
    return vels
end

function InitVels(configuration::Configuration,balancingMPI::BalancingMPI,velocity::NoVelocity)
    n_total=sum(configuration.moles_length)
    vels=zeros(n_total,3)
    index=1
    for comm_mole_i in balancingMPI.mole_index_comms
        for mole_i in comm_mole_i
            molecule=configuration.moles[mole_i]
            mole_length=size(molecule.vels)[1]
            vels[index:mole_length+index-1,1:3]=molecule.vels
            index=index+mole_length
        end
    end
    @info "Initialize Velocities Size $(n_total)"
    return vels
end

function InitForces(configuration::Configuration)
    n_total=sum(configuration.moles_length)
    forces=zeros(n_total,3)
    @info "Initialize Forces Size $(n_total)"
    return forces
end

function InitEnergies(configuration::Configuration)
    n_total=sum(configuration.moles_length)
    energies=zeros(n_total)
    @info "Initialize Energy Size $(n_total)"
    return energies
end


module CropPlanner

import HTTP
import JSON
import DataFrames
using DataFramesMeta
using Pipe
import CSV
using Dates
using Unitful, Unitful.DefaultSymbols
using UnitfulUS

include("addearliestharvest!.jl")
include("addseedingdateoftransplant!.jl")
include("convertstringarray.jl")
include("converttodate.jl")
include("converttodays.jl")
include("getbasketprice.jl")
include("listavailablecropondate.jl")
include("tosiunits.jl")

export listavailablecropondate

vegdic = JSON.parse(String(read("data/vegetable_name_dictionary.json")))

# Create table of range of harvest_dates
crop_calendar = CSV.read("data/crop plan - ecoumene_crop_calendar.csv", DataFrames.DataFrame)
crop_calendar[!,:vegetable] = [get(vegdic, v, v) for v in  crop_calendar[!,:vegetable]]

rename!(crop_calendar, [ i => crop_calendar[1,i] for i in 3:ncol(crop_calendar)])
crop_calendar1 = crop_calendar[3:end,:]

possible_harvest_weeks = @pipe crop_calendar1 |>
    @where(_, :activity .== "harvest") |>
    stack(_, 3:ncol(_)) |>
    dropmissing |>
    @transform(_, week = [i.value for i in Week.(Date.(:variable, DateFormat("m/d/y")))]) |>
    @select(_, :vegetable, :activity, :week, date = :variable)

CSV.write("data/possible_harvest_weeks.csv", possible_harvest_weeks)

# Load tables
crop_notes = CSV.read("data/crop plan - jm_crop_notes.csv", DataFrames.DataFrame)
crop_notes[!,:vegetable] = [get(vegdic, v, v) for v in  crop_notes[!,:vegetable] ]
CSV.write("data/crop plan - jm_crop_notes.csv",crop_notes)

garden_plan = CSV.read("data/crop plan - jm_garden_plan.csv", DataFrames.DataFrame)
garden_plan[!,:vegetable] = [get(vegdic, v, v) for v in  garden_plan[!,:vegetable] ]
CSV.write("data/crop plan - jm_garden_plan.csv",garden_plan)

transplant_table = CSV.read("data/crop plan - jm_transplant_table.csv", DataFrames.DataFrame)
transplant_table[!,:vegetable] = [get(vegdic, v, v) for v in  transplant_table[!,:vegetable] ]
CSV.write("data/crop plan - jm_transplant_table.csv",transplant_table)

yields_and_sales = CSV.read("data/crop plan - jm_yields_and_sales.csv", DataFrames.DataFrame)
yields_and_sales[!,:vegetable] = [get(vegdic, v, v) for v in  yields_and_sales[!,:vegetable] ]
CSV.write("data/crop plan - jm_yields_and_sales.csv",yields_and_sales)

possible_harvest_weeks = CSV.read("data/possible_harvest_weeks.csv", DataFrame)

all_vegdf = DataFrame(vegetable = unique(outerjoin(crop_notes, garden_plan, transplant_table, yields_and_sales, possible_harvest_weeks, on = :vegetable,makeunique = true)[:,:vegetable]))

crop_df = outerjoin(crop_notes, garden_plan, transplant_table, yields_and_sales, on = :vegetable,makeunique = true)

crop_df_cln = @transform(crop_df,
    intensive_spacing_between_rows_in = tosiunits.(:intensive_spacing_between_rows_in),
    intensive_spacing_within_rows_in = tosiunits.(:intensive_spacing_within_rows_in),
    preferred_cultivars = convertstringarray.(:preferred_cultivars),
    dates_of_seedings = convertstringarray.(:dates_of_seeding_of_transplant),
    dates_of_transplanting = convertstringarray.(:dates_of_transplanting),
    period_of_harvest_days = convertodays.(:period_of_harvest_days),
    days_in_garden_min = convertodays.(:days_in_garden_min),
    days_in_garden_max = convertodays.(:days_in_garden_max),
    date_in = converttodate.(:date_in),
    date_out = converttodate.(:date_out)
    )

crop_df_cln[ismissing.(crop_df_cln[!,:period_of_harvest_days]), :period_of_harvest_days] .= Day(0)

addseedingdateoftransplant!(crop_df_cln)
addearliestharvest!(crop_df_cln)

show(names(crop_df_cln))
sort!(crop_df_cln, :earliest_harvest)

# Create possible_harvest_weeks for JM info
jm_harvest_range = @pipe crop_df_cln |>
    groupby(_, :vegetable) |>
    @combine(_,
        min_harvest = minimum(:earliest_harvest),
        max_harvest = maximum(:date_out)
    )

CSV.write("tmp.csv", jm_harvest_range)

end


"""
    mpi_point_to_point(data, comm = mpi_comm(), root = mpi_root)

Declare `data` as mpi-distributed and set communication strategy to point-to-point.

Sets up the [`MPIData`](@ref) structure with [`MPIPointToPoint`](@ref) strategy.
"""
function mpi_point_to_point(data, comm = mpi_comm(), root = mpi_root)
    MPI.Initialized() || error("MPI needs to be initialised first.")
    np = MPI.Comm_size(comm)
    id = MPI.Comm_rank(comm)
    s = MPIPointToPoint(eltype(storage(data)), np, id, comm)
    return MPIData(data, comm, root, s)
end

"""
    MPIPointToPoint{N,A}

Point-to-point communication strategy. Uses circular communication using `MPI.Send` and
`MPI.Recv!`.

# Constructor

* `MPIPointToPoint(::Type{P}, np, id, comm)`: Construct an instance with pair type `P` on
  `np` processes with current rank `id`.

"""
struct MPIPointToPoint{P,N} <: DistributeStrategy
    np::Int32
    id::Int32
    comm::MPI.Comm
    datatype::MPI.Datatype
    buffers::NTuple{N,Vector{P}}
end
function MPIPointToPoint(::Type{Pair{K,V}}, np, id, comm) where {K,V}
    P = Pair{K,V}
    datatype = MPI.Datatype(P)
    return MPIPointToPoint{P,np}(np, id, comm, datatype, ntuple(_ -> P[], np))
end
"""
    recvbuff(s::MPIPointToPoint)

Get the receive buffer.
"""
function recvbuff(s::MPIPointToPoint)
    return s.buffers[s.id + 1]
end
"""
    sendbuff(s::MPIPointToPoint, id)

Get the send buffer associated with `id`.
"""
function sendbuff(s::MPIPointToPoint, id)
    return s.buffers[id + 1]
end
"""
    receive!(target, s::MPIPointToPoint, id)

Recieve from rank with `id` and move recieved values to `target`.
"""
function receive!(target, s::MPIPointToPoint{P}, id) where P
    status = MPI.Probe(id, 0, s.comm)
    count = MPI.Get_count(status, s.datatype)
    resize!(recvbuff(s), count)
    rb = recvbuff(s)
    MPI.Recv!(MPI.Buffer(rb, length(rb), s.datatype), id, 0, s.comm)
    for (key, val) in recvbuff(s)
        deposit!(target, key, val, nothing)
    end
    return target
end
"""
    send!(s::MPIPointToPoint{P})

Send the contents of the send buffers to all other ranks.
"""
function send!(s::MPIPointToPoint{<:Any,N}) where {N}
    for id in 0:N-1
        id == s.id && continue
        sb = sendbuff(s, id)
        MPI.Send(MPI.Buffer(sb, length(sb), s.datatype), id, 0, s.comm)
    end
end

function mpi_combine_walkers!(target, source, s::MPIPointToPoint{<:Any,N}) where {N}
    foreach(empty!, s.buffers)

    # sort source into send buffers, put appropriate values into target.
    for (key, val) in pairs(source)
        tr = targetrank(key, s.np)
        if tr == s.id
            deposit!(target, key, val, nothing)
        else
            push!(sendbuff(s, tr), key => val)
        end
    end
    # Receive from lower ranks.
    for id in 0:(s.id - 1)
        receive!(target, s, id)
    end
    # Send to all ranks.
    send!(s)
    # Receive from higher ranks.
    for id in (s.id + 1):(N - 1)
        receive!(target, s, id)
    end
    MPI.Barrier(s.comm)
    return target
end


using JuMP
struct DummyOptimizer <: MOI.AbstractOptimizer end
MOI.is_empty(::DummyOptimizer) = true
MOI.supports_constraint(::DummyOptimizer, ::Type{<:MOI.AbstractFunction}, ::Type{<:MOI.AbstractSet}) = true
function MOI.optimize!(::DummyOptimizer, model::MOI.ModelLike)
    return MOI.Utilities.identity_index_map(model), false
end
MOI.get(::DummyOptimizer, ::MOI.TerminationStatus) = MOI.OTHER_ERROR
MOI.get(::DummyOptimizer, ::MOI.RawStatusString) = "Dummy status"


using GRIB

keylist = Vector{String}()
grib_file = "/home/neo/Projekte/Predly/weather/gfs_workflow/download/data/gfs.20210109/06/gfs.t06z.pgrb2.0p25.f002"

parameters = [7,11,586,422,423,424,426,428,430,433,434,435,437,438,439,440,441,442,443,451,453,455,456,457,458,472,473,474,476,479,480,481,482,483,495,496,498,588]

unneeded_identifiers = ["values", "codedValues", "latLonValues", "latitudes", "longitudes", "distinctLatitudes", "distinctLongitudes"]

f = GribFile(grib_file)

function getValue(f, band::Int64, lat::Float64, lon::Float64)

    start_time = time()

    for (index, message) in enumerate(f)

        band_time = time()
        data_array = []
        parameter_name = message["parameterName"]
        parameter_unit = message["parameterUnits"]
        parameter_ecmf = message["nameECMF"]
        units_ecmf     = message["unitsECMF"]
        name           = message["name"]
        units          = message["units"]
        level          = message["level"]
        type_of_level  = message["typeOfLevel"]
        short_name     = message["shortName"]

        if index == band

            Nearest(message) do near
                lons, lats, values, distances = find(near, message, 8, 50)
                seek(f, 0)
                println("Took ", time()-start_time)
                println(lons,lats,values,distances)
                return lons, lats, values, distances
            end
        end
    end

end

getValue(f, 450, 50.0, 8.25)

destroy(f)


using HDF5
using Merlin

include("tagset.jl")
include("ner.jl")
include("model.jl")

const wordembeds_file = ".data/glove.6B.100d.h5"
#const datapath = joinpath(dirname(@__FILE__), ".data")

# training
ner = NER()
train(ner, ".data/eng.train", ".data/eng.testb")
#save("NER.jld2", Dict("a"=>seg))

# decoding
#seg = load("NER.jld2")
#println(seg["a"].char2id)
#seg = Merlin.load("NER.merlin")


@testset "filter.jl" begin
    println("             > filter.jl")
    m_out  = 1024
    m_in   = 1024

    state = BinomialState(128, m_out, m_in)
    model = BinomialGridModel(
        m_out,
        1:5,
        LinRange(0.05,0.95,5),
        LinRange(0.1,2,5),
        LinRange(0.05,2,5),
        LinRange(0.05,2,5)
    )

    fstate = NestedParticleState(state, model)
    filter = NestedParticleFilter(12)
    obs    = BinomialObservation(0.3f0, 0.1f0)

    if true
        println("")
        println("Benchmarking one filter update step: should take less than 10ms")
        display(@benchmark CUDA.@sync update!($fstate, $obs, $filter))
        println("")
        println("")
    end
end


""" Struct for 2D spectra specialized for EXPRES and its additional data products.
    Drafted, but not yet folded into examples, as still using BasicSpectra2D instead.
    For now extra info is showing up in metadata.
    """
struct Spectra2DEXPRES{T1<:Real,T2<:Real,T3<:Real,
     AA1<:AbstractArray{T1,2},AA2<:AbstractArray{T2,2},AA3<:AbstractArray{T3,2},
     InstT<:AbstractInstrument
    } <: AbstractSpectra2D
    λ::AA1
    flux::AA2
    var::AA3
    blaze::AA2
    continuum::AA2
    tellurics::AA2
    inst::InstT
    metadata::Dict{Symbol,Any} # Dict{Symbol,Any}
end

function Spectra2DEXPRES(λ::A1, flux::A2, var::A3, inst::InstT; blaze::A2, continuum::A2, tellurics::A2,
     metadata::Dict{Symbol,Any} = Dict{Symbol,Any}() ) where {
     T1<:Real, T2<:Real, T3<:Real, A1<:AbstractArray{T1,2}, A2<:AbstractArray{T2,2}, A3<:AbstractArray{T3,2},
     InstT<:AbstractInstrument  }
    @assert size(λ) == size(flux)
    @assert size(λ) == size(var)
    @assert 1 <= size(λ,1) <= max_pixel_in_order(inst)-min_pixel_in_order(inst)+1
    @assert 1 <= size(λ,2) <= max_order(inst)-min_order(inst)+1
    Spectra2DEXPRES{eltype(λ),eltype(flux),eltype(var),typeof(λ),typeof(flux),typeof(var),typeof(inst)}(λ,flux,var,blaze,continuum,tellurics,inst,metadata)
end


@testset "Interpolation" begin
    println("INTERPOLATION TEST")

    test_set = [
        (input="T/polar", freqrange=(1kHz, 12kHz), max_mag=30.0, mae_mag=5.0, max_phi=180.0, mae_phi=60.0),
        (input="T/polar_noXO", freqrange=(1kHz, 12kHz), max_mag=30.0, mae_mag=5.0, max_phi=180.0, mae_phi=60.0),

        (input="M/polar", freqrange=(200Hz, 12kHz), max_mag=35.0, mae_mag=3.0, max_phi=180.0, mae_phi=35.0),
        (input="M/polar_noXO", freqrange=(200Hz, 12kHz), max_mag=35.0, mae_mag=3.0, max_phi=180.0, mae_phi=35.0),

        (input="W/polar", freqrange=(0Hz, 1kHz), max_mag=2.0, mae_mag=0.5, max_phi=20.0, mae_phi=3.0),
        (input="W/polar_noXO", freqrange=(0Hz, 1kHz), max_mag=2.0, mae_mag=0.5, max_phi=20.0, mae_phi=3.0),
    ]
    # NOTE: errors are quite large, mainly for higher frequencies
    for e in test_set
        println("  \"$(e.input)\"")
        input_dir = joinpath(test_dir, e.input)
        print("\tA:")
        @test test_interpolation(input_dir, e.freqrange, compare_magdb, e.max_mag, e.mae_mag)
        print("\tΦ:")
        @test test_interpolation(input_dir, e.freqrange, compare_phasedeg, e.max_phi, e.mae_phi)
    end
end


#####
##### StepRange
#####

# Show we could interop with `StepRange`, which is another contender for a TimeSpans alternative
function Base.StepRange(span::AlignedSpans.AlignedSpan)
    # the rounding here is not ideal
    t = Nanosecond(round(Int, TimeSpans.nanoseconds_per_sample(span.sample_rate)))
    return (span.first_index * t):t:(span.last_index * t)
end

function AlignedSpan(r::StepRange{T,S}) where {T<:Period,S<:Period}
    sample_rate = TimeSpans.NS_IN_SEC / Dates.value(convert(Nanosecond, step(r)))
    i = first(r) / step(r)
    j = last(r) / step(r)
    return AlignedSpan(sample_rate, Int(i), Int(j))
end

@testset "StepRange roundtripping" begin
    for sample_rate in [1.0, 0.5, 100.0, 128.33]
        for (i, j) in [1 => 10, 5 => 20, 3 => 6, 78 => 79]
            span = AlignedSpan(sample_rate, i, j)
            r = StepRange(span)
            span2 = AlignedSpan(r)
            @test span.first_index == span2.first_index
            @test span.last_index == span2.last_index
            @test span.sample_rate ≈ span2.sample_rate rtol = 1e-7 # annoyingly imprecise...
        end
    end
end

#####
##### TimeSpans
#####

@testset "TimeSpans roundtripping" begin
    for sample_rate in [1.0, 0.5, 100.0, 128.33]
        # AlignedSpan -> TimeSpan -> AlignedSpan
        for (i, j) in [1 => 10, 5 => 20, 3 => 6, 78 => 79]
            for mode in (RoundSpanDown, RoundInward, ConstantSamplesRoundingMode(RoundDown))
                as = AlignedSpan(sample_rate, i, j)
                ts = TimeSpan(as)

                # We don't perfectly roundtrip with sample rate 128.33 for two of the rounding modes.
                # However, if we widen the timespan by 1 nanosecond, we can roundtrip in that case.
                # So let's test that, at least.
                if sample_rate == 128.33
                    if mode == RoundInward
                        ts = TimeSpan(max(Nanosecond(0), start(as) - Nanosecond(1)),
                                      stop(as))
                    elseif mode == ConstantSamplesRoundingMode(RoundDown)
                        ts = TimeSpan(start(as), stop(as) + Nanosecond(1))
                    end
                end

                rt = AlignedSpan(sample_rate, ts, mode)
                @test as == rt
            end
        end
    end
end

@test TimeSpans.istimespan(AlignedSpan(1, 1, 1))

#####
##### Onda
#####

@testset "Samples indexing" begin
    for sample_rate in [1.0, 0.5, 100.0, 128.33]
        samples = make_test_samples(sample_rate)
        for (i, j) in [1 => 10, 5 => 20, 3 => 6, 78 => 79]
            @test samples[:, AlignedSpan(sample_rate, i, j)] == samples[:, i:j]
            @test_throws ArgumentError samples[:, AlignedSpan(sample_rate + 1, i, j)]
        end
    end
end

#####
##### JSON/Arrow
#####

spans = [AlignedSpan(1.0, 5, 10), AlignedSpan(111.345, 500, 10000)]
@testset "JSON3 roundtripping" begin
    @test JSON3.read(JSON3.write(spans), Vector{AlignedSpan}) == spans
end

@testset "Arrow roundtripping" begin
    @test Arrow.Table(Arrow.tobuffer((; spans))).spans == spans
end


include("default_params.jl")
include("model_def.jl")

ks_errors = DataFrame(DG=[],DG_limit=[],Order_1=[],QBDRAP=[])
l1_errors = DataFrame(DG=[],DG_limit=[],Order_1=[],QBDRAP=[])
l1_cell_probs_errors = DataFrame(DG=[],DG_limit=[],Order_1=[],QBDRAP=[])
ks_error_row = zeros(4)
l1_error_row = zeros(4)
l1_cell_probs_row = zeros(4)

sim_cdf = CSV.read(
    (@__DIR__)*"/data/sims_cdf_evaluated.csv",
    DataFrame,
)

for o in orders
    cols = [1;3;4]
    for (c,m) in enumerate(string.(approx_types[[1;3;4]]).*string.([1;3;4]))
        approx_data = CSV.read(
            (@__DIR__)*"/data/first_return_cdf_approximations/order_"*string(o)*"_model_"*m*".csv",
            DataFrame,
        )
        ks_error = max(
            maximum(abs.(approx_data.phase_2-sim_cdf.phase_2)),
            maximum(abs.(approx_data.phase_4-sim_cdf.phase_4)),
        )
        ks_error_row[cols[c]] = ks_error
        l1_cdf = sum(abs.(approx_data.phase_2[1:end-1]+approx_data.phase_2[2:end]-
                sim_cdf.phase_2[1:end-1]-sim_cdf.phase_2[2:end]).*diff(approx_data.x))/2 +
            sum(abs.(approx_data.phase_4[1:end-1]+approx_data.phase_4[2:end]-
                sim_cdf.phase_4[1:end-1]-sim_cdf.phase_4[2:end]).*diff(approx_data.x))/2
        l1_error_row[cols[c]] = l1_cdf
        cell_edges_idx = findall(mod.(approx_data.x,0.4).==0.0)
        sim_cell_probs = DataFrame(
            x=[approx_data.x[1]; approx_data.x[cell_edges_idx[2:end]] - approx_data.x[cell_edges_idx[1:end-1]]],
            phase_2=[sim_cdf.phase_2[1]; sim_cdf.phase_2[cell_edges_idx[2:end]] - sim_cdf.phase_2[cell_edges_idx[1:end-1]]],
            phase_4=[sim_cdf.phase_4[1]; sim_cdf.phase_4[cell_edges_idx[2:end]] - sim_cdf.phase_4[cell_edges_idx[1:end-1]]],
        )
        approx_cell_probs = DataFrame(
            x=[approx_data.x[1]; approx_data.x[cell_edges_idx[2:end]] - approx_data.x[cell_edges_idx[1:end-1]]],
            phase_2=[approx_data.phase_2[1]; approx_data.phase_2[cell_edges_idx[2:end]] - approx_data.phase_2[cell_edges_idx[1:end-1]]],
            phase_4=[approx_data.phase_4[1]; approx_data.phase_4[cell_edges_idx[2:end]] - approx_data.phase_4[cell_edges_idx[1:end-1]]],
        )
        l1_cell_probs_row[cols[c]] = sum(abs.(approx_cell_probs.phase_2-sim_cell_probs.phase_2)) +
            sum(abs.(approx_cell_probs.phase_4-sim_cell_probs.phase_4))
    end
    push!(ks_errors,ks_error_row)
    push!(l1_errors,l1_error_row)
    push!(l1_cell_probs_errors,l1_cell_probs_row)
end

CSV.write((@__DIR__)*"/data/errors/ks_error.csv",ks_errors)
CSV.write((@__DIR__)*"/data/errors/l1_error.csv",l1_errors)
CSV.write((@__DIR__)*"/data/errors/l1_cell_probs_error.csv",l1_cell_probs_errors)

ks_errors_lwr = DataFrame(DG=[],DG_limit=[],Order_1=[],QBDRAP=[])
l1_errors_lwr = DataFrame(DG=[],DG_limit=[],Order_1=[],QBDRAP=[])
l1_cell_probs_errors_lwr = DataFrame(DG=[],DG_limit=[],Order_1=[],QBDRAP=[])
ks_error_row_lwr = zeros(4)
l1_error_row_lwr = zeros(4)
l1_cell_probs_error_row_lwr = zeros(4)

ks_errors_upr = DataFrame(DG=[],DG_limit=[],Order_1=[],QBDRAP=[])
l1_errors_upr = DataFrame(DG=[],DG_limit=[],Order_1=[],QBDRAP=[])
l1_cell_probs_errors_upr = DataFrame(DG=[],DG_limit=[],Order_1=[],QBDRAP=[])
ks_error_row_upr = zeros(4)
l1_error_row_upr = zeros(4)
l1_cell_probs_error_row_upr = zeros(4)

ks_error_boot = zeros(n_boot)
l1_error_boot = zeros(n_boot)
l1_cell_probs_error_boot = zeros(n_boot)
for o in orders
    cols = [1;3;4]
    for (c,m) in enumerate(string.(approx_types[[1;3;4]]).*string.([1;3;4]))
        approx_data = CSV.read(
            (@__DIR__)*"/data/first_return_cdf_approximations/order_"*string(o)*"_model_"*m*".csv",
            DataFrame,
        )
        cell_edges_idx = findall(mod.(approx_data.x,0.4).==0.0)
        approx_cell_probs = DataFrame(
                x=[approx_data.x[1]; approx_data.x[cell_edges_idx[2:end]] - approx_data.x[cell_edges_idx[1:end-1]]],
                phase_2=[approx_data.phase_2[1]; approx_data.phase_2[cell_edges_idx[2:end]] - approx_data.phase_2[cell_edges_idx[1:end-1]]],
                phase_4=[approx_data.phase_4[1]; approx_data.phase_4[cell_edges_idx[2:end]] - approx_data.phase_4[cell_edges_idx[1:end-1]]],
            )
        for n in 1:n_boot
            sim_cdf = CSV.read(
                (@__DIR__)*"/data/bootstrap_first_return/sample_"*string(n)*".csv",
                DataFrame,
            )
            ks_error = max(
                maximum(abs.(approx_data.phase_2-sim_cdf.phase_2)),
                maximum(abs.(approx_data.phase_4-sim_cdf.phase_4)),
            )
            ks_error_boot[n] = ks_error
            l1_cdf = sum(abs.(approx_data.phase_2[1:end-1]+approx_data.phase_2[2:end]-
                    sim_cdf.phase_2[1:end-1]-sim_cdf.phase_2[2:end]).*diff(approx_data.x))/2 +
                sum(abs.(approx_data.phase_4[1:end-1]+approx_data.phase_4[2:end]-
                    sim_cdf.phase_4[1:end-1]-sim_cdf.phase_4[2:end]).*diff(approx_data.x))/2
            l1_error_boot[n] = l1_cdf
            sim_cell_probs = DataFrame(
                x=[approx_data.x[1]; approx_data.x[cell_edges_idx[2:end]] - approx_data.x[cell_edges_idx[1:end-1]]],
                phase_2=[sim_cdf.phase_2[1]; sim_cdf.phase_2[cell_edges_idx[2:end]] - sim_cdf.phase_2[cell_edges_idx[1:end-1]]],
                phase_4=[sim_cdf.phase_4[1]; sim_cdf.phase_4[cell_edges_idx[2:end]] - sim_cdf.phase_4[cell_edges_idx[1:end-1]]],
            )
            l1_cell_probs_error_boot[n] = sum(abs.(approx_cell_probs.phase_2-sim_cell_probs.phase_2)) +
                sum(abs.(approx_cell_probs.phase_4-sim_cell_probs.phase_4))
        end
        ks_error_row_lwr[cols[c]] = quantile(ks_error_boot,qtiles[1])
        l1_error_row_lwr[cols[c]] = quantile(l1_error_boot,qtiles[1])
        l1_cell_probs_error_row_lwr[cols[c]] = quantile(l1_cell_probs_error_boot,qtiles[1])
        ks_error_row_upr[cols[c]] = quantile(ks_error_boot,qtiles[2])
        l1_error_row_upr[cols[c]] = quantile(l1_error_boot,qtiles[2])
        l1_cell_probs_error_row_upr[cols[c]] = quantile(l1_cell_probs_error_boot,qtiles[2])
    end
    push!(ks_errors_lwr,ks_error_row_lwr)
    push!(l1_errors_lwr,l1_error_row_lwr)
    push!(l1_cell_probs_errors_lwr,l1_cell_probs_error_row_lwr)
    push!(ks_errors_upr,ks_error_row_upr)
    push!(l1_errors_upr,l1_error_row_upr)
    push!(l1_cell_probs_errors_upr,l1_cell_probs_error_row_upr)
end

CSV.write((@__DIR__)*"/data/errors/ks_error_lwr.csv",ks_errors_lwr)
CSV.write((@__DIR__)*"/data/errors/l1_error_lwr.csv",l1_errors_lwr)
CSV.write((@__DIR__)*"/data/errors/l1_cell_probs_error_lwr.csv",l1_cell_probs_errors_lwr)
CSV.write((@__DIR__)*"/data/errors/ks_error_upr.csv",ks_errors_upr)
CSV.write((@__DIR__)*"/data/errors/l1_error_upr.csv",l1_errors_upr)
CSV.write((@__DIR__)*"/data/errors/l1_cell_probs_error_upr.csv",l1_cell_probs_errors_upr)


using Metalhead, Test
using Flux
using Flux: Zygote

# PRETRAINED_MODELS = [(VGG19, false), ResNet50, GoogLeNet, DenseNet121, SqueezeNet]
PRETRAINED_MODELS = []

function gradtest(model, input)
  y, pb = Zygote.pullback(() -> model(input), Flux.params(model))
  gs = pb(ones(Float32, size(y)))

  # if we make it to here with no error, success!
  return true
end

@testset "AlexNet" begin
  model = AlexNet()
  @test size(model(rand(Float32, 256, 256, 3, 2))) == (1000, 2)
  @test_throws ArgumentError AlexNet(pretrain = true)
  @test_skip gradtest(model, rand(Float32, 256, 256, 3, 2))
end

@testset "VGG" begin
  @testset "$model(BN=$bn)" for model in [VGG11, VGG13, VGG16, VGG19], bn in [true, false]
    imsize = (224, 224)
    m = model(batchnorm = bn)

    @test size(m(rand(Float32, imsize..., 3, 2))) == (1000, 2)
    if (model, bn) in PRETRAINED_MODELS
      @test (model(batchnorm = bn, pretrain = true); true)
    else
      @test_throws ArgumentError model(batchnorm = bn, pretrain = true)
    end
    @test_skip gradtest(m, rand(Float32, imsize..., 3, 2))
  end
end

@testset "ResNet" begin
  @testset for model in [ResNet18, ResNet34, ResNet50, ResNet101, ResNet152]
    m = model()

    @test size(m(rand(Float32, 256, 256, 3, 2))) == (1000, 2)
    if model in PRETRAINED_MODELS
      @test (model(pretrain = true); true)
    else
      @test_throws ArgumentError model(pretrain = true)
    end
    @test_skip gradtest(m, rand(Float32, 256, 256, 3, 2))
  end

  @testset "Shortcut C" begin
    m = Metalhead.resnet(Metalhead.basicblock, :C;
                         channel_config = [1, 1],
                         block_config = [2, 2, 2, 2])

    @test size(m(rand(Float32, 256, 256, 3, 2))) == (1000, 2)
  end
end

@testset "GoogLeNet" begin
  m = GoogLeNet()
  @test size(m(rand(Float32, 224, 224, 3, 2))) == (1000, 2)
  @test_throws ArgumentError (GoogLeNet(pretrain = true); true)
  @test_skip gradtest(m, rand(Float32, 224, 224, 3, 2))
end

@testset "Inception3" begin
  m = Inception3()
  @test size(m(rand(Float32, 299, 299, 3, 2))) == (1000, 2)
  @test_throws ArgumentError Inception3(pretrain = true)
  @test_skip gradtest(m, rand(Float32, 299, 299, 3, 2))
end

@testset "SqueezeNet" begin
  m = SqueezeNet()
  @test size(m(rand(Float32, 227, 227, 3, 2))) == (1000, 2)
  @test_throws ArgumentError (SqueezeNet(pretrain = true); true)
  @test_skip gradtest(m, rand(Float32, 227, 227, 3, 2))
end

@testset "DenseNet" begin
  @testset for model in [DenseNet121, DenseNet161, DenseNet169, DenseNet201]
    m = model()

    @test size(m(rand(Float32, 224, 224, 3, 2))) == (1000, 2)
    if model in PRETRAINED_MODELS
      @test (model(pretrain = true); true)
    else
      @test_throws ArgumentError model(pretrain = true)
    end
    @test_skip gradtest(m, rand(Float32, 224, 224, 3, 2))
  end
end


export decimation, masterTrigger, currentFrame, ramWriterMode, connectADC, startADC, stopADC, samplesPerPeriod, periodsPerFrame,
     numSlowDACChan, setSlowDACLUT, enableSlowDAC, slowDACStepsPerRotation, samplesPerSlowDACStep, prepareSlowDAC,
     currentWP, slowDACInterpolation, numSlowADCChan, numLostStepsSlowADC, bufferSize, keepAliveReset, triggerMode,
     slowDACStepsPerFrame, enableDACLUT, ADCPerformanceData, RPPerformance, RPStatus, RPInfo, startPipelinedData, PerformanceData, numChan

struct ADCPerformanceData
  deltaRead::UInt64
  deltaSend::UInt64
end

struct RPStatus
  overwritten::Bool
  corrupted::Bool
  stepsLost::Bool
  adcEnabled::Bool
  dacEnabled::Bool
end

struct PerformanceData
  wpRead::UInt64
  adc::ADCPerformanceData
  dac::DACPerformanceData
  status::RPStatus
end

struct RPPerformance
  data::Vector{PerformanceData}
end

struct RPInfo
  performances::Vector{RPPerformance}
end

function RPInfo()
  return RPInfo([RPPerformance([])])
end

decimation(rp::RedPitaya) = query(rp,"RP:ADC:DECimation?", Int64)
function decimation(rp::RedPitaya, dec)
  rp.decimation = Int64(dec)
  send(rp, string("RP:ADC:DECimation ", rp.decimation))
end

numChan(rp::RedPitaya) = 2

numSlowDACChan(rp::RedPitaya) = query(rp,"RP:ADC:SlowDAC?", Int64)
function numSlowDACChan(rp::RedPitaya, value)
  if value <= 0 || value > 4
    error("Num slow DAC channels needs to be between 1 and 4!")
  end
  send(rp, string("RP:ADC:SlowDAC ", Int64(value)))
end

function setSlowDACLUT(rp::RedPitaya, lut::Array)
  lutFloat32 = map(Float32, lut)
  send(rp, string("RP:ADC:SlowDACLUT"))
  @debug "Writing slow DAC LUT"
  write(rp.dataSocket, lutFloat32)
end

function enableDACLUT(rp::RedPitaya, lut::Array)
  lutBool = map(Bool, lut)
  send(rp, string("RP:ADC:EnableDACLUT"))
  @debug "Writing enable DAC LUT"
  write(rp.dataSocket, lutBool)
end

function enableSlowDAC(rp::RedPitaya, enable::Bool, numFrames::Int64=0,
            ffRampUpTime::Float64=0.4, ffRampUpFraction::Float64=0.8)
  enableI = Int32(enable)
  return query(rp, string("RP:ADC:SlowDACEnable ", enableI,
              ",", numFrames, ",", ffRampUpTime, ",", ffRampUpFraction), Int64)
end

function slowDACInterpolation(rp::RedPitaya, enable::Bool)
  enableI = Int32(enable)
  send(rp, string("RP:ADC:SlowDACInterpolation ", enableI))
end

numSlowADCChan(rp::RedPitaya) = query(rp,"RP:ADC:SlowADC?", Int64)
function numSlowADCChan(rp::RedPitaya, value)
  send(rp, string("RP:ADC:SlowADC ", Int64(value)))
end

numLostStepsSlowADC(rp::RedPitaya) = query(rp,"RP:ADC:SlowDACLostSteps?", Int64)

function samplesPerPeriod(rp::RedPitaya)
  return rp.samplesPerPeriod
end
function samplesPerPeriod(rp::RedPitaya, value)
  rp.samplesPerPeriod = value
  samplesPerSlowDACStep(rp, value)
end

function periodsPerFrame(rp::RedPitaya)
  return rp.periodsPerFrame
end
function periodsPerFrame(rp::RedPitaya, value)
  rp.periodsPerFrame = value
end

samplesPerSlowDACStep(rp::RedPitaya) = query(rp,"RP:ADC:SlowDAC:SamplesPerStep?", Int64)
function samplesPerSlowDACStep(rp::RedPitaya, value)
  send(rp, string("RP:ADC:SlowDAC:SamplesPerStep ", value))
end

slowDACStepsPerRotation(rp::RedPitaya) = query(rp,"RP:ADC:SlowDAC:StepsPerRotation?", Int64)
function slowDACStepsPerRotation(rp::RedPitaya, value)
  send(rp, string("RP:ADC:SlowDAC:StepsPerRotation ", value))
end

function prepareSlowDAC(rp::RedPitaya, samplesPerStep, stepsPerRotation, numOfChan)
  numSlowDACChan(rp, numOfChan)
  samplesPerSlowDACStep(rp, samplesPerStep)
  slowDACStepsPerRotation(rp, stepsPerRotation)
end

function slowDACStepsPerFrame(rp::RedPitaya, stepsPerFrame)
  samplesPerFrame = rp.periodsPerFrame * rp.samplesPerPeriod
  samplesPerStep = div(samplesPerFrame, stepsPerFrame)
  samplesPerSlowDACStep(rp, samplesPerStep)
  slowDACStepsPerRotation(rp, stepsPerFrame) # Sets PDMClockDivider
end

function currentFrame(rp::RedPitaya)
  return Int64(floor(currentWP(rp) / (rp.samplesPerPeriod * rp.periodsPerFrame)))
end

function currentPeriod(rp::RedPitaya)
  return Int64(floor(currentWP(rp) / (rp.samplesPerPeriod)))
end

currentWP(rp::RedPitaya) = query(rp,"RP:ADC:WP:CURRENT?", Int64)
bufferSize(rp::RedPitaya) = query(rp,"RP:ADC:BUFFER:SIZE?", Int64)

function masterTrigger(rp::RedPitaya, val::Bool)
  valStr = val ? "ON" : "OFF"
  send(rp, string("RP:MasterTrigger ", valStr))
end
masterTrigger(rp::RedPitaya) = occursin("ON", query(rp,"RP:MasterTrigger?"))

function keepAliveReset(rp::RedPitaya, val::Bool)
  valStr = val ? "ON" : "OFF"
  send(rp, string("RP:KeepAliveReset ", valStr))
end
keepAliveReset(rp::RedPitaya) = occursin("ON", query(rp,"RP:KeepAliveReset?"))

# "TRIGGERED" or "CONTINUOUS"
function ramWriterMode(rp::RedPitaya, mode::String)
  send(rp, string("RP:RamWriterMode ", mode))
end

# "INTERNAL" or "EXTERNAL"
function triggerMode(rp::RedPitaya, mode::String)
  send(rp, string("RP:Trigger:Mode ", mode))
end

connectADC(rp::RedPitaya) = send(rp, "RP:ADC:ACQCONNect")
function startADC(rp::RedPitaya)
  send(rp, "RP:ADC:ACQSTATUS ON")
end
stopADC(rp::RedPitaya) = send(rp, "RP:ADC:ACQSTATUS OFF")

wasOverwritten(rp::RedPitaya) = query(rp, "RP:STATus:OVERwritten?", Bool)
wasCorrupted(rp::RedPitaya) = query(rp, "RP:STATus:CORRupted?", Bool)

function serverStatus(rp::RedPitaya)
  send(rp, "RP:STATus?")
  return readServerStatus(rp)
end

function readServerStatus(rp::RedPitaya)
  statusRaw = read!(rp.dataSocket, Array{Int8}(undef, 1))[1]
  status = RPStatus((statusRaw & 1) != 0, # overwritten
   (statusRaw & (1 << 1)) != 0, # corrupted
   (statusRaw & (1 << 2)) != 0, # stepsLost
   (statusRaw & (1 << 3)) != 0, # adcEnabled
   (statusRaw & (1 << 4)) != 0) # dacEnabled
  return status
end

function performanceData(rp::RedPitaya)
  send(rp, "RP:PERF?")
  return readPerformanceData(rp)
end

function readPerformanceData(rp::RedPitaya)
  adc = readADCPerformanceData(rp)
  dac = readDACPerformanceData(rp)
  return (adc, dac)
end

function readADCPerformanceData(rp::RedPitaya)
  perf = read!(rp.dataSocket, Array{UInt64}(undef, 2))
  return ADCPerformanceData(perf[1], perf[2])
end

# Low level read. One has to take care that the numFrames are available
function readSamples_(rp::RedPitaya, reqWP, numSamples)
  command = string("RP:ADC:DATA? ",Int64(reqWP),",",Int64(numSamples))
  send(rp, command)

  @debug "read data ..."
  u = read!(rp.dataSocket, Array{Int16}(undef, 2 * Int64(numSamples)))
  @debug "read data!"
  return u
end

# Low level read, reads samples, error and perf. Values need to be already requested
function readSamplesChunk_(rp::RedPitaya, reqWP::Int64, numSamples::Int64, into=nothing)
  @debug "read samples chunk ..."
  if into === nothing
    into = Array{Int16}(undef, 2 * Int64(numSamples))
  end
  data = read!(rp.dataSocket, into)
  status = readServerStatus(rp)
  (adc, dac) = readPerformanceData(rp)
  @debug "read samples chunk ..."
  perf = PerformanceData(UInt64(reqWP), adc, dac, status)
  return (data, perf)
end

# Low level read, that includes performance and error data
function readDetailedSamples_(rp::RedPitaya, reqWP::Int64, numSamples::Int64)
  command = string("RP:ADC:DATA:DETAILED? ",Int64(reqWP),",",Int64(numSamples))
  send(rp, command)
  return readSamplesChunk_(rp, reqWP, numSamples)
end

function readSamplesIntermediate_(rp::RedPitaya, reqWP::Int64, numSamples::Int64)
  data = readSamples_(rp, reqWP, numSamples)
  status = serverStatus(rp)
  (adc, dac) = performanceData(rp)
  perf = PerformanceData(UInt64(reqWP), adc, dac, status)
  return (data, perf)
end

function readSamplesOld_(rp::RedPitaya, reqWP::Int64, numSamples::Int64)
  data = readSamples_(rp, reqWP, numSamples)
  overwritten = wasOverwritten(rp)
  corrupted = wasCorrupted(rp)
  status = RPStatus(overwritten, corrupted, false, true, true)
  (adc, dac) = performanceData(rp)
  perf = PerformanceData(UInt64(reqWP), adc, dac, status)
  return (data, perf)
end

function startPipelinedData(rp::RedPitaya, reqWP::Int64, numSamples::Int64, chunkSize::Int64)
  command = string("RP:ADC:DATA:PIPELINED? ", reqWP, ",", numSamples, ",", chunkSize)
  send(rp, command)
end

# Low level read. One has to take care that the numFrames are available
function readDataSlow_(rp::RedPitaya, startFrame, numFrames)
  numPeriods = rp.periodsPerFrame
  numChan = numSlowADCChan(rp)

  command = string("RP:ADC:SLOW:FRAMES:DATA ",Int64(startFrame),",",Int64(numFrames))
  #send(rp, command)

  @debug "read data ..."
  #u = read!(rp.dataSocket, Array{Float32}(undef, numChan * numFrames * numPeriods))
  @debug "read data!"
  #return reshape(u, numChan, numPeriods, numFrames)
  return zeros(Float32, numChan, numPeriods, numFrames)
end


using RecordedArrays: getdim

@testset "MCIndices" begin
    im = MCIndices(([1, 2, 3], [1, 3, 4], [2, 4, 5]))
    @test size(im) == (3, 3, 3)
    for i in ndims(im)
        @test getdim(im, i) == im.indices[i]
    end
    for (i, ind) in enumerate(im)
        Is = Base._ind2sub(im, i)
        @test map(getindex, im.indices, Is) == ind
    end
    for i in ndims(im)
        push!(getdim(im, i), 6)
        @test im.indices[i][end] == 6
        @test length(im.indices[i]) == 4
    end
end

@testset "DOKSparseArray" begin
    s = DOKSparseArray(Dict{Tuple{},Int}(() => 1), Size())
    @test ArrayInterface.parent_type(s) == Dict{Tuple{},Int}
    @test s[] == 1
    s[] = 2
    @test get(s, (), 1) == 2
end


using Oceananigans.Utils: with_tracers
using Oceananigans.TurbulenceClosures: with_tracers

@inline zeroforcing(args...) = 0

const assumed_field_locations = Dict(:tracer => (Cell, Cell, Cell),
                                          :u => (Face, Cell, Cell),
                                          :v => (Cell, Face, Cell),
                                          :w => (Face, Cell, Face))

"""
    ModelForcing(; u=zeroforcing, v=zeroforcing, w=zeroforcing, tracer_forcings...)

Return a named tuple of forcing functions for each solution field.

Example
=======

julia> u_forcing = SimpleForcing((x, y, z, t) -> exp(z) * cos(t))

julia> model = IncompressibleModel(forcing=ModelForcing(u=u_forcing))
"""
function ModelForcing(; u=zeroforcing, v=zeroforcing, w=zeroforcing, tracer_forcings...)
    u = for_field_name(:u, u)
    v = for_field_name(:v, v)
    w = for_field_name(:w, w)

    # Re-build tracer forcings
    tracer_names = Tuple(f.first for f in tracer_forcings)
    tracer_functions = Tuple(for_field_name(f.first, f.second) for f in tracer_forcings)
    tracer_forcings = NamedTuple{tracer_names}(tracer_functions)

    return merge((u=u, v=v, w=w), tracer_forcings)
end

for_field_name(field_name, forcing_function) = forcing_function # Fallback for ordinary functions

""" Returns a SimpleForcing object at `field_name`'s assumed location. """
function for_field_name(field_name, f::SimpleForcing)

    # Deduce the location of `field_name`.
    X, Y, Z = field_name ∈ (:u, :v, :w) ? assumed_field_locations[field_name] :
                                          assumed_field_locations[:tracer]

    return SimpleForcing{X, Y, Z}(field_name, f.forcing, f.parameters, f.field_in_signature)
end

default_tracer_forcing(args...) = zeroforcing

ModelForcing(tracers, proposal_forcing) = with_tracers(tracers, proposal_forcing, default_tracer_forcing,
                                                       with_velocities=true)