text:
module MultilayerQG

export
  fwdtransform!,
  invtransform!,
  streamfunctionfrompv!,
  pvfromstreamfunction!,
  updatevars!,

  set_q!,
  set_ψ!,
  energies,
  fluxes

using
  FFTW,
  CUDA,
  LinearAlgebra,
  StaticArrays,
  Reexport

@reexport using FourierFlows

using LinearAlgebra: mul!, ldiv!
using FFTW: rfft, irfft
using FourierFlows: parsevalsum, parsevalsum2, superzeros, plan_flows_rfft

nothingfunction(args...) = nothing

"""
    Problem(; parameters...)

Construct a multi-layer QG problem.
"""
function Problem(nlayers::Int,                        # number of fluid layers
                     dev = CPU();
              # Numerical parameters
                      nx = 128,
                      Lx = 2π,
                      ny = nx,
                      Ly = Lx,
                      dt = 0.01,
              # Physical parameters
                      f₀ = 1.0,                       # Coriolis parameter
                       β = 0.0,                       # y-gradient of Coriolis parameter
                       g = 1.0,                       # gravitational constant
                       U = zeros(nlayers),            # imposed zonal flow U(y) in each layer
                       H = 1/nlayers * ones(nlayers), # rest fluid height of each layer
                       ρ = Array{Float64}(1:nlayers), # density of each layer
                     eta = nothing,                   # topographic PV
              # Bottom Drag and/or (hyper)-viscosity
                       μ = 0,
                       ν = 0,
                      nν = 1,
              # Timestepper and equation options
                 stepper = "RK4",
                  calcFq = nothingfunction,
              stochastic = false,
                  linear = false,
                       T = Float64)

   # topographic PV
   eta === nothing && (eta = zeros(dev, T, (nx, ny)))

   grid = TwoDGrid(dev, nx, Lx, ny, Ly; T=T)
   params = Params(nlayers, g, f₀, β, ρ, H, U, eta, μ, ν, nν, grid, calcFq=calcFq, dev=dev)
   vars = calcFq == nothingfunction ? Vars(dev, grid, params) : (stochastic ? StochasticForcedVars(dev, grid, params) : ForcedVars(dev, grid, params))
   eqn = linear ? LinearEquation(dev, params, grid) : Equation(dev, params, grid)

  FourierFlows.Problem(eqn, stepper, dt, grid, vars, params, dev)
end

abstract type BarotropicParams <: AbstractParams end

struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams
  # prescribed params
   nlayers :: Int        # Number of fluid layers
         g :: T          # Gravitational constant
        f₀ :: T          # Constant planetary vorticity
         β :: T          # Planetary vorticity y-gradient
         ρ :: Aphys3D    # Array with density of each fluid layer
         H :: Aphys3D    # Array with rest height of each fluid layer
         U :: Aphys3D    # Array with imposed constant zonal flow U(y) in each fluid layer
       eta :: Aphys2D    # Array containing topographic PV
         μ :: T          # Linear bottom drag
         ν :: T          # Viscosity coefficient
        nν :: Int        # Hyperviscous order (nν=1 is plain old viscosity)
   calcFq! :: Function   # Function that calculates the forcing on QGPV q

  # derived params
        g′ :: Aphys1D    # Array with the reduced gravity constants for each fluid interface
        Qx :: Aphys3D    # Array containing x-gradient of PV due to eta in each fluid layer
        Qy :: Aphys3D    # Array containing y-gradient of PV due to β, U, and eta in each fluid layer
         S :: Atrans4D   # Array containing coeffients for getting PV from  streamfunction
       S⁻¹ :: Atrans4D   # Array containing coeffients for inverting PV to streamfunction
  rfftplan :: Trfft      # rfft plan for FFTs
end

struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: BarotropicParams
  # prescribed params
         β :: T          # Planetary vorticity y-gradient
         U :: Aphys3D    # Imposed constant zonal flow U(y)
       eta :: Aphys2D    # Array containing topographic PV
         μ :: T          # Linear bottom drag
         ν :: T          # Viscosity coefficient
        nν :: Int        # Hyperviscous order (nν=1 is plain old viscosity)
   calcFq! :: Function   # Function that calculates the forcing on QGPV q

  # derived params
        Qx :: Aphys3D    # Array containing x-gradient of PV due to eta
        Qy :: Aphys3D    # Array containing meridional PV gradient due to β, U, and eta
  rfftplan :: Trfft      # rfft plan for FFTs
end

function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 1}) where TU
  T = eltype(grid)
  if length(U) == nlayers
    U_2D = zeros(dev, T, (1, nlayers))
    U_2D[:] = U
    U_2D = repeat(U_2D, outer=(grid.ny, 1))
  else
    U_2D = zeros(dev, T, (grid.ny, 1))
    U_2D[:] = U
  end
  U_3D = zeros(dev, T, (1, grid.ny, nlayers))
  @views U_3D[1, :, :] = U_2D
  return U_3D
end

function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 2}) where TU
  T = eltype(grid)
  U_3D = zeros(dev, T, (1, grid.ny, nlayers))
  @views U_3D[1, :, :] = U
  return U_3D
end

function convert_U_to_U3D(dev, nlayers, grid, U::Number)
  T = eltype(grid)
  A = ArrayType(dev)
  U_3D = reshape(repeat([T(U)], outer=(grid.ny, 1)), (1, grid.ny, nlayers))
  return A(U_3D)
end

function Params(nlayers, g, f₀, β, ρ, H, U, eta, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE, dev::Device=CPU()) where TU

  T = eltype(grid)
  A = ArrayType(dev)

   ny, nx = grid.ny , grid.nx
  nkr, nl = grid.nkr, grid.nl
   kr, l  = grid.kr , grid.l

    U = convert_U_to_U3D(dev, nlayers, grid, U)

  Uyy = real.(ifft(-l.^2 .* fft(U)))
  Uyy = repeat(Uyy, outer=(nx, 1, 1))

  etah = rfft(A(eta))
  etax = irfft(im * kr .* etah, nx)
  etay = irfft(im * l  .* etah, nx)

  Qx = zeros(dev, T, (nx, ny, nlayers))
  @views @. Qx[:, :, nlayers] += etax

  Qy = zeros(dev, T, (nx, ny, nlayers))
  Qy = T(β) .- Uyy  # T(β) is needed to ensure that Qy remains same type as U
  @views @. Qy[:, :, nlayers] += etay

  rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort)

  if nlayers==1
    return SingleLayerParams(T(β), U, eta, T(μ), T(ν), nν, calcFq, Qx, Qy, rfftplanlayered)

  else # if nlayers≥2

    ρ = reshape(T.(ρ), (1,  1, nlayers))
    H = reshape(T.(H), (1,  1, nlayers))

    g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface

    Fm = @. T(f₀^2 / (g′ * H[2:nlayers]))
    Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1]))

    typeofSkl = SArray{Tuple{nlayers, nlayers}, T, 2, nlayers^2} # StaticArrays of type T and dims = (nlayers x nlayers)

    S = Array{typeofSkl, 2}(undef, (nkr, nl))
    calcS!(S, Fp, Fm, nlayers, grid)

    S⁻¹ = Array{typeofSkl, 2}(undef, (nkr, nl))
    calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid)

    S, S⁻¹, Fp, Fm  = A(S), A(S⁻¹), A(Fp), A(Fm)     # convert to appropriate ArrayType

    CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1])
    for j = 2:nlayers-1
      CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] - Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j])
    end
    CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers])

    return Params(nlayers, T(g), T(f₀), T(β), A(ρ), A(H), U, eta, T(μ), T(ν), nν, calcFq, A(g′), Qx, Qy, S, S⁻¹, rfftplanlayered)
  end
end

numberoflayers(params) = params.nlayers
numberoflayers(::SingleLayerParams) = 1

# ---------
# Equations
# ---------

function hyperdissipation(dev, params, grid)
  T = eltype(grid)
  L = ArrayType(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params)))
  @. L = - params.ν * grid.Krsq^params.nν
  @views @. L[1, 1, :] = 0
  return L
end

function LinearEquation(dev, params, grid)
  L = hyperdissipation(dev, params, grid)
  return FourierFlows.Equation(L, calcNlinear!, grid)
end

function Equation(dev, params, grid)
  L = hyperdissipation(dev, params, grid)
  return FourierFlows.Equation(L, calcN!, grid)
end

# ----
# Vars
# ----

struct Vars{Aphys, Atrans, F, P} <: AbstractVars
        q :: Aphys
        ψ :: Aphys
        u :: Aphys
        v :: Aphys
       qh :: Atrans
       ψh :: Atrans
       uh :: Atrans
       vh :: Atrans
      Fqh :: F
  prevsol :: P
end

const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing}
const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray}

"""
    Vars(dev, grid, params)

Returns the vars for unforced multi-layer QG problem with `grid` and `params`.
"""
function Vars(dev::Dev, grid, params) where Dev
  T = eltype(grid)
  nlayers = numberoflayers(params)

  @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v
  @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh

  return Vars(q, ψ, u, v, qh, ψh, uh, vh, nothing, nothing)
end

"""
    ForcedVars(dev, grid, params)

Returns the vars for forced multi-layer QG problem with `grid` and `params`.
"""
function ForcedVars(dev::Dev, grid, params) where Dev
  T = eltype(grid)
  nlayers = numberoflayers(params)

  @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v
  @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh

  return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, nothing)
end

"""
    StochasticForcedVars(dev, rid, params)

Returns the vars for forced multi-layer QG problem with `grid` and `params`.
"""
function StochasticForcedVars(dev::Dev, grid, params) where Dev
  T = eltype(grid)
  nlayers = numberoflayers(params)

  @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v
  @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh prevsol

  return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, prevsol)
end

fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var)
invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh)

function streamfunctionfrompv!(ψh, qh, params, grid)
  for j=1:grid.nl, i=1:grid.nkr
    CUDA.@allowscalar @views ψh[i, j, :] .= params.S⁻¹[i, j] * qh[i, j, :]
  end
end

function pvfromstreamfunction!(qh, ψh, params, grid)
  for j=1:grid.nl, i=1:grid.nkr
    CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :]
  end
end

function streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid)
  @. ψh = -grid.invKrsq * qh
end

function pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid)
  @. qh = -grid.Krsq * ψh
end

"""
    calcS!(S, Fp, Fm, nlayers, grid)

Constructs the array S, which consists of nlayer x nlayer static arrays S_kl that relate
the q's and ψ's at every wavenumber: q̂_{k, l} = S_kl * ψ̂_{k, l}.
"""
function calcS!(S, Fp, Fm, nlayers, grid)
  F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp))
  for n=1:grid.nl, m=1:grid.nkr
    CUDA.@allowscalar k² = grid.Krsq[m, n]
    Skl = SMatrix{nlayers, nlayers}(- k² * I + F)
    S[m, n] = Skl
  end
  return nothing
end

"""
    calcS⁻¹!(S, Fp, Fm, nlayers, grid)

Constructs the array S⁻¹, which consists of nlayer x nlayer static arrays (S_kl)⁻¹ that
relate the q's and ψ's at every wavenumber: ψ̂_{k, l} = (S_kl)⁻¹ * q̂_{k, l}.
"""
function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid)
  T = eltype(grid)
  F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp))
  for n=1:grid.nl, m=1:grid.nkr
    CUDA.@allowscalar k² = grid.Krsq[m, n] == 0 ? 1 : grid.Krsq[m, n]
    Skl = - k² * I + F
    S⁻¹[m, n] = SMatrix{nlayers, nlayers}(I / Skl)
  end
  S⁻¹[1, 1] = SMatrix{nlayers, nlayers}(zeros(T, (nlayers, nlayers)))
  return nothing
end

# -------
# Solvers
# -------

function calcN!(N, sol, t, clock, vars, params, grid)
  nlayers = numberoflayers(params)
  calcN_advection!(N, sol, vars, params, grid)
  @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers]   # bottom linear drag
  addforcing!(N, sol, t, clock, vars, params, grid)
  return nothing
end

function calcNlinear!(N, sol, t, clock, vars, params, grid)
  nlayers = numberoflayers(params)
  calcN_linearadvection!(N, sol, vars, params, grid)
  @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers]   # bottom linear drag
  addforcing!(N, sol, t, clock, vars, params, grid)
  return nothing
end

"""
    calcN_advection!(N, sol, vars, params, grid)

Calculates the advection term.
"""
function calcN_advection!(N, sol, vars, params, grid)
  @. vars.qh = sol

  streamfunctionfrompv!(vars.ψh, vars.qh, params, grid)

  @. vars.uh = -im * grid.l  * vars.ψh
  @. vars.vh =  im * grid.kr * vars.ψh

  invtransform!(vars.u, vars.uh, params)
  @. vars.u += params.U                    # add the imposed zonal flow U

  uQx, uQxh = vars.q, vars.uh              # use vars.q and vars.uh as scratch variables
  @. uQx  = vars.u * params.Qx             # (U+u)*∂Q/∂x
  fwdtransform!(uQxh, uQx, params)
  @. N = - uQxh                            # -\hat{(U+u)*∂Q/∂x}

  invtransform!(vars.v, vars.vh, params)

  vQy, vQyh = vars.q, vars.vh              # use vars.q and vars.vh as scratch variables
  @. vQy = vars.v * params.Qy              # v*∂Q/∂y
  fwdtransform!(vQyh, vQy, params)
  @. N -= vQyh                             # -\hat{v*∂Q/∂y}

  invtransform!(vars.q, vars.qh, params)

  uq , vq  = vars.u , vars.v               # use vars.u and vars.v as scratch variables
  uqh, vqh = vars.uh, vars.vh              # use vars.uh and vars.vh as scratch variables
  @. uq *= vars.q                          # (U+u)*q
  @. vq *= vars.q                          # v*q

  fwdtransform!(uqh, uq, params)
  fwdtransform!(vqh, vq, params)

  @. N -= im * grid.kr * uqh + im * grid.l * vqh    # -\hat{∂[(U+u)q]/∂x} - \hat{∂[vq]/∂y}

  return nothing
