text:
using OrdinaryDiffEq
using Trixi
using Plots

v1 = 0.1 # MACH: v1 = 0.38 -> 10.16%, v1 = 0.1 -> 2.67? %
tspan = (0.0, 35) # ≈ 20 für stabilen Zustand
CFL = 0.8
cells_per_dimension = (64, 64)

visualization = VisualizationCallback(interval=200,variable_names=["v1","v2"], plot_creator=Trixi.save_plot)# clims=(-0.5,0.5),

function WR2_initial_condition_constant(x, t, equations::CompressibleEulerEquations2D)
    rho = 1.0
    rho_v1 = v1
    rho_v2 = 0
    rho_e = 25.0
    return SVector(rho, rho_v1, rho_v2, rho_e)
  end
initial_condition = WR2_initial_condition_constant

surface_flux = FluxPlusDissipation(flux_chandrashekar, DissipationLocalLaxFriedrichs(max_abs_speed_naive)) # flux_lax_friedrichs #
volume_flux  = flux_chandrashekar
volume_integral = VolumeIntegralFluxDifferencing(volume_flux)
basis = LobattoLegendreBasis(4)
solver = DGSEM(basis, surface_flux, volume_integral)

###############################################################################
# semidiscretization of the compressible Euler equations

equations = CompressibleEulerEquations2D(1.4)

# mapping O-mesh
function mapping(xi_, eta_)

    ξ = xi_
    η = eta_

    size = 2    # gebietsgröße
    d = 2       # ≈ durchmesser
    x = size * (d + ξ ) * cos(π * (η + 1))
    y = size * (d + ξ ) * sin(π * (η + 1))

    return SVector(x, y)
end

function boundary_condition_stream(u_inner, orientation, direction, x, t,
    surface_flux_function,
    equations::CompressibleEulerEquations2D)
    # Far Field Conditions
    # if (x[1] < 1)
    #     # if t < π/2
    #     #     u_boundary = SVector(1, 0.38 , 0, 10) #+ 1* sin(π*t/10)*10
    #     # else
    #     #     u_boundary = SVector(1, 0.38, 0, 10)
    #     # end
    # else
        u_boundary = SVector(1, v1, 0, 25)
    # end
    # Calculate boundary flux
    if direction in (2, 4) # u_inner is "left" of boundary, u_boundary is "right" of boundary
        flux = surface_flux_function(u_inner, u_boundary, orientation, equations)
    else # direction == 4 # u_boundary is "left" of boundary, u_inner is "right" of boundary
        flux = surface_flux_function(u_boundary, u_inner, orientation, equations)
    end

    return flux
end

function boundary_condition_constant_farfield( u_inner, orientation, direction, x, t,
    surface_flux_function,
    equations::CompressibleEulerEquations2D)
    # Far Field Conditions
    u_boundary = SVector(1, v1, 0, 10)

    # Calculate boundary flux
    if direction in (2, 4) # u_inner is "left" of boundary, u_boundary is "right" of boundary
        flux = surface_flux_function(u_inner, u_boundary, orientation, equations)
    else # direction == 4 # u_boundary is "left" of boundary, u_inner is "right" of boundary
        flux = surface_flux_function(u_boundary, u_inner, orientation, equations)
    end

    return flux
end

function boundary_condition_freeslip(u_inner, orientation, direction, x, t,
    surface_flux_function,
    equations::CompressibleEulerEquations2D)
    # Freeslip wall Conditions
    rho, v1, v2, p = cons2prim(u_inner, equations)
    u_boundary = prim2cons(SVector(rho, -v1, -v2, p), equations)

    # Calculate boundary flux
    if direction == 2 # u_inner is "left" of boundary, u_boundary is "right" of boundary
        flux = surface_flux_function(u_inner, u_boundary, orientation, equations)
    else # direction == 4 # u_boundary is "left" of boundary, u_inner is "right" of boundary
        flux = surface_flux_function(u_boundary, u_inner, orientation, equations)
    end

    return flux
end

mesh = CurvedMesh(cells_per_dimension, mapping,periodicity=(false,true))

#mesh = CurvedMesh(cells_per_dimension, coordinates_min, coordinates_max, periodicity=(false, true))

# Implement boundary conditions
# x neg and x_pos are not periodic
# boundary_conditions = boundary_condition_periodic
boundary_conditions = (x_neg=boundary_condition_freeslip,
                       x_pos=boundary_condition_stream,
                       y_neg=boundary_condition_periodic,
                       y_pos=boundary_condition_periodic)

semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver,boundary_conditions=boundary_conditions)

###############################################################################
# ODE solvers, callbacks etc.

summary_callback = SummaryCallback()

save_solution = SaveSolutionCallback(interval=10,
                                     save_initial_solution=true,
                                     save_final_solution=true,
                                     solution_variables=cons2prim)

stepsize_callback = StepsizeCallback(cfl=CFL)
analysis_interval = 100
analysis_callback = AnalysisCallback(semi, interval=analysis_interval, save_analysis=true,
                                                 extra_analysis_errors=(:conservation_error,),
                                                 extra_analysis_integrals=(entropy, energy_total,
                                                                           energy_kinetic, energy_internal)

                                                )

callbacks = CallbackSet(summary_callback, stepsize_callback, analysis_callback, visualization)

###############################################################################
# run the simulation

ode = semidiscretize(semi, tspan)

sol = solve(ode, CarpenterKennedy2N54(williamson_condition=false),
            dt=1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
            save_everystep=false, callback=callbacks);
summary_callback() # print the timer summary

# pd=PlotData2D(sol)
# b=plot(pd["rho"])
# plot!(getmesh(pd))

# plot(sol)
# # savefig(plotd, "test.png")

max_stars_repo_path: src/WR2/MS4 ZylinderFlow.jl

max_issues_repo_path: src/WR2/MS4 ZylinderFlow.jl

max_forks_repo_path: src/WR2/MS4 ZylinderFlow.jl


text:
"""
    loadElser(fname)

Load (and symmetrize) benchmark phase retrieval data

# Implementation
Loads data from
https://github.com/veitelser/phase-retrieval-benchmarks
provided file name (+path from cwd)

Symmetrize the data (make it twice is large; ignore optimization
due to real-valuedness).

# Examples
```julia-repl
julia> using FPRS;
julia> sqrtI, nAtoms, suppSize = loadElser("../data/data100E")
```
"""
function loadElser(fname::String)
    fIn = readdlm(fname,'\t')
    nAtoms = parse(Int16, fname[end-3:end-1])
    suppSize = 8 * nAtoms
    M = 128
    sqrtI = zeros(M,M)
    sqrtI[:,1:64] = fIn[:,:]
    sqrtI[1,end:-1:66] = fIn[1,2:end]
    sqrtI[2:end, end:-1:66] = fIn[end:-1:2,2:end]
    return sqrtI, nAtoms, suppSize
end

function loadElserPure(fname::String)
    fInOriginal = readdlm(fname,'\t')
    fIn = sqrt.(fInOriginal) # to match c++ and julia FFTW normalization constants
    fIn = transpose(fIn) # To match julia's rfft format
    M = 128
    Mh = 64
    fmag = zeros(Mh, M)
    fmag[1:64,:] = fIn
    return fmag
end

ext: jl

max_stars_repo_name: tsipenyuk/FPRS

max_issues_repo_name: tsipenyuk/FPRS

max_forks_repo_name: tsipenyuk/FPRS


text:
using MAT
using LinearAlgebra
using TimerOutputs
using SparseArrays
using SuiteSparse

to = TimerOutput()

function testfun(x)
	afile = Array{Any}(undef, 8)
	afile[8] = matopen("/home/ducky/code/progetto-calcolo/data/matrix/Hook_1498.mat")
	afile[1] = matopen("/home/ducky/code/progetto-calcolo/data/matrix/G3_circuit.mat")
	afile[2] = matopen("/home/ducky/code/progetto-calcolo/data/matrix/nd24k.mat")
	afile[3] = matopen("/home/ducky/code/progetto-calcolo/data/matrix/bundle_adj.mat")
	afile[4] = matopen("/home/ducky/code/progetto-calcolo/data/matrix/ifiss_mat.mat")
	afile[5] = matopen("/home/ducky/code/progetto-calcolo/data/matrix/TSC_OPF_1047.mat")
	afile[6] = matopen("/home/ducky/code/progetto-calcolo/data/matrix/ns3Da.mat")
	afile[7] = matopen("/home/ducky/code/progetto-calcolo/data/matrix/GT01R.mat")
	file = afile[x]

		to = TimerOutput()
		prob = read(file, "Problem")
		A= prob["A"]
		# A = sparse(Asparse)
		xe = ones(size(A,2),1)

		b= A*xe

		println("Begin Factor")
		@timeit to "Factorization Cost" begin
				A1 = factorize(A)
		end

		@timeit to "FactoResolver" begin
			xF = A1\b
		end
		res_normF = norm(xe-xF)/norm(xe)

		@timeit to "LinearResolver" begin
			xM = A\b
		end

		res_norm = norm(xM-xe)/norm(xe)
		res_normF = norm(xM-xe)/norm(xe)
		print("Printing the norm for ", prob["name"], " file!\n")
		print(res_norm, "\n",
		res_normF,
		"\nPrinting timings\n")
		show(to)
		print("\n")
		println(typeof(A), " tipo di A")
		println(typeof(A1), " tipo di A1")

end

max_stars_repo_name: illoxian/mldivide-comparator

max_issues_repo_name: illoxian/mldivide-comparator

max_forks_repo_name: illoxian/mldivide-comparator


text:
using Plots
using BenchmarkTools