end

"""
    calcN_linearadvection!(N, sol, vars, params, grid)

Calculates the advection term of the linearized equations.
"""
function calcN_linearadvection!(N, sol, vars, params, grid)
  @. vars.qh = sol

  streamfunctionfrompv!(vars.ψh, vars.qh, params, grid)

  @. vars.uh = -im * grid.l  * vars.ψh
  @. vars.vh =  im * grid.kr * vars.ψh

  invtransform!(vars.u, vars.uh, params)
  @. vars.u += params.U                    # add the imposed zonal flow U
  uQx, uQxh = vars.q, vars.uh              # use vars.q and vars.uh as scratch variables
  @. uQx  = vars.u * params.Qx             # (U+u)*∂Q/∂x
  fwdtransform!(uQxh, uQx, params)
  @. N = - uQxh                            # -\hat{(U+u)*∂Q/∂x}

  invtransform!(vars.v, vars.vh, params)

  vQy, vQyh = vars.q, vars.vh              # use vars.q and vars.vh as scratch variables

  @. vQy = vars.v * params.Qy              # v*∂Q/∂y
  fwdtransform!(vQyh, vQy, params)
  @. N -= vQyh                             # -\hat{v*∂Q/∂y}

  invtransform!(vars.q, vars.qh, params)

  @. vars.u  = params.U
  Uq , Uqh  = vars.u , vars.uh             # use vars.u and vars.uh as scratch variables
  @. Uq *= vars.q                          # U*q

  fwdtransform!(Uqh, Uq, params)

  @. N -= im * grid.kr * Uqh               # -\hat{∂[U*q]/∂x}

  return nothing
end

addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing

function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid)
  params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid)
  @. N += vars.Fqh
  return nothing
end

# ----------------
# Helper functions
# ----------------

"""
    updatevars!(vars, params, grid, sol)
    updatevars!(prob)

Update all problem variables using `sol`.
"""
function updatevars!(vars, params, grid, sol)
  @. vars.qh = sol
  streamfunctionfrompv!(vars.ψh, vars.qh, params, grid)
  @. vars.uh = -im * grid.l  * vars.ψh
  @. vars.vh =  im * grid.kr * vars.ψh

  invtransform!(vars.q, deepcopy(vars.qh), params)
  invtransform!(vars.ψ, deepcopy(vars.ψh), params)
  invtransform!(vars.u, deepcopy(vars.uh), params)
  invtransform!(vars.v, deepcopy(vars.vh), params)

  return nothing
end

updatevars!(prob) = updatevars!(prob.vars, prob.params, prob.grid, prob.sol)

"""
    set_q!(sol, params, vars, grid, q)
    set_q!(prob)

Set the solution `prob.sol` as the transform of `q` and updates variables.
"""
function set_q!(sol, params, vars, grid, q)
  A = typeof(vars.q)
  fwdtransform!(vars.qh, A(q), params)
  @. vars.qh[1, 1, :] = 0
  @. sol = vars.qh
  updatevars!(vars, params, grid, sol)

  return nothing
end

function set_q!(sol, params::SingleLayerParams, vars, grid, q::AbstractArray{T, 2}) where T
  A = typeof(vars.q[:, :, 1])
  q_3D = vars.q
  @views q_3D[:, :, 1] = A(q)
  set_q!(sol, params, vars, grid, q_3D)

  return nothing
end

set_q!(prob, q) = set_q!(prob.sol, prob.params, prob.vars, prob.grid, q)

"""
    set_ψ!(params, vars, grid, sol, ψ)
    set_ψ!(prob)

Set the solution `prob.sol` to correspond to the transform of streamfunction `ψ` and
updates variables.
"""
function set_ψ!(sol, params, vars, grid, ψ)
  A = typeof(vars.ψ)
  fwdtransform!(vars.ψh, A(ψ), params)
  pvfromstreamfunction!(vars.qh, vars.ψh, params, grid)
  invtransform!(vars.q, vars.qh, params)

  set_q!(sol, params, vars, grid, vars.q)

  return nothing
end

function set_ψ!(sol, params::SingleLayerParams, vars, grid, ψ::AbstractArray{T, 2}) where T
  A = typeof(vars.ψ[:, :, 1])
  ψ_3D = vars.ψ
  @views ψ_3D[:, :, 1] = A(ψ)

  set_ψ!(sol, params, vars, grid, ψ_3D)

  return nothing
end

set_ψ!(prob, ψ) = set_ψ!(prob.sol, prob.params, prob.vars, prob.grid, ψ)

"""
    energies(vars, params, grid, sol)
    energies(prob)

Returns the kinetic energy of each fluid layer KE``_1, ...,`` KE``_{n}``, and the
potential energy of each fluid interface PE``_{3/2}, ...,`` PE``_{n-1/2}``, where ``n``
is the number of layers in the fluid. (When ``n=1``, only the kinetic energy is returned.)

The kinetic energy at the ``j``-th fluid layer is
```math
\\textrm{KE}_j = \\frac{H_j}{H} \\int \\frac1{2} |\\boldsymbol{\\nabla} \\psi_j|^2 \\frac{\\mathrm{d}^2 \\boldsymbol{x}}{L_x L_y} \\ , \\quad j = 1, \\dots, n \\ ,
```
while the potential energy that corresponds to the interface ``j+1/2`` (i.e., interface between the ``j``-th and ``(j+1)``-th fluid layer) is
```math
\\textrm{PE}_{j+1/2} = \\int \\frac1{2} \\frac{f_0^2}{g'_{j+1/2}} (\\psi_j - \\psi_{j+1})^2 \\frac{\\mathrm{d}^2 \\boldsymbol{x}}{L_x L_y} \\ , \\quad j = 1, \\dots, n-1 \\ .
```
"""
function energies(vars, params, grid, sol)
  nlayers = numberoflayers(params)
  KE, PE = zeros(nlayers), zeros(nlayers-1)

  @. vars.qh = sol
  streamfunctionfrompv!(vars.ψh, vars.qh, params, grid)

  abs²∇𝐮h = vars.uh        # use vars.uh as scratch variable
  @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh)

  for j = 1:nlayers
    CUDA.@allowscalar KE[j] = 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H)
  end

  for j = 1:nlayers-1
    CUDA.@allowscalar PE[j] = 1 / (2 * grid.Lx * grid.Ly) * params.f₀^2 / params.g′[j] * parsevalsum(abs2.(vars.ψh[:, :, j+1] .- vars.ψh[:, :, j]), grid)
  end

  return KE, PE
end

function energies(vars, params::SingleLayerParams, grid, sol)
  @. vars.qh = sol
  streamfunctionfrompv!(vars.ψh, vars.qh, params, grid)

  abs²∇𝐮h = vars.uh        # use vars.uh as scratch variable
  @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh)

  return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h, grid)
end

energies(prob) = energies(prob.vars, prob.params, prob.grid, prob.sol)

"""
    fluxes(vars, params, grid, sol)
    fluxes(prob)

Returns the lateral eddy fluxes within each fluid layer, lateralfluxes``_1,...,``lateralfluxes``_n``
and also the vertical eddy fluxes at each fluid interface
verticalfluxes``_{3/2}``, ...,`` verticalfluxes``_{n-1/2}, where ``n`` is the number of layers in the fluid.
(When ``n=1``, only the lateral fluxes are returned.)

The lateral eddy fluxes whithin the ``j``-th fluid layer are
```math
\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j \\, \\upsilon_j \\, \\partial_y u_j
\\frac{\\mathrm{d}^2 \\boldsymbol{x}}{L_x L_y} \\ , \\quad j = 1, \\dots, n \\ ,
```
while the vertical eddy fluxes at the ``j+1/2``-th fluid interface  (i.e., interface between
the ``j``-th and ``(j+1)``-th fluid layer) are
```math
\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f_0^2}{g'_{j+1/2} H} (U_j - U_{j+1}) \\,
\\upsilon_{j+1} \\, \\psi_{j} \\frac{\\mathrm{d}^2 \\boldsymbol{x}}{L_x L_y} \\ , \\quad
j = 1 , \\dots , n-1.
```
"""
function fluxes(vars, params, grid, sol)
  nlayers = numberoflayers(params)

  lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1)

  updatevars!(vars, params, grid, sol)

  ∂u∂yh = vars.uh           # use vars.uh as scratch variable
  ∂u∂y  = vars.u            # use vars.u  as scratch variable

  @. ∂u∂yh = im * grid.l * vars.uh
  invtransform!(∂u∂y, ∂u∂yh, params)

  lateralfluxes = (sum(@. params.H * params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :]
  lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H))

  for j = 1:nlayers-1
    CUDA.@allowscalar verticalfluxes[j] = sum(@views @. params.f₀^2 / params.g′[j] * (params.U[: ,:, j] - params.U[:, :, j+1]) * vars.v[:, :, j+1] * vars.ψ[:, :, j] ; dims=(1, 2))[1]
    CUDA.@allowscalar verticalfluxes[j] *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H))
  end

  return lateralfluxes, verticalfluxes
end

function fluxes(vars, params::SingleLayerParams, grid, sol)
  updatevars!(vars, params, grid, sol)

  ∂u∂yh = vars.uh           # use vars.uh as scratch variable
  ∂u∂y  = vars.u            # use vars.u  as scratch variable

  @. ∂u∂yh = im * grid.l * vars.uh
  invtransform!(∂u∂y, ∂u∂yh, params)

  lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :]
  lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly)

  return lateralfluxes
end

fluxes(prob) = fluxes(prob.vars, prob.params, prob.grid, prob.sol)

end # module


text:
module RANSACBenchmark

using LinearAlgebra

using RANSAC
using StaticArrays
using Random
using BenchmarkTools: @benchmark

export  benchmarkcloud1,
        benchmarkcloud2

export  benchmarksamplefunction

# abbreviations
const SV = SVector
nSV(v...) = normalize(SVector(v...))
const URN = Union{Real,Nothing}

include("samplers.jl")
include("samplebenchmark.jl")
include("exampleclouds.jl")

end # module

ext: jl


text:
module PlotEM1D
using MCMC_Driver, TransD_GP, PyPlot, JLD, StatsBase, Statistics, CSEM1Dkernels, LinearAlgebra

function get_posterior(opt_in::TransD_GP.Options, opt_EM::MCMC_Driver.EMoptions;
    fdataname="",
    burnin=300000,
    nbins=50,
    rho=[0.],
    alpha=0.025,
    alphacolor = "blue",
    alphalwidth = 0.3,
    qp1=0.01,
    qp2=0.99,
    figsize = (10,6),
    fontsize = 16,
    convfigsize = (8,6),
    kfigsize = (9,6),
    convfontsize = 14,
    cmapcdf = "RdYlBu_r",
    cmappdf = "inferno_r",
    showtitle=false,
    titlestr="",
    topadjust=0.9,
    fixviewextent=[0. 0. 0. 0.],
    rhobgcolor="yellow",
    rhofgcolor="k",
    nlldelta = 100,
    nxbins = 20,
    nftbins = 20,
    nxftfontsize = 16,
    pdfnormalize=false)
    @assert length(fixviewextent) == 4

    M, n, x_ft, iter, misfit, T0loc = get_posterior(opt_in, fdataname=fdataname)

    from = findfirst(iter.>burnin)
    f1 = figure(figsize=convfigsize)
    s1 = subplot(311)
    plot(iter,n)
    s1[:set_ylim](opt_in.nmin, opt_in.nmax)
    ylabel("# train")
    s2 = subplot(312, sharex=s1)
    plot(iter, misfit)
    ylabel("-log L")
    s3 = subplot(313, sharex=s1)
    plot(iter, T0loc)
    s3[:set_yticks](collect(minimum(T0loc):2:maximum(T0loc)))
    s2[:set_ylim](mean(misfit[from:end])-nlldelta, mean(misfit[from:end])+nlldelta)
    s3[:set_xlim](burnin+1, iter[end])
    ylabel("T=1 location")
    xlabel("iterations")
    MCMC_Driver.nicenup(gcf(), fsize=convfontsize)
    if showtitle
        titlestr == "" && (titlestr = fdataname)
        f1[:suptitle](titlestr, fontsize=convfontsize)
        f1[:subplots_adjust](top=topadjust)
    end

    xall = opt_in.xall
    f,ax = plt[:subplots](1,2, sharex=true, sharey=true, figsize=figsize)
    rhomin, rhomax = 0.0, 0.0
    for (i,mm) in enumerate(M)
        if i >= from
            # ax[1][:step](mm, xall[:], linewidth=alphalwidth, color=alphacolor, alpha=alpha)
            #plot(x_ft[i][1:n[i],2], x_ft[i][1:n[i],1], ".", markersize=20)
            # title("model $(iter[i])")
        end
        rhomin_mm = min(mm...)
        rhomax_mm = max(mm...)
        rhomin_mm < rhomin && (rhomin = rhomin_mm)
        rhomax_mm > rhomax && (rhomax = rhomax_mm)
    end

    # if rho != [0.]
    #     ax[1][:step](log10.(rho[2:end]), opt_EM.z[2:end], linewidth=3, color=rhobgcolor)
    #     ax[1][:step](log10.(rho[2:end]), opt_EM.z[2:end], linewidth=2, color=rhofgcolor)
    # end
    # ax[1][:grid]()
    # ax[1][:invert_yaxis]()

    edges = LinRange(rhomin, rhomax, nbins+1)
    himage = zeros(Float64, length(M[1]), nbins)
    cumhimage = zeros(Float64, length(M[1]), nbins)
    CI = zeros(Float64, length(M[1]), 2)
    z = opt_EM.z
    for ilayer=1:length(M[1])
        himage[ilayer,:] = (fit(Histogram, [m[ilayer] for m in M[from:end]], edges).weights)
        himage[ilayer,:] = himage[ilayer,:]/sum(himage[ilayer,:])/(diff(edges)[1])
        cumhimage[ilayer,:] = cumsum(himage[ilayer,:])/sum(himage[ilayer,:])
        pdfnormalize && (himage[ilayer,:] = himage[ilayer,:]/maximum(himage[ilayer,:]))
        CI[ilayer,:] = [quantile([m[ilayer] for m in M[from:end]],(qp1, qp2))...]
    end
    im1 = ax[1][:imshow](himage, extent=[edges[1],edges[end],xall[end],xall[1]], aspect="auto",
            cmap=cmappdf)
    cb1 = colorbar(im1, ax=ax[1])
    cb1[:ax][:set_xlabel]("pdf")
    ax[1][:grid]()
    im2 = ax[2][:imshow](cumhimage, extent=[edges[1],edges[end],xall[end],xall[1]], aspect="auto",
            cmap=cmapcdf)
    ax[2][:grid]()
    cb2 = colorbar(im2, ax=ax[2])
    cb2[:ax][:set_xlabel]("cdf")
    if rho != [0.]
        ax[1][:step](log10.(rho[2:end]), opt_EM.z[2:end], linewidth=3, color=rhobgcolor)
        ax[1][:step](log10.(rho[2:end]), opt_EM.z[2:end], linewidth=2, color=rhofgcolor)
        ax[2][:step](log10.(rho[2:end]), opt_EM.z[2:end], linewidth=3, color=rhobgcolor)
        ax[2][:step](log10.(rho[2:end]), opt_EM.z[2:end], linewidth=2, color=rhofgcolor)
    end
    ax[1][:plot](CI, xall[:], linewidth=2, color="g")
    ax[1][:plot](CI, xall[:], linewidth=2, color="k", linestyle="--")
    ax[2][:plot](CI, xall[:], linewidth=2, color="g")
    ax[2][:plot](CI, xall[:], linewidth=2, color="k", linestyle="--")
    # ax[2][:plot](CI, xall[:], "--k")
    ylim(opt_EM.zRx, xall[end])
    gca()[:invert_yaxis]()
    ax[1][:set_xlabel](L"\log_{10} \rho")
    ax[1][:set_ylabel]("depth (m)")
    ax[2][:set_xlabel](L"\log_{10} \rho")
    if !(fixviewextent == [0. 0. 0. 0.])
        ax[1][:set_xlim](fixviewextent[1:2])
        ax[1][:set_ylim](fixviewextent[3:4])
        ax[1][:invert_yaxis]()
    end
    MCMC_Driver.nicenup(f, fsize=fontsize)
    if showtitle
        titlestr == "" && (titlestr = fdataname)
        f[:suptitle](titlestr, fontsize=fontsize)
        f[:subplots_adjust](top=topadjust)
    end

    x, ft = zeros(sum(n[from:end])), zeros(sum(n[from:end]))
    nlast = 0
    for (i, xt) in enumerate(x_ft[from:end])
        x[nlast+1:nlast+n[from-1+i]]  = xt[1:n[from-1+i],1]
        ft[nlast+1:nlast+n[from-1+i]] = xt[1:n[from-1+i],2]
        nlast = nlast+n[from-1+i]
    end
    f2, ax2 = plt[:subplots](2,2, figsize=kfigsize)
    edgesx = LinRange(opt_in.xbounds[1], opt_in.xbounds[2], nxbins+1)
    edgesrho = LinRange(opt_in.fbounds[1], opt_in.fbounds[2], nftbins+1)
    h = fit(Histogram, (x, ft), (edgesx, edgesrho)).weights
    im3 = ax2[1][:imshow](h, extent=[edgesrho[1],edgesrho[end],edgesx[end],edgesx[1]], aspect="auto", vmin=0.0, vmax=2*max(h...))
    ax2[1][:set_ylabel]("depth m")
    ax2[1][:set_xlabel](L"\log_{10}\rho")
    rhist = sum(h,dims=1)[:]
    rhist = rhist/sum(rhist)/diff(edgesrho)[1]
    ax2[2][:bar](0.5*(edgesrho[2:end]+edgesrho[1:end-1]), rhist, width=diff(edgesrho)[1])
    ax2[2][:plot]([edgesrho[1], edgesrho[end]], 1/(opt_in.fbounds[end]-opt_in.fbounds[1])*[1,1],"--k")
    ax2[2][:set_ylim](0, 3*maximum(rhist))
    ax2[2][:set_xlabel](L"\log_{10}\rho")
    ax2[2][:set_ylabel]("pdf")
    xhist = sum(h,dims=2)[:]
    xhist = xhist/sum(xhist)/diff(edgesx)[1]
    ax2[3][:barh](0.5*(edgesx[2:end]+edgesx[1:end-1]), xhist, height=diff(edgesx)[1])
    ax2[3][:plot](1/(opt_in.xbounds[end]-opt_in.xbounds[1])*[1,1], [edgesx[1], edgesx[end]],"--k")
    ax2[3][:set_xlim](0, 3*maximum(xhist))
    ax2[3][:set_ylim](ax2[1][:get_ylim]())
    ax2[3][:set_ylabel]("depth m")
    ax2[3][:set_xlabel]("pdf")
    # ax2[3][:xaxis][:set_label_position]("top")
    ax2[2][:get_shared_y_axes]()[:join](ax2[1], ax2[3])
    ax2[2][:get_shared_x_axes]()[:join](ax2[1], ax2[2])
    ax2[2][:set_xlim](ax2[1][:get_xlim]())
    k = fit(Histogram, n[from:end], (opt_in.nmin-0.5:opt_in.nmax+0.5)).weights
    k = k/sum(k)
    @show mean(k)
    ax2[4][:bar](opt_in.nmin:opt_in.nmax, k, width=1)
    ax2[4][:plot]([opt_in.nmin, opt_in.nmax],1/(opt_in.nmax-opt_in.nmin+1)*[1,1],"--k")
    ax2[4][:set_xlabel]("# training points")
    ax2[4][:set_ylim](0, 3*max(k...))
    ax2[4][:set_ylabel]("pdf")
    ax2[4][:set_xlim](opt_in.nmin, opt_in.nmax)
    MCMC_Driver.nicenup(gcf(), fsize=nxftfontsize)
    savefig(fdataname*"_k.png", dpi=300)

    figure(f1[:number])
    savefig(fdataname*"_conv.png", dpi=300)
    figure(f[:number])
    savefig(fdataname*"_post.png", dpi=300)
    return ax[1], im1, cb1, ax[2], im2, cb2
end

function get_posterior(opt_in::TransD_GP.Options; fdataname="", decimate=1)

    M = TransD_GP.history(opt_in, stat=:fstar)
    n = TransD_GP.history(opt_in, stat=:nodes)
    x_ft = TransD_GP.history(opt_in, stat=:x_ftrain)
    iter = TransD_GP.history(opt_in, stat=:iter)
    misfit = load("misfit_T0_"*fdataname*".jld", "misfit")
    T0loc = load("misfit_T0_"*fdataname*".jld", "T0loc")

    return M[1:decimate:end], n[1:decimate:end],
    x_ft[1:decimate:end], iter[1:decimate:end], misfit[1:decimate:end], T0loc[1:decimate:end]
end

function get_field(mfstar::AbstractArray, d::AbstractArray,
    opt::TransD_GP.Options, opt_EM::MCMC_Driver.EMoptions;
    plotstackerrors=true, plotMLerrors=true, sfactor=1.0)
    @assert length(opt_EM.rho) == length(mfstar) + opt_EM.nfixed

    opt_EM.rho[opt_EM.nfixed+1:end] = 10.0.^mfstar

    Er = CSEM1Dkernels.getCSEM1DanisoHED(opt_EM.freqs,
    opt_EM.rRx, opt_EM.zRx, opt_EM.zTx, opt_EM.z, [opt_EM.rho opt_EM.rho],
    0, RxAzim = opt_EM.RxAzim, TxDip = opt_EM.TxDip)

    σ = zeros(size(Er))
    if plotMLerrors || plotstackerrors
        fig, ax = plt[:subplots](1,2, sharex=true, figsize=(12,8))
    end
    for lfreq = 1:length(opt_EM.freqs)
        nr = sum(.!isnan.(d[:,lfreq]))
        select = .!isnan.(d[:,lfreq])
        r = Er[select,lfreq] - d[select,lfreq]
        σ[:,lfreq] = sqrt(0.5/nr)*norm(r./abs.(d[select,lfreq]), 2)*abs.(d[:,lfreq])
        if plotstackerrors
            ax[1][:errorbar](opt_EM.rRx[select], abs.(d[select,lfreq]), yerr = 2*sfactor*opt_EM.sd[select,lfreq], linestyle="none", marker="x",
            elinewidth=0, capsize=3)
            ax[2][:errorbar](opt_EM.rRx[select], unwrap(angle.(d[select,lfreq])), yerr = 2*sfactor*opt_EM.sd[select,lfreq]./abs.(d[select,lfreq]),
            linestyle="none", marker="x", elinewidth=0, capsize=3)
        end
        if plotMLerrors
            ax[1][:errorbar](opt_EM.rRx[select], abs.(Er[select,lfreq]), yerr = 2*σ[select,lfreq],
            elinewidth=2, color="k")
            ax[2][:errorbar](opt_EM.rRx[select], unwrap(angle.(Er[select,lfreq])), yerr = 2*σ[select,lfreq]./abs.(d[select,lfreq]),
            elinewidth=2, color="k")
        end
    end
    ax[1][:set_yscale]("log")
    ax[1][:grid]()
    ax[2][:grid]()
    return Er, σ
end

function plot_post_fields(d::AbstractArray,
    opt::TransD_GP.Options,
    opt_EM::MCMC_Driver.EMoptions,
    fdataname::String;
    sfactor=1.0, decimate=1000, burnin=0, alpha=0.1, figsize=(10,7), fig2size=(12,2), fontsize=14,
    msize=10, showMLnoise=false, showresids=false, nresidbins=11)

    @assert burnin != 0
    M, = get_posterior(opt, fdataname=fdataname)
    from, = divrem(burnin, opt.save_freq)
    M = M[from+1:decimate:end]
    cols = 2
    showMLnoise && (cols=3; figsize=(15,7))
    fig, ax = plt[:subplots](1,cols, sharex=true, figsize=figsize)
    σ = zeros(size(d))
    cmap = ["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf"]
    f3, ax3 = plt[:subplots](2,1, sharex=true, figsize=figsize)
    if showresids
        fig2, ax2 = plt[:subplots](1,length(opt_EM.freqs), figsize=fig2size, sharex=true, sharey=true)
        residedges = LinRange(-5.5,5.5,nresidbins)
        residbins = 0.5*(residedges[1:end-1]+residedges[2:end])
        histresid = zeros(length(residbins))
    end
    for (i, m) in enumerate(M)
        @assert length(opt_EM.rho) == length(m) + opt_EM.nfixed
        opt_EM.rho[opt_EM.nfixed+1:end] = 10.0.^m
        Er = CSEM1Dkernels.getCSEM1DanisoHED(opt_EM.freqs,
        opt_EM.rRx, opt_EM.zRx, opt_EM.zTx, opt_EM.z, [opt_EM.rho opt_EM.rho],
        0, RxAzim = opt_EM.RxAzim, TxDip = opt_EM.TxDip)

        for lfreq = 1:length(opt_EM.freqs)
            nr = sum(.!isnan.(d[:,lfreq]))
            select = .!isnan.(d[:,lfreq])
            r = Er[select,lfreq] - d[select,lfreq]
            σ[:,lfreq] = sqrt(0.5/nr)*norm(r./abs.(d[select,lfreq]), 2)*abs.(d[:,lfreq])
            ax3[1][:plot](opt_EM.rRx[select], real(r./σ[select,lfreq]), color=cmap[lfreq], alpha=alpha, ".", markersize=msize)
            ax3[2][:plot](opt_EM.rRx[select], imag(r./σ[select,lfreq]), color=cmap[lfreq], alpha=alpha, ".", markersize=msize)
            if showresids
                histresid[:] += fit(Histogram, real(r./σ[select,lfreq]), residedges).weights +
                                fit(Histogram, imag(r./σ[select,lfreq]), residedges).weights
                # histi[:] += fit(Histogram, imag(r./σ[select,lfreq]), residedges).weights
            end
            ax[1][:plot](opt_EM.rRx[select], abs.(Er[select,lfreq]), color=cmap[lfreq], alpha=alpha)
            ax[2][:plot](opt_EM.rRx[select], unwrap(angle.(Er[select,lfreq])), color=cmap[lfreq], alpha=alpha)
            if i == length(M)
                 ax[1][:errorbar](opt_EM.rRx[select], abs.(d[select,lfreq]), yerr = 2*sfactor*opt_EM.sd[select,lfreq],
                 linestyle="none", elinewidth=0, capsize=3, label=string(opt_EM.freqs[lfreq])*" Hz")
                ax[2][:errorbar](opt_EM.rRx[select], unwrap(angle.(d[select,lfreq])), yerr = 2*sfactor*opt_EM.sd[select,lfreq]./abs.(d[select,lfreq]),
                linestyle="none", elinewidth=0, capsize=3)
                showMLnoise && ax[3][:plot](opt_EM.rRx[select], sfactor*opt_EM.sd[select,lfreq], "x", markersize=msize)

                if showresids
                    histresid = histresid/sum(histresid)/(residbins[2]-residbins[1])
                    # histi = histi/sum(histi)/(residbins[2]-residbins[1])
                    ax2[lfreq][:plot](residbins, histresid, "-r")
                    # ax2[lfreq][:plot](residbins, histi, "-b")
                    ax2[lfreq][:plot](-6:.1:6, 1/sqrt(2*pi)*exp.(-0.5*(-6:.1:6).^2),"--k")
                    ax2[lfreq][:set_title]("$(opt_EM.freqs[lfreq]) Hz")
                    ax2[lfreq][:set_xlabel]("residuals")
                    ax2[lfreq][:set_ylim](0,0.55)
                    figure(fig2[:number])
                    MCMC_Driver.nicenup(gcf(),fsize=fontsize)
                end
            end
            showMLnoise && ax[3][:plot](opt_EM.rRx[select], σ[select,lfreq], ".", color=cmap[lfreq], markersize=msize, alpha=alpha)
        end
    end
    ax[1][:set_yscale]("log")
    ax[1][:grid]()
    ax[1][:set_xlabel]("offset m")
    ax[1][:set_ylabel]("Amplitude "* L"\frac{V}{Am^2}")
    ax[2][:grid]()
    ax[2][:set_xlabel]("offset m")
    ax[2][:set_ylabel]("Phase rad")
    if showMLnoise
        ax[3][:set_xlabel]("offset m")
        ax[3][:set_ylabel](L"\^{\sigma_{|amp|}}")
        ax[3][:grid]()
        ax[3][:set_yscale]("log")
        # @info("here23")
        # ax[3][:get_shared_y_axes]()[:join](ax[1], ax[3])
        # ax[3][:autoscale]()
    end

    figure(fig[:number])
    MCMC_Driver.nicenup(gcf(),fsize=fontsize)
    savefig(fdataname*"_post_fields.png", dpi=300)
    figure(fig2[:number])
    savefig(fdataname*"_post_resids.png", dpi=300)
    nothing