function vfsolveplot(vnew, kgrid, tolerance, imax)
    α = 0.3
    δ = 0.1
    β = 0.9

    display(plot(kgrid, vnew, label = "V0"))
    # display(plot(kgrid, kgrid, label = "k'0"))
    sleep(2)
    v = vnew .+ 2*tolerance
    kprime = zeros(n)
    i = 1
    while maximum(abs.(v - vnew)) > tolerance && i <= imax
        v = vnew;
        c = kgrid.^α + (1-δ)*kgrid .- kgrid'
        c[c .< 0] .= 0
        (vnew, cartesianindex) = findmax(log.(c) .+ β*v', dims = 2);
        kprimeindex = getindex.(cartesianindex, 2)
        kprime = kgrid[kprimeindex]
        i % 5 == 0 ? display(plot!(kgrid, vnew, label = "V$i")) : nothing
        # i % 5 == 0 ? display(plot!(kgrid, kprime, label = "k'$i")) : nothing
        sleep(1)
        i += 1;
    end
    return (v = vnew, kprime = kprime, kprimeindex = kprimeindex)
end

kupper = 2
klower = 0.001
n = 100
kgrid = collect(range(klower, stop = kupper, length = n))
(v, kprime, kprimeindex) = vfsolveplot(zeros(n), kgrid, 0.001, 1000);


text:
#states constructors

struct State_fixed{T<:AbstractVector}
    st::T
    ope::Op_fixed
    State_fixed(st,ope) = length(st) != binomial(dim(ope),nume(ope)) ? error("lenght of vector does not match dimension") : new{typeof(st)}(st/sqrt(st'*st),ope)
end
#
# struct State_fixed{T<:AbstractVector}
#     st::T
#     ope::Op
#     nume::Int64
# end

st(s::State_fixed) = s.st

ope(s::State_fixed) = s.ope

nume(s::State_fixed) = nume(s.ope)

dim(s::State_fixed) = dim(s.ope)

basis(s::State_fixed) = basis(s.ope)

typ(s::State_fixed) = eltype(s.st)

function rhosp(sta::State_fixed)
    precis = 15
    n = dim(ope(sta))
    num = nume(sta)
    rhospv = spzeros(typ(sta),n,n)
    estate = st(sta)
    o = ope(sta)
    for i in 1:n
        for j in 1:n
            rhospv[i,j] = round(estate'*ada(o, i, j)*estate, digits = precis)
        end
    end
    return rhospv
end

# This functions takes stat, an array
# in the full 2^n space, into the reduced
# space.
function fixed_state(stat::AbstractVector, nume::Int64)
    n = Int(log(2,length(stat)))
    posn = ffilter(n, nume)
    reduced = stat[setdiff(1:end,posn)]
    return reduced
end

function fixed_state(sta::State, nume::Int64)
    stat = st(sta)
    n = Int(log(2,length(stat)))
    posn = ffilter(n, nume)
    reduced = stat[setdiff(1:end,posn)]
    return State_fixed(reduced,Op_fixed(n,nume))
end

# This functions takes stat, an array
# in the reduced space, into the full 2^n
# space. Index is the index from basis_m
# and d the dimension
function unfixed_state(stat::AbstractVector, d::Int64, m::Int64)
    index = basis_m(d,m)[2]
    l = size(stat)[1]
    state2n = zeros(2^d)
    for (i,v) in enumerate(index)
        state2n[Int(v)]=stat[i]
    end
    return state2n
end

function unfixed_state(sta::State_fixed)
    d = dim(ope(sta))
    index = basis_m(d,nume(sta))[2]
    stat = st(sta)
    l = size(stat)[1]
    state2n = zeros(2^d)
    for (i,v) in enumerate(index)
        state2n[Int(v)]=stat[i]
    end
    return State(state2n,Op(d))
end


text:
using Gadfly, Romeo, GLFW, GLAbstraction, Reactive, ModernGL, GLWindow, Color, Distributions
function clear!(x::Vector{RenderObject})
	while !isempty(x)
		value = pop!(x)
		delete!(value)
	end
end

const taskqueue = Any[]
function viz(func::Function, args...; custumizations...)
	signals = map(args) do x
		isa(x, Signal) ? x : Input(x)
	end
	function execute(x...)
		func(x...)
	end
	visobj = lift(execute, signals...)
	returnobj = visualize(visobj.value; custumizations...)
	lift(visobj) do vis
		if isempty(taskqueue)
			push!(taskqueue, @async visualize(vis; custumizations...))
		elseif length(taskqueue) == 1
			push!(taskqueue, vis)
		elseif length(taskqueue) == 2
			if istaskdone(first(taskqueue))
				empty!(taskqueue)
				push!(taskqueue, @async visualize(vis; custumizations...))
			else
				taskqueue[2] = vis
			end
		end
		returnobj
	end
	returnobj
end

function init_romeo()
	sourcecode_area = lift(Romeo.ROOT_SCREEN.area) do x
		Rectangle(0, 0, div(x.w, 7)*3, x.h)
	end
	visualize_area = lift(Romeo.ROOT_SCREEN.area) do x
		Rectangle(div(x.w,7)*3, 0, div(x.w, 7)*3, x.h)
	end
	search_area = lift(visualize_area) do x
		Rectangle(x.x, x.y, x.w, div(x.h,10))
	end
	edit_area = lift(Romeo.ROOT_SCREEN.area) do x
		Rectangle(div(x.w, 7)*6, 0, div(x.w, 7), x.h)
	end

	sourcecode_screen   = Screen(Romeo.ROOT_SCREEN, area=sourcecode_area)
	visualize_screen    = Screen(Romeo.ROOT_SCREEN, area=visualize_area)
	search_screen       = Screen(visualize_screen,  area=search_area)
	edit_screen         = Screen(Romeo.ROOT_SCREEN, area=edit_area)

	w_height = lift(Romeo.ROOT_SCREEN.area) do x
		x.h
	end
	source_offset = lift(w_height) do x
		translationmatrix(Vec3(30,x-30,0))
	end
	w_height_search = lift(search_screen.area) do x
		x.h
	end
	search_offset = lift(w_height_search) do x
		translationmatrix(Vec3(30,x-30,0))
	end
	D = Dict{Symbol, Any}(
		:a => Input(1f0),
		:b => Input(10f0)
	)
	signals = edit(D, screen=sourcecode_screen)
	barplot = viz(plot, sin, D[:a], D[:b], screen=visualize_screen)

	append!(sourcecode_screen.renderlist, signals)
	push!(visualize_screen.renderlist, barplot)
end

init_romeo()
glClearColor(1,1,1,1)
while Romeo.ROOT_SCREEN.inputs[:open].value
	Romeo.renderloop(Romeo.ROOT_SCREEN)
	sleep(0.0001)
end
GLFW.Terminate()


text:
using Flora.RNN
using MLDatasets # FashionMNIST

# load full training set
train_x, train_y = FashionMNIST.traindata() # 60_000

picks = [
    [0, 3, 9], # 0 T-Shirt, 3 Dress, 9 Ankle boot
    [1, 2, 5],
]

m = chain(3)
for pick in picks
    X, Y = prepare(train_x, train_y, pick)
    p = train!(:fashion, m, pick, X, Y; n=3)
    show_accuracy(m, pick, X, Y)
    save(p)
end

#=
train! n=3
pick = [0, 3, 9]
mean(argmax(m(x)) .== argmax(y)) = 0.9375
pick = [1, 2, 5]
mean(argmax(m(x)) .== argmax(y)) = 0.9864444444444445
=#


text:
export AbstractPredictor,
    AbstractPredictorCache,
    cache,
    setup!,
    update!,
    predict!

"""
    AbstractPredictor

The path tracking problem can be considered a initial value problem.
A predictor make use of this fact and predicts for a given pair ``(x₀,t)``
with ``H(x₀,t)=0`` a new ``x`` such that ``H(x, t + Δt)=0``.

The differential equation is
```math
Hₓ(x(t), t)ẋ(t)+H_t(x(t), t) = 0
```
which follows from the fact ``d/dt H(x(t),t) ≡ 0 ∀ t∈[0,1]`` and the total derivative
of ``H`` w.r.t. ``t``.
"""
abstract type AbstractPredictor end
abstract type AbstractPredictorCache end

"""
    cache(::AbstractPredictor, ::HomotopyWithCache, x, ẋ, t)::AbstractPredictorCache

Construct a cache to avoid allocations.
"""
cache(p::AbstractPredictor, args...) = throw(MethodError(cache, tuple(p, args...)))

"""
    predict!(xnext, ::AbstractPredictor, ::AbstractPredictorCache, H::HomotopyWithCache, x, t, Δt, ẋ)

Perform a prediction step for the value of `x` with step size `Δt`. `ẋ` is the derivate of `x(t)` at `t`.
"""
function predict! end

"""
    update!(cache::AbstractStatefulPredictorCache, H, x, ẋ, t, fac)

Update the cache. `x` is the new path value at `t` and `ẋ` is the derivative at `t`.
`fac` is a factorization of the Jacobian at `(x,t)`.
"""
update!(::AbstractPredictorCache, H, x, ẋ, t, fac) = nothing

"""
    setup!(cache::AbstractStatefulPredictorCache, H, x, ẋ, t, Jac)

Setup the cache. `x` is the new path value at `t` and `ẋ` is the derivative at `t`.
`fac` is a factorization of the Jacobian at `(x,t)`. This falls back to calling `update`.
"""
setup!(C::AbstractPredictorCache, H, x, ẋ, t, Jac) = update!(C, H, x, ẋ, t, Jac)

"""
    reset!(cache, x, t)

Reset the cache. `x` is the path value at `t`.
"""
function reset!(::AbstractPredictorCache, x, t)
    nothing
end

"""
    order(::AbstractPredictor)

The order of a predictor is defined as the order of asymptotic error.
"""
function order end

# HELPERS

"""
    minus_ẋ!(out, H, x, t, J, dt)

Evaluate `Hₓ(x(t), t)⁻¹∂H∂t(x(t), t)` and store the result in `out`. `A` needs
to be able to store the Jacobian of `H`.
"""
function minus_ẋ!(ẋ, H, x, t, Jac::Jacobian, dt)
    jacobian_and_dt!(Jac.J, dt, H, x, t)
    updated_jacobian!(Jac)
    solve!(ẋ, Jac, dt)
end

include("predictors/null_predictor.jl")
include("predictors/euler.jl")
include("predictors/rk3.jl")
include("predictors/rk4.jl")
include("predictors/heun.jl")
include("predictors/midpoint.jl")
include("predictors/ralston.jl")
include("predictors/pade21.jl")


text:
struct ECE{B<:AbstractBinningAlgorithm,D} <: CalibrationErrorEstimator
    """Binning algorithm."""
    binning::B
    """Distance function."""
    distance::D
end

"""
    ECE(binning[, distance = TotalVariation()])

Create an estimator of the expected calibration error (ECE) with the given `binning`
algorithm and `distance` function.
"""
ECE(binning::AbstractBinningAlgorithm) = ECE(binning, TotalVariation())

# estimate ECE
function _calibrationerror(ece::ECE, predictions::AbstractMatrix{<:Real},
                           labels::AbstractVector{<:Integer})
    @unpack binning, distance = ece

    # bin predictions and labels
    bins = perform(binning, predictions, labels)

    # compute output type of estimates in each bin
    T = typeof(bin_eval_end(distance,
                            zero(result_type(distance, predictions, Int[])), 1))

    # compute weighted sum of distances
    d = zero(T)
    for bin in values(bins)
        d += bin_eval(bin, distance)
    end

    # normalize result
    inv(length(labels)) * d
end


text:
using Hexagons

immutable HexagonalBinGeometry <: Gadfly.GeometryElement
    default_statistic::Gadfly.StatisticElement
    tag::Symbol

    function HexagonalBinGeometry(
            default_statistic::Gadfly.StatisticElement;
            tag::Symbol=empty_tag)
        new(default_statistic, tag)
    end

    function HexagonalBinGeometry(; xbincount=200, ybincount=200, tag::Symbol=empty_tag)
        new(Gadfly.Stat.hexbin(xbincount=xbincount, ybincount=ybincount), tag)
    end
end

const hexbin = HexagonalBinGeometry

function default_statistic(geom::HexagonalBinGeometry)
    geom.default_statistic
end

function element_aesthetics(geom::HexagonalBinGeometry)
    [:x, :y, :xsize, :ysize, :color]
end

function render(geom::HexagonalBinGeometry, theme::Gadfly.Theme, aes::Gadfly.Aesthetics)
    default_aes = Gadfly.Aesthetics()
    default_aes.color = PooledDataArray(RGB{Float32}[theme.default_color])
    default_aes.xsize = [1.0]
    default_aes.ysize = [1.0]
    aes = inherit(aes, default_aes)

    Gadfly.assert_aesthetics_defined("Geom.hexbin", aes, :x, :y)
    Gadfly.assert_aesthetics_equal_length("Geom.hexbin", aes, :x, :y)

    n = length(aes.x)
    visibility = Bool[!isna(c) for c in takestrict(cycle(aes.color), n)]
    xs = aes.x[visibility]
    ys = aes.y[visibility]
    xsizes = collect(eltype(aes.xsize), takestrict(cycle(aes.xsize), n))[visibility]
    ysizes = collect(eltype(aes.ysize), takestrict(cycle(aes.ysize), n))[visibility]
    cs = collect(eltype(aes.color), takestrict(cycle(aes.color), n))[visibility]
    n = length(xs)

    return compose!(
        context(),
        Compose.polygon([hexpoints(xs[i], ys[i], xsizes[i], ysizes[i]) for i in 1:n], geom.tag),
        linewidth(0.1mm), # pad the hexagons so they ovelap a little
        fill(cs),
        stroke(cs),
        svgclass("geometry"))
end


text:
export RedPitayaDAQParams, RedPitayaDAQ, disconnect, setSlowDAC, getSlowADC, connectToServer,
       setTxParamsAll, disconnect
using RedPitayaDAQServer

@enum RPTriggerMode begin
  INTERNAL
  EXTERNAL
end

Base.@kwdef mutable struct RedPitayaDAQParams <: DAQParams
  "All configured channels of this DAQ device."
  channels::Dict{String, DAQChannelParams}

  "IPs of the Red Pitayas"
  ips::Vector{String}
  "Trigger mode of the Red Pitayas. Default: `EXTERNAL`."
  triggerMode::RPTriggerMode = EXTERNAL
  "Time to wait after a reset has been issued."
  resetWaittime::typeof(1.0u"s") = 45u"s"
  calibFFCurrentToVolt::Vector{Float32}
  calibIntToVolt::Array{Float32}
  ffRampUpFraction::Float32 = 1.0 # TODO RampUp + RampDown, could be a Union of Float or Vector{Float} and then if Vector [1] is up and [2] down
  ffRampUpTime::Float32 = 0.1 # and then the actual ramping could be a param of RedPitayaDAQ
  passPDMToFastDAC::Vector{Bool}
end

Base.@kwdef struct RedPitayaTxChannelParams <: TxChannelParams
  channelIdx::Int64
  limitPeak::typeof(1.0u"V")
  sinkImpedance::SinkImpedance = SINK_HIGH
  allowedWaveforms::Vector{Waveform} = [WAVEFORM_SINE]
  feedback::Union{DAQFeedback, Nothing} = nothing
  calibration::Union{typeof(1.0u"V/T"), Nothing} = nothing
end

Base.@kwdef struct RedPitayaLUTChannelParams <: DAQChannelParams
  channelIdx::Int64
  calibration::Union{typeof(1.0u"V/T"), typeof(1.0u"V/A"), Nothing} = nothing
end

"Create the params struct from a dict. Typically called during scanner instantiation."
function RedPitayaDAQParams(dict::Dict{String, Any})
  return createDAQParams(RedPitayaDAQParams, dict)
end

function createDAQChannels(::Type{RedPitayaDAQParams}, dict::Dict{String, Any})
  # TODO This is mostly copied from createDAQChannels, maybe manage to get rid of the duplication
  channels = Dict{String, DAQChannelParams}()
  for (key, value) in dict
    splattingDict = Dict{Symbol, Any}()
    if value["type"] == "tx"
      splattingDict[:channelIdx] = value["channel"]
      splattingDict[:limitPeak] = uparse(value["limitPeak"])

      if haskey(value, "sinkImpedance")
        splattingDict[:sinkImpedance] = value["sinkImpedance"] == "FIFTY_OHM" ? SINK_FIFTY_OHM : SINK_HIGH
      end

      if haskey(value, "allowedWaveforms")
        splattingDict[:allowedWaveforms] = toWaveform.(value["allowedWaveforms"])
      end

      if haskey(value, "feedback")
        channelID=value["feedback"]["channelID"]
        calibration=uparse(value["feedback"]["calibration"])

        splattingDict[:feedback] = DAQFeedback(channelID=channelID, calibration=calibration)
      end

      if haskey(value, "calibration")
        splattingDict[:calibration] = uparse.(value["calibration"])
      end

      if haskey(value, "passPDMToFastDAC")
        splattingDict[:passPDMToFastDAC] = value["passPDMToFastDAC"]
      end

      channels[key] = RedPitayaTxChannelParams(;splattingDict...)
    elseif value["type"] == "rx"
      channels[key] = DAQRxChannelParams(channelIdx=value["channel"])
    elseif value["type"] == "txSlow"
      calib = nothing
      if haskey(value, "calibration")
        calib = uparse.(value["calibration"])
      end
      channels[key] = RedPitayaLUTChannelParams(channelIdx=value["channel"], calibration = calib)
    end
  end

  return channels
end

Base.@kwdef mutable struct RedPitayaDAQ <: AbstractDAQ
  "Unique device ID for this device as defined in the configuration."
  deviceID::String
  "Parameter struct for this devices read from the configuration."
  params::RedPitayaDAQParams
  "Flag if the device is optional."
	optional::Bool = false
  "Flag if the device is present."
  present::Bool = false
  "Vector of dependencies for this device."
  dependencies::Dict{String, Union{Device, Missing}}

  "Reference to the Red Pitaya cluster"
  rpc::Union{RedPitayaCluster, Nothing} = nothing
  rpv::Union{RedPitayaClusterView, Nothing} = nothing

  rxChanIDs::Vector{String} = []
  refChanIDs::Vector{String} = []
  acqSeq::Union{Vector{AbstractSequence}, Nothing} = nothing
  samplesPerStep::Int32 = 0
  decimation::Int32 = 64
  samplingPoints::Int = 1
  sampleAverages::Int = 1
  acqPeriodsPerFrame::Int = 1
  acqPeriodsPerPatch::Int = 1
  acqNumFrames::Int = 1
  acqNumFrameAverages::Int = 1
  acqNumAverages::Int = 1
end

function init(daq::RedPitayaDAQ)
  @info "Initializing Red Pitaya DAQ with ID `$(daq.deviceID)`."

  # Restart the DAQ if necessary
  try
    daq.rpc = RedPitayaCluster(daq.params.ips)
  catch e
    if hasDependency(daq, SurveillanceUnit)
      su = dependency(daq, SurveillanceUnit)
      if hasResetDAQ(su)
        @info "Connection to DAQ could not be established! Restart (wait $(daq.resetWaittime) seconds...)!"
        resetDAQ(su)
        sleep(daq.resetWaittime)
        daq.rpc = RedPitayaCluster(daq.params.ips)
      else
        rethrow()
      end
    else
      @error "Error with Red Pitaya occured and the DAQ does not have access to a surveillance "*
             "unit for resetting it. Please check closely if this should be the case."
      rethrow()
    end
  end

  try
    daq.params.calibIntToVolt = reshape(daq.params.calibIntToVolt, 2, :)
  catch e
    @error e
  end

  #setACQParams(daq)
  masterTrigger(daq.rpc, false)
  triggerMode(daq.rpc, string(daq.params.triggerMode))
  ramWriterMode(daq.rpc, "TRIGGERED")
  modeDAC(daq.rpc, "STANDARD")
  #masterTrigger(daq.rpc, true)

  daq.present = true
end

neededDependencies(::RedPitayaDAQ) = []
optionalDependencies(::RedPitayaDAQ) = [TxDAQController, SurveillanceUnit]

Base.close(daq::RedPitayaDAQ) = daq.rpc

#### Sequence ####
function setSequenceParams(daq::RedPitayaDAQ, luts::Vector{Union{Nothing, Array{Float64}}}, enableLuts::Vector{Union{Nothing, Array{Bool}}})
  if length(luts) != length(daq.rpc)
    throw(DimensionMismatch("$(length(luts)) LUTs do not match $(length(daq.rpc)) RedPitayas"))
  end
  if length(enableLuts) != length(daq.rpc)
    throw(DimensionMismatch("$(length(enableLuts)) enableLUTs do not match $(length(daq.rpc)) RedPitayas"))
  end
  # Restrict to sequences of equal length, not a requirement of RedPitayaDAQServer, but of MPIMeasurements for simplicityacqSeq
  sizes = map(x-> size(x, 2) , filter(!isnothing, luts))
  if minimum(sizes) != maximum(sizes)
    throw(DimensionMismatch("LUTs do not have equal amount of steps"))
  end

  @info "Set sequence params"

  stepsPerRepetition = div(daq.acqPeriodsPerFrame, daq.acqPeriodsPerPatch)
  samplesPerSlowDACStep(daq.rpc, div(samplesPerPeriod(daq.rpc) * periodsPerFrame(daq.rpc), stepsPerRepetition))
  daq.samplesPerStep = samplesPerSlowDACStep(daq.rpc)
  clearSequence(daq.rpc)

  acqSeq = []
  for (i, rp) in enumerate(daq.rpc)
    lut = luts[i]
    enableLUT = enableLuts[i]
    if !isnothing(lut)
      numSlowDACChan(rp, size(lut, 1))
      @show lut
      #TODO IMPLEMENT SHORTER RAMP DOWN TIMING FOR SYSTEM MATRIX
      rpSeq = ArbitrarySequence(lut, enableLUT, stepsPerRepetition, daq.acqNumFrames*daq.acqNumFrameAverages,
                      computeRamping(daq.rpc, size(lut, 2), daq.params.ffRampUpTime, daq.params.ffRampUpFraction))
      push!(acqSeq, rpSeq)
      appendSequence(rp, rpSeq)
      # TODO enableLuts not yet implemented
    else
      # TODO What to do in this case, see maybe fill with zeros in other setSequenceParams
      # PauseSequence()
    end
  end

  daq.acqSeq = isempty(acqSeq) ? nothing : acqSeq

end
function setSequenceParams(daq::RedPitayaDAQ, sequence::Sequence)
  luts = Array{Union{Nothing, Array{Float64}}}(nothing, length(daq.rpc))
  enableLuts = Array{Union{Nothing, Array{Bool}}}(nothing, length(daq.rpc))

  lutChannels = [channel for channel in daq.params.channels if channel[2] isa RedPitayaLUTChannelParams]
  seqChannels = acyclicElectricalTxChannels(sequence)
  channelMapping = []
  for channel in seqChannels
    index = findfirst(x-> id(channel) == x[1], lutChannels)
    if !isnothing(index)
      push!(channelMapping, (lutChannels[index][2], channel))
    else
      throw(ScannerConfigurationError("No txSlow Channel defined for Field channel $(id(channel))"))
    end
  end

  @show channelMapping
  for rp in 1:length(daq.rpc)
    start = (rp - 1) * 4 + 1
    currentPossibleChannels = collect(start:start+3)
    currentMapping = [(lut, seq) for (lut, seq) in channelMapping if lut.channelIdx in currentPossibleChannels]
    if !isempty(currentMapping)
      lut = createLUT(start, currentMapping)
      luts[rp] = lut
    end
  end
  daq.acqPeriodsPerPatch = acqNumPeriodsPerPatch(sequence)
  @show luts
  setSequenceParams(daq, luts, enableLuts)
end

function createLUT(start, channelMapping)
  channelMapping = sort(channelMapping, by = x -> x[1].channelIdx)
  lutValues = []
  lutIdx = []
  for (lutChannel, seqChannel) in channelMapping
    tempValues = MPIFiles.values(seqChannel)
    if !isnothing(lutChannel.calibration)
      tempValues = tempValues.*lutChannel.calibration
    end
    tempValues = ustrip.(u"V", tempValues)
    push!(lutValues, tempValues)
    push!(lutIdx, lutChannel.channelIdx)
  end

  # Idx from 1 to 4
  lutIdx = (lutIdx.-start).+1
  # Fill skipped channels with 0.0, assumption: size of all lutValues is equal
  lut = zeros(Float32, maximum(lutIdx), size(lutValues[1], 1))
  for (i, lutIndex) in enumerate(lutIdx)
    lut[lutIndex, :] = lutValues[i]
  end
  return lut
end

function prepareSequence(daq::RedPitayaDAQ, sequence::Sequence)
  if !isnothing(daq.acqSeq)
    @info "Preparing sequence"
    success = RedPitayaDAQServer.prepareSequence(daq.rpc)
    if !success
      @warn "Failed to prepare sequence"
    end
  end
end

function endSequence(daq::RedPitayaDAQ, endFrame)
  sampPerFrame = samplesPerPeriod(daq.rpc) * periodsPerFrame(daq.rpc)
  endSample = (endFrame + 1) * sampPerFrame
  wp = currentWP(daq.rpc)
  # Wait for sequence to finish
  numQueries = 0
    while wp < endSample
      sampleDiff = endSample - wp
      waitTime = (sampleDiff / (125e6/daq.decimation))
      sleep(waitTime) # Queries are expensive, try to sleep to minimize amount of queries
      numQueries += 1
      wp = currentWP(daq.rpc)
  end
  stopTx(daq)
end

function getFrameTiming(daq::RedPitayaDAQ)
  # TODO How to signal end of sequences without any LUTs
  startSample = RedPitayaDAQServer.start(daq.acqSeq[1]) * daq.samplesPerStep
  startFrame = div(startSample, samplesPerPeriod(daq.rpc) * periodsPerFrame(daq.rpc))
  endFrame = div((length(daq.acqSeq[1]) * daq.samplesPerStep), samplesPerPeriod(daq.rpc) * periodsPerFrame(daq.rpc))
  return startFrame, endFrame
end

#### Producer/Consumer ####
mutable struct RedPitayaAsyncBuffer <: AsyncBuffer
  samples::Union{Matrix{Int16}, Nothing}
  performance::Vector{Vector{PerformanceData}}
end
AsyncBuffer(daq::RedPitayaDAQ) = RedPitayaAsyncBuffer(nothing, Vector{Vector{PerformanceData}}(undef, 1))

channelType(daq::RedPitayaDAQ) = SampleChunk

function updateAsyncBuffer!(buffer::RedPitayaAsyncBuffer, chunk)
  samples = chunk.samples
  perfs = chunk.performance
  push!(buffer.performance, perfs)
  if !isnothing(buffer.samples)
    buffer.samples = hcat(buffer.samples, samples)
  else
    buffer.samples = samples
  end
  for (i, p) in enumerate(perfs)
    if p.status.overwritten || p.status.corrupted
        @warn "RedPitaya $i lost data"
    end
end
end

function frameAverageBufferSize(daq::RedPitayaDAQ, frameAverages)
  return samplesPerPeriod(daq.rpc), length(daq.rxChanIDs), periodsPerFrame(daq.rpc), frameAverages
end

function startProducer(channel::Channel, daq::RedPitayaDAQ, numFrames)
  startFrame, endFrame = getFrameTiming(daq)
  startTx(daq)

  samplesPerFrame = samplesPerPeriod(daq.rpc) * periodsPerFrame(daq.rpc)
  startSample = startFrame * samplesPerFrame
  samplesToRead = samplesPerFrame * numFrames
  chunkSize = Int(ceil(0.1 * (125e6/daq.decimation)))

  rpu = daq.rpc
  if !isnothing(daq.rpv)
    rpu = daq.rpv
  end

  # Start pipeline
  @info "Pipeline started"
  try
    readPipelinedSamples(rpu, startSample, samplesToRead, channel, chunkSize = chunkSize)
  catch e
    @error e
    # TODO disconnect and reconnect to recover from open pipeline
    @info "Attempting reconnect to reset pipeline"
    daq.rpc = RedPitayaCluster(daq.params.ips)
    daq.rpv = nothing
  end
  @info "Pipeline finished"
  return endFrame
end

function convertSamplesToFrames!(buffer::RedPitayaAsyncBuffer, daq::RedPitayaDAQ)
  unusedSamples = buffer.samples
  samples = buffer.samples
  frames = nothing
  samplesPerFrame = samplesPerPeriod(daq.rpc) * periodsPerFrame(daq.rpc)
  samplesInBuffer = size(samples)[2]
  framesInBuffer = div(samplesInBuffer, samplesPerFrame)

  if framesInBuffer > 0
      samplesToConvert = view(samples, :, 1:(samplesPerFrame * framesInBuffer))
      chan = numChan(daq.rpc)
      rpu = daq.rpc
      if !isnothing(daq.rpv)
        rpu = daq.rpv
        chan = numChan(daq.rpv)
      end
      frames = convertSamplesToFrames(rpu, samplesToConvert, chan, samplesPerPeriod(daq.rpc), periodsPerFrame(daq.rpc), framesInBuffer, daq.acqNumAverages, 1)
      if (samplesPerFrame * framesInBuffer) + 1 <= samplesInBuffer
          unusedSamples = samples[:, (samplesPerFrame * framesInBuffer) + 1:samplesInBuffer]
      else
        unusedSamples = nothing
      end
  end

  buffer.samples = unusedSamples
  return frames
end

function retrieveMeasAndRef!(buffer::RedPitayaAsyncBuffer, daq::RedPitayaDAQ)
  frames = convertSamplesToFrames!(buffer, daq)
  uMeas = nothing
  uRef = nothing
  if !isnothing(frames)
    rxIds = channelIdx(daq, daq.rxChanIDs)
    refIds = channelIdx(daq, daq.refChanIDs)
    # Map channel index to their respective index in the view
    if !isnothing(daq.rpv)
      rxIds = map(x->clusterToView(daq.rpv, x), rxIds)
      refIds = map(x->clusterToView(daq.rpv, x), rxIds)
    end
    uMeas = frames[:,rxIds,:,:]
    uRef = frames[:, refIds,:,:]
  end
  return uMeas, uRef
end

#### Tx and Rx
function setup(daq::RedPitayaDAQ, sequence::Sequence)
  setupRx(daq, sequence)
  setupTx(daq, sequence)
end

function setupTx(daq::RedPitayaDAQ, sequence::Sequence)
  @info "Setup tx"
  periodicChannels = periodicElectricalTxChannels(sequence)

  if any([length(component.amplitude) > 1 for channel in periodicChannels for component in channel.components])
    error("The Red Pitaya DAQ cannot work with more than one period in a frame or frequency sweeps yet.")
  end

  # Iterate over sequence(!) channels
  for channel in periodicChannels
    channelIdx_ = channelIdx(daq, id(channel)) # Get index from scanner(!) channel

    offsetVolts = offset(channel)*calibration(daq, id(channel))
    offsetDAC(daq.rpc, channelIdx_, ustrip(u"V", offsetVolts))
    #jumpSharpnessDAC(daq.rpc, channelIdx_, daq.params.jumpSharpness) # TODO: Can we determine this somehow from the sequence?

    for (idx, component) in enumerate(components(channel))
      freq = ustrip(u"Hz", txBaseFrequency(sequence)) / divider(component)
      frequencyDAC(daq.rpc, channelIdx_, idx, freq)
    end

    # In the Red Pitaya, the signal type can only be set per channel
    waveform_ = unique([waveform(component) for component in components(channel)])
    if length(waveform_) == 1
      if !isWaveformAllowed(daq, id(channel), waveform_[1])
        throw(SequenceConfigurationError("The channel of sequence `$(name(sequence))` with the ID `$(id(channel))` "*
                                       "defines a waveforms of $waveform_, but the scanner channel does not allow this."))
      end
      waveform_ = uppercase(fromWaveform(waveform_[1]))
      signalTypeDAC(daq.rpc, channelIdx_, waveform_)
    else
      throw(SequenceConfigurationError("The channel of sequence `$(name(sequence))` with the ID `$(id(channel))` "*
                                       "defines different waveforms in its components. This is not supported "*
                                       "by the Red Pitaya."))
    end
  end

  pass = isempty(daq.params.passPDMToFastDAC) ? [false for i = 1:length(daq.rpc)] : daq.params.passPDMToFastDAC
  @show pass
  passPDMToFastDAC(daq.rpc, pass)

  #setSequenceParams(daq, sequence) # This might need to be removed for calibration measurement time savings
end

function setupRx(daq::RedPitayaDAQ)
  @info "Setup rx"
  decimation(daq.rpc, daq.decimation)
  samplesPerPeriod(daq.rpc, daq.samplingPoints * daq.acqNumAverages)
  periodsPerFrame(daq.rpc, daq.acqPeriodsPerFrame)
  #numSlowADCChan(daq.rpc, 4) # Not used as far as I know
end
function setupRx(daq::RedPitayaDAQ, sequence::Sequence)
  @assert txBaseFrequency(sequence) == 125.0u"MHz" "The base frequency is fixed for the Red Pitaya "*
  "and must thus be 125 MHz and not $(txBaseFrequency(sequence))."

  # The decimation can only be a power of 2 beginning with 8
  decimation_ = upreferred(txBaseFrequency(sequence)/rxSamplingRate(sequence))
  if decimation_ in [2^n for n in 3:8]
    daq.decimation = decimation_
  else
    throw(ScannerConfigurationError("The decimation derived from the rx bandwidth of $(rxBandwidth(sequence)) and "*
      "the base frequency of $(txBaseFrequency(sequence)) has a value of $decimation_ "*
      "but has to be a power of 2"))
  end

  daq.acqNumFrames = acqNumFrames(sequence)
  daq.acqNumFrameAverages = acqNumFrameAverages(sequence)
  daq.acqNumAverages = acqNumAverages(sequence)
  daq.samplingPoints = rxNumSamplesPerPeriod(sequence)
  daq.acqPeriodsPerFrame = acqNumPeriodsPerFrame(sequence)

  daq.rxChanIDs = []
  for channel in rxChannels(sequence)
    push!(daq.rxChanIDs, id(channel))
  end

  # TODO possibly move some of this into abstract daq
  daq.refChanIDs = []
  txChannels = [channel[2] for channel in daq.params.channels if channel[2] isa TxChannelParams]
  daq.refChanIDs = unique([tx.feedback.channelID for tx in txChannels if !isnothing(tx.feedback)])

  # Construct view to save bandwidth
  rxIDs = sort(union(channelIdx(daq, daq.rxChanIDs), channelIdx(daq, daq.refChanIDs)))
  selection = [false for i = 1:length(daq.rpc)]
  for i in map(x->div(x -1, 2) + 1, rxIDs)
    selection[i] = true
  end
  daq.rpv = RedPitayaClusterView(daq.rpc, selection)

  setupRx(daq)
end
function setupRx(daq::RedPitayaDAQ, decimation, numSamplesPerPeriod, numPeriodsPerFrame, numFrames; numFrameAverage=1, numAverages=1)
  daq.decimation = decimation
  daq.samplingPoints = numSamplesPerPeriod
  daq.acqPeriodsPerFrame = numPeriodsPerFrame
  daq.acqNumFrames = numFrames
  daq.acqNumFrameAverages = numFrameAverage
  daq.acqNumAverages = numAverages
  daq.rpv = nothing
  setupRx(daq)
end

# Starts both tx and rx in the case of the Red Pitaya since both are entangled by the master trigger.
function startTx(daq::RedPitayaDAQ)
  startADC(daq.rpc)
  masterTrigger(daq.rpc, true)
  @info "Started tx"
end

function stopTx(daq::RedPitayaDAQ)
  #setTxParams(daq, zeros(ComplexF64, numTxChannels(daq),numTxChannels(daq)))
  stopADC(daq.rpc)
  masterTrigger(daq.rpc, false)
  @info "Stopped tx"
  #RedPitayaDAQServer.disconnect(daq.rpc)
end

function prepareControl(daq::RedPitayaDAQ)
  clearSequence(daq.rpc)
end

function prepareTx(daq::RedPitayaDAQ, sequence::Sequence)
  stopTx(daq)
  @info "Preparing amplitude and phase"
  allAmps  = Dict{String, Vector{typeof(1.0u"V")}}()
  allPhases = Dict{String, Vector{typeof(1.0u"rad")}}()
  for channel in periodicElectricalTxChannels(sequence)
    name = id(channel)
    amps = []
    phases = []
    for comp in components(channel)
      # Lengths check == 1 happens in setupTx already
      amp = amplitude(comp)
      if dimension(amp) == dimension(1.0u"T")
        amp = (amp * calibration(daq, name))
      end
      push!(amps, amp)
      push!(phases, phase(comp))
    end

    allAmps[name] = amps

    allPhases[name] = phases

  end
  setTxParams(daq, allAmps, allPhases)
end

"""
Set the amplitude and phase for all the selected channels.

Note: `amplitudes` and `phases` are defined as a dictionary of
vectors, since every channel referenced by the dict's key could
have a different amount of components.
"""
function setTxParams(daq::RedPitayaDAQ, amplitudes::Dict{String, Vector{Union{Float32, Nothing}}}, phases::Dict{String, Vector{Union{Float32, Nothing}}})
  for (channelID, components_) in amplitudes
    channelVoltage = 0
    for amplitude_ in components_
      if !isnothing(amplitude_)
        channelVoltage += amplitude_
      end
    end

    if channelVoltage >= ustrip(u"V", limitPeak(daq, channelID))
      error("This should never happen!!! \nTx voltage on channel with ID `$channelID` is above the limit with a voltage of $channelVoltage.")
    end
  end

  for (channelID, components_) in phases
    for (componentIdx, phase_) in enumerate(components_)
      if !isnothing(phase_)
        phaseDAC(daq.rpc, channelIdx(daq, channelID), componentIdx, phase_)
      end
    end
  end

  for (channelID, components_) in amplitudes
    for (componentIdx, amplitude_) in enumerate(components_)
      if !isnothing(amplitude_)
        amplitudeDAC(daq.rpc, channelIdx(daq, channelID), componentIdx, amplitude_)
      end
    end
  end
end

function setTxParams(daq::RedPitayaDAQ, amplitudes::Dict{String, Vector{typeof(1.0u"V")}}, phases::Dict{String, Vector{typeof(1.0u"rad")}}; convolute=true)
  # Determine the worst case voltage per channel
  # Note: this would actually need a fourier synthesis with the given signal type,
  # but I don't think this is necessary
  for (channelID, components_) in amplitudes
    channelVoltage = 0
    for amplitude_ in components_
      channelVoltage += ustrip(u"V", amplitude_)
    end

    if channelVoltage >= ustrip(u"V", limitPeak(daq, channelID))
      error("This should never happen!!! \nTx voltage on channel with ID `$channelID` is above the limit with a voltage of $channelVoltage.")
    end
  end

  for (channelID, components_) in phases
    for (componentIdx, phase_) in enumerate(components_)
      phaseDAC(daq.rpc, channelIdx(daq, channelID), componentIdx, ustrip(u"rad", phase_))
    end
  end

  for (channelID, components_) in amplitudes
    for (componentIdx, amplitude_) in enumerate(components_)
        amplitudeDAC(daq.rpc, channelIdx(daq, channelID), componentIdx, ustrip(u"V", amplitude_))
    end
  end
end

currentFrame(daq::RedPitayaDAQ) = RedPitayaDAQServer.currentFrame(daq.rpc)
currentPeriod(daq::RedPitayaDAQ) = RedPitayaDAQServer.currentPeriod(daq.rpc)

function readData(daq::RedPitayaDAQ, startFrame::Integer, numFrames::Integer, numBlockAverages::Integer=1)
  u = RedPitayaDAQServer.readData(daq.rpc, startFrame, numFrames, numBlockAverages, 1, useCalibration = true)

  @info "size u in readData: $(size(u))"
  uMeas = u[:,channelIdx(daq, daq.rxChanIDs),:,:]u"V"
  uRef = u[:,channelIdx(daq, daq.refChanIDs),:,:]u"V"

  # lostSteps = numLostStepsSlowADC(master(daq.rpc))
  # if lostSteps > 0
  #   @error "WE LOST $lostSteps SLOW DAC STEPS!"
  # end

  @debug size(uMeas) size(uRef)

  return uMeas, uRef
end

function readDataPeriods(daq::RedPitayaDAQ, numPeriods, startPeriod, acqNumAverages)
  u = RedPitayaDAQServer.readDataPeriods(daq.rpc, startPeriod, numPeriods, acqNumAverages, useCalibration = true)

  uMeas = u[:,channelIdx(daq, daq.rxChanIDs),:]
  uRef = u[:,channelIdx(daq, daq.refChanIDs),:]

  return uMeas, uRef
end

numTxChannelsTotal(daq::RedPitayaDAQ) = numChan(daq.rpc)
numRxChannelsTotal(daq::RedPitayaDAQ) = numChan(daq.rpc)
numTxChannelsActive(daq::RedPitayaDAQ) = numChan(daq.rpc) #TODO: Currently, all available channels are active
numRxChannelsActive(daq::RedPitayaDAQ) = numRxChannelsReference(daq)+numRxChannelsMeasurement(daq)
numRxChannelsReference(daq::RedPitayaDAQ) = length(daq.refChanIDs)
numRxChannelsMeasurement(daq::RedPitayaDAQ) = length(daq.rxChanIDs)
numComponentsMax(daq::RedPitayaDAQ) = 4
canPostpone(daq::RedPitayaDAQ) = true
canConvolute(daq::RedPitayaDAQ) = false

######## OLD #########
function disconnect(daq::RedPitayaDAQ)
  RedPitayaDAQServer.disconnect(daq.rpc)
end


text:
include("postprocessing.jl")
include("anderson.jl")
include("dielectric_functions.jl")
using Plots.PlotMeasures

function setup_plots()
    gr()
    default(size=tuple(Int.(ceil.(0.75 .* [600, 400]))...),
            guidefontsize=10, grid=false)
end

function dump_latex_tables()
    write(joinpath(@__DIR__, "table_bulk.tex"),  table_bulk(df,  latex=true))
    write(joinpath(@__DIR__, "table_mixed.tex"), table_mixed(df, latex=true))
end

function plot_modes(case="Al_repeat_40", colorbar=false)
    h5f = h5open(joinpath(@__DIR__, case, "None_optimal_damping.hdf5"))

    # Read modes and average over X axes
    xlen = size(h5f["arnoldi/X"], 1)
    modes = [
        sum(h5f["arnoldi/X"][:, :, :,   i], dims=1)[1, :, :] / xlen
        for i in 1:size(h5f["arnoldi/X"], 4)
    ]
    λs = h5f["arnoldi/λ"][:]

    data = [
            modes[end-0];
            fill(0, 10, size(modes[end], 2));
            modes[end-1];
            fill(0, 10, size(modes[end], 2));
            modes[end-2];
            fill(0, 10, size(modes[end], 2));
            modes[end-3];
            fill(0, 10, size(modes[end], 2));
            modes[end-4];
           ]

    p = heatmap(data, c=:RdBu_11, aspect_ratio=1, showaxis=false, ticks=false,
                margin=0mm, clim=(-0.006, 0.006), grid=false, legend=colorbar)
    yaxis!(p, raw"$y$ direction")
    xaxis!(p, raw"$z$ direction")

    # Draw black border
    w = size(modes[end], 1) + 0.5
    h = size(modes[end], 2) + 0.5
    xs = [0.5, 0.5, h,   h, 0.5]
    ys = [0.5,   w, w, 0.5, 0.5]
    for a in 0:4
        plot!(p, xs, ys .+ a .* (w .+ 9.5), color=:black, legend=false, lw=0.5)
    end

    # Determine atomic (y, z) coordinates
    system, _, repeat = split(case, "_")
    lattice, atoms = load_system(system, parse(Int, repeat))
    positions = vcat(last.(atoms)...)
    cleanup(x) = x > 0.98 ? x - 1 : x
    atom_yzs = [(cleanup(pos[3]) * size(modes[end], 2),
                 cleanup(pos[2]) * size(modes[end], 1)) for pos in positions]
    for a in 0:4
        scatter!(p, first.(atom_yzs) .+ 1,
                 last.(atom_yzs) .+ a .* (w .+ 9.5),
                 m=:circle, color=:black, markersize=1.4, markeralpha=0.5,
                 markerstrokealpha=0.5)
    end

    p
end

function plot_ldos_density_2d(case; clim=(0, 0.15), aspect_ratio=2)
    h5f = h5open(joinpath(@__DIR__, case, "ldos.hdf5"))
    stemp = @sprintf "%.3f" read(h5f, "temperature_scf")

    # Average over X axes
    xlen = size(h5f["ldos_$stemp"], 1)
    data = [
        sum(h5f["ldos_$stemp"][:, :, :], dims=1)[1, :, :] / xlen;
        fill(0, 10, size(h5f["ldos_$stemp"], 3));
        sum(h5f["ρ_real"][:, :, :], dims=1)[1, :, :] / xlen;
    ]

    p = heatmap(data, c=cgrad(:devon, rev=true), aspect_ratio=aspect_ratio,
                showaxis=false, ticks=false, size=(450, 200),
                margin=0mm, grid=false, clim=clim)
    yaxis!(p, raw"$y$ direction")
    xaxis!(p, raw"$z$ direction")

    # Draw black border
    w = size(h5f["ldos_$stemp"], 2) + 0.5
    h = size(h5f["ldos_$stemp"], 3) + 0.5
    xs = [0.5, 0.5, h,   h, 0.5]
    ys = [0.5,   w, w, 0.5, 0.5]
    plot!(p, xs, ys, color=:black, legend=false, lw=0.5)
    plot!(p, xs, ys .+ w .+ 9.5,
          color=:black, legend=false, lw=0.5)

    p
end

function plot_ldos_density_1d(case; boxes=[], temperatures=nothing, flip=false,
                              legend=:topright, letter=nothing)
    h5f = h5open(joinpath(@__DIR__, case, "ldos.hdf5"))
    isnothing(temperatures) && (temperatures = [read(h5f, "temperature_scf")])

    # Average over X axes
    xlen = size(h5f["ρ_real"], 1)
    ylen = size(h5f["ρ_real"], 2)
    rho_avg  = sum(sum(h5f["ρ_real"][:, :, :], dims=1), dims=2)[1, 1, :] / xlen / ylen;

    bmin = Inf
    bmax = -Inf
    ldos_avg = Vector{Any}(undef, length(temperatures))
    for (i, temp) in enumerate(temperatures)
        stemp = @sprintf "%.3f" temp
        ldos_avg[i] = sum(sum(h5f["ldos_$stemp"][:, :, :],
                              dims=1), dims=2)[1, 1, :] / xlen / ylen;
        bmin = min(bmin, min(minimum(ldos_avg[i]), minimum(rho_avg)))
        bmax = max(bmax, max(maximum(ldos_avg[i]), maximum(rho_avg)))
    end

    # Determine atomic positions for tick placement
    system, _, repeat = split(case, "_")
    lattice, atoms = load_system(system, parse(Int, repeat))
    positions = vcat(last.(atoms)...)
    ticks = [pos[3] * size(h5f["ρ_real"], 3) for pos in positions]

    p = plot(;legend=legend, xticks=ticks, xformatter=x -> "", grid=false, )
    for (s, e) in boxes
        c=:grey85
        plot!(p,
              [s,    s,    e,    e,    s] .* length(rho_avg),
              [bmin, bmax, bmax, bmin, bmin],
              fill=(0, c), label="", line=c)
    end
    lss = append!([:solid, :dot, :dash, :dashdot], fill(:solid, length(temperatures)))
    cols = [3, :orangered1, 4, 5, 6, 7, 8, 9]
    for (i, temp) in reverse(collect(enumerate(temperatures)))
        stemp = @sprintf "%.3f" temp
        label = "LDOS" * (length(temperatures) > 1 ? " (T=$stemp)" : "")
        plot!(p, ldos_avg[i], label=label, color=cols[i], ls=lss[i], lw=1.8)
    end
    plot!(p, rho_avg, label="ρ", color=1, ls=:solid, lw=3.5)
    xaxis!(p, "z", xflip=flip)

    if !isnothing(letter)
        start = flip ? length(rho_avg) - 5 : 3
        annotate!(p, start, 1.05bmax, Plots.text(letter, 10, :left))
    end

    p
end

const NAMEMAP = Dict(
    "RestaSiO2"  => raw"Dielectric (SiO₂)",
    "RestaGaAs"  => raw"Dielectric (GaAs)",
    "HybridGaAs" => "LDOS+Dielectric",
    "HybridVac"  => "LDOS",
    "Custom"     => "Localised",
    "TF"         => "TFW",
)

function plot_convergence(case; namemap=NAMEMAP, iteration=50, tol=1e-10,
                          preconditioners=["None", "Kerker"],
                          ylabel="Total energy absolute error",
                          show_alpha=false, ymax=1e3, letter=nothing)
    style = Dict(
        # both
        "None"          => (color=4, linestyle=:dot,   mark=:cross, markerstrokecolor=4),
        "Kerker"        => (color=2, linestyle=:solid, mark=:xcross, markerstrokecolor=2),
        # bulk
        "RestaSiO2"     => (color=0, linestyle=:solid, mark=:utriangle, markerstrokecolor=0),
        "RestaGaAs"     => (color=1, linestyle=:dash,  mark=:dtriangle, markerstrokecolor=1),
        # compare
        "HybridVac"     => (color=1, linestyle=:solid, mark=:utriangle, markerstrokecolor=1),
        "TF"            => (color=0, linestyle=:dash,  mark=:dtriangle, markerstrokecolor=0),
    )

    # First pass: For DFTK runs collect mimimal total energy
    directory = joinpath(@__DIR__, case)
    mintotal = collect_minimal_total_energy(directory, style)

    # Add DFTK errors
    errors = Dict{String, Vector{Float64}}()
    for mixkey in preconditioners
        resultfile = joinpath(directory, mixkey * ".json")
        isfile(resultfile) || continue

        energies = open(JSON.parse, resultfile, "r")["energies"]
        energies === nothing && continue
        values = abs.(energies .- mintotal)
        idx = findfirst(values .< 1e-16)
        isnothing(idx) || (values = values[1:idx-1])
        isempty(values) && continue

        errors[mixkey] = values
    end

    # Add Quantum-Espresso errors
    qe_fn = joinpath(directory, "TF/espresso.pwo")
    if isfile(qe_fn) && "TF" in preconditioners
        errors["TF"] = [p.accuracy for p in parse_pwo(qe_fn)]
    end

    # If no data return
    isempty(errors) && return plot()

    maxiter = 0
    p = plot(legend=:topright, grid=false)
    for mixkey in preconditioners
        dampfile = joinpath(directory, mixkey * ".optimal_damping.json")
        αopt = isfile(dampfile) ? open(JSON.parse, dampfile, "r")["optimal_damping"] : NaN
        mixkey == "TF" && (αopt = 0.7)  # QE calculations all use α=0.7

        maxiter = max(maxiter, length(errors[mixkey]))
        label = get(NAMEMAP, mixkey, mixkey)
        show_alpha && (label=label * " (α=$αopt)")
        plot!(p, errors[mixkey]; label=label, markersize=4, style[mixkey]..., lw=1.5)
    end
    xlims!(p, (0, min(iteration, maxiter)))
    ylims!(p, (tol, ymax))
    yaxis!(p, ylabel, :log10)
    xaxis!(p, "Iteration")
    if !isnothing(letter)
        pos = ymax / 10 * 9
        plot!(p, [1.1, 12], [pos, pos], label="", line=(0.3, :white), lw=22)
        annotate!(p, 1.2, pos, Plots.text(letter, 10, :left))
    end

    p
end

function main()
    setup_plots()

    for (letter, sys) in [("(b) Al+vacuum", "AlVac_repeat_10"), ("(a) Al", "Al_repeat_12"),
                          ("(c) GaAs", "GaAs_repeat_20")]
        savefig(plot_anderson("$(sys)_anderson", letter=letter), "$(sys)_anderson.pdf")
    end

    # Plots of dielectric functions
    savefig(plot_dielectric_functions(), "dielectric_functions.pdf")
    pgfplotsx(); savefig(plot_modes(), "aluminium_modes.pdf"); gr()

    # LDOS / density plots
    # savefig(plot_ldos_density_2d("AlSiO2H_repeat_20"),              "AlSiO2_ldos.pdf")
    # savefig(plot_ldos_density_2d("AlVac_repeat_20", clim=(0, 0.2)), "AlVac_ldos.pdf")
    savefig(plot_ldos_density_1d("AlSiO2H_repeat_20", boxes=[(0.395, 0.995),],
                                 temperatures=[0.001, 0.032], flip=true,
                                 letter="(b) Al+SiO₂"),
            "AlSiO2_ldos_1d.pdf")
    savefig(plot_ldos_density_1d("AlVac_repeat_20", boxes=[(0, 0.516), (0.99, 1.0)],
                                 temperatures=[0.001, 0.032], letter="(a) Al+vacuum"),
            "AlVac_ldos_1d.pdf")

    # Convergence plots
    ldos_vs_tf_precs = ["None", "HybridVac", "Kerker", "TF"]
    savefig(plot_convergence("AlVac_repeat_20", preconditioners=ldos_vs_tf_precs,
                             ylabel="Estimated SCF error", letter="(a) Al+vacuum"),
            "AlVac_repeat_20.pdf")
    savefig(plot_convergence("AlSiO2H_repeat_20", preconditioners=ldos_vs_tf_precs,
                             ylabel="Estimated SCF error", letter="(b) Al+SiO₂"),
            "AlSiO2H_repeat_20.pdf")

    bulk_precs = ["None", "RestaSiO2", "RestaGaAs", "Kerker"]
    savefig(plot_convergence("SiO2_repeat_39", preconditioners=bulk_precs, letter="(a) SiO₂"),
            "SiO2_repeat_39.pdf")
    savefig(plot_convergence("GaAs_repeat_40", preconditioners=bulk_precs, letter="(b) GaAs"),
            "GaAs_repeat_40.pdf")
    savefig(plot_convergence("Al_repeat_40", preconditioners=bulk_precs, letter="(c) Al"),
            "Al_repeat_40.pdf")
end

(abspath(PROGRAM_FILE) == @__FILE__) && main()


text:
using Plots, RDatasets, Query, DataFrames, StatPlots

gr()

iris = dataset("datasets", "iris")

function SplitDFbySpecies(iris, species)
    df = @from i in iris begin
         @where i.Species == species
         @select {i.PetalLength, i.PetalWidth}
         @collect DataFrame
    end
    return df
end

function dfmean(df::DataFrame)
    return [mean(i) for i in df.columns]
end

function dfcov(df)
    cols = df.columns
    ncols = size(cols)[1]
    ndatas = size(cols[1])[1]
    array = zeros(ncols, ncols)
    mean = dfmean(df)
    for i in 1:ncols
        for j in 1:ncols
            array[i, j] = 1/ndatas * (cols[i] .- mean[i])' * (cols[j] .- mean[j])
        end
    end
    return array
end

function lintf(df, S)
    groups = Array(df[:Species])
    petals = df[[:PetalLength, :PetalWidth]]

    colnames = names(petals)
    values = Array(petals)
    values = values * S
    tf_df = DataFrame()
    tf_df[colnames[1]] = values[:, 1]
    tf_df[colnames[2]] = values[:, 2]
    tf_df[:Species] = groups

    return tf_df
end

function swap(array)
    temp = zeros(2, 2)
    temp[:, 1] = array[:, 2]
    temp[:, 2] = array[:, 1]

    return temp
end

@df iris scatter(:PetalLength, :PetalWidth, group=:Species,
                 m=(0.5, [:+ :h :star7], 12), bg=RGB(.2,.2,.2))

scatter!(xlabel="Petal Length", ylabel="Petal Width", title="No Decorrelation",
         xlims=(1, 7), ylims=(-1, 4))
savefig("../figures/no_decorrelation.png")

coviris = dfcov(iris[[:PetalLength, :PetalWidth]])
S = -swap(eigvecs(coviris))

decorrelated = lintf(iris, S)
covdecorrelated = dfcov(decorrelated[[:PetalLength, :PetalWidth]])

@df decorrelated scatter(:PetalLength, :PetalWidth, group=:Species,
                         m=(0.5, [:+ :h :star7], 12), bg=RGB(.2,.2,.2))
scatter!(xlabel="Petal Length", ylabel="Petal Width", title="Decorrelated",
         xlims=(1, 7), ylims=(-2, 3))
savefig("../figures/decorrelation.png")


text:
# This code is part of QuantumCircuits.
#
# (C) Copyright Rafał Pracht 2022.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

using Test

using QuantumCircuits
using QuantumCircuits.Circuit: getCode

using QuantumCircuits.Circuit # TODO

# Expected circuits
expc = QCircuit(2)
add!(expc, QuantumCircuits.Circuit.X, 0)
add!(expc, QuantumCircuits.Circuit.H, 1)
add!(expc, QuantumCircuits.Circuit.CX, 1, 0)

# Get the circuit from the only expression
acirc = QCircuit(2)
acirc.x(0)
acirc.h(1)
acirc.cx(1, 0)
acirc == expc

# capture the local variable
i=0
j=1
acirc = QCircuit(2)
acirc.x(i)
acirc.h(j)
acirc.cx(j, i)
@test acirc == expc
@test acirc.qubits == 2

######################################################################################
expc = QCircuit(5)
n = expc.qubits-1
add!(expc, QuantumCircuits.Circuit.H, 0:n)
foreach(i -> add!(expc, QuantumCircuits.Circuit.CX, i, i+1), 0:n-1)

# With loops
acirc = QCircuit(5)
n = acirc.qubits - 1
foreach(i -> acirc.h(i), 0:n)
foreach(i -> acirc.cx(i, i+1), 0:n-1)
@test acirc == expc

######################################################################################

# Expected circuits
expc = QCircuit(2)
add!(expc, QuantumCircuits.Circuit.Sx, 0)
add!(expc, QuantumCircuits.Circuit.Y, 1)
add!(expc, QuantumCircuits.Circuit.Z, 0)
add!(expc, QuantumCircuits.Circuit.S, 1)
add!(expc, QuantumCircuits.Circuit.Sd, 0)
add!(expc, QuantumCircuits.Circuit.T, 1)
add!(expc, QuantumCircuits.Circuit.Td, 0)
add!(expc, QuantumCircuits.Circuit.U, 1, π/5, π/6, π/7)
add!(expc, QuantumCircuits.Circuit.U3, 0, π/5, π/6, π/7)
add!(expc, QuantumCircuits.Circuit.Rx, 1, π/2)
add!(expc, QuantumCircuits.Circuit.Ry, 0, π/3)
add!(expc, QuantumCircuits.Circuit.Rz, 1, π/4)
add!(expc, QuantumCircuits.Circuit.Rzx, 0, 1, π/2)
add!(expc, QuantumCircuits.Circuit.U4, 0, 1, [i for i in 1:15])

# @gate
acirc = QCircuit(2)
acirc.sx(0)
acirc.y(1)
acirc.z(0)
acirc.s(1)
acirc.sdg(0)
acirc.t(1)
acirc.tdg(0)
acirc.u(1, π/5, π/6, π/7)
acirc.u3(0, π/5, π/6, π/7)
acirc.rx(1, π/2)
acirc.ry(0, π/3)
acirc.rz(1, π/4)
acirc.rzx(0, 1, π/2)
acirc.u4(0, 1, [i for i in 1:15])
@test acirc == expc

# Test if we can property createt with random parameters
acirc = QCircuit(2)
acirc.u3(0)
acirc.u4(0, 1)
@test length(getCode(acirc)) == 2
@test typeof(getCode(acirc)[1]) == QuantumCircuits.Circuit.U3
@test typeof(getCode(acirc)[2]) == QuantumCircuits.Circuit.U4


text:
using Base:
    IOError, UV_EEXIST, UV_ESRCH,
    Process

using Base.Filesystem:
    File, open, JL_O_CREAT, JL_O_RDWR, JL_O_RDONLY, JL_O_EXCL,
    samefile

using Random

## SECTION - Becnhmarks compositions
function compositions_benchmark(; clear_results=false)

    comps = Dict{Any,Any}()
    clear_results && rm(datadir("composition_results"); recursive=true, force=true)
    mkpath(datadir("composition_results"))
    number_of_compositions = 0
    symbols_dict = Dict{String,Int64}()
    n_compositions_files = length(readdir(datadir("compositions")))
    # multiplyin by 2 because for each dom_size used in training, we use 2 for testing
    #configs_channel = Channel{Int}(n_compositions_files*2);
    #search_space_SA = SharedVector{Int}(n_compositions_files*2)

    function aux(file_name)
        if startswith(file_name, "con=")
            json = JSON.parsefile(joinpath(datadir("compositions"), file_name))
            counter = 1
            while (haskey(json, string(counter)))
                @info "composition №: $counter, worker: $(myid())"

                dom_size = json["params"]["domains_size"]

                fn1, fn2 = generate_file_name(json, counter, symbols_dict, dom_size)
                path1, path2 = joinpath(datadir("composition_results"),fn1),
                                joinpath(datadir("composition_results"),fn2)

                if isfile(path1) && isfile(path2)
                    @warn "The results for this composition already exist"
                else
                    path1, path2 = tryopen_exclusive(path1), tryopen_exclusive(path2)

                    n_symbols, memoize, population, generations, icn_iterations,
                    partial_search_limit, icn_time, maths, temp_concept, metric,
                    comp, selection_rate, dom_size, search, complete_search_limit,
                    solutions_limit, param = extract_data_from_json(
                        json, counter
                    )

                    @warn "retrieved data for comp: " maths

                    concept = Constraints.concept(BENCHED_CONSTRAINTS[temp_concept])

                    @warn "initial dom_size" dom_size
                    for i in [1,100-dom_size]
                        dom_size = dom_size + i
                        @warn "testing against dom_size" dom_size

                        solutions, non_sltns, _ = search_space(
                            #search_space_SA,
                            dom_size,
                            concept,
                            param;
                            #leave search param flexible by default,
                            complete_search_limit,
                            solutions_limit,
                        )

                        #write_configs(configs_channel)

                        @info "search space retrieved"

                        icn_composition_string = comp[findfirst("function", comp)[1]:end]
                        icn_composition = eval(Meta.parse(icn_composition_string))
                        timed_loss = @timed loss(
                            solutions,
                            non_sltns,
                            icn_composition,
                            eval(Meta.parse(metric)),
                            dom_size,
                            param;
                            samples=solutions_limit,
                        )

                        results = timed_loss.value[1]
                        sol_length = timed_loss.value[2]
                        normalised_results = normalise(Symbol(metric), results, dom_size)
                        push!(comps, "dom_size" => dom_size)
                        push!(comps, "icn_time" => icn_time)
                        push!(comps, "search" => search)
                        push!(comps, "concept" => temp_concept)
                        push!(comps, "complete_search_limit" => complete_search_limit)
                        push!(comps, "memoize" => memoize)
                        push!(comps, "sampling" => solutions_limit)
                        push!(comps, "population" => population)
                        push!(comps, "generations" => generations)
                        push!(comps, "icn_iterations" => icn_iterations)
                        push!(comps, "partial_search_limit" => partial_search_limit)
                        push!(comps, "maths" => maths)
                        push!(comps, "composition" => comp)
                        push!(comps, "composition_number" => counter)
                        push!(comps, "selection_rate" => selection_rate)
                        push!(comps, "time" => timed_loss.time / sol_length)
                        push!(comps, "accuracy" => results)
                        push!(comps, "normalised" => normalised_results)
                        try
                        push!(comps, "mean" => mean(normalised_results))
                        catch e
                            @info "########################## DEBUG ##########################"
                            println(normalised_results)
                            @info "########################## DEBUG ##########################"
                            println(results)
                            @info "########################## DEBUG ##########################"
                            println(timed_loss)
                            @info "########################## DEBUG ##########################"
                        end
                        push!(comps, "median" => median(normalised_results))
                        push!(comps, "std" => std(normalised_results; corrected=false))
                        push!(comps, "var" => var(normalised_results; corrected=false))
                        push!(comps, "cov" => cov(normalised_results; corrected=false))
                        push!(comps, "symbols_count" => n_symbols)
                        i ==1 ? export_compositions(comps, path1) : export_compositions(comps, path2)
                        # export_csv(
                        #     comps, joinpath(datadir("composition_results"), "results.csv")
                        # )
                    end
                end
                counter += 1
            end
            number_of_compositions += counter - 1
            @info "number of compositions processed: " number_of_compositions
        end
    end

    #pmap(aux, shuffle(readdir(datadir("compositions"))))
    map(aux, Random.shuffle(readdir(datadir("compositions"))))

    export_symbols_dict(symbols_dict)

    return nothing
end

## SECTION - Normalisation

# Dispatcher for normalised results
normalise(metric, results, dom_size) = normalise(results, dom_size, Val(metric))

# Division by #variable (=dom_size) * dom_size
normalise(results, dom_size, ::Val{:manhattan}) = results / dom_size^2

# Division by #variable (=dom_size)
normalise(results, dom_size, ::Val{:hamming}) = results / dom_size

## SECTION - Exports
function export_compositions(comps, path)

    if (!isnothing(path))
        write(path, JSON.json(comps, 2))
    end

    return nothing
end

function export_symbols_dict(symbols_dict::Dict{String,Int64})
    path = joinpath(datadir("composition_results"), "symbols_dict.csv")
    return CSV.write(path, DataFrame(symbols_dict))
end

# This function must be very bad performance wise, please tell me if you have better solutions :-)
function export_csv(comps, path)
    temp_dict = Dict{String,String}()
    delete!(comps, "composition")
    for column_name in collect(keys(comps))
        if isequal(column_name, "accuracy") || isequal(column_name, "normalised")
            value_str = "["
            for value in comps[column_name]
                value_str = string(value_str, string(value), ", ")
            end
            value_str = string(value_str[1:(end - 2)], "]")
            temp_dict[column_name] = value_str
        else
            temp_dict[column_name] = string(comps[column_name])
        end
    end

    df = DataFrame(temp_dict)

    if isfile(path)
        CSV.write(path, df; append=true)
    else
        CSV.write(path, df)
    end
end

## SECTION - Generation of elements

function generate_param(param, dom_size)
    # Generate an appropriate parameter for the concept if relevant
    param = if isnothing(param)
        nothing
    elseif param == 1
        rand(1:dom_size)
    else
        rand(1:dom_size, param)
    end
    return param
end

function generate_file_name(file, counter, symbols_dict, dom_size)
    concept = file["params"]["concept"][1]
    symbols = file[string(counter)]["symbols"]

    # sort transformation symbols to avoid duplicate compositions
    sort!(symbols[1])

    file_name = concept
    for symbol in reverse(symbols)[1:(end - 1)]
        if symbol[1] ∉ keys(symbols_dict)
            push!(symbols_dict, symbol[1] => length(symbols_dict) + 1)
        end
        file_name = string(file_name, "_", symbols_dict[symbol[1]])
    end

    for transformation in reverse(symbols[1])
        if transformation ∉ keys(symbols_dict)
            push!(symbols_dict, transformation => length(symbols_dict) + 1)
        end
        file_name = string(file_name, "_", symbols_dict[transformation])
    end
    path1 =  string(file_name,"_dom_size_" ,string(dom_size+1),".json")
    path2 =  string(file_name, "_dom_size_100",".json")
    return path1, path2
end

## SECTION - Other utility functions

function extract_data_from_json(file, counter)
    symbols_count = length(file[string(counter)]["symbols"][1]) + 3
    icn_time = file["icn_time"]
    memoize = file["params"]["memoize"]
    population = file["params"]["population"]
    generations = file["params"]["generations"]
    icn_iterations = file["params"]["icn_iterations"]
    search = file["params"]["search"]
    partial_search_limit = file["params"]["partial_search_limit"]
    concept = Symbol(file["params"]["concept"][1])
    maths = file[string(counter)]["maths"]
    metric = file["params"]["metric"]
    comp = file[string(counter)]["Julia"]
    selection_rate = file[string(counter)]["selection_rate"]
    dom_size = file["params"]["domains_size"]
    search = eval(Meta.parse(":" * file["params"]["search"]))
    complete_search_limit = file["params"]["complete_search_limit"]
    solutions_limit = file["params"]["sampling"]
    param = generate_param(file["params"]["concept"][2], dom_size)

    return symbols_count,
    memoize,
    population,
    generations,
    icn_iterations,
    partial_search_limit,
    icn_time,
    maths,
    concept,
    metric,
    comp,
    selection_rate,
    dom_size,
    search,
    complete_search_limit,
    solutions_limit,
    param
end

# Calculate the loss of a composition
function loss(solutions, non_sltns, composition, metric, dom_size, param; samples=nothing)
    @info "started calculating loss"
    @warn "solutions: " solutions
    @warn "non_sltns: " non_sltns
    l = length(solutions)
    X = if isnothing(samples) || length(non_sltns) < samples
        l += length(non_sltns)
        Iterators.flatten((solutions, non_sltns))
    else
        l += samples
        Iterators.flatten((solutions, rand(non_sltns, samples)))
    end

    @warn "X: " X

    result = map(
        x -> abs(Base.invokelatest(composition, x; param, dom_size) - metric(x, solutions)),
        X,
    )

    @info "loss calculated"

    for m in methods(composition)
        Base.delete_method(m)
    end
    return result, l
end

function count_compositions()
    total_compositions = 0
    unique_compositions = 0
    existing_comps = []
    for file_name in cd(readdir, joinpath(datadir("compositions")))
        if startswith(file_name, "con=")
            json = JSON.parsefile(joinpath(datadir("compositions"), file_name))
            unique_counter = 1
            counter = 1
            symbol_number = 1
            while (haskey(json, string(counter)))
                path = joinpath(
                    datadir("composition_results"),
                    #generate_file_name(json, counter, symbols_dict),
                )
                if path ∉ existing_comps
                    unique_counter += 1
                    push!(existing_comps, path)
                end
                counter += 1
            end
            total_compositions += counter - 1
            unique_compositions += unique_counter
        end
        @info "current number of compositions: " total_compositions
        @info "current number of unique compositions: " unique_compositions
    end
    @info "total number of compositions: " total_compositions
    @info "total number of unique compositions: " unique_compositions
end

## SECTION - Unused

# relative standard deviation
rsd(results) = std(results; corrected=false) / mean(results)

# Helper function to open files in O_EXCL mode
function tryopen_exclusive(path::String, mode::Integer = 0o666)
    try
        return open(path, JL_O_RDWR | JL_O_CREAT | JL_O_EXCL, mode)
    catch ex
        (isa(ex, IOError) && ex.code == UV_EEXIST) || rethrow(ex)
    end
    return nothing
end

# TODO:: create exclusive
# TODO:: cleanup script
# TODO:: add dom_size to comp benchmark file name


text:
module Types

using DemeNet: Certificate, Contract, Intent, Consensus, AbstractID, ID, DemeID, Deme
using PeaceVote: AbstractVote, AbstractProposal, AbstractBraid
using Base: UUID

import Base.Dict

struct Vote <: AbstractVote
    pid::Int ### One gets it from a BraidChain loking into a sealed proposal
    vote::Int ### Number or a message
end

struct Proposal <: AbstractProposal
    msg::AbstractString
    options::Vector{T} where T<:AbstractString
end

import Base.==
==(a::Proposal,b::Proposal) = a.msg==b.msg && a.options==b.options

struct Braid <: AbstractBraid
    index::Union{Nothing,Int} ### latest index of the ledger
    hash::Union{Nothing,BigInt} ### hash of the ledger up to the latest index
    ids::Vector{ID} ### the new ids for the public keys
end

function Dict(p::Proposal)
    dict = Dict()
    dict["msg"] = p.msg
    dict["options"] = p.options
    return dict
end

function Proposal(dict::Dict)
    msg = dict["msg"]
    options = dict["options"]
    return Proposal(msg,options)
end

function Dict(v::Vote)
    dict = Dict()
    dict["pid"] = v.pid
    dict["vote"] = v.vote
    return dict
end

function Vote(dict::Dict)
    pid = dict["pid"]
    vote = dict["vote"]
    return Vote(pid,vote)
end

function Dict(braid::Braid)
    dict = Dict()
    dict["ids"] = [string(i.id,base=16) for i in braid.ids]
    return dict
end

function Braid(dict::Dict)
    ids = ID[ID(parse(BigInt,i,base=16)) for i in dict["ids"]] ### Any until I fix SynchronicBallot
    return Braid(nothing,nothing,ids)
end

end


text:
export RandomCograph, is_cograph, CographRepresentation

"""
`RandomCograph(vtcs)` creates a random cograph with vertices
in `vtcs`.

`RandomCograph(n)` creates a random cograph with vertex set `1:n`.
"""
function RandomCograph(vtcs::Vector{T}) where {T}
    G = SimpleGraph{T}()

    # basis cases
    if length(vtcs) == 0
        return G
    end
    if length(vtcs) == 1
        add!(G, vtcs[1])
        return G
    end

    # split the vertex set
    alist = Vector{T}()
    blist = Vector{T}()
    while length(alist) == 0 || length(blist) == 0
        alist = Vector{T}()
        blist = Vector{T}()
        for v in vtcs
            if rand() < 0.5
                push!(alist, v)
            else
                push!(blist, v)
            end
        end
    end

    # recurse
    G = RandomCograph(alist)
    H = RandomCograph(blist)

    # copy H into G
    for v in blist
        add!(G, v)
    end
    for e in elist(H)
        add!(G, e[1], e[2])
    end

    # final coin flip
    if rand() < 0.5
        for a in alist
            for b in blist
                add!(G, a, b)
            end
        end
    end
    cache_save(G, :name, "Cograph")
    return G
end

RandomCograph(n::Int) = RandomCograph(collect(1:n))
RandomCograph(A::BitSet) = RandomCograph(collect(A))
RandomCograph(A::Set) = RandomCograph(collect(A))

"""
`is_cograph(G)` returns `true` is `G` is a cograph, i.e., a
complement reducible graph.
"""
function is_cograph(G::SimpleGraph{T}) where {T}
    if NV(G) <= 3    # all graphs with 3 or fewer vertices are cographs
        return true
    end
    if is_connected(G)
        GG = G'
        if is_connected(GG)
            return false
        end
        return is_cograph(GG)
    end

    clist = parts(components(G))
    for A in clist
        H = induce(G, A)
        if !is_cograph(H)
            return false
        end
    end
    return true
end

"""
`CographRepresentation(G::SimpleGraph)` returns a `String` that
presents a construction of `G` using disjoint union `+` and join `*`.
Throws an error if `G` is not a cograph.
"""
function CographRepresentation(G::SimpleGraph)
    if cache_check(G, :CographRepresentation)
        return cache_recall(G, :CographRepresentation)
    end

    if NV(G) == 0
        return "()"
    end
    if NV(G) == 1
        v = first(G.V)
        return "$v"
    end

    if is_connected(G)
        comps = collect(parts(components(G')))
        ncomps = length(comps)
        if ncomps == 1
            error("This is not a cograph")
        end
        result = ""
        for k = 1:ncomps
            H = induce(G, comps[k])
            if NV(H) > 1
                result *= "(" * CographRepresentation(H) * ")"
            else
                result *= CographRepresentation(H)
            end
            if k < ncomps
                result *= "∨"
            end
        end
    else
        comps = collect(parts(components(G)))
        ncomps = length(comps)
        result = ""
        for k = 1:ncomps
            H = induce(G, comps[k])
            if NV(H) > 1
                result *= "(" * CographRepresentation(H) * ")"
            else
                result *= CographRepresentation(H)
            end
            if k < ncomps
                result *= "+"
            end
        end
    end

    cache_save(G, :CographRepresentation, result)
    return result
end


text:
const MPIInteger = Union{Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64, UInt64}
const MPIFloatingPoint = Union{Float32, Float64}
const MPIComplex = Union{ComplexF32, ComplexF64}

const MPIDatatype = Union{Char,
                            Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64,
                            UInt64,
                            Float32, Float64, ComplexF32, ComplexF64}
MPIBuffertype{T} = Union{Ptr{T}, Array{T}, SubArray{T}, Ref{T}}

MPIBuffertypeOrConst{T} = Union{MPIBuffertype{T}, SentinelPtr}

Base.cconvert(::Type{MPIPtr}, x::Union{Ptr{T}, Array{T}, Ref{T}}) where T = Base.cconvert(Ptr{T}, x)
function Base.cconvert(::Type{MPIPtr}, x::SubArray{T}) where T
    Base.cconvert(Ptr{T}, x)
end
function Base.unsafe_convert(::Type{MPIPtr}, x::MPIBuffertype{T}) where T
    ptr = Base.unsafe_convert(Ptr{T}, x)
    reinterpret(MPIPtr, ptr)
end
function Base.cconvert(::Type{MPIPtr}, ::Nothing)
    reinterpret(MPIPtr, C_NULL)
end

macro assert_minlength(buffer, count)
    quote
        if $(esc(buffer)) isa AbstractArray
            @assert length($(esc(buffer))) >= $(esc(count))
        end
    end
end

"""
    MPI.MPIPtr

A pointer to an MPI buffer. This type is used only as part of the implicit conversion in
`ccall`: a Julia object can be passed to MPI by defining methods for
`Base.cconvert(::Type{MPIPtr}, ...)`/`Base.unsafe_convert(::Type{MPIPtr}, ...)`.

Currently supported are:
 - `Ptr`
 - `Ref`
 - `Array`
 - `SubArray`
 - `CuArray` if CuArrays.jl is loaded.

Additionally, certain sentinel values can be used, e.g. `MPI_IN_PLACE` or `MPI_BOTTOM`.
"""
MPIPtr

"""
    MPI.Buffer

An MPI buffer for communication operations.

# Fields
$(DocStringExtensions.FIELDS)

# Usage

    Buffer(data, count::Integer, datatype::Datatype)

Generic constructor.

    Buffer(data)

Construct a `Buffer` backed by `data`, automatically determining the appropriate `count`
and `datatype`. Methods are provided for

 - `Ref`
 - `Array`
 - `CuArray` if CuArrays.jl is loaded
 - `SubArray`s of an `Array` or `CuArray` where the layout is contiguous, sequential or
   blocked.

"""
struct Buffer{A}
    """a Julia object referencing a region of memory to be used for communication. It is
    required that the object can be `cconvert`ed to an [`MPIPtr`](@ref)."""
    data::A

    """the number of elements of `datatype` in the buffer. Note that this may not
    correspond to the number of elements in the array if derived types are used."""
    count::Cint

    """the [`MPI.Datatype`](@ref) stored in the buffer."""
    datatype::Datatype
end
Buffer(buf::Buffer) = buf
Buffer(data, count::Integer, datatype::Datatype) = Buffer(data, Cint(count), datatype)

function Buffer(arr::Array)
    Buffer(arr, Cint(length(arr)), Datatype(eltype(arr)))
end
function Buffer(ref::Ref)
    Buffer(ref, Cint(1), Datatype(eltype(ref)))
end

# SubArray
function Buffer(sub::Base.FastContiguousSubArray)
    Buffer(sub, Cint(length(sub)), Datatype(eltype(sub)))
end
function Buffer(sub::Base.FastSubArray)
    datatype = Types.create_vector(length(sub), 1, sub.stride1,
                                   Datatype(eltype(sub); commit=false))
    Types.commit!(datatype)
    Buffer(sub, Cint(1), datatype)
end
function Buffer(sub::SubArray{T,N,P,I,false}) where {T,N,P,I<:Tuple{Vararg{Union{Base.ScalarIndex, Base.Slice, AbstractUnitRange}}}}
    datatype = Types.create_subarray(size(parent(sub)),
                                     map(length, sub.indices),
                                     map(i -> first(i)-1, sub.indices),
                                     Datatype(eltype(sub), commit=false))
    Types.commit!(datatype)
    Buffer(parent(sub), Cint(1), datatype)
end

"""
    Buffer_send(data)

Construct a [`Buffer`](@ref) object for a send operation from `data`, allowing cases where
`isbits(data)`.
"""
Buffer_send(data) = isbits(data) ? Buffer(Ref(data)) : Buffer(data)

const BUFFER_NULL = Buffer(C_NULL, 0, DATATYPE_NULL)


text:
module Loess

include("loess1d.jl")

mutable struct MultivariateLoessModel{T <: AbstractFloat}

    xs::AbstractMatrix{T}       # An n by m predictor matrix containing n observations from m predictors
    ys::AbstractMatrix{T}       # A length n response vector
    bs::Matrix{T}               # Least squares coefficients
    verts::Dict{Vector{T}, Int} # kd-tree vertexes mapped to indexes
    kdtree::KDTree{T}

end

"""
    loess(xs, ys, normalize=true, span=0.75, degreee=2)

Fit a loess model.

Args:
  `xs`: A `n` by `m` matrix with `n` observations from `m` independent predictors
  `ys`: A length `n` response vector.
  `normalize`: Normalize the scale of each predicitor. (default true when `m > 1`)
  `span`: The degree of smoothing, typically in [0,1]. Smaller values result in smaller
      local context in fitting.
  `degree`: Polynomial degree.

Returns:
  A fit `LoessModel`.

"""
function loess(xs::AbstractMatrix{T}, ys::AbstractMatrix{T};
               normalize::Bool=true,
               span::AbstractFloat=0.75,
               degree::Integer=2) where T<:AbstractFloat

    if size(xs, 1) != size(ys, 1)
        throw(DimensionMismatch("Predictor and response arrays must of the same length"))
    end

    n, m = size(xs)
    q = ceil(Int, (span * n))

    # TODO: We need to keep track of how we are normalizing so we can
    # correctly apply predict to unnormalized data. We should have a normalize
    # function that just returns a vector of scaling factors.
    if normalize && m > 1
        xs = tnormalize!(copy(xs))
    end

    kdtree = KDTree(xs, 0.05 * span)
    verts = Array{T}(undef, length(kdtree.verts), m)

    # map verticies to their index in the bs coefficient matrix
    verts = Dict{Vector{T}, Int}()
    for (k, vert) in enumerate(kdtree.verts)
        verts[vert] = k
    end

    # Fit each vertex
    ds = Array{T}(undef, n) # distances
    perm = collect(1:n)
    bs = Array{T}(undef, length(kdtree.verts), 1 + degree * m)

    # TODO: higher degree fitting
    us = Array{T}(undef, q, 1 + degree * m)
    vs = Array{T}(undef, q)
    for (vert, k) in verts
        # reset perm
        for i in 1:n
            perm[i] = i
        end

        # distance to each point
        for i in 1:n
            ds[i] = euclidean(vec(vert), vec(xs[i,:]))
        end

        # copy the q nearest points to vert into X
        partialsort!(perm, q, by=i -> ds[i])
        dmax = maximum([ds[perm[i]] for i = 1:q])

        for i in 1:q
            pi = perm[i]
            w = tricubic(ds[pi] / dmax)
            us[i,1] = w
            for j in 1:m
                x = xs[pi, j]
                wxl = w
                for l in 1:degree
                    wxl *= x
                    us[i, 1 + (j-1)*degree + l] = wxl # w*x^l
                end
            end
            vs[i] = ys[pi] * w
        end
        bs[k,:] = us \ vs
    end

    MultivariateLoessModel{T}(xs, ys, bs, verts, kdtree)
end

loess(xs::AbstractVector{T}, ys::AbstractMatrix{T}; kwargs...) where {T<:AbstractFloat} =
    loess(reshape(xs, (length(xs), 1)), ys; kwargs...)

# Predict response values from a trained loess model and predictor observations.
#
# Loess (or at least most implementations, including this one), does not perform,
# extrapolation, so none of the predictor observations can exceed the range of
# values in the data used for training.
#
# Args:
#   zs/z: An n' by m matrix of n' observations from m predictors. Where m is the
#       same as the data used in training.
#
# Returns:
#   A length n' vector of predicted response values.
#
function predict(model::MultivariateLoessModel{T}, z::T) where T <: AbstractFloat
    predict(model, T[z])
end

function predict(model::MultivariateLoessModel{T}, zs::AbstractVector{T}) where T <: AbstractFloat
    m = size(model.xs, 2)

    # in the univariate case, interpret a non-singleton zs as vector of
    # ponits, not one point
    if m == 1 && length(zs) > 1
        return predict(model, reshape(zs, (length(zs), 1)))
    end

    if length(zs) != m
        error("$(m)-dimensional model applied to length $(length(zs)) vector")
    end

    adjacent_verts = traverse(model.kdtree, zs)

    if m == 1
        @assert(length(adjacent_verts) == 2)
        z = zs[1]
        u = (z - adjacent_verts[1][1]) /
        (adjacent_verts[2][1] - adjacent_verts[1][1])

        y1 = evalpoly(zs, model.bs[model.verts[[adjacent_verts[1][1]]],:])
        y2 = evalpoly(zs, model.bs[model.verts[[adjacent_verts[2][1]]],:])
        return (1.0 - u) * y1 + u * y2
    else
        error("Multivariate blending not yet implemented")
        # TODO:
        #   1. Univariate linear interpolation between adjacent verticies.
        #   2. Blend these estimates. (I'm not sure how this is done.)
    end
end

function predict(model::MultivariateLoessModel{T}, zs::AbstractMatrix{T}) where T <: AbstractFloat
    ys = Array{T}(undef, size(zs, 1))
    for i in 1:size(zs, 1)
        # the vec() here is not necessary on 0.5 anymore
        ys[i] = predict(model, vec(zs[i,:]))
    end
    ys
end

end


text:
@testset "Gabor wavelet" begin
# there are two methods and one struct for calculating Gabor wavelets
# calculate Gabor wavelet given ω
N=4
ω=zeros(2^N)
σ=1.0
ξ=0.0
γ = Scat.gabor1d(ω, σ, ξ)
@test length(γ)==length(ω)
@test γ==ones(length(ω))

# calculate Gabor wavelet without ω
N=4
σ=1.0/sqrt(2.0)
ξ=0.0
ω=FFTW.fftfreq(N)
γ = Scat.gabor1d(N, σ, ξ)
@test length(γ)==length(ω)
@test exp.(-ω.^2)≈γ

# test inner constructor for GaborWavelet struct
N=4
σ=1.0
ξ=0.0
j=1
γ = GaborWavelet(N, σ, ξ, j)
@test length(γ.γ)==N
@test γ.γ == Scat.gabor1d(N, σ, ξ)
end

@testset "Gaussian filter" begin
# there are two methods and one struct for calculating Gaussian lowpass filters
# calculate Gaussian lowpass filter given ω
N=4
ω=zeros(2^N)
σ=1.0
G = Scat.gauss1d(ω, σ)
@test length(G)==length(ω)
@test G==ones(length(ω))

# calculate Gaussian lowpass filter without ω
N=4
σ=1.0/sqrt(2.0)
ω=collect(0.0:N-1)./N
G = Scat.gauss1d(N, σ)
@test length(G)==length(ω)
@test exp.(-ω.^2)≈G

# test inner constructor for GaussianFilter struct
N=4
σ=1.0
ξ=0.0
j=1
G = GaussianFilter(N, σ)
@test length(G.ϕ)==N
@test G.ϕ == Scat.gauss1d(N, σ)
end

@testset "Morlet wavelet" begin
# there are two methods and one struct for calculating Morlet wavelets
# calculate Morlet wavelet given Gabor wavelet and Gaussian lowpass filter
N=4
σ=1.0/sqrt(2.0)
ξ=0.0
G=ones(N)
γ=ones(N)
ψ=Scat.morlet1d(γ, G)
@test length(ψ)==N
@test ψ==zeros(N)

# calculate Morlet wavelet without ω
N=4
σ=1.0/sqrt(2.0)
ξ=0.0
ψ=Scat.morlet1d(N, σ, ξ)
@test length(G)==N

# test both inner constructors for MorletWavelet struct
N=4
σ=1.0
ξ=0.0
j=1
# get GaborWavelet and GaussianFilter structs to build MorletWavelet
G = GaussianFilter(N, σ)
γ = GaborWavelet(N, σ, ξ, j)
ψ = MorletWavelet(γ, G)
@test length(ψ.ψ)==N
@test ψ.ψ == Scat.morlet1d(γ.γ, G.ϕ)
end


text:
using Printf
using Statistics
using Random
using Oceananigans
using Oceananigans.Units
using GLMakie

using Oceananigans.TurbulenceClosures: FluxTapering

gradient = "y"
filename = "coarse_baroclinic_adjustment_" * gradient

# Architecture
architecture = CPU()

# Domain
Lz = 1kilometers     # depth [m]
Ny = 20
Nz = 20
save_fields_interval = 1hour
stop_time = 30days
Δt = 20minutes

grid = LatitudeLongitudeGrid(architecture;
                             topology = (Bounded, Bounded, Bounded),
                             size = (Ny, Ny, Nz),
                             longitude = (-5, 5),
                             latitude = (40, 50),
                             z = (-Lz, 0),
                             halo = (3, 3, 3))

coriolis = HydrostaticSphericalCoriolis()

Δy = 1000kilometers / Ny
@show κh = νh = Δy^4 / 10days
vertical_closure = VerticalScalarDiffusivity(ν=1e-2, κ=1e-4)
horizontal_closure = HorizontalScalarBiharmonicDiffusivity(ν=νh)

gerdes_koberle_willebrand_tapering = FluxTapering(1e-1)
gent_mcwilliams_diffusivity = IsopycnalSkewSymmetricDiffusivity(κ_skew = 1e3,
                                                                κ_symmetric = (b=0, c=1e3),
                                                                slope_limiter = gerdes_koberle_willebrand_tapering)

closures = (vertical_closure, horizontal_closure, gent_mcwilliams_diffusivity)

@info "Building a model..."

model = HydrostaticFreeSurfaceModel(grid = grid,
                                    coriolis = coriolis,
                                    buoyancy = BuoyancyTracer(),
                                    closure = closures,
                                    tracers = (:b, :c),
                                    momentum_advection = VectorInvariant(),
                                    tracer_advection = WENO5(),
                                    free_surface = ImplicitFreeSurface())

@info "Built $model."

"""
Linear ramp from 0 to 1 between -Δy/2 and +Δy/2.

For example:

y < y₀           => ramp = 0
y₀ < y < y₀ + Δy => ramp = y / Δy
y > y₀ + Δy      => ramp = 1
"""
function ramp(λ, y, Δ)
    gradient == "x" && return min(max(0, λ / Δ + 1/2), 1)
    gradient == "y" && return min(max(0, (y - 45) / Δ + 1/2), 1)
end

# Parameters
N² = 4e-6 # [s⁻²] buoyancy frequency / stratification
M² = 8e-8 # [s⁻²] horizontal buoyancy gradient

Δy = 1 # degree
Δz = 100

Δc = 100kilometers * 2Δy
Δb = 100kilometers * Δy * M²
ϵb = 1e-2 * Δb # noise amplitude

bᵢ(λ, y, z) = N² * z + Δb * ramp(λ, y, Δy)
cᵢ(λ, y, z) = exp(-y^2 / 2Δc^2) * exp(-(z + Lz/4)^2 / 2Δz^2)

set!(model, b=bᵢ, c=cᵢ)

#####
##### Simulation building
#####

simulation = Simulation(model; Δt, stop_time)

# add timestep wizard callback
wizard = TimeStepWizard(cfl=0.1, max_change=1.1, max_Δt=Δt)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(20))

# add progress callback
wall_clock = [time_ns()]

function print_progress(sim)
    @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): (%6.3e, %6.3e, %6.3e) m/s, next Δt: %s\n",
            100 * (sim.model.clock.time / sim.stop_time),
            sim.model.clock.iteration,
            prettytime(sim.model.clock.time),
            prettytime(1e-9 * (time_ns() - wall_clock[1])),
            maximum(abs, sim.model.velocities.u),
            maximum(abs, sim.model.velocities.v),
            maximum(abs, sim.model.velocities.w),
            prettytime(sim.Δt))

    wall_clock[1] = time_ns()

    return nothing
end

simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(20))

simulation.output_writers[:fields] = JLD2OutputWriter(model, merge(model.velocities, model.tracers),
                                                      schedule = TimeInterval(save_fields_interval),
                                                      filename = filename * "_fields",
                                                      overwrite_existing = true)

@info "Running the simulation..."

run!(simulation)

@info "Simulation completed in " * prettytime(simulation.run_wall_time)

#####
##### Visualize
#####

fig = Figure(resolution = (1400, 700))

filepath = filename * "_fields.jld2"

ut = FieldTimeSeries(filepath, "u")
bt = FieldTimeSeries(filepath, "b")
ct = FieldTimeSeries(filepath, "c")

# Build coordinates, rescaling the vertical coordinate
x, y, z = nodes((Center, Center, Center), grid)

#####
##### Plot buoyancy...
#####

times = bt.times
Nt = length(times)

if gradient == "y" # average in x
    un(n) = interior(mean(ut[n], dims=1), 1, :, :)
    bn(n) = interior(mean(bt[n], dims=1), 1, :, :)
    cn(n) = interior(mean(ct[n], dims=1), 1, :, :)
else # average in y
    un(n) = interior(mean(ut[n], dims=2), :, 1, :)
    bn(n) = interior(mean(bt[n], dims=2), :, 1, :)
    cn(n) = interior(mean(ct[n], dims=2), :, 1, :)
end

@show min_c = 0
@show max_c = 1
@show max_u = 4 * maximum(abs, un(Nt))
min_u = - max_u

axu = Axis(fig[2, 1], xlabel="$gradient (deg)", ylabel="z (m)", title="Zonal velocity")
axc = Axis(fig[3, 1], xlabel="$gradient (deg)", ylabel="z (m)", title="Tracer concentration")
slider = Slider(fig[4, 1:2], range=1:Nt, startvalue=1)
n = slider.value

u = @lift un($n)
b = @lift bn($n)
c = @lift cn($n)

hm = heatmap!(axu, y, z, u, colorrange=(min_u, max_u), colormap=:balance)
contour!(axu, y, z, b, levels = 25, color=:black, linewidth=2)
cb = Colorbar(fig[2, 2], hm)

hm = heatmap!(axc, y, z, c, colorrange=(0, 0.5), colormap=:speed)
contour!(axc, y, z, b, levels = 25, color=:black, linewidth=2)
cb = Colorbar(fig[3, 2], hm)

title_str = @lift "Baroclinic adjustment with GM at t = " * prettytime(times[$n])
ax_t = fig[1, 1:2] = Label(fig, title_str)

display(fig)

record(fig, filename * ".mp4", 1:Nt, framerate=8) do i
    @info "Plotting frame $i of $Nt"
    n[] = i
end


text:
#!/usr/bin/env julia

"""
Tests for the [`Primes`](@ref) common module.
"""
module TestPrimes

include(joinpath("..", "src", "Primes.jl"))

using Test

@testset "Primes" begin
    @testset "prime_factorization" begin
        @test Primes.prime_factorization(2) == [(2, 1)]
        @test Primes.prime_factorization(3) == [(3, 1)]
        @test Primes.prime_factorization(4) == [(2, 2)]
        @test Primes.prime_factorization(5) == [(5, 1)]
        @test Primes.prime_factorization(6) == [(2, 1), (3, 1)]
        @test Primes.prime_factorization(12) == [(2, 2), (3, 1)]
        @test Primes.prime_factorization(600) == [(2, 3), (3, 1), (5, 2)]
        @test Primes.prime_factorization(2903) == [(2903, 1)]
        @test(
            Primes.prime_factorization(61740)
                == [(2, 2), (3, 2), (5, 1), (7, 3)]
        )
        @test(
            Primes.prime_factorization(4928693)
                == [(7, 1), (11, 3), (23, 2)]
        )
        @test(
            Primes.prime_factorization(169165232)
                == [(2, 4), (17, 1), (313, 1), (1987, 1)]
        )
        @test(
            Primes.prime_factorization(74435892358158)
                == [(2, 1), (3, 3), (29, 1), (3049, 2), (5113, 1)]
        )
    end
end

end

max_stars_repo_name: curtislb/ProjectEuler

max_issues_repo_name: curtislb/ProjectEuler

max_forks_repo_name: curtislb/ProjectEuler


text:
# Testing

using Test

mycomb = Hotcomb((a=1:2, b=1:4,c='a':'c'))
@test keys(mycomb) == (:a, :b, :c)
@test value(mycomb) == [1:2, 1:4, 'a':'c']
@test size(mycomb) == (24,3)
@test collect(mycomb) == mycomb[:,:] # Seems to be working
@test mycomb[13,2] == mycomb[13,:b]
@test mycomb[5,:] == [1, 2, 'b']

@test_throws "Symbol :d not found" mycomb[3,:d]

###############################################################################
# Benchmarking

using BenchmarkTools, InteractiveUtils

mycomb = Hotcomb([7,7,7,7,7,7,7])
value(mycomb)
size(mycomb)

mycomb_collect = mycomb[]

mycomb[1000,:]
@benchmark mycomb[1000,:]

mycomb_collect[1000,:]
@benchmark mycomb_collect[1000,:]

varinfo() # Accessing mycomb as Hotcomb is more time intensive but uses less memory than generating full collection