end

end


text:
using LinearAlgebra, Plots, ForwardDiff, Ipopt, MathOptInterface
using Interpolations
# using SNOPT7

const MOI = MathOptInterface

include("direct_variational_methods_kuka_particle.jl")
include("/home/taylor/.julia/dev/SNOPT7/src/SNOPT7.jl")
struct Problem <: MOI.AbstractNLPEvaluator
    nq
    nu
    nγ
    ns

    p_dy
    p_sd
    p_cm

    p_dy_jac
    p_sd_jac
    p_cm_jac

    jac_sparsity

    N
    enable_hessian::Bool
end

function Problem(nq,nu,nγ,ns,N)
    shift = nq+nu+nγ+ns

    p_dy = (N-1)*nq
    p_sd = (N+1)*nγ
    p_cm = N-1

    p_dy_jac = (N-1)*nq*(shift + 2nq)
    p_sd_jac = (N+1)*nγ*nq
    p_cm_jac = (N-1)*(nγ + ns + nq)

    sparsity_dy = sparsity_dynamics()
    sparsity_sd = sparsity_signed_distance(p_dy)
    sparsity_cm = sparsity_complementarity(p_dy+p_sd)

    sparsity = []
    append!(sparsity,sparsity_dy)
    append!(sparsity,sparsity_sd)
    append!(sparsity,sparsity_cm)

    Problem(nq,nu,nγ,ns,p_dy,p_sd,p_cm,p_dy_jac,p_sd_jac,p_cm_jac,sparsity,N,false)
end

MOI.features_available(prob::Problem) = [:Grad, :Jac]
MOI.initialize(prob::Problem, features) = nothing
MOI.jacobian_structure(prob::Problem) = prob.jac_sparsity
MOI.hessian_lagrangian_structure(prob::Problem) = []

function MOI.eval_objective(prob::Problem, Z)
    return cost(Z)
end

function MOI.eval_objective_gradient(prob::Problem, grad_f, Z)
    ∇cost!(grad_f, Z)
end

function MOI.eval_constraint(prob::Problem, g, Z)
    g_dy = view(g,1:prob.p_dy)
    g_sd = view(g,prob.p_dy .+ (1:prob.p_sd))
    g_cm = view(g,(prob.p_dy + prob.p_sd) .+ (1:prob.p_cm))
    dynamics_constraints!(g_dy,Z)
    signed_distance_constraints!(g_sd,Z)
    complementarity_constraints!(g_cm,Z)
end

function MOI.eval_constraint_jacobian(prob::Problem, jac, Z)
    jac_dy = view(jac,1:prob.p_dy_jac)
    jac_sd = view(jac,prob.p_dy_jac .+ (1:prob.p_sd_jac))
    jac_cm = view(jac,(prob.p_dy_jac + prob.p_sd_jac) .+ (1:prob.p_cm_jac))

    ∇dynamics_constraints!(jac_dy,Z)
    ∇signed_distance_constraints!(jac_sd,Z)
    ∇complementarity_constraints!(jac_cm,Z)
end

MOI.eval_hessian_lagrangian(prob::Problem, H, x, σ, μ) = nothing
NN
prob = Problem(nq,nu,nγ,ns,N)

Q0 = [q0,line_trajectory(q1,qN,N)...]
# Q0 = copy(Q)
# Q0 = [q0 for k = 1:N+1]
# Q0 = [q0,[[Q_kuka[k];q0[8:10]] for k = 2:N+1]...]

# U0 = copy(U_hold)
U0 = copy(U_hold_rand)
# U0 = copy(U)
# U0 = copy(Q_kuka)
Γ0 = [0.01*rand(nγ) for k = 1:N-1]
S0 = [ones(ns) for k = 1:N-1]

# Q0 = copy(Q)
# U0 = copy(U)
# Γ0 = copy(Γ)
# S0 = copy(S)

ZZ = packZ(Q0,U0,Γ0,S0)

con_bnds = [(prob.p_dy,:equality),(prob.p_sd,:inequality),(prob.p_cm,:inequality)]

Z_low, Z_up = primal_bounds(q0,q1,qN)
c_low, c_up = constraint_bounds(con_bnds)

nlp_bounds = MOI.NLPBoundsPair.(c_low,c_up)
block_data = MOI.NLPBlockData(nlp_bounds,prob,true)

# solver = Ipopt.Optimizer()
solver = SNOPT7.Optimizer()

Z = MOI.add_variables(solver,NN)

for i = 1:NN
    zi = MOI.SingleVariable(Z[i])
    MOI.add_constraint(solver, zi, MOI.LessThan(Z_up[i]))
    MOI.add_constraint(solver, zi, MOI.GreaterThan(Z_low[i]))
    MOI.set(solver, MOI.VariablePrimalStart(), Z[i], ZZ[i])
end

# Solve the problem
MOI.set(solver, MOI.NLPBlock(), block_data)
MOI.set(solver, MOI.ObjectiveSense(), MOI.MIN_SENSE)
@time MOI.optimize!(solver)

# Get the solution
res = MOI.get(solver, MOI.VariablePrimal(), Z)

Q,U,Γ,S = unpackZ(res)

Q_kuka = [Q[k][1:7] for k = 1:N+1]
Q_ee = [Array(end_effector_function(Q[k][1:7])) for k = 1:N+1]
Q_p = [Q[k][8:10] for k = 1:N+1]
#
# Xblk = zeros(3,N+1)
# for k = 1:N+1
#     Xblk[:,k] = Q_p[k]
# end
#
# Eblk = zeros(3,N+1)
# for k = 1:N+1
#     Eblk[:,k] = Q_ee[k]
# end
#
# Ublk = zeros(nu,N-1)
# for k = 1:N-1
#     Ublk[:,k] = U[k]
# end
#
# Γblk = zeros(N-1)
# for k = 1:N-1
#     Γblk[k] = Γ[k][1]
# end
# Sblk = zeros(N-1)
# for k = 1:N-1
#     Sblk[k] = S[k][1]
# end
#
# p = plot(tspan,Xblk[1,2:end],width=2,label="x",legend=:topleft,title="Ball trajectory",xlabel="time (s)",ylabel="position (m)")
# p = plot!(tspan,Xblk[2,2:end],width=2,label="y")
# p = plot!(tspan,Xblk[3,2:end],width=2,label="z")
# savefig(p,joinpath(pwd(),"ball_in_cup_kuka/","ball_trajectory.png"))
#
# p = plot(tspan,Eblk[1,2:end],width=2,label="x",legend=:topleft,title="End effector trajectory",xlabel="time (s)",ylabel="position (m)")
# p = plot!(tspan,Eblk[2,2:end],width=2,label="y")
# p = plot!(tspan,Eblk[3,2:end],width=2,label="z")
# savefig(p,joinpath(pwd(),"ball_in_cup_kuka/","ee_trajectory.png"))
#
# plot(Ublk')
# tspan = range(0,stop=(h*(N-1)),length=N)
#
# p = plot(tspan,[Γblk;Γblk[end]],width=2,linetype=:steppost,title="Contact Force",label="",xlabel="time (s)",ylabel="force (N)")
# savefig(p,joinpath(pwd(),"ball_in_cup_kuka/","force.png"))
#
# p = plot(tspan,[Sblk;Sblk[end]],width=2,linetype=:steppost,title="max slack = $(maximum(S))",label="",xlabel="time (s)",ylabel="slack",ylim=(-1,1))
# savefig(p,joinpath(pwd(),"ball_in_cup_kuka/","slack.png"))
# maximum(S)
#
# maximum(S)

using MeshCatMechanisms
using MeshCat
using GeometryTypes
using CoordinateTransformations

vis = MeshCatMechanisms.Visualizer()
open(vis)
mvis = MechanismVisualizer(mechanismkuka, visualskuka, vis[:base])

function MeshCat.setanimation!(mvis::MechanismVisualizer,
                      times::AbstractVector{<:Real},
                      configurations::AbstractVector{<:AbstractVector{<:Real}},
                      config_p::AbstractVector{<:AbstractVector{<:Real}};
                      fps::Integer=30,
                      play::Bool=true,
                      repetitions::Integer=1)

    q0 = copy(configuration(MeshCatMechanisms.state(mvis)))
    @assert axes(times) == axes(configurations)
    @assert axes(times) == axes(config_p)

    interpolated_configurations = Interpolations.interpolate((times,), configurations, Gridded(Linear()))
    interpolated_configurations2 = Interpolations.interpolate((times,), config_p, Gridded(Linear()))

    animation = MeshCatMechanisms.Animation()
    num_frames = floor(Int, (times[end] - first(times)) * fps)

    setobject!(mvis["load"],HyperSphere(Point3f0(0), convert(Float32,0.05)) ,MeshPhongMaterial(color=RGBA(1, 0, 0, 1.0)))
    settransform!(mvis["load"], Translation(interpolated_configurations2(0.0)...))

    for frame in 0:num_frames
        time = first(times) + frame / fps
        set_configuration!(MeshCatMechanisms.state(mvis), interpolated_configurations(time))

        atframe(animation, MeshCatMechanisms.visualizer(mvis), frame) do frame_visualizer
            # println(interpolated_configurations2(time))
            settransform!(frame_visualizer["load"], Translation(interpolated_configurations2(time)...))
            MeshCatMechanisms._render_state!(MechanismVisualizer(MeshCatMechanisms.state(mvis), frame_visualizer))
        end
    end
    setanimation!(MeshCatMechanisms.visualizer(mvis), animation, play=play, repetitions=repetitions)
    set_configuration!(MeshCatMechanisms.state(mvis), q0)
end

setanimation!(mvis,t,Q_kuka[2:end],Q_p[2:end])

function output_traj(Q,idx=2:N+1,filename=joinpath(pwd(),"kuka.txt"))
    f = open(filename,"w")
    for k = idx
        q1, q2, q3, q4, q5, q6, q7 = Q[k][1:7]
        str = "$q1 $q2 $q3 $q4 $q5 $q6 $q7"
        if k != N
            str *= " "
        end
        write(f,str)
    end

    close(f)
end

output_traj(Q)
output_traj(U,collect(1:N-1),joinpath(pwd(),"kuka_control.txt"))


text:
import NLopt, ForwardDiff, DiffResults

import Base: setproperty!, getproperty
function auto_diff(f)
    (x, g) -> if !isempty(g)
        res = DiffResults.DiffResult(0.0, g)
        ForwardDiff.gradient!(res, f, x)
        res.value
    else
        f(x)
    end
end

function combine_gradient(f, gf)
    (x, g) -> if !isempty(g)
        g .= gf(x)
        f(x)
    else
        f(x)
    end
end

function wrap_function(f, g)
    if g === :Auto || g === :AUTO
        auto_diff(f)
    else
        combine_gradient(f, g)
    end
end

function create_nlopt(algo, lb, ub; options...)
    opt = NLopt.Opt(algo, length(lb))
    opt.lower_bounds = lb
    opt.upper_bounds = ub
    foreach(options) do (p, v)
        setproperty!(opt, p, v)
    end
    return opt
end


text:
# regularize/z-test.jl

using Test: @test

@test Aodwt(:test)
@test diff_map(:test)


text:
module SpringCollab2020StaticPuffer

using ..Ahorn, Maple

@mapdef Entity "SpringCollab2020/StaticPuffer" StaticPuffer(x::Integer, y::Integer, right::Bool=false)

const placements = Ahorn.PlacementDict(
    "Static Puffer (Right) (Spring Collab 2020)" => Ahorn.EntityPlacement(
        StaticPuffer,
        "point",
        Dict{String, Any}(
            "right" => true
        )
    ),
    "Static Puffer (Left) (Spring Collab 2020)" => Ahorn.EntityPlacement(
        StaticPuffer,
        "point",
        Dict{String, Any}(
            "right" => false
        )
    )
)

sprite = "objects/puffer/idle00"

function Ahorn.selection(entity::StaticPuffer)
    x, y = Ahorn.position(entity)
    scaleX = get(entity, "right", false) ? 1 : -1

    return Ahorn.getSpriteRectangle(sprite, x, y, sx=scaleX)
end

function Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::StaticPuffer, room::Maple.Room)
    scaleX = get(entity, "right", false) ? 1 : -1

    Ahorn.drawSprite(ctx, sprite, 0, 0, sx=scaleX)
end

end


text:
# ---
# title: 1631. Path With Minimum Effort
# id: problem1631
# author: Indigo
# date: 2021-01-30
# difficulty: Medium
# categories: Binary Search, Depth-first Search, Union Find, Graph
# link: <https://leetcode.com/problems/path-with-minimum-effort/description/>
# hidden: true
# ---
#
# You are a hiker preparing for an upcoming hike. You are given `heights`, a 2D
# array of size `rows x columns`, where `heights[row][col]` represents the
# height of cell `(row, col)`. You are situated in the top-left cell, `(0, 0)`,
# and you hope to travel to the bottom-right cell, `(rows-1, columns-1)` (i.e.,
# **0-indexed** ). You can move **up** , **down** , **left** , or **right** ,
# and you wish to find a route that requires the minimum **effort**.
#
# A route's **effort** is the **maximum absolute difference** **** in heights
# between two consecutive cells of the route.
#
# Return _the minimum **effort** required to travel from the top-left cell to
# the bottom-right cell._
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/10/04/ex1.png)
#
#
#
#     Input: heights = [[1,2,2],[3,8,2],[5,3,5]]
#     Output: 2
#     Explanation: The route of [1,3,5,3,5] has a maximum absolute difference of 2 in consecutive cells.
#     This is better than the route of [1,2,2,2,5], where the maximum absolute difference is 3.
#
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2020/10/04/ex2.png)
#
#
#
#     Input: heights = [[1,2,3],[3,8,4],[5,3,5]]
#     Output: 1
#     Explanation: The route of [1,2,3,4,5] has a maximum absolute difference of 1 in consecutive cells, which is better than route [1,3,5,3,5].
#
#
# **Example 3:**
#
# ![](https://assets.leetcode.com/uploads/2020/10/04/ex3.png)
#
#
#
#     Input: heights = [[1,2,1,1,1],[1,2,1,2,1],[1,2,1,2,1],[1,2,1,2,1],[1,1,1,2,1]]
#     Output: 0
#     Explanation: This route does not require any effort.
#
#
#
#
# **Constraints:**
#
#   * `rows == heights.length`
#   * `columns == heights[i].length`
#   * `1 <= rows, columns <= 100`
#   * `1 <= heights[i][j] <= 106`
#
#
## @lc code=start
using LeetCode

function minimum_effort_path(heights::Vector{Vector{Int}})
    m, n = length(heights), length(heights[1])
    ds = IntDisjointSets(m * n)
    map = SortedMultiDict{Int, Pair{Int, Int}}()
    for i in 1:m-1
        for j in 1:n-1
            insert!(map, abs(heights[i][j] - heights[i+1][j]), Pair((i - 1) * n + j, i * n + j))
            insert!(map, abs(heights[i][j] - heights[i][j+1]), Pair((i - 1) * n + j, (i - 1) * n + j + 1))
        end
    end
    for i in 1:m-1
        insert!(map, abs(heights[i][n] - heights[i+1][n]), Pair(i * n, (i + 1) * n))
    end
    for j in 1:n-1
        insert!(map, abs(heights[m][j] - heights[m][j+1]), Pair((m-1) * n + j, (m-1) * n + j + 1))
    end
    for (dist, edge) in map
        union!(ds, edge.first, edge.second)
        if in_same_set(ds, 1, m * n)
            return dist
        end
    end
    return 0
end
## @lc code=end


text:
#using GLOQ
#x = rand(Float64,8)
#X = [0 x[1]+x[2]*im x[3]+x[4]*im;
#     0 x[5]         x[6]+x[7]*im;
#     0 0            x[8]]
#H = 0.5*(X+adjoint(X))
HK = [ 0.0 0.0 1.0;
       0.0 0.0 0.0;
       1.0 0.0 0.0]
HS = [ 0.0 1.0 0.0;
      -1.0 0.0 0.0;
       0.0 0.0 0.0]
H = HK + 1.0im.*HS
L1 = [0.0 0.25 0.0;
      0.0 0.0  0.5;
      0.0 0.0  0.0]
L2 = [0.0 0.0 0.0;
      0.0 0.5 0.0;
      0.0 0.0 1.0]
L_list = (L1,L2)
LH,LD = GLOQ.make_lindblad_operator_complex(H,L_list)
LK,LS,LD2 = GLOQ.make_lindblad_operator(HK,HS,L_list)
L = LH+LD
println("Test whether the real-valued and complex-valued operator match each other")
@test (LD==LD2)&&( LH==-1.0im*(LK+1.0im*LS) )

println("\n Lindblad operator assembled\n")

tspan_length = 21
t_span = range(0.0,pi,length=tspan_length)
t_span2 = zeros(Float64,tspan_length)
dt = pi/(tspan_length-1);
for i = 1:tspan_length
      t_span2[i] = (i-1)*dt
end
u1 = [1.0;0;0]
rho1_history = GLOQ.exponential_solver_complex(u1,L,t_span;initial_type = "states")
u2 = [0;1.0;0]
rho2 = (u2*u2')[:]
rho2_history = GLOQ.exponential_solver_complex(rho2,L,t_span2;initial_type = "density")
u3 = [1.0/sqrt(2.0);-im/sqrt(2.0);0]
rho3_history = GLOQ.exponential_solver_complex(u3,L,t_span;initial_type = "states")
u4 = [0.0;1.0/sqrt(2.0);-im/sqrt(2.0)]
rho4_history = GLOQ.exponential_solver_complex(u4,L,t_span;initial_type = "states")
u5 = [1.0/sqrt(2.0);0.0;-im/sqrt(2.0)]
rho5_history = GLOQ.exponential_solver_complex(u5,L,t_span;initial_type = "states")

population1_history = GLOQ.get_population(rho1_history)
population2_history = GLOQ.get_population(rho2_history)
population3_history = GLOQ.get_population(rho3_history)
population4_history = GLOQ.get_population(rho4_history)
population5_history = GLOQ.get_population(rho5_history)

#################################################
# test real solvers
rho1_history_u,rho1_history_v = GLOQ.exponential_solver(real(u1),-imag(u1),LK,LS,LD2,t_span;initial_type = "states")
rho2_u,rho2_v = GLOQ.convert_state_to_density(real(u2),-imag(u2))
rho2_history_u,rho2_history_v = GLOQ.exponential_solver(rho2_u,rho2_v,LK,LS,LD2,t_span;initial_type = "density")
rho3_history_u,rho3_history_v = GLOQ.exponential_solver(real(u3),-imag(u3),-imag(LH),real(LH),LD,t_span;initial_type = "states")
rho4_u,rho4_v = GLOQ.convert_state_to_density(real(u4),-imag(u4))
rho4_history_u,rho4_history_v = GLOQ.exponential_solver(rho4_u,rho4_v,LK,LS,LD2,t_span;initial_type = "density")
rho5_history_u,rho5_history_v = GLOQ.exponential_solver(real(u5),-imag(u5),LK,LS,LD2,t_span;initial_type = "states")

@test ( (norm(rho1_history_u-1.0im*rho1_history_v-rho1_history,Inf)<1e-14)&&
        (norm(rho2_history_u-1.0im*rho2_history_v-rho2_history,Inf)<1e-14)&&
        (norm(rho3_history_u-1.0im*rho3_history_v-rho3_history,Inf)<1e-14)&&
        (norm(rho4_history_u-1.0im*rho4_history_v-rho4_history,Inf)<1e-14)&&
        (norm(rho5_history_u-1.0im*rho5_history_v-rho5_history,Inf)<1e-14) )

population1_history2 = GLOQ.get_population(rho1_history_u)
population2_history2 = GLOQ.get_population(rho2_history_u)
population3_history2 = GLOQ.get_population(rho3_history_u)
population4_history2 = GLOQ.get_population(rho4_history_u)
population5_history2 = GLOQ.get_population(rho5_history_u)

@test ( (norm(population1_history-population1_history2,Inf)<1e-14)&&
        (norm(population2_history-population2_history2,Inf)<1e-14)&&
        (norm(population3_history-population3_history2,Inf)<1e-14)&&
        (norm(population4_history-population4_history2,Inf)<1e-14)&&
        (norm(population5_history-population5_history2,Inf)<1e-14) )

# check the interface with DifferentialEquations

rho3 = (u3*u3')[:]
problem3 = GLOQ.LindbladODEProblemComplex(rho3,L,t_span[end];initial_type="density")
rho3_history_diffeq = solve(problem3,saveat=t_span)
difference = maximum(abs.(rho3_history-rho3_history_diffeq))
@test(difference<5e-7)

N_density = length(u3)^2;
problem3_real = GLOQ.LindbladODEProblem(real(u3),-imag(u3),LK,LS,LD,t_span[end];initial_type="states")
rho3_history_diffeq_real = solve(problem3_real,saveat=t_span)

println(maximum(abs.(real(rho3_history)-rho3_history_diffeq_real[1:N_density,:])))
println(maximum(abs.(imag(rho3_history)+rho3_history_diffeq_real[N_density+1:end,:])))
@test( (maximum(abs.(real(rho3_history)-rho3_history_diffeq_real[1:N_density,:]))<5e-5)&&
       (maximum(abs.(imag(rho3_history)+rho3_history_diffeq_real[N_density+1:end,:]))<5e-5) )


text:
export Maskedtrain1st!, Standardtrain1st!

using Flux
using Flux: onehotbatch
using Juno
using Base.Iterators
using StatsBase: mean
using Zygote
using Zygote: @nograd
using LinearAlgebra
using CUDA

function perturb(αs::AbstractArray)
    counter = [ones(length(a)) for a in αs]
    while sum(sum(counter)) > 0
        rn = rand([i for i in 1:length(counter) if sum(counter[i]) > 0])
        row = rand([i for i in 1:length(counter[rn]) if counter[rn][i] == 1])
        inds = findall(softmax(αs[rn][row]) .> 0)
        if length(inds) <= 1
            counter[rn][row] = 0
        else
            perturbs = [copy(αs) for i in inds]
            for i in 1:length(inds)
                perturbs[i] = deepcopy(αs)
                perturbs[i][rn][row][inds[i]] = -Inf32
            end
            return (rn, row, inds, perturbs)
        end
    end
    return (-1, -1, [], [])
end

function Standardtrain1st!(accuracy, loss, model, train, opt, losses=[0.0,0.0]; cbepoch = () -> (), cbbatch = () -> ())
    local train_loss
    w = all_ws_sansbn(model)
    for train_batch in TrainCuIterator(train)
        gsw = gradient(w) do
            train_loss = loss(model, train_batch...)
            return train_loss
        end
        losses[1] = train_loss
        CUDA.reclaim()
        GC.gc()
        Flux.Optimise.update!(opt, w, gsw)
        cbbatch()
    end
    cbepoch()
end

function Maskedtrain1st!(accuracy, loss, model, train, val, opt, losses=[0.0,0.0]; cbepoch = () -> (), cbbatch = () -> ())
    w = all_ws_sansbn(model)
    for _ in 1:1
        rn, row, inds, perturbs = perturb([model.normal_αs, model.reduce_αs])
        if rn == -1
            continue
        end
        vals = [accuracy(model, val, pert = pert) for pert in perturbs]
        to_remove = sortperm(vals)[1:length(vals)-1]
        if rn == 1
            model.normal_αs[row][to_remove] .= -Inf32
            display((rn, row, softmax(model.normal_αs[row])))
        else
            model.reduce_αs[row][to_remove] .= -Inf32
            display((rn. row, softmax(model.reduce_αs[row])))
        end
    end
    cbbatch()
    for train_batch in CuIterator(train)
        gsw = gradient(w) do
            train_loss = loss(model, train_batch...)
            return train_loss
        end
        losses[1] = train_loss
        CUDA.reclaim()
        GC.gc()
        Flux.Optimise.update!(opt, w, gsw)
        cbbatch()
    end

    cbepoch()
end

max_stars_repo_name: SuReLI/DifferentiableNAS

max_issues_repo_name: SuReLI/DifferentiableNAS

max_forks_repo_name: SuReLI/DifferentiableNAS


text:
module ImageReconstruction

using Base.Threads
using Interpolations
using FFTW
using ImageTransformations
export radon, iradon, sinogram

"""
    radon(I::AbstractMatrix, θ::AbstractRange, t::AbstractRange)

Radon transform of a image `I` producing a sinogram from view angles `θ` in
radians and detector sampling `t`.
"""
function radon(I::AbstractMatrix, θ::AbstractRange, t::AbstractRange)
    P = zeros(eltype(I), length(t), length(θ))
    ax1, ax2 = axes(I)

    nax1, nax2 = length(ax1), length(ax2)
    for j in ax2, i in ax1
        x = j - nax2 / 2 + 0.5
        y = i - nax1 / 2 + 0.5
        @inbounds for (k, θₖ) in enumerate(θ)
            t′ = x * cos(θₖ) + y * sin(θₖ)

            a = convert(Int, round((t′ - minimum(t)) / step(t) + 1))

            (a < 1 || a > length(t)) && continue
            α = abs(t′ - t[a])

            I′ = I[i, j]
            P[a, k] += (1 - α) * I′

            (a > length(t) + 1) && continue
            P[a+1, k] += α * I′
        end
    end

    P
end

function _ramp_spatial(N::Int, τ, Npad::Int = N)
    @assert Npad ≥ N
    N2 = N ÷ 2
    hval(n) = n == 0 ? 1 / (4*τ^2) : - mod(n, 2)/(π * n * τ)^2
    [i ≤ N ? hval(i - N2 - 1) : 0. for i = 1:Npad]
end

"""
	iradon(P::AbstractMatrix, θ::AbstractRange, t::AbstractRange)

Inverse radon transform of a sinogram `P` with view angles `θ` in radians and
detector sampling `t` producing an image on a 128x128 matrix.
"""
function iradon(P::AbstractMatrix, θ::AbstractRange, t::AbstractRange)
    pixels = 128

    N = length(t)
    K = length(θ)
    Npad = nextpow(2, 2 * N - 1)
    τ = step(t)
    ramp = fft(_ramp_spatial(N, τ, Npad))
    i = N ÷ 2 + 1
    j = i + N - 1

    T = eltype(P)
    I = [zeros(T, pixels, pixels) for _ = 1:nthreads()]
    P′ = [Vector{Complex{T}}(undef, Npad) for _ = 1:nthreads()]
    Q = [Vector{T}(undef, N) for _ = 1:nthreads()]

    l = SpinLock()
    @inbounds @threads for (k, θₖ) in collect(enumerate(θ))
        id = threadid()
        Pid, Qid, Iid = P′[id], Q[id], I[id]

        # filter projection
        Pid[1:N] .= P[:, k]
        Pid[N+1:end] .= 0

        # Need to prevent multiple thread execution during fft/ifft.
        # double free occurs otherwise.
        # https://github.com/JuliaMath/FFTW.jl/issues/134
        lock(l)
        fft!(Pid)
        Pid .*= ramp
        ifft!(Pid)
        unlock(l)
        Qid .= τ .* real.(@view Pid[i:j])

        Qₖ = LinearInterpolation(t, Qid)

        # backproject
        for c in CartesianIndices(first(I))
            x = c.I[2] - pixels ÷ 2 + 0.5
            y = c.I[1] - pixels ÷ 2 + 0.5
            x^2 + y^2 ≥ pixels^2 / 4 && continue
            t′ = x * cos(θₖ) + y * sin(θₖ)
            Iid[c] += Qₖ(t′)
        end
    end
    sum(I) .* π ./ K
end

"""
    sinogram(I::AbstractMatrix, n::AbstractRange)

Generate a sinogram `I` with `n` views  and producing an image on a 128x128 matrix.
"""
function sinogram(I::AbstractMatrix, n::Int64)

    angles = LinRange(1, 180, n)

    #Add padding (zeros) to avoid data loss while rotating
    image_size = 128
    image_diagonal = sqrt(image_size^2 + image_size^2)
    padding_amount = ceil(Int64,image_diagonal - image_size) + 2
    padded_image = zeros(image_size+padding_amount, image_size+padding_amount)
    padded_image[ceil(Int,padding_amount/2):(ceil(Int,padding_amount/2)+image_size-1),
                 ceil(Int,padding_amount/2):(ceil(Int,padding_amount/2)+image_size-1)] = I

    # Rotate the image for each angle and sum up the columns of each projected image to create sinogram
    #n = length(angles)
    sinogram_image = zeros(size(padded_image,2), n)
    for i = 1:n
        #temp_image = imrotate(padded_image, 90-angles(i), 'bilinear', 'crop')
        temp_image = imrotate(padded_image, 2*(90-angles[i])/pi)
        sinogram_image[:,i] .= sum(temp_image)
    end
    return sinogram_image
end

end # module


text:
"""
    DeathDistribution

An abstract type used to form an assumption of how deaths occur throughout a
    year. See `Balducci()`, `Uniform()`, and `Constant()` for concrete
    assumption types.
"""
abstract type DeathDistribution end

"""
    Balducci()

A `DeathDistribution` type that assumes a decreasing force of mortality
over the year.
"""
struct Balducci <: DeathDistribution end

"""
    Uniform()

A `DeathDistribution` type that assumes an increasing force of mortality
over the year.
"""
struct Uniform <: DeathDistribution end

"""
    Constant()

A `DeathDistribution` type that assumes a constant force of mortality
over the year.
"""
struct Constant <: DeathDistribution end


text:
using DelimitedFiles

function calculate_cost(position, position_vec)
    diff_vec = map(x->abs(x - position), position_vec)
    return sum(diff_vec)
end

function test_calculate_cost()
    test_list = [16,1,2,0,4,2,7,1,2,14]
    @assert calculate_cost(2, test_list) == 37
    @assert calculate_cost(1, test_list) == 41
    @assert calculate_cost(3, test_list) == 39
    @assert calculate_cost(10, test_list) == 71
end

test_calculate_cost()

# read input file, note single quotes around the delimiter
input = readdlm("input.txt", ',',Int)
# convert to a column vector
input = vec(input)

min_cost = typemax(Int)
min_cost_position = 0
for i in minimum(input):maximum(input)
    cost = calculate_cost(i, input)
    if cost < min_cost
        global min_cost = cost
        global min_cost_position = i
    end
end

println("Best position ", min_cost_position, " cost ", min_cost)


text:
function generate(rng::AbstractRNG,
                  ::Type{EventObservations},
                  sim::Simulation{S, M},
                  delay_infection::UnivariateDistribution{D},
                  delay_removal::UnivariateDistribution{D};
                  seq_len::Union{Nothing, Int64}=nothing,
                  force::Bool = false) where {
                    S <: Union{SEIR, SIR},
                    M <: ILM,
                    D <: Distributions.ValueSupport}
  infection = fill(NaN, individuals(sim.events))
  removal = fill(NaN, individuals(sim.events))
  if force
    for i in findall(sim.events.infection .!== NaN)
      if sim.events.infection[i] == -Inf
        infection[i] = -Inf
        if sim.events.removal[i] == -Inf
          removal[i] = -Inf
        else
          removal[i] = sim.events.removal[i] + rand(rng, delay_removal)
        end
      else
        if isnan(sim.events.removal[i])
          infection[i] = sim.events.infection[i] + rand(rng, delay_infection)
        elseif sim.events.removal[i] > -Inf
          infection_delay_ub = sim.events.removal[i] - sim.events.infection[i]
          infection[i] = sim.events.infection[i] +
            rand(rng, Truncated(delay_infection, 0.0, infection_delay_ub))
          removal[i] = sim.events.removal[i] + rand(rng, delay_removal)
        end
      end
      @debug "Infection observation of i = $i at t = $(round(infection[i], digits=3)) (actual infection onset at t = $(round(sim.events.infection[i], digits=3)))"
      @debug "Removal observation of i = $i at t = $(round(removal[i], digits=3)) (actual removal at t = $(round(sim.events.removal[i], digits=3)))"
    end
  else
    for i in findall(sim.events.infection .!== NaN)
      if sim.events.infection[i] == -Inf
        infection[i] = -Inf
        if sim.events.removal == -Inf
          removal[i] = -Inf
        else
          removal[i] = sim.events.removal[i] + rand(rng, delay_removal)
        end
      else
        infection_delay = rand(rng, delay_infection)
        if sim.events.removal[i] === NaN
          infection[i] = sim.events.infection[i] + infection_delay
        else
          if infection_delay + sim.events.infection[i] < sim.events.removal[i]
            infection[i] = sim.events.infection[i] + infection_delay
            removal[i] = sim.events.removal[i] + rand(rng, delay_removal)
          end
        end
      end
      @debug "Infection observation of i = $i t = $(round(infection[i], digits=3)) (actual infection onset at t = $(round(sim.events.infection[i], digits=3)))"
      @debug "Removal observation of i = $i at t = $(round(removal[i], digits=3)) (actual removal at t = $(round(sim.events.removal[i], digits=3)))"
    end
  end
  if M == PhyloILM
    seq_len == nothing && error("Must specify a sequence length with $(M)s")
    tree, event_nodes = generate(Tree,
                                sim.events,
                                infection,
                                sim.transmission_network)
    seq_full = simulate(RNASeq, tree, sim.substitution_model, seq_len)
    seq_obs = Union{Nothing, RNASeq}[isnothing(event_nodes[i]) ? nothing : seq_full[event_nodes[i]] for i = eachindex(infection)]
    return EventObservations{S}(infection, removal, seq_obs)
  elseif M == TNILM
    return EventObservations{S}(infection, removal)
  end
end

function generate(rng::AbstractRNG,
                  ::Type{EventObservations},
                  sim::Simulation{S, M},
                  delay_infection::UnivariateDistribution{D};
                  seq_len::Union{Nothing, Int64}=nothing) where {
                    S <: Union{SEI, SI},
                    M <: ILM,
                    D <: Distributions.ValueSupport}
  infection = fill(NaN, individuals(sim.events))
  @simd for i in findall(.!isnan.(sim.events.infection))
    if sim.events.infection[i] == -Inf
      infection[i] = -Inf
    else
      infection[i] = sim.events.infection[i] + rand(rng, delay_infection)
    end
    @debug "Infection observation of i = $i at t = $(round(infection[i], digits=3))) (actual infection onset at t = $(round(sim.events.infection[i], digits=3)))"
  end
  if M == PhyloILM
    seq_len == nothing && error("Must specify a sequence length with $(M)s")
    tree, event_nodes = generate(Tree,
                                sim.events,
                                infection,
                                sim.transmission_network)
    seq_full = simulate(RNASeq, tree, sim.substitution_model, seq_len)
    seq_obs = Union{Nothing, RNASeq}[isnothing(event_nodes[i]) ? nothing : seq_full[event_nodes[i]] for i = eachindex(infection)]
    return EventObservations{S}(infection, seq_obs)
  elseif M == TNILM
    return EventObservations{S}(infection)
  end
end

observe(rng::AbstractRNG, x::Simulation, y::UnivariateDistribution, z::UnivariateDistribution; kwargs...) = generate(rng, EventObservations, x, y, z; kwargs...)
observe(rng::AbstractRNG, x::Simulation, y::UnivariateDistribution; kwargs...) = generate(rng, EventObservations, x, y; kwargs...)

observe(x::Simulation, y::UnivariateDistribution, z::UnivariateDistribution; kwargs...) = observe(default_rng(), x,y,z; kwargs...)
observe(x::Simulation, y::UnivariateDistribution; kwargs...) = observe(default_rng(), x, y; kwargs...)


text:
module ShellStructureTopo

include("topo.jl")

end # module


text:
function _dft_filesize_old_state(e::CustomEvent, file::String)
    get!(e, file, 0.0)
end

function _dft_filesize_new_state(::CustomEvent, file::String)
    filesize(file)
end

function _dft_filesize_update!(e::CustomEvent, file::String)
    new_size = new_state(e, file)
    setindex!(e, new_size, file)
    return e
end

function _dft_filesize_trigger(::CustomEvent, old_size, new_size)
    return old_size < new_size
end

function FileSizeEvent()
    return CustomEvent{String, Float64}(;
        old_state = _dft_filesize_old_state,
        new_state = _dft_filesize_new_state,
        update! = _dft_filesize_update!,
        trigger = _dft_filesize_trigger,
    )
end


text:
function entropies()

    # H[q(α)] + H[q(β)] +  H[q(f)]

end


text:
using CUDA.CUDNN:
    cudnnRNNForward,
    cudnnRNNForward!,
    cudnnRNNBackwardData_v8,
    cudnnRNNBackwardWeights_v8,
    cudnnRNNDescriptor,
    cudnnRNNDescriptor_t,
    cudnnSetRNNDescriptor_v8,
    cudnnGetRNNWeightSpaceSize,
    cudnnGetRNNTempSpaceSizes,
    cudnnRNNAlgo_t,
        CUDNN_RNN_ALGO_STANDARD,        # 0, robust performance across a wide range of network parameters
        CUDNN_RNN_ALGO_PERSIST_STATIC,  # 1, fast when the first dimension of the input tensor is small (meaning, a small minibatch), cc>=6.0
        CUDNN_RNN_ALGO_PERSIST_DYNAMIC, # 2, similar to static, optimize using the specific parameters of the network and active GPU, cc>=6.0
        CUDNN_RNN_ALGO_COUNT,           # 3
    cudnnRNNMode_t,
        CUDNN_RNN_RELU, # 0, /* basic RNN cell type with ReLu activation */
        CUDNN_RNN_TANH, # 1, /* basic RNN cell type with tanh activation */
        CUDNN_LSTM,     # 2, /* LSTM with optional recurrent projection and clipping */
        CUDNN_GRU,      # 3, /* Using h' = tanh(r * Uh(t-1) + Wx) and h = (1 - z) * h' + z * h(t-1); */
    cudnnRNNBiasMode_t,
        CUDNN_RNN_NO_BIAS,         # 0, /* rnn cell formulas do not use biases */
        CUDNN_RNN_SINGLE_INP_BIAS, # 1, /* rnn cell formulas use one input bias in input GEMM */
        CUDNN_RNN_DOUBLE_BIAS,     # 2, /* default, rnn cell formulas use two bias vectors */
        CUDNN_RNN_SINGLE_REC_BIAS, # 3  /* rnn cell formulas use one recurrent bias in recurrent GEMM */
    cudnnDirectionMode_t,
        CUDNN_UNIDIRECTIONAL, # 0, /* single direction network */
        CUDNN_BIDIRECTIONAL,  # 1, /* output concatination at each layer */
    cudnnRNNInputMode_t,
        CUDNN_LINEAR_INPUT, # 0, /* adjustable weight matrix in first layer input GEMM */
        CUDNN_SKIP_INPUT,   # 1, /* fixed identity matrix in the first layer input GEMM */
    cudnnMathType_t,
        CUDNN_DEFAULT_MATH,                    # 0,
        CUDNN_TENSOR_OP_MATH,                  # 1,
        CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION, # 2,
        CUDNN_FMA_MATH,                        # 3,
    #/* For auxFlags in cudnnSetRNNDescriptor_v8() and cudnnSetRNNPaddingMode() */
        CUDNN_RNN_PADDED_IO_DISABLED, # 0
        CUDNN_RNN_PADDED_IO_ENABLED,  # (1U << 0)
    cudnnForwardMode_t,
        CUDNN_FWD_MODE_INFERENCE, # 0
        CUDNN_FWD_MODE_TRAINING,  # 1
    cudnnRNNDataDescriptor_t,
    cudnnSetRNNDataDescriptor,
    cudnnRNNDataLayout_t,
        CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED,   # 0, /* padded, outer stride from one time-step to the next */
        CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_PACKED,     # 1, /* sequence length sorted and packed as in basic RNN api */
        CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED, # 2, /* padded, outer stride from one batch to the next */
    cudnnWgradMode_t,
        CUDNN_WGRAD_MODE_ADD, # 0, /* add partial gradients to wgrad output buffers */
        CUDNN_WGRAD_MODE_SET, # 1, /* write partial gradients to wgrad output buffers */
    cudnnTensorDescriptor,
    cudnnDropoutDescriptor,
    cudnnDataType,
    math_mode

@testset "cudnn/rnn" begin
    X,H,B,T = 8,8,4,2
    w = CUDA.randn(10000)
    x = CUDA.randn(X,B,T)
    hx1 = CUDA.randn(H,B,1)
    cx1 = CUDA.randn(H,B,1)

    function rnntest(;
        hx = nothing,
        cx = nothing,
        hy = nothing,
        cy = nothing,
        layout::cudnnRNNDataLayout_t = CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED,
        seqLengthArray::Union{Nothing,Vector{Cint}} = nothing,
        fwdMode::cudnnForwardMode_t = CUDNN_FWD_MODE_INFERENCE,
        # descriptor keywords
        hiddenSize::Integer = H,
        algo::cudnnRNNAlgo_t = CUDNN_RNN_ALGO_STANDARD,
        cellMode::cudnnRNNMode_t = CUDNN_LSTM,
        biasMode::cudnnRNNBiasMode_t = CUDNN_RNN_DOUBLE_BIAS,
        dirMode::cudnnDirectionMode_t = CUDNN_UNIDIRECTIONAL,
        inputMode::cudnnRNNInputMode_t = CUDNN_LINEAR_INPUT,
        mathPrec::DataType = eltype(x),
        mathType::cudnnMathType_t = math_mode(),
        inputSize::Integer = size(x,1),
        projSize::Integer = hiddenSize,
        numLayers::Integer = 1,
        dropout::Real = 0,
        auxFlags::Integer = CUDNN_RNN_PADDED_IO_ENABLED,
    )
        d = cudnnRNNDescriptor(algo, cellMode, biasMode, dirMode, inputMode,
                               cudnnDataType(eltype(x)), cudnnDataType(mathPrec), mathType,
                               Int32(inputSize), Int32(hiddenSize), Int32(projSize),
                               Int32(numLayers), cudnnDropoutDescriptor(Cfloat(dropout)),
                               UInt32(auxFlags))
        y = cudnnRNNForward(w, x; hx, cx, hy, cy, layout, seqLengthArray, fwdMode,
                            hiddenSize, algo, cellMode, biasMode, dirMode, inputMode,
                            mathPrec, mathType, inputSize, projSize, numLayers, dropout,
                            auxFlags)
        _y = copy(y)
        _hy = (hy === nothing ? hy : copy(hy[]))
        _cy = (cy === nothing ? cy : copy(cy[]))
        @test _y ≈ cudnnRNNForward!(y, w, x; hx, cx, hy, cy, layout, seqLengthArray, fwdMode,
                                    hiddenSize, algo, cellMode, biasMode, dirMode, inputMode,
                                    mathPrec, mathType, inputSize, projSize, numLayers, dropout, auxFlags)
        (_hy === hy === nothing || _hy ≈ hy[])
        (_cy === cy === nothing || _cy ≈ cy[])
        @test _y ≈ cudnnRNNForward(w, x, d; hx, cx, hy, cy, layout, seqLengthArray, fwdMode)
        (_hy === hy === nothing || _hy ≈ hy[])
        (_cy === cy === nothing || _cy ≈ cy[])
        @test _y ≈ cudnnRNNForward!(y, w, x, d; hx, cx, hy, cy, layout, seqLengthArray, fwdMode)
        (_hy === hy === nothing || _hy ≈ hy[])
        (_cy === cy === nothing || _cy ≈ cy[])
    end

    rnntest()
    rnntest(hx=hx1)
    rnntest(cx=cx1)
    rnntest(hy=Ref{Any}())
    rnntest(cy=Ref{Any}())
    rnntest(layout=CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_PACKED)
    rnntest(layout=CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED)
    rnntest(seqLengthArray=Cint[1,2,1,2])
    rnntest(fwdMode=CUDNN_FWD_MODE_TRAINING)
    rnntest(hiddenSize=16)
    # XXX: it's unclear which devices support this algorithm
    if capability(device()) >= v"6.1"
        rnntest(algo=CUDNN_RNN_ALGO_PERSIST_STATIC)
    end
    #rnntest(algo=CUDNN_RNN_ALGO_PERSIST_DYNAMIC) # causes segfault
    rnntest(cellMode=CUDNN_RNN_RELU)
    rnntest(cellMode=CUDNN_RNN_TANH)
    rnntest(cellMode=CUDNN_GRU)
    rnntest(biasMode=CUDNN_RNN_NO_BIAS)
    rnntest(biasMode=CUDNN_RNN_SINGLE_INP_BIAS)
    rnntest(biasMode=CUDNN_RNN_SINGLE_REC_BIAS)
    rnntest(dirMode=CUDNN_BIDIRECTIONAL)
    rnntest(inputMode=CUDNN_SKIP_INPUT)
    rnntest(mathPrec=Float32) # only possible option for F32 input
    rnntest(mathType=CUDNN_DEFAULT_MATH)
    rnntest(mathType=CUDNN_TENSOR_OP_MATH)
    rnntest(mathType=CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION)
    rnntest(projSize=4)
    rnntest(numLayers=2)
    rnntest(dropout=0.5)
    rnntest(auxFlags=CUDNN_RNN_PADDED_IO_DISABLED)
    rnntest(auxFlags=CUDNN_RNN_PADDED_IO_ENABLED)
end


text:
# Copyright (c) 2021, R. Dong, C. Goodbreak, H. Harrington, G. Pogudin
# Copyright (c) 2020, A. Ovchinnikov, A. Pillay, G. Pogudin, T. Scanlon

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

#------------------------------------------------------------------------------

"""
    linear_compartment_model(graph, inputs, outputs, leaks)

Input: defines a linear compartment model with nodes numbered from 1 to `n` by
- `graph` - and array of integer arrays representing the adjacency lists of the graph
- `inputs` - array of input nodes
- `outputs` - array of output nodes
- `leaks` - array of sink nodes

Output:
- the corresponding ODE system in the notation of https://doi.org/10.1007/s11538-015-0098-0
"""
function linear_compartment_model(
    graph::Vector{Vector{Int}},
    inputs::Vector{Int},
    outputs::Vector{Int},
    leaks::Vector{Int}
)
    n = length(graph)
    x_vars_names = ["x$i" for i in 1:n]
    y_vars_names = ["y$i" for i in outputs]
    u_vars_names = ["u$i" for i in inputs]
    edges_vars_names = Array{String, 1}()
    for i in 1:n
        for j in graph[i]
            push!(edges_vars_names, "a_$(j)_$(i)")
        end
    end
    for s in leaks
        push!(edges_vars_names, "a_0_$(s)")
    end

    R, vars = StructuralIdentifiability.Nemo.PolynomialRing(
        StructuralIdentifiability.Nemo.QQ,
        vcat(x_vars_names, y_vars_names, u_vars_names, edges_vars_names)
    )
    x_vars = @view vars[1:n]
    x_equations = Dict{fmpq_mpoly, Union{fmpq_mpoly, Generic.Frac{fmpq_mpoly}}}(x => R(0) for x in x_vars)
    for i in 1:n
        for j in graph[i]
            rate = str_to_var("a_$(j)_$(i)", R)
            x_equations[x_vars[j]] += x_vars[i] * rate
            x_equations[x_vars[i]] -= x_vars[i] * rate
        end
        if i in leaks
            rate = str_to_var("a_0_$(i)", R)
            x_equations[x_vars[i]] += -x_vars[i] * rate
        end
        if i in inputs
            x_equations[x_vars[i]] += str_to_var("u$i", R)
        end
    end

    y_equations = Dict{fmpq_mpoly, Union{fmpq_mpoly, Generic.Frac{fmpq_mpoly}}}(str_to_var("y$i", R) => str_to_var("x$i", R) for i in outputs)

    return ODE{fmpq_mpoly}(x_equations, y_equations, Array{fmpq_mpoly}([str_to_var("u$i", R) for i in inputs]))
end

#------------------------------------------------------------------------------


text:
function chebyshev_gauss_lobatto(n::S,domain = [1.0,-1.0]) where {S<:Integer}

  # These nodes a just the Chebyshev extrema by a different name.

  # Construct the nodes on the [-1.0,1.0] interval

  if n <= 0
    error("The number of nodes must be positive.")
  end

  if n == 1
    nodes   = [0.0]
    weights = [1.0*pi]
  else
    nodes    = zeros(n)
    nodes[1] = -1.0
    nodes[n] = 1.0

    weights    = zeros(n)
    weights[1] = pi/(2*(n-1))
    weights[n] = pi/(2*(n-1))

    for i = 2:div(n,2)
      x = cos(pi*(i-1)/(n-1))
      nodes[i]       = -x
      nodes[end-i+1] = x
      weights[i]       = pi/(n-1)
      weights[end-i+1] = pi/(n-1)
    end

    if isodd(n)
      nodes[round(Int,(n+1)/2)]   = 0.0
      weights[round(Int,(n+1)/2)] = pi/(n-1)
    end
  end

  # Scale the nodes to the desired domain

  nodes = scale_nodes(nodes,domain)

  return nodes, weights

end


text:
module GroundDataAA
"""
Module for the download and processing of atmospheric data gathered by measuring stations located in Alto Adige, Italy.
"""
#=
Alto Adige:
    Link:
        Stazioni Meteo (JSON):
            http://dati.retecivica.bz.it/services/meteo/v1/stations
        Dati Sensori Meteo (JSON):
            http://dati.retecivica.bz.it/services/meteo/v1/sensors
        Stazioni QA (JSON):
            http://dati.retecivica.bz.it/services/airquality/stations
        Dati sensori QA (JSON):
            http://dati.retecivica.bz.it/services/airquality/timeseries
=#

using CSV
using DataFrames
using Dates
using HTTP
using JSONTables

export getData,
       getRegionAttributes, getRegionIds, getRegionStationsInfo

"""
    getRegionAttributes( [ type::Symbol=:METEO ] )

Obtain the names of the columns of the region's dataframe required by `GroundData.createMap`'s `attributes` parameter to create `GroundData.standardize`'s `map` parameter.
"""
function getRegionAttributes( type::Symbol=:METEO )
    return type == :METEO ? [ :DESC_I, :UNIT, :VALUE, nothing, :DATE, :LONG, :LAT, :ALT, nothing, nothing, :rmh ] :
        type == :AIRQUALITY ? [ :MCODE, nothing, :VALUE, nothing, :DATE, :LONG, :LAT, nothing, :FLAGS, nothing ] :
            throw(DomainError(type, "`type` must be either `:METEO` OR `:AIRQUALITY`."))
end

"""
    getRegionAttributes( [ type::Symbol=:METEO ] )

Obtain the names of the columns of the dataframe required for `GroundData.standardize`'s `bridge` parameter.
"""
function getRegionIds( type::Symbol=:METEO )
    if type != :METEO && type != :AIRQUALITY
        throw(DomainError(type, "`type` must be either `:METEO` OR `:AIRQUALITY`."))
    end
    return :SCODE
end

"""
    getRegionStationInfo( [ type::Symbol=:METEO  ] )

Obtain the names of the columns of the region's stations dataframe required by `GroundData.createMap`'s `attributes` parameter to be used in `GroundData.generateUuidsTable`.
"""
function getRegionStationsInfo( type::Symbol=:METEO )
    if type != :METEO && type != :AIRQUALITY
        throw(DomainError(type, "`type` must be either `:METEO` OR `:AIRQUALITY`."))
    end
    return [ :SCODE, :NAME_I, :LONG, :LAT ]
end

"""
    getData(; <keyword arguments> )

Obtain data of category `type` and source of category `kind`.

# Arguments
- `type::Symbol=:METEO`: defines the type of data to be downloaded, it must either be `:METEO` or `:AIRQUALITY`.
- `kind::Symbol=:STATIONS`: defines if the data to be downloaded has to regard information on the stations or their actual measurements, it must either be `:STATIONS` or `:SENSORS`.
"""
function getData(; type::Symbol=:METEO, kind::Symbol=:STATIONS )
 # The URL changes based on the `type` and `kind` of the data to be retrieved
    opt1 = type == :METEO ? "meteo/v1" :
        type == :AIRQUALITY ? "airquality" :
        throw(DomainError(type, "`type` must either be `:METEO` or `:AIRQUALITY`."))
    opt2 = kind == :STATIONS ? "stations" :
        kind == :SENSORS ? ( type == :METEO ? "sensors" : "timeseries" ) :
        throw(DomainError(kind, "`kind` must either be `:STATIONS` or `:SENSORS`."))
 # Data obtained from the page in raw string form
    page = String( HTTP.get( "http://dati.retecivica.bz.it/services/$opt1/$opt2" ).body )
 # Stations' data requires specific processing to be transformed in a dataframe
    if kind == :STATIONS
        if type == :METEO
            chars = " : "
            div = "\r\n\t\t},\r\n\t\t"
            lim = 11
        else
            chars = ":"
            div = "}\r\n,"
            lim = 5
        end
        features = split( page , "\"features\"$chars"  )[2]
        stations = split( features, "\"properties\"$chars" )
        stations = [ split( station, div )[1] for station in stations ]
        page = "[" * join( stations[2:end], "," )[1:end-lim] * "]"
    end
    data = DataFrame( jsontable(page) )
    if type == :METEO && kind == :SENSORS
        insertcols!( data, :rmh => "0m" )
        transform!( data, [:DATE] => ByRow( x -> ismissing(x) ? missing : DateTime( x[1:19], "yyyy-mm-ddTH:M:S" ) ) => :DATE )
    end
    return data
end

#   ressta = getData( type=:METEO, source=:STATIONS )
#   ressen = getData( type=:METEO, source=:SENSORS )
#   ressta = getData( type=:AIRQUALITY, source=:STATIONS )
#   ressen = getData( type=:AIRQUALITY, source=:SENSORS )

end # module

max_stars_repo_name: DavideFavaro/Tirocinio

max_issues_repo_name: DavideFavaro/Tirocinio

max_forks_repo_name: DavideFavaro/Tirocinio


text:
using MathOptInterface, Test

const CleverDicts = MathOptInterface.Utilities.CleverDicts

# Note: `MyKey` is just for testing. You wouldn't want to use it in practice
# because the key type of the dictionary isn't a concrete type.
struct MyKey{X}
    MyKey(x) = new{Int64(x)}()
end
CleverDicts.key_to_index(::MyKey{X}) where {X} = X
CleverDicts.index_to_key(::Type{MyKey}, index::Int64) = MyKey(index)

@testset "CleverDict" begin
    @testset "MyKey type" begin
        d = CleverDicts.CleverDict{MyKey,String}()
        key = CleverDicts.add_item(d, "first")
        @test key == MyKey(1)
        @test d[MyKey(1)] == "first"
    end

    @testset "Abstract Value" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,Any}()
        key = CleverDicts.add_item(d, :a)
        @test key == MathOptInterface.VariableIndex(1)
        @test d[MathOptInterface.VariableIndex(1)] == :a
        key = CleverDicts.add_item(d, "b")
        @test key == MathOptInterface.VariableIndex(2)
        @test d[MathOptInterface.VariableIndex(2)] == "b"
        for (k, v) in d
            if k.value == 1
                @test v == :a
            else
                @test k.value == 2
                @test v == "b"
            end
        end
    end

    @testset "get/set" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,String}()
        key = CleverDicts.add_item(d, "first")
        @test key == MathOptInterface.VariableIndex(1)
        @test get(d, key, nothing) == "first"
        @test get(d, MathOptInterface.VariableIndex(2), nothing) === nothing
        @test Dict(key => "first") == d
        @test Dict(key => "second") != d
        sizehint!(d, 1)
        @test d[key] == "first"
        @test haskey(d, key) == true
        @test_throws KeyError d[MathOptInterface.VariableIndex(2)]
        delete!(d, key)
        sizehint!(d, 2)
        @test_throws KeyError d[key]
        # set index is valid now
        # @test_throws KeyError d[key] = "key"
        @test haskey(d, key) == false
        key2 = CleverDicts.add_item(d, "second")
        @test key2 == MathOptInterface.VariableIndex(2)
        @test d[key2] == "second"
        d[key2] = "third"
        @test d[key2] == "third"
        @test get(d, key, nothing) === nothing
        @test get(d, key2, nothing) === "third"
        @test Dict(key2 => "second") != d
        @test Dict(key2 => "third") == d

        empty!(d)

        key = CleverDicts.add_item(d, "first")
        @test key == MathOptInterface.VariableIndex(1)
        @test d[key] == "first"
        d[key] = "zeroth"
        @test d[key] == "zeroth"
        @test haskey(d, key) == true
        @test_throws KeyError d[MathOptInterface.VariableIndex(2)]
        delete!(d, key)
        @test_throws KeyError d[key]
        # set index is valid now
        # @test_throws KeyError d[key] = "key"
        @test haskey(d, key) == false
        key2 = CleverDicts.add_item(d, "second")
        @test key2 == MathOptInterface.VariableIndex(2)
        @test d[key2] == "second"
    end

    @testset "LinearIndex" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,String}()
        key = CleverDicts.add_item(d, "first")
        @test d[CleverDicts.LinearIndex(1)] == "first"
        key2 = CleverDicts.add_item(d, "second")
        @test d[CleverDicts.LinearIndex(2)] == "second"
        @test length(d) == 2
        delete!(d, key)
        @test d[CleverDicts.LinearIndex(1)] == "second"
        @test_throws KeyError d[CleverDicts.LinearIndex(2)]
        @test length(d) == 1
    end

    @testset "keys/values" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,String}()
        key = CleverDicts.add_item(d, "first")
        key2 = CleverDicts.add_item(d, "second")
        @test collect(keys(d)) == [
            MathOptInterface.VariableIndex(1),
            MathOptInterface.VariableIndex(2),
        ]
        @test collect(values(d)) == ["first", "second"]
        delete!(d, key)
        key3 = CleverDicts.add_item(d, "third")
        @test collect(keys(d)) == [
            MathOptInterface.VariableIndex(2),
            MathOptInterface.VariableIndex(3),
        ]
        @test collect(values(d)) == ["second", "third"]
    end

    @testset "iterate" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,String}()
        key = CleverDicts.add_item(d, "first")
        key2 = CleverDicts.add_item(d, "second")
        my_keys = MathOptInterface.VariableIndex[]
        my_values = String[]
        for (k, v) in d
            push!(my_keys, k)
            push!(my_values, v)
        end
        @test my_keys == [
            MathOptInterface.VariableIndex(1),
            MathOptInterface.VariableIndex(2),
        ]
        @test my_values == ["first", "second"]
        delete!(d, key)
        key3 = CleverDicts.add_item(d, "third")
        my_keys = MathOptInterface.VariableIndex[]
        my_values = String[]
        for (k, v) in d
            push!(my_keys, k)
            push!(my_values, v)
        end
        @test my_keys == [
            MathOptInterface.VariableIndex(2),
            MathOptInterface.VariableIndex(3),
        ]
        @test my_values == ["second", "third"]
    end

    @testset "iterate ii" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,String}()
        key = CleverDicts.add_item(d, "first")
        key2 = CleverDicts.add_item(d, "second")
        my_keys = MathOptInterface.VariableIndex[]
        my_values = String[]
        for (k, v) in d
            push!(my_keys, k)
            push!(my_values, v)
        end
        @test my_keys == [
            MathOptInterface.VariableIndex(1),
            MathOptInterface.VariableIndex(2),
        ]
        @test my_values == ["first", "second"]
        delete!(d, key)
        @test d[CleverDicts.LinearIndex(1)] == "second"
        key3 = CleverDicts.add_item(d, "third")
        my_keys = MathOptInterface.VariableIndex[]
        my_values = String[]
        for (k, v) in d
            push!(my_keys, k)
            push!(my_values, v)
        end
        @test my_keys == [
            MathOptInterface.VariableIndex(2),
            MathOptInterface.VariableIndex(3),
        ]
        @test my_values == ["second", "third"]
    end

    @testset "iterate iii" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,String}()
        y = 0
        for (k, v) in d
            y += 1
        end
        @test y == 0
    end

    @testset "haskey" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,String}()
        @test !haskey(d, 1)
        k = CleverDicts.add_item(d, "a")
        @test haskey(d, k)
        j = CleverDicts.add_item(d, "b")
        @test haskey(d, j)
        delete!(d, k)
        @test !haskey(d, k)
        @test haskey(d, j)
    end

    @testset "haskey" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,String}()
        @test isempty(d) == true
        k = CleverDicts.add_item(d, "a")
        @test isempty(d) == false
        delete!(d, k)
        @test isempty(d) == true
        j = CleverDicts.add_item(d, "b")
        @test isempty(d) == false
    end

    @testset "delete!" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,String}()
        @test length(d) == 0
        @test delete!(d, MathOptInterface.VariableIndex(0)) == nothing
        k1 = CleverDicts.add_item(d, "a")
        k2 = CleverDicts.add_item(d, "b")
        d[CleverDicts.LinearIndex(2)] == "b"
        delete!(d, k1)
        d[CleverDicts.LinearIndex(1)] == "b"
        k3 = CleverDicts.add_item(d, "c")
        @test d[k3] == "c"
        @test d[CleverDicts.LinearIndex(1)] == "b"
        @test d[CleverDicts.LinearIndex(2)] == "c"
    end

    @testset "Dense Operations" begin
        # inverse_hash :: Int64 -> Int
        inverse_hash(x::Int64) = Int(x * 2)
        # hash :: Int -> Int64
        hash(x::Int) = Int64(div(x, 2))
        d = CleverDicts.CleverDict{Int,Float64}(hash, inverse_hash, 3)

        d[4] = 0.25
        @test !haskey(d, 2)
        @test haskey(d, 4)
        @test !haskey(d, 6)
        @test length(d) == 1
        @test collect(d) == [4 => 0.25]
        @test d[4] == 0.25

        d[2] = 1.5
        @test haskey(d, 2)
        @test haskey(d, 4)
        @test !haskey(d, 6)
        @test length(d) == 2
        @test collect(d) == [2 => 1.5, 4 => 0.25]
        @test d[2] == 1.5
        @test d[4] == 0.25

        d[6] = 0.75
        @test haskey(d, 2)
        @test haskey(d, 4)
        @test haskey(d, 6)
        @test length(d) == 3
        @test collect(d) == [2 => 1.5, 4 => 0.25, 6 => 0.75]
        @test d[2] == 1.5
        @test d[4] == 0.25
        @test d[6] == 0.75
    end

    @testset "Negative index" begin
        d = CleverDicts.CleverDict{MathOptInterface.VariableIndex,String}()
        d[MathOptInterface.VariableIndex(-3)] = "a"
        @test d[MathOptInterface.VariableIndex(-3)] == "a"
        @test_throws ErrorException CleverDicts.add_item(d, "b")
        d[MathOptInterface.VariableIndex(0)] = "b"
        @test d[MathOptInterface.VariableIndex(-3)] == "a"
        @test d[MathOptInterface.VariableIndex(0)] == "b"
        @test_throws ErrorException CleverDicts.add_item(d, "c")
    end

    @testset "convert" begin
        vals = [MathOptInterface.VariableIndex(-i) for i in 1:10]
        d = Dict(MathOptInterface.VariableIndex(i) => vals[i] for i in 1:10)
        T = MathOptInterface.Utilities.DenseVariableDict{
            MathOptInterface.VariableIndex,
        }
        c = convert(T, d)
        @test c isa T
        @test c.is_dense
        @test c.last_index == 10
        @test c.vector == vals
        @test c === convert(T, c)
    end
end


text:
## -------------------------------------------------------------------
# Yield corr
let
    method = FBA_Z_FIX_ADJ_MIN_VG_MIN_COST

    fontsize = 13
    common_params = (;
        thickness_scaling = 1.3,
        xguidefontsize = fontsize, yguidefontsize = fontsize,
        dpi = 1000
    )

    p = plot(;
        xlabel = "D",
        ylabel = "biom. adj. factor",
        common_params...
    )

    all_vals = []
    for (iJR, Data) in [
        (ChK.iJR904, ChK.KayserData),
        (ChN.iJR904, ChN.NanchenData),
        (ChF.iJR904, ChF.FolsomData),
        # (ChH.iJR904, ChH.HeerdenData),
    ]
        src = nameof(Data)
        marker = (6, source_markers[Data])

        DAT = ChE.load_DAT(src)
        EXPS = DAT[:EXPS]
        LP_DAT = ChE.load_LP_DAT(src)

        @info("Doing", src, method)

        D = DAT[:exp, :flx, "D", EXPS]
        exp_vals = LP_DAT[method, :Sfac, EXPS]

        scatter!(p, abs.(D), abs.(exp_vals);
            label = "",
            color = :black,
            marker, alpha = 0.8,
        )

    end

    plot!(;legend = :bottomright)
    mysavefig(p, "biom_adj_study")
end