include("Header.jl")

function getDataType(data)
    if     data == 0x5054 return "topography"
    elseif data == 0x5245 return "error"
    elseif data == 0x4850 return "phase"
    else return "none" end
end

function getAdRange(data)
    # ADレンジ定数（ユニポーラ0～1V）(使われていないらしい)
    if     data == 0x00000001 @assert false
    # ADレンジ定数（ユニポーラ0～2.5V）(使われていないらしい)
    elseif data == 0x00000002 @assert false
    # ADレンジ定数（ユニポーラ0～5V）(使われていないらしい)
    elseif data == 0x00000004 @assert false
    # ADレンジ定数（バイポーラ±1V）
    elseif data == 0x00010000 return 2.0
    # ADレンジ定数（バイポーラ±2.5V）
    elseif data == 0x00020000 return 5.0
    # ADレンジ定数（バイポーラ±5V）
    elseif data == 0x00040000 return 10.0
    # ADレンジ定数（バイポーラ±80V, データを編集した場合に仮想的にこれを使う。実際にバイポーラ80VでAD変換したわけではない。また、分解能は16ビットにする）
    elseif data == 0x00080000 return 160.0
    # 何も当てはまらない
    else @assert false end
    return nothing
end

function getOperatorName(io, operatorNameSize)
    ret = []
    for i in 1:operatorNameSize
        push!(ret, Int64(read(io, UInt8)))
    end
    return ret
end

function getComment(io, commentOffsetSize, commentSize)
    for i in 1:commentOffsetSize
        read(io, Bool)
    end
    ret = []
    for i in 1:commentSize
        push!(ret, Int64(read(io, UInt8)))
    end
    return ret
end

function readDate(io::IOStream)
    year              = Int64(read(io, Int16))
    month             = Int64(read(io, UInt8))
    day               = Int64(read(io, UInt8))
    hour              = Int64(read(io, UInt8))
    minute            = Int64(read(io, UInt8))
    second            = Int64(read(io, UInt8))
    return DateTime(year, month, day, hour, minute, second)
end

function readHeaderV0(io::IOStream)
    fileVersion       = 0
    dataType1ch       = getDataType(read(io, Int16))
    dataType2ch       = getDataType(read(io, Int16))
    fileHeaderSize    = Int64(read(io, Int32))
    frameHeaderSize   = Int64(read(io, Int32))
    operatorNameSize  = Int64(read(io, Int32))
    commentOffsetSize = Int64(read(io, Int32))
    commentSize       = Int64(read(io, Int32))
    pixelX            = Int64(read(io, Int16))
    pixelY            = Int64(read(io, Int16))
    scanningRangeX    = Int64(read(io, Int16))
    scanningRangeY    = Int64(read(io, Int16))
    frameRate         = Float64(read(io, Float32))
    piezoExtensionZ   = Float64(read(io, Float32))
    piezoGainZ        = Float64(read(io, Float32))
    adRange           = getAdRange(read(io, UInt32))
    AdResolution      = (2.0)^Int64(read(io, Int32))
    isAveraged        = read(io, Bool)
    averageWindow     = Int64(read(io, Int32))
    legacy            = Int64(read(io, Int16)) # ダミー
    day               = readDate(io)
    roundingDegree    = Int64(read(io, UInt8))
    maxRangeX         = Float64(read(io, Float32))
    maxRangeY         = Float64(read(io, Float32))
    booked1           = Int64(read(io, Int32))
    booked2           = Int64(read(io, Int32))
    booked3           = Int64(read(io, Int32))
    initFrame         = Int64(read(io, Int32))
    numFrames         = Int64(read(io, Int32))
    machineId         = Int64(read(io, Int32))
    fileId            = Int64(read(io, Int16))
    operatorName      = getOperatorName(io, operatorNameSize)
    sensorSensitivity = Float64(read(io, Float32))
    phaseSensitivity  = Float64(read(io, Float32))
    scannigDirection  = Int64(read(io, Int32))
    comment           = getComment(io, commentOffsetSize, commentSize)

    return HeaderV0(fileVersion, dataType1ch, dataType2ch, fileHeaderSize, frameHeaderSize, operatorNameSize, commentOffsetSize, commentSize, pixelX, pixelY, scanningRangeX, scanningRangeY, frameRate, piezoExtensionZ, piezoGainZ, adRange, AdResolution, isAveraged, averageWindow, day, roundingDegree, maxRangeX, maxRangeY, booked1, booked2, booked3, initFrame, numFrames, machineId, fileId, operatorName, sensorSensitivity, phaseSensitivity, scannigDirection, comment)
end


function deactivate!(plank, loss)
    @inbounds plank.ac .*= (1.0 .- loss)
end

##### grazing and grazing loss
function grazing!(plank, arch::Architecture, plk, p)
    ##### calculate grazing loss
    calc_loss!(plk.DOC.data, plk.POC.data, plank, plank.ac, plank.xi, plank.yi, plank.zi,
               plank.graz, p.grazFracC, arch)

    ##### deactivate grazed individuals
    deactivate!(plank, plank.graz)

    return nothing
end

##### mortality and mortality loss
function mortality!(plank, arch::Architecture, plk, p)
    ##### calculate mortality loss
    calc_loss!(plk.DOC.data, plk.POC.data, plank, plank.ac, plank.xi, plank.yi, plank.zi,
               plank.mort, p.mortFracC, arch)

    ##### deactivate dead individuals
    deactivate!(plank, plank.mort)

    return nothing
end

@kernel function get_tind_kernel!(idx, con, con_ind, de_ind)
    i = @index(Global, Linear)
    if con[i] == 1.0
        idx[i] = de_ind[con_ind[i]]
    end
end
function get_tind!(idx, con, con_ind, de_ind, arch)
    kernel! = get_tind_kernel!(device(arch), 256, (size(idx,1)))
    event = kernel!(idx, con, con_ind, de_ind)
    wait(device(arch), event)
    return nothing
end

##### copy ready to divide individuals to inactive rows
@kernel function copy_daughter_individuals_kernel!(plank, con, idx)
    i = @index(Global, Linear)
    if (con[i] == 1.0) & (idx[i] ≠ 0)
        # @print("index: $(idx[i]), $i \n")
        @inbounds plank.x[idx[i]]    = plank.x[i]
        @inbounds plank.y[idx[i]]    = plank.y[i]
        @inbounds plank.z[idx[i]]    = plank.z[i]
        @inbounds plank.Sz[idx[i]]   = plank.Sz[i]
        @inbounds plank.Bm[idx[i]]   = plank.Bm[i]
        @inbounds plank.Chl[idx[i]]  = plank.Chl[i]
        @inbounds plank.gen[idx[i]]  = plank.gen[i]
        @inbounds plank.ac[idx[i]]   = plank.ac[i]
        @inbounds plank.dvid[idx[i]] = plank.dvid[i]
        @inbounds plank.graz[idx[i]] = plank.graz[i]
        @inbounds plank.mort[idx[i]] = plank.mort[i]
    end
end
function copy_daughter_individuals!(plank, con, idx::AbstractArray{Int64,1}, arch)
    kernel! = copy_daughter_individuals_kernel!(device(arch), 256, (size(plank.ac,1)))
    event = kernel!(plank, con, idx)
    wait(device(arch), event)
    return nothing
end

##### cell division
@kernel function divide_to_half_kernel!(plank)
    i = @index(Global)
    @inbounds plank.Sz[i]  *= (2.0 - plank.dvid[i]) / 2
    @inbounds plank.Bm[i]  *= (2.0 - plank.dvid[i]) / 2
    @inbounds plank.Chl[i] *= (2.0 - plank.dvid[i]) / 2
    @inbounds plank.gen[i] += plank.dvid[i]
    @inbounds plank.age[i] *= (1.0 - plank.dvid[i])
    @inbounds plank.iS[i]   = plank.iS[i] * (1.0 - plank.dvid[i]) + plank.Sz[i] * plank.dvid[i]
    @inbounds plank.Th[i]  *= (1.0 - plank.dvid[i])
end
function divide_to_half!(plank, arch)
    kernel! = divide_to_half_kernel!(device(arch), 256, (size(plank.ac,1)))
    event = kernel!(plank)
    wait(device(arch), event)
    return nothing
end
function divide!(plank, deactive_ind, arch::Architecture)
    con_ind = cumsum(plank.dvid)
    get_tind!(plank.idx, plank.dvid, Int.(con_ind), deactive_ind, arch)
    copy_daughter_individuals!(plank, plank.dvid, Int.(plank.idx), arch)
    divide_to_half!(plank, arch)
    # println("time step ends")
    return nothing
end


"""
$(DocStringExtensions.README)
"""
module StochasticDiffEq

using DocStringExtensions
  import RandomNumbers: Xorshifts

  using Reexport
  @reexport using DiffEqBase

  import OrdinaryDiffEq

  using UnPack, RecursiveArrayTools, DataStructures
  using DiffEqNoiseProcess, Random, ArrayInterface
  using NLsolve, ForwardDiff, StaticArrays, MuladdMacro, FiniteDiff, Base.Threads

  import DiffEqBase: ODE_DEFAULT_NORM, ODE_DEFAULT_ISOUTOFDOMAIN,
         ODE_DEFAULT_PROG_MESSAGE, ODE_DEFAULT_UNSTABLE_CHECK

  using DiffEqBase: DiffEqArrayOperator

  using DiffEqBase: TimeGradientWrapper, UJacobianWrapper, TimeDerivativeWrapper, UDerivativeWrapper

  import RecursiveArrayTools: chain

  using Logging, SparseArrays

  using LinearAlgebra, Random, FillArrays

  import ForwardDiff.Dual

  import DiffEqBase: step!, initialize!, DEAlgorithm,
                     AbstractSDEAlgorithm, AbstractRODEAlgorithm, DEIntegrator, AbstractDiffEqInterpolation,
                     DECache, AbstractSDEIntegrator, AbstractRODEIntegrator, AbstractContinuousCallback,
                     Tableau

  # Integrator Interface
  import DiffEqBase: resize!,deleteat!,addat!,full_cache,user_cache,u_cache,du_cache,
                     rand_cache,ratenoise_cache,
                     resize_non_user_cache!,deleteat_non_user_cache!,addat_non_user_cache!,
                     terminate!,get_du, get_dt,get_proposed_dt,set_proposed_dt!,
                     u_modified!,savevalues!,add_tstop!,add_saveat!,set_reltol!,
                     set_abstol!, postamble!, last_step_failed, has_Wfact, has_jac

  using DiffEqBase: check_error!, is_diagonal_noise, @..

  using DiffEqBase: nlsolvefail, isnewton, set_new_W!, get_W, iipnlsolve, oopnlsolve, _vec, _reshape

  using DiffEqBase: NLSolver

  using DiffEqBase: FastConvergence, Convergence, SlowConvergence, VerySlowConvergence, Divergence

  import DiffEqBase: calculate_residuals, calculate_residuals!, nlsolve_f, unwrap_cache, islinear

  import DiffEqBase: iip_get_uf, oop_get_uf, build_jac_config

  using SparseDiffTools: forwarddiff_color_jacobian!, ForwardColorJacCache

  const CompiledFloats = Union{Float32,Float64}

  import DiffEqJump
  import DiffEqJump: JumpProblem

  import Base.Threads
  @static if VERSION < v"1.3"
    seed_multiplier() = Threads.threadid()
  else
    seed_multiplier() = 1
  end

  include("misc_utils.jl")
  include("algorithms.jl")
  include("options_type.jl")
  include("interp_func.jl")
  include("caches/cache_types.jl")
  include("caches/basic_method_caches.jl")
  include("caches/explicit_3s_mil_methods.jl")
  include("caches/lamba_caches.jl")
  include("caches/iif_caches.jl")
  include("caches/sdirk_caches.jl")
  include("caches/implicit_split_step_caches.jl")
  include("caches/sra_caches.jl")
  include("caches/rossler_caches.jl")
  include("caches/srk_weak_caches.jl")
  include("caches/kencarp_caches.jl")
  include("caches/predcorr_caches.jl")
  include("caches/SROCK_caches.jl")
  include("caches/tau_caches.jl")
  include("caches/dynamical_caches.jl")
  include("integrators/type.jl")
  include("dense.jl")
  include("alg_utils.jl")
  include("integrators/stepsize_controllers.jl")
  include("integrators/integrator_utils.jl")
  include("cache_utils.jl")
  include("integrators/integrator_interface.jl")
  include("iterator_interface.jl")
  include("composite_solution.jl")
  include("solve.jl")
  include("initdt.jl")
  include("perform_step/low_order.jl")
  include("perform_step/explicit_3s_mil_methods.jl")
  include("perform_step/lamba.jl")
  include("perform_step/iif.jl")
  include("perform_step/sri.jl")
  include("perform_step/sra.jl")
  include("perform_step/srk_weak.jl")
  include("perform_step/sdirk.jl")
  include("perform_step/implicit_split_step.jl")
  include("perform_step/kencarp.jl")
  include("perform_step/predcorr.jl")
  include("perform_step/split.jl")
  include("perform_step/composite.jl")
  include("perform_step/SROCK_perform_step.jl")
  include("perform_step/tau_leaping.jl")
  include("perform_step/dynamical.jl")
  include("tableaus.jl")
  include("SROCK_tableaus.jl")
  include("iterated_integrals.jl")
  include("SROCK_utils.jl")
  include("composite_algs.jl")
  include("weak_utils.jl")

  export StochasticDiffEqAlgorithm, StochasticDiffEqAdaptiveAlgorithm,
          StochasticCompositeAlgorithm

  export EM, LambaEM, PCEuler, RKMil, SRA, SRI, SRIW1,
         SRA1, SRA2, SRA3,
         SOSRA, SOSRA2, RKMilCommute, RKMilGeneral,
         SRIW2, SOSRI, SOSRI2, SKenCarp,
         SROCK1, SROCK2, SROCKEM, SKSROCK, TangXiaoSROCK2, KomBurSROCK2, SROCKC2,
         WangLi3SMil_A, WangLi3SMil_B, WangLi3SMil_C, WangLi3SMil_D, WangLi3SMil_E, WangLi3SMil_F,
         AutoSOSRI2, AutoSOSRA2,
         DRI1, DRI1NM, RI1, RI3, RI5, RI6, RDI1WM, RDI2WM, RDI3WM, RDI4WM,
         RS1, RS2,
         PL1WM, PL1WMA,
         NON, COM, NON2

  export SIEA, SMEA, SIEB, SMEB

  export EulerHeun, LambaEulerHeun

  export SimplifiedEM

  export SplitEM, IIF1M, IIF2M, IIF1Mil

  export ImplicitEM, ImplicitEulerHeun, ISSEM, ISSEulerHeun,
         ImplicitRKMil, STrapezoid, SImplicitMidpoint

  export TauLeaping, CaoTauLeaping

  export BAOAB

  export StochasticDiffEqRODEAlgorithm, StochasticDiffEqRODEAdaptiveAlgorithm,
         StochasticDiffEqRODECompositeAlgorithm

  export RandomEM

  export IteratedIntegralApprox, IICommutative, IIWiktorsson

  #General Functions
  export solve, init, solve!, step!

  #Misc Tools
  export checkSRIOrder, checkSRAOrder,  checkRIOrder, checkRSOrder,
         checkNONOrder,
         constructSRIW1, constructSRA1,
         constructDRI1, constructRI1, constructRI3, constructRI5, constructRI6,
         constructRDI1WM, constructRDI2WM, constructRDI3WM, constructRDI4WM,
         constructRS1, constructRS2,
         constructNON, constructNON2

end # module


module Molly

using BioStructures
using CellListMap
using ChainRulesCore
import Chemfiles # Imported to avoid clashing names
using Colors
using Combinatorics
using CUDA
using Distances
using Distributions
using EzXML
using ForwardDiff
using KernelDensity
using NearestNeighbors
import NNlib
using ProgressMeter
using Reexport
using Requires
using Unitful
using Zygote

@reexport using AtomsBase
@reexport using StaticArrays
@reexport using Unitful

using Base.Threads
using LinearAlgebra
using SparseArrays

include("types.jl")
include("cutoffs.jl")
include("spatial.jl")
include("forces.jl")
include("simulators.jl")
include("coupling.jl")
include("neighbors.jl")
include("loggers.jl")
include("analysis.jl")
include("chain_rules.jl")
include("zygote.jl")
include("setup.jl")

function __init__()
    @require GLMakie="e9467ef8-e4e7-5192-8a1a-b1aee30e663a" include("makie.jl")
end

end


###############################################################################
#
# Steps of non-steady state stomatal controls
#
###############################################################################
#=
0. Initialize plant and environment with minimal stomatal opening
1. Update transpiration rate
2. Update canopy temperature
3. Update plant hydraulic profile every 10 minutes
4. Calculate assimilation rate
5. Calculate dH₂O/dt and dCO₂/dt for the environment
6. Calculate dGs/dt for the given environment
7. Update stomatal conductance
8. Goto step 1

This function will be splited into many functions later!
=#

function test_diurnal(
            node::SPACMono{FT},
            Δt::FT=FT(10)
) where {FT<:AbstractFloat}
    # 0.1 create variables required
    @unpack n_canopy, wl_set = node;
    canopy_rt = node.canopy_rt;
    fraction_sl::Array{FT,1} = repeat(canopy_rt.lidf, outer=[ length(canopy_rt.lazitab) ]) / length(canopy_rt.lazitab);
    n_sl = length(canopy_rt.lidf) * length(canopy_rt.lazitab);

    # change node.in_rad to alter direct and diffuse fraction
    # sum of direct  energy in W:
    # e_dir = sum( (wl_set.swl[2:end] .- wl_set.swl[1:end-1]) .* node.in_rad.E_direct  ) / 1000
    # sum of diffuse energy in W:
    # e_dif = sum( (wl_set.swl[2:end] .- wl_set.swl[1:end-1]) .* node.in_rad.E_diffuse ) / 1000
    # target e_dir is tar_dir, do it this way:
    # node.in_rad.E_direct  .*= tar_dir / e_dir
    # target e_dif is tar_dif, do it this way:
    # node.in_rad.E_diffuse .*= tar_dif / e_dif

    # 0.2 A diurnal cycle for radiation and Tair
    Tmean  = FT(295.15)
    DeltaT = 3
    omega  = 2 * π / (24*3600)
    t      = range(0, stop=24*3600*2, step=Δt)
    phi_t  = omega .* t .- π .* ones(FT, size(t)) ./ 2
    PARr_t = zeros(FT, size(t))
    Tair_t = zeros(FT, size(t))
    Tlef_t = zeros(FT, size(t))
    N      = length(PARr_t)

    # 0.3 customize dirnal cycles of temperatures and PAR
    for i_tim in 1:N
        PARr_t[i_tim] = max( sin(phi_t[i_tim]), FT(0.0) )
        Tair_t[i_tim] = Tmean + DeltaT * sin(phi_t[i_tim] - π/3)
        Tlef_t[i_tim] = Tmean + DeltaT * sin(phi_t[i_tim] - π/3)
    end

    # 0.4 initialize the canopy temperature
    for i_can in 1:n_canopy
        # Assume RH = 50%
        photosystem_temperature_dependence!(node.plant_ps[i_can].ps.PSM, node.envirs[i_can], Tlef_t[1]);
        node.plant_ps[i_can].ps.p_H₂O_sat = saturation_vapor_pressure(Tlef_t[1]);
        node.envirs[i_can].p_H₂O = node.envirs[i_can].p_H₂O_sat / 2;
        # latent heat flux as well
        node.plant_ps[i_can].LV = latent_heat_vapor(Tlef_t[i_can]) * 1000 / 18;
        # vc dependecy as well
        temperature_effects!(node.plant_hs.leaves[i_can], Tlef_t[i_can]);
    end

    # run dirnal cycle
    # NPP per ground area
    Agpp_t = zeros(FT, size(t));
    # NPP per ground area
    Asum_t = zeros(FT, size(t));
    # trasnpiration rate per ground area
    Esum_t = zeros(FT, size(t));
    for i_tim in 1:N
        # 0. recalculate the canopy RT stuff
        zenith = zenith_angle(node.latitude, 180 + t[i_tim]/86400 );
        zenith = min(88, zenith);
        node.angles.tts = zenith;
        canopy_geometry!(canopy_rt, node.angles, node.can_opt, node.rt_con);
        canopy_matrices!(node.leaves_rt, node.can_opt);
        short_wave!(canopy_rt, node.can_opt, node.can_rad, node.in_rad, node.soil_opt, node.rt_con);
        canopy_fluxes!(canopy_rt, node.can_opt, node.can_rad, node.in_rad, node.soil_opt, node.leaves_rt, wl_set, node.rt_con);
        thermal_fluxes!(node.leaves_rt, node.can_opt, node.can_rad, canopy_rt, node.soil_opt, [FT(400.0)], wl_set);
        SIF_fluxes!(node.leaves_rt, node.can_opt, node.can_rad, canopy_rt, node.soil_opt, wl_set, node.rt_con, node.rt_dim);

        # 2. Update transpiration rate
        a_gpp = FT(0);
        a_sum = FT(0);
        e_sum = FT(0);
        for i_can in 1:n_canopy
            rt_layer = node.n_canopy + 1 - i_can;

            # set the diffuse par to all the leaves
            #node.plant_ps[i_can].APAR           .= PARr_t[i_tim]  * node.can_rad.absPAR_shadeCab[rt_layer] * FT(1e6);
            #node.plant_ps[i_can].APAR[1:end-1] .+= PARr_t[i_tim] .* reshape(node.can_rad.absPAR_sunCab[:,:,rt_layer],(:,1))[:,1] .* FT(1e6);

            # calculate the fraction of sunlit and shaded leaves
            f_view = (node.can_opt.Ps[rt_layer] + node.can_opt.Ps[rt_layer+1]) / 2;
            #la_new = [f_view .* fraction_sl; 1-f_view];
            #node.plant_ps[i_can].LAIx .= la_new;
            ##=
            for i_leaf in 1:n_sl
                node.plant_ps[i_can].APAR[i_leaf] = PARr_t[i_tim] * node.can_rad.absPAR_sunCab[(rt_layer-1)*n_sl + i_leaf] * FT(1e6);
                node.plant_ps[i_can].LAIx[i_leaf] = f_view * fraction_sl[i_leaf];
            end
            node.plant_ps[i_can].APAR[end] = PARr_t[i_tim] * node.can_rad.absPAR_shadeCab[rt_layer] * FT(1e6);
            node.plant_ps[i_can].LAIx[end] = 1 - f_view;
            # =#

            # update leaf temperature
            # update the parameters from stomatal model
            node.plant_ps[i_can].T = Tlef_t[i_tim];
            update_leaf_TP!(node.plant_ps[i_can], node.plant_hs.leaves[i_can], node.envirs[i_can]);
            # uncomment this for Sperry, Eller, and Gentine models
            #update_leaf_AK!(node.plant_ps[i_can], node.plant_hs.leaves[i_can], node.envirs[i_can]);

            #canopyi.t_list[1:end-1] .= reshape(node.can_rad.T_sun3D[:,:,rt_layer],(:,1))[:,1]
            #canopyi.t_list[end]      = node.can_rad.T_shade[rt_layer]
            # Assume RH = 50%
            node.envirs[i_can].t         = Tair_t[i_tim];
            node.envirs[i_can].p_H₂O_sat = saturation_vapor_pressure(Tair_t[i_tim]);
            node.envirs[i_can].p_H₂O     = node.envirs[i_can].p_H₂O_sat / 2;

            # update the flow rates
            g_i_can = FT(0);
            a_i_can = FT(0);
            e_i_can = FT(0);
            for i_leaf in 1:(n_sl+1)
                g_i_can += node.plant_ps[i_can].Ag[i_leaf] * node.plant_ps[i_can].LAIx[i_leaf] * node.plant_ps[i_can].LA;
                a_i_can += node.plant_ps[i_can].An[i_leaf] * node.plant_ps[i_can].LAIx[i_leaf] * node.plant_ps[i_can].LA;
                e_i_can += node.plant_ps[i_can].g_lw[i_leaf] * (node.plant_ps[i_can].p_sat - node.envirs[i_can].p_H₂O) / node.envirs[i_can].P_AIR * node.plant_ps[i_can].LAIx[i_leaf] *
                           node.plant_ps[i_can].LA;
            end
            #g_i_can = sum( node.plant_ps[i_can].Ag .* node.plant_ps[i_can].LAIx ) * node.plant_ps[i_can].LA;
            #a_i_can = sum( node.plant_ps[i_can].An .* node.plant_ps[i_can].LAIx ) * node.plant_ps[i_can].LA;
            #e_i_can = sum( node.plant_ps[i_can].g_lw .* (node.plant_ps[i_can].p_sat - node.envirs[i_can].p_H₂O) ./ node.envirs[i_can].P_AIR .* node.plant_ps[i_can].LAIx ) * node.plant_ps[i_can].LA;

            # calculate the photosynthetic rates
            gas_exchange!(node.plant_ps[i_can], node.envirs[i_can], GswDrive());
            gsw_control!(node.plant_ps[i_can], node.envirs[i_can]);

            # use the ball-berry model here for now as the ∂A/∂E and ∂A/∂Θ functions are not yet ready
            gsw_ss = stomatal_conductance(node.stomata_model, node.plant_ps[i_can], node.envirs[i_can], FT(1));

            # assume τ = 10 minutes
            for i_leaf in 1:(n_sl+1)
                node.plant_ps[i_can].g_sw[i_leaf] += (gsw_ss[i_leaf] - node.plant_ps[i_can].g_sw[i_leaf]) / 600 * Δt;
            end
            #node.plant_ps[i_can].g_sw .+= (gsw_ss .- node.plant_ps[i_can].g_sw) / 600 * Δt;

            a_gpp += g_i_can;
            a_sum += a_i_can;
            e_sum += e_i_can;
        end

        Agpp_t[i_tim] = a_gpp / node.ga;
        Asum_t[i_tim] = a_sum / node.ga;
        Esum_t[i_tim] = e_sum / node.ga;

        # update pressure profile every 10 minutes
        # TO BE DONE within PlantHydraulics

        _par = PARr_t[i_tim]  * node.can_rad.absPAR_shadeCab[1] * FT(1e6);
    end

    #=
    # plot the results
    figure(1)
    clf()
    plot(t ./ 3600, Agpp_t, "g-");
    plot(t ./ 3600, Asum_t, "k-");

    figure(2)
    clf()
    plot(t ./ 3600, Esum_t, "k-");
    =#

    return nothing
end


using Random
using DABenchmarks
using LaTeXStrings

Random.seed!(0)

using Distributed
machine = [("MACHINENAME", 2)]
addprocs(machine, dir = "/home/USER", exename="/path/to/julia", exeflags="--project")
# addprocs(2, exeflags="--project") # For local workers

@everywhere using StochasticPrograms
@everywhere using ProgressiveHedgingSolvers
@everywhere using HydroModels
@everywhere using Gurobi

function prepare_dbenchmark(nscenarios::Integer, nsamples::Integer; timeout::Integer = 1000)
    ph = ProgressiveHedgingSolver(()->GurobiSolver(OutputFlag=0), execution = :synchronous, penalty = :adaptive, ζ = 0.01, α = 0.75, log = false, τ = 5e-6)
    solvers = [ph]
    solvernames = ["Synchronous progressive-hedging"]
    # Create Day-ahead benchmark
    return Scaling(solvernames, solvers, nscenarios, nsamples; timeout = timeout)
end
@info "Running progressive-hedging benchmarks on 2 workers..."
dab = prepare_dbenchmark(1000, 5); warmup_benchmarks!(dab); run_benchmarks!(dab); save_results!(dab, "ph_2.json")
exit()


"""
    AOGCM1D(ndays=60)

Atmosphere-Ocean, coupled, one-dimensional column model.

```
outputs,parameters=AOGCM1D(10)
p = dirname(pathof(AirSeaFluxes))
include(joinpath(p,"recipes_plots.jl"))
p1,p2,p3=plot_final(outputs,parameters)
display(p3)
```
"""
function AOGCM1D(ndays=60)

cpl=1;
moist=1;

dt=60; #seconds
#ndays=60; # duration of simulation in days

n=Int(round(ndays*24*60*60/dt)); # number of time steps

H=50; # atmospheric layer thickness
dz=1; # oceanic layer thickness
vlevo=50; # number of oceanic vertical levels
vleva=50; # number of atmospheric vertical levels
rhow0=1029; # kg/m^3
cp=4184; #J/K/kg
cpa=1005; #J/K/kg
sb = 5.670e-8;
d2k=273.15;
Av=2500000;
gamma_blk=0.01;
cvapor_fac=640380;
cvapor_exp=5107.4;
cvapor_exp_ice=5897.8;
saltsat=0.980;
p0=101300; #reference pressure
gravity_mks=9.81;
humid_fac=0.606;
tref=25; #reference temperature
sref=25;
SO=sref; #salinity not implemented
alpha=2e-4; #expansion coefficients for temperature
beta=7.4e-4; #expansion coefficients for salt
emissivity=0.97; # ocean emissivity
kappa=1e-5; # m^2/kg air absorption coefficient
Rgas=287.05;
karman=0.4;

ZA=H*[1:vleva;];
ZO=zeros(vlevo+1);
ZO[2:vlevo+1]=-[1:vlevo;]*dz;
kT0=ones(vlevo)*1e-4; # SST vertical mixing
kT=kT0; # SST vertical mixing

Aimp=0.5; # Atmospheric implicitness
Oimp=0.5; # Oceanic implicitness

R=0.62; g1=0.6; g2=20;

f=R*exp.(ZO/g1)+(1-R)*exp.(ZO/g2);
f[ZO.<-200].=0;
f=-diff(f,dims=1);

albedo=0.06;
lat=30;
lon=0;
t=[1:n;]*dt/3600;
jstart=70;# julian day, borial summer
jadd=round.((t.-t[1])/24);
jd=jstart.+jadd;

AOI=angle_of_incidence.(lat,lon,jd,t); AOI[AOI.<0].=0;
SW=(1362*(1-albedo).*f'.*AOI);
#mean(SW(:,1))
#SW=(repmat((sin(2*pi/n*(1:n)*ndays+pi/2)+1)/2*1362*(1-albedo),vlev,1).*
#repmat(f',1,n))'.*cos(lat/180*pi);

TO=zeros(n+1,vlevo);
UO=zeros(n+1,vlevo);
SST=zeros(n+1);

UA=zeros(n+1,vleva);
qA=zeros(n+1,vleva);
THETA=zeros(n+1,vleva);
rhoa=zeros(n+1,vleva);
TA=zeros(n+1,vleva)

UOst=zeros(vlevo); TOst=zeros(vlevo);
UAst=zeros(vleva); qAst=zeros(vleva); THETAst=zeros(vleva);

############################################################################
# Initial conditions
############################################################################
#U0=ones(n,1)*5;
#THETA0=(sin(2*pi/n*(1:n)*ndays/5-pi/2)+21);
TA[1,:]=25 .-0.007.*H.*[1:vleva;];
P=zeros(vleva);
P[1]=101300;
P[2:vleva]=P[1].*exp.(-cumsum(H./(Rgas.*(TA[1,2:end].+d2k)./gravity_mks)));
rhoa[1,:]=P./(Rgas.*(TA[1,:].+d2k));
THETA[1,:]=(TA[1,:].+d2k).*(P[1]./P).^(2/7).-d2k;
THETA[1,1:20].=THETA[1,1];
THETA0=ones(n,1)*THETA[1,end].+1;#+gamma_blk*H+0.01;#+0.51;#+SW*kappa/cpa*dt;
T0=ones(n+1,vlevo).*26;#repmat((16+(vlev-1:-1:0)/vlev*8),n+1,1);
TO[:,:].=T0;
TO[:,1]=T0[:,1]; TO[1,:]=T0[1,1].+15 .*([vlevo:-1:1;]/vlevo.-1);
SST.=TO[:,1];
ssq = saltsat*cvapor_fac*exp(-cvapor_exp./(TA[1]+d2k))./rhoa[1,1];
RH=0.8;
q0=0.5*ssq;
#Tv=(THETA(i-1,:)+d2k).*(1+humid_fac*q(i-1));
qA[1,:]=saltsat*cvapor_fac*exp.(-cvapor_exp./(TA[1,:].+d2k))./rhoa[1,:].*RH;
UA[1,:]=(0.2/karman).*(log.((ZA.-H/2)./1e-4));
#U(1,:)=U(1,end);
U0=(0.2/karman).*(log((ZA[end]+H/2)/1e-4)).*ones(1,n+1);
atemp=THETA[1,1]+d2k; aqh=qA[1,1];
speed=abs(UA[1,1]); sst=SST[1];
tmp0=bulkformulae(atemp,aqh,speed,sst,H/2,H/2,H/2,H/2,rhoa[1,1]);
huol0=tmp0["huol"]
TAup=270;
#U0=(sin(2*pi/n*(1:n)*ndays/5+2*pi/n*3600/dt*sf)*8+U(1,end));
#UO(:)=0.;
#U0(3600*24*5/dt:end)=U(1,end); one pertubation
############################################################################

LH=zeros(1,n); SH=zeros(1,n); TAU=zeros(1,n); TAUO=zeros(1,n);
LW=zeros(1,n); E=zeros(1,n); ssq=zeros(1,n);
xi=zeros(1,n); ce=zeros(1,n); ch=zeros(1,n);
ustar=zeros(1,n); qstar=zeros(1,n); tstar=zeros(1,n);
psimh=zeros(1,n); psixh=zeros(1,n); Ri=zeros(n,vleva);
rd=zeros(1,n); rh=zeros(1,n); re=zeros(1,n);
Qnet=zeros(1,n);
KAm=zeros(n,vleva);KAt=zeros(n,vleva);
KOm=zeros(n,vlevo);KOt=zeros(n,vlevo);

for i=2:n+1
    if rem(i,1440*60/dt)==0
        println(string("Day  ",round(i/(1440*60/dt),digits=5)))
    end

    ########################################################################
    # surface
    ########################################################################
    tmp=bulkformulae(TA[i-1,1]+d2k,qA[i-1,1],abs(UA[i-1,1]-UO[i-1,1]),SST[i-1,1],
        H/2,H/2,H/2,10,rhoa[i-1,1]);

    hl=tmp["hl"]
    hs=tmp["hs"]
    ce[i-1]=tmp["ce"]
    ch[i-1]=tmp["ch"]
    ssq[i-1]=tmp["ssq"]
    xi[i-1]=tmp["huol"]
    rd[i-1]=tmp["rd"]
    rh[i-1]=tmp["rh"]
    re[i-1]=tmp["re"]
    ustar[i-1]=tmp["ustar"]
    qstar[i-1]=tmp["qstar"]
    tstar[i-1]=tmp["tstar"]
    psimh[i-1]=tmp["psimh"]
    psixh[i-1]=tmp["psixh"]

    LH[i-1]=-hl;
    SH[i-1]=-hs;
    E[i-1]=-hl/Av;
    ϵa=0.725+0.17*log10(sum(qA[1,:]/1000*1.2*H)*100)
    LW[i-1]=emissivity*sb*((SST[i-1]+d2k).^4 .- ϵa.*(TAup).^4);
    TAU[i-1]=rhoa[i-1,1].*ustar[i-1].^2 .*sign(ustar[i-1]);
    Qnet[i-1]=SW[i-1,1]-LH[i-1]-SH[i-1]-LW[i-1];
    ########################################################################
    # Atmospheric vertical diffusion
    ########################################################################

    Tv=(THETA[i-1,:].+d2k).*(1 .+humid_fac*qA[i-1,:]);
    KAm[i-1,:], KAt[i-1,:]= holtslag(Tv,qA[i-1,1],UA[i-1,:],
                 UA[i-1,:].*0,LH[i-1],SH[i-1],ustar[i-1],[0; ZA]);
    kam=KAm[i-1,:];
    kat=KAt[i-1,:];

    ########################################################################
    # Oceanic vertical diffusion
    ########################################################################

    rhow=rhow0.*(1 .-alpha.*(TO[i-1,:].-tref).+beta.*(SO.-sref));
    B=gravity_mks.*(rhow0.-rhow)./rhow0;
    Uo=(vlevo:-1:1)/vlevo/10; Vo=zeros(size(Uo)); FWflux=zeros(size(Qnet));
    KOm[i-1,:],KOt[i-1,:],_=
        kpp(TO[i-1,:],TO[i-1,:].*0,UO[i,:],Vo,rhow,Qnet[i-1]
             ,0,ustar[i-1],ZO,lat);
    kot=KOt[i-1,:];
    kom=KOm[i-1,:];

    ########################################################################
    # Atmospheric explicit time stepping
    ########################################################################

    TAUS=sign(UA[i-1,1]-UO[i-1,1]);

    UAst[1]=UA[i-1,1]-TAU[i-1]*TAUS/(H*rhoa[i-1,1])*dt-(1-Aimp)*(
        kam[1].*(UA[i-1,1]-UA[i-1,2])/H^2*dt);
    UAst[end]=UA[i-1,end]-(1-Aimp)*(
        kam[end]*(UA[i-1,end]-U0[i-1])/H^2*dt);
    UAst[2:end-1]=UA[i-1,2:end-1]+(1-Aimp)*(
        kam[1:end-2].*(UA[i-1,1:end-2]-UA[i-1,2:end-1])/H^2*dt-
        kam[2:end-1].*(UA[i-1,2:end-1]-UA[i-1,3:end  ])/H^2*dt);

    qAst[1]=qA[i-1,1]-E[i-1]/(H*rhow[1])*dt-
        (1-Aimp)*(kat[1]*(qA[i-1,1]-qA[i-1,2])/H^2*dt);
    qAst[end]=qA[i-1,end]-(1-Aimp)*(kat[end]*(qA[i-1,end]-q0)/H^2*dt);
    qAst[2:end-1]=qA[i-1,2:end-1]+(1-Aimp)*(
        kat[1:end-2].*(qA[i-1,1:end-2]-qA[i-1,2:end-1])/H^2*dt-
        kat[2:end-1].*(qA[i-1,2:end-1]-qA[i-1,3:end  ])/H^2*dt);

    THETAst[1]=THETA[i-1,1]+(SH[i-1])/cpa/(H*rhoa[i-1,1])*dt-
         (1-Aimp)*(kat[1]*(THETA[i-1,1]-THETA[i-1,2])/H^2*dt);
    THETAst[end]=THETA[i-1,end]-
         (1-Aimp)*(kat[end]*(THETA[i-1,end]-THETA0[i-1])/H^2*dt);
    THETAst[2:end-1]=THETA[i-1,2:end-1]+(1-Aimp)*(
         kat[1:end-2].*(THETA[i-1,1:end-2]-THETA[i-1,2:end-1])/H^2*dt-
         kat[2:end-1].*(THETA[i-1,2:end-1]-THETA[i-1,3:end  ])/H^2*dt);

    ########################################################################
    # Oceanic explicit time stepping
    ########################################################################

    if cpl==1

        TOst[1]=TO[i-1,1]+(SW[i-1,1]-LH[i-1]-SH[i-1]-LW[i-1])/cp/(dz*rhow[1])*dt-
            (1-Oimp)*(kot[1]*(TO[i-1,1]-TO[i-1,2])/dz^2*dt);
        TOst[end]=TO[i-1,end]+SW[i-1,end]/cp/(dz*rhow[end])*dt+
            (1-Oimp)*(kot[end-1].*(TO[i-1,end-1]-TO[i-1,end])/dz^2*dt);
        TOst[2:end-1]=TO[i-1,2:end-1]+(1-Oimp)*(
            kot[1:end-2].*(TO[i-1,1:end-2]-TO[i-1,2:end-1])/dz^2*dt-
            kot[2:end-1].*(TO[i-1,2:end-1]-TO[i-1,3:end  ])/dz^2*dt)+
            SW[i-1,2:end-1]./cp./(dz*rhow[2:end-1])*dt;

        UOst[1]=UO[i-1,1]+
             (TAU[i-1])*TAUS/(rhow[1])/dz*dt-
             (1-Oimp)*(kom[1]*(UO[i-1,1]-UO[i-1,2])/dz^2*dt);
        UOst[end]=UO[i-1,end]+
             (1-Oimp)*(kom[end-1].*(UO[i-1,end-1]-UO[i-1,end])/dz^2*dt);
        UOst[2:end-1]=UO[i-1,2:end-1]+(1-Oimp)*(
             kom[1:end-2].*(UO[i-1,1:end-2]-UO[i-1,2:end-1])/dz^2*dt-
             kom[2:end-1].*(UO[i-1,2:end-1]-UO[i-1,3:end  ])/dz^2*dt);

    end

    if Aimp==0
        UA[i,:]=UAst; qA[i,:]=qAst; THETA[i,:]=THETAst;
    else
        du=-Aimp*kam[1:end-1].*dt./H^2; dl=-Aimp*kam[2:end].*dt./H^2;
        d=(1 .+Aimp.*dt./H^2 .*([kam[1:end-1]; 0]+[0; kam[2:end]]));
        A=inv(Tridiagonal(dl,d,du));
        UA[i,:]=A*UAst;
        du=-Aimp*kat[1:end-1].*dt./H^2; dl=-Aimp*kat[2:end].*dt./H^2;
        d=(1 .+Aimp.*dt./H^2 .*([kat[1:end-1]; 0]+[0; kat[2:end]]));
        A=inv(Tridiagonal(dl,d,du));
        qA[i,:]=A*qAst; THETA[i,:]=A*THETAst;
    end

    if Oimp==0
        UO[i,:]=UOst; TO[i,:]=TOst;
    else
        du=-Oimp*kom[1:end-1].*dt./dz^2; dl=-Oimp*kom[2:end].*dt./dz^2;
        d=(1 .+Oimp.*dt./dz^2 .*([kom[1:end-1]; 0]+[0; kom[2:end]]));
        A=inv(Tridiagonal(dl,d,du));
        UO[i,:]=A*UOst;
        du=-Oimp*kot[1:end-1].*dt./dz^2; dl=-Oimp*kot[2:end].*dt./dz^2;
        d=(1 .+Oimp.*dt./dz^2 .*([kot[1:end-1]; 0]+[0; kot[2:end]]));
        A=inv(Tridiagonal(dl,d,du));
        TO[i,:]=A*TOst;
    end
    SST[i]=TO[i,1]

    TA[i,:]=(THETA[i,:].+d2k).*(P./P[1]).^(2/7).-d2k;
    rhoa[i,:]=P./(Rgas.*(TA[i,:].+d2k));

    if moist==1
        qsat=saltsat*cvapor_fac*exp.(-cvapor_exp./(TA[i,:].+d2k))./rhoa[i-1,:];
        dq=(qA[i,:].-qsat).*((qA[i,:].-qsat).>0);
        dq=min.(dq,qA[i,:]);
        dTHETA=Av.*dq./cp; #[J/kg][j-1 kg K]
        #println([i maximum(dTHETA) maximum(dq)])
        ind=qA[i,:].>qsat;
        qA[i,ind]=qA[i,ind]-dq[ind];
        #THETA[i,:]=THETA[i,: ]+dTHETA;
    end

end

outputs=Dict("U0" => U0,"UA" => UA,"THETA" => THETA,"SST" => SST,"TO" => TO,
"xi" => xi,"ustar" => ustar,"psimh" => psimh,"psixh" => psixh,
"rd" => rd,"rh" => rh,"re" => re,"SW" => SW,"LW" => LW,"SH" => SH,"LH" => LH)

parameters=Dict("dt"=>dt, "n"=>n, "ndays"=>ndays, "cp"=>cp, "dz"=>dz, "rhow0"=>rhow0)

return outputs,parameters

end


using Base.Broadcast: BroadcastFunction, broadcasted, materialize

ChainRulesCore.@non_differentiable Base.getindex(m::AbstractAttenMask, I::Integer...)
ChainRulesCore.@non_differentiable Base.getindex(m::MaskIndexer, I::Integer...)
ChainRulesCore.@non_differentiable Base.getindex(m::AbstractAttenMask, I::Tuple)
ChainRulesCore.@non_differentiable Base.getindex(m::MaskIndexer, I::Tuple)
ChainRulesCore.@non_differentiable (::Type{<:AbstractAttenMask})(args...)
ChainRulesCore.@non_differentiable (::Type{<:AbstractAttenMaskOp})(args...)
ChainRulesCore.@non_differentiable getmask(arg...)

function ChainRulesCore.rrule(::typeof(apply_mask), op::NaiveAttenMaskOp, mask, score)
    m = as_bool(randomness(mask)) ? getmask(mask, score) : mask
    naive_apply_mask_pullback(Ȳ) = (NoTangent(), NoTangent(), NoTangent(), @thunk unthunk(Ȳ) .* m)
    return score .* m, naive_apply_mask_pullback
end

function ChainRulesCore.rrule(config::RuleConfig, ::typeof(apply_broadcast_mask), f, mask, score, scale)
    tmp = getmask(mask, score, scale)
    Y, back = rrule_via_ad(config, broadcast, f, tmp, score)
    function fallback_apply_broadcast_mask_pullback(Ȳ)
        Ȳs = (back∘unthunk)(Ȳ)
        return (NoTangent(), Ȳs[2], NoTangent(), Ȳs[4], NoTangent())
    end
    Y, fallback_apply_broadcast_mask_pullback
end

function ChainRulesCore.rrule(config::RuleConfig, ::typeof(apply_broadcast_mask), ::typeof(+), mask, score, scale)
    apply_broadcast_mask_pullback(Ȳ) = (NoTangent(), NoTangent(), NoTangent(), Ȳ, NoTangent())
    apply_broadcast_mask(+, mask, score, scale), apply_broadcast_mask_pullback
end

function ChainRulesCore.rrule(config::RuleConfig, ::typeof(apply_broadcast_mask), ::typeof(*), mask, score, scale)
    rnm = as_bool(randomness(mask))
    m = rnm ? getmask(mask, score, scale) : mask
    function apply_broadcast_mask_pullback(Ȳ)
        thk = @thunk begin
            tmp = unthunk(Ȳ)
            rnm ? tmp .* m : @. tmp * (mask * scale)
        end
        return (NoTangent(), NoTangent(), NoTangent(), thk, NoTangent())
    end
    fwd = rnm ? score .* m : apply_broadcast_mask(*, mask, score, scale)
    fwd, apply_broadcast_mask_pullback
end

@init @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" begin
    Zygote.unbroadcast(x::AbstractAttenMask, _) = nothing

    function Zygote._pullback(ctx::Zygote.AContext, ::typeof(Broadcast.broadcasted), ::MaskStyle{M}, f, args...) where M
        return y, ∇broadcasted = Zygote._pullback(ctx, Broadcast.broadcasted, M(), f, args...)
    end
end


module SurrogateModels

"""
    regression(X,y,m)

Output a regression matrix given inputs

 - `X`: ``n × k`` matrix with `n` training points in `k` dimensional feature space

 - `y`: ``n × 1`` vector of responses (######)

 - `m`: Order of polynomial (Presently only `1` and `2` is supported)

"""
function regression(x,y,m)

    return(x)
end

end
# vim:filetype=julia:foldmethod=syntax


using Gadfly, Colors, BlackBoxOptimizationBenchmarking
using Distributed, Statistics
##

run_bench = true

nprocs()< 2 && addprocs(2)

## Setup

bbob_dir = dirname(pathof(BlackBoxOptimizationBenchmarking))

@everywhere begin

    using BlackBoxOptimizationBenchmarking
    const BBOB = BlackBoxOptimizationBenchmarking
    bbob_dir = dirname(pathof(BBOB))
    include(joinpath(bbob_dir,"../scripts/optimizers_interface.jl"))

    dimensions = [3 6 12]
    Ntrials = 1
    Δf = 1e-6
    run_lengths = round.(Int,range(20, stop=3_000, length=15))
    #run_lengths = round.(Int,linspace(20,20_000,15))
    funcs = BBOB.list_functions()
end

c = opt -> Chain(opt,NelderMead(),0.9)

optimizers = [
    c(SAMIN(;verbosity=0)),
    NelderMead(),
    OptimRestart(NelderMead()),
    c(SimulatedAnnealing()),
    c(BlackBoxOptimMethod(:adaptive_de_rand_1_bin_radiuslimited)),
    c(BlackBoxOptimMethod(:adaptive_de_rand_1_bin)),
    c(NLoptOptimMethod(:GN_ISRES)),
    c(NLoptOptimMethod(:GN_ESCH)),
    c(NLoptOptimMethod(:GD_STOGO)),
    c(BlackBoxOptimMethod(:xnes)),
    c(BlackBoxOptimMethod(:generating_set_search)),
    c(BlackBoxOptimMethod(:de_rand_2_bin)),
    #BlackBoxOptimMethod(:resampling_memetic_search),#poor performance
    PyMinimize("Nelder-Mead"),
    #PyCMA(),
]

opt_strings = map(string,optimizers)

## run benchmark

if run_bench

    results = pmap(opt->BBOB.benchmark(opt, funcs, run_lengths, Ntrials, dimensions, Δf), optimizers)
    error("done")
    @info("Done, saving data and making plots.")

    # save output
    outdir = joinpath(bbob_dir,"..","data")

    writedlm(joinpath(outdir,"mean_succ.txt"),mean_succ)
    writedlm(joinpath(outdir,"mean_dist.txt"),mean_dist)
    writedlm(joinpath(outdir,"mean_fmin.txt"),mean_fmin)
    writedlm(joinpath(outdir,"runtime.txt"),runtime)
else
    # load output
    outdir = joinpath(bbob_dir,"..","data")

    s = (length(optimizers),length(funcs),length(run_lengths),length(dimensions))
    mean_succ = reshape( readdlm(joinpath(outdir,"mean_succ.txt")), s)
    mean_dist = reshape( readdlm(joinpath(outdir,"mean_dist.txt")), s)
    mean_fmin = reshape( readdlm(joinpath(outdir,"mean_fmin.txt")), s)
    runtime = readdlm(joinpath(outdir,"runtime.txt"))
end

## make plots

outdir = joinpath(bbob_dir,"..","data","plots")
cols = Colors.distinguishable_colors(size(mean_succ,1)+1,colorant"white")[2:end]
line_style = i -> mod(i,2)==0 ? [:solid] : [:dash]

#all together

idx = sortperm([mean(mean_succ[i,:,end,:]) for i=1:length(optimizers)],rev=true)

p = plot(
    [layer(
        x=collect(run_lengths),y=mean(mean_succ[i,:,:,:],dims=(1,3)),Geom.line,
        Theme(default_color=cols[i], line_style=line_style(i), line_width=2pt)
    ) for i in 1:size(mean_succ,1)]...,
    Coord.cartesian(ymax=1),
    Guide.title("All functions"), Guide.xlabel("Run Length"), Guide.ylabel("Success rate"),
    Guide.manual_color_key("", opt_strings[idx], cols[idx])
)

#

#draw(PDF(joinpath(outdir,"mean_succ.pdf"),22cm,14cm),p)
#draw(PNG(joinpath(outdir,"mean_succ.png"),dpi=150,16cm,10cm),p)
draw(SVG(joinpath(outdir,"mean_succ.svg"),16cm,10cm),p)

# per dimension
for (k,D) in zip(1:length(dimensions), dimensions)

    idx = sortperm([mean(mean_succ[i,:,end,k]) for i=1:length(optimizers)],rev=true)

    p = plot(
        [layer(
            x=collect(run_lengths),y=mean(mean_succ[i,:,:,k],dims=1),Geom.line,
            Theme(default_color=cols[i], line_style=line_style(i), line_width=2pt)
        ) for i in 1:size(mean_succ,1)]...,
        Coord.cartesian(ymax=1),
        Guide.title("All functions, D: $D"), Guide.xlabel("Run Length"), Guide.ylabel("Success rate"),
        Guide.manual_color_key("", opt_strings[idx], cols[idx])
    )
    #draw(PDF(joinpath(outdir,"per_dimension","mean_succ_$(D).pdf"),22cm,14cm),p)
    #draw(PNG(joinpath(outdir,"per_dimension","mean_succ_$(D).png"),dpi=150,16cm,10cm),p)
    draw(SVG(joinpath(outdir,"per_dimension","mean_succ_$(D).svg"),16cm,10cm),p)

end

## f min

idx = sortperm([exp(mean(log.(mean_fmin[i,:,end,:].+eps()))) for i=1:length(optimizers)],rev=true)

# all together
p = plot(
    [layer(
        x=collect(run_lengths),y=exp.(mean(log.(mean_fmin[i,:,:,:].+eps()),dims=(1,3))),Geom.line,
        Theme(default_color=cols[i],line_width=2pt)
    )
    for i in 1:size(mean_succ,1)]...,
    layer(yintercept=[Δf],Geom.hline,Theme(default_color=colorant"gray")),
    Guide.title("All functions"), Guide.xlabel("Run Length"), Guide.ylabel("fmin"),
    Guide.manual_color_key("", opt_strings[idx], cols[idx]),
#    Scale.x_log10,
    Scale.y_log10,
)
#draw(PDF(joinpath(outdir,"mean_fmin.pdf"),22cm,14cm),p)
#draw(PNG(joinpath(outdir,"mean_fmin.png"),dpi=150,16cm,10cm),p)
draw(SVG(joinpath(outdir,"mean_fmin.svg"),16cm,10cm),p)

#per dimension
for (k,D) in zip(1:length(dimensions), dimensions)

    idx = sortperm([exp.(mean(log.(mean_fmin[i,:,end,k].+eps()))) for i=1:length(optimizers)],rev=true)

    p = plot(
        [layer(
            x=collect(run_lengths)/D,y=exp.(mean(log.(mean_fmin[i,:,:,k].+eps()),dims=1)),Geom.line,
            Theme(default_color=cols[i],line_width=2pt)
        ) for i in 1:size(mean_succ,1)]...,
        layer(yintercept=[Δf],Geom.hline,Theme(default_color=colorant"gray")),
        Guide.title("All functions"), Guide.xlabel("Run Length / D"), Guide.ylabel("fmin"),
        Guide.manual_color_key("", opt_strings[idx], cols[idx]),
    #    Scale.x_log10,
        Scale.y_log10,
    )

    #draw(PDF(joinpath(outdir,"per_dimension","mean_fmin_$(D).pdf"),20cm,12cm),p)
    #draw(PNG(joinpath(outdir,"per_dimension","mean_fmin_$(D).png"),dpi=150,16cm,10cm),p)
    draw(SVG(joinpath(outdir,"per_dimension","mean_fmin_$(D).svg"),16cm,10cm),p)
end

## distance to minimizer

for (k,D) in zip(1:length(dimensions), dimensions)

    p = plot(
        [layer(x=collect(run_lengths)/D,y=median(mean_dist[i,:,:,k],dims=1),Geom.line,Theme(default_color=cols[i])) for i in 1:size(mean_succ,1)]...,
        Guide.title("All functions"), Guide.xlabel("Run Length / D"), Guide.ylabel("Distance to xmin"),
        Guide.manual_color_key("", opt_strings, cols),
    #    Scale.x_log10,
        Scale.y_log10,
    )

#    #draw(PDF(joinpath(outdir,"per_dimension","mean_dist_$(D).pdf"),20cm,12cm),p)
    draw(SVG(joinpath(outdir,"per_dimension","mean_dist_$(D).svg"),16cm,10cm),p)
end

## runtime

idx = sortperm(mean(runtime,dims=2)[:],rev=true)

p = plot(
    layer(
        y=opt_strings[idx], x=(mean(runtime,dims=2)/Base.minimum(mean(runtime,dims=2)))[idx],
        color = cols[idx],
        Geom.bar(orientation=:horizontal)
    ),
    Guide.title("All functions"), Guide.ylabel(""),
    Guide.xlabel("Relative Run Time",orientation=:horizontal),
    Scale.x_log10,
#    Scale.y_log10,
)

#draw(PDF(joinpath(outdir,"runtime.pdf"),20cm,12cm),p)
#draw(PNG(joinpath(outdir,"runtime.png"),dpi=150,16cm,10cm),p)
draw(SVG(joinpath(outdir,"runtime.svg"),16cm,10cm),p)
nothing

# convert svg to png (Gadfly's png export is broken at the moment)
# for f in *.svg; do
#     svgexport "$f" "${f%.svg}.png" 1.5x
# done

##


module TestExamplesMPIP4estMesh

using Test
using Trixi

include("test_trixi.jl")

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
const EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "p4est_2d_dgsem")

@testset "P4estMesh MPI" begin

# Run basic tests
@testset "Examples 2D" begin
  # Linear scalar advection
  @trixi_testset "elixir_advection_basic.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_basic.jl"),
      # Expected errors are exactly the same as with TreeMesh!
      l2   = [8.311947673061856e-6],
      linf = [6.627000273229378e-5])
  end

  @trixi_testset "elixir_advection_unstructured_flag.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_unstructured_flag.jl"),
      l2   = [0.0005379687442422346],
      linf = [0.007438525029884735])
  end

  @trixi_testset "elixir_advection_restart.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_restart.jl"),
      l2   = [4.507575525876275e-6],
      linf = [6.21489667023134e-5])
  end
end

end # P4estMesh MPI

end # module


using GLAbstraction, GeometryTypes, ModernGL, Compat, FileIO, GLWindow, FixedSizeArrays, FixedPointNumbers, ColorTypes
using Base.Test

function is_ci()
	get(ENV, "TRAVIS", "") == "true" || get(ENV, "APPVEYOR", "") == "true" || get(ENV, "CI", "") == "true"
end

include("macro_test.jl")

if !is_ci() # only do test if not CI... this is for automated testing environments which fail for OpenGL stuff, but I'd like to test if at least including works

window = create_glcontext("test", resolution=(500,500))

include("accessors.jl")
include("uniforms.jl")
include("texture.jl")

# Test for creating a GLBuffer with a 1D Julia Array
# You need to supply the cardinality, as it can't be inferred
# indexbuffer is a shortcut for GLBuffer(GLUInt[0,1,2,2,3,0], 1, buffertype = GL_ELEMENT_ARRAY_BUFFER)
indexes = indexbuffer(GLuint[0,1,2])
# Test for creating a GLBuffer with a 1D Julia Array of Vectors
#v = Vec2f[Vec2f(0.0, 0.5), Vec2f(0.5, -0.5), Vec2f(-0.5,-0.5)]

v = Vec2f0[Vec2f0(0.0, 0.5), Vec2f0(0.5, -0.5), Vec2f0(-0.5,-0.5)]

verts = GLBuffer(v)
@test size(verts, 1) == 3
@test size(verts, 2) == 1
# lets define some uniforms
# uniforms are shader variables, which are supposed to stay the same for an entire draw call

const triangle = RenderObject(
	Dict(
		:vertex => verts,
		:name_doesnt_matter_for_indexes => indexes
	),
	LazyShader(load("test.vert"), load("test.frag"))
)

postrender!(triangle, render, triangle.vertexarray)

glClearColor(0,0,0,1)
while isopen(window)
  	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
	render(triangle)
	swapbuffers(window)
	pollevents()
	sleep(0.01)
end

end


using ACE, ACEbase, StaticArrays
using Printf, LinearAlgebra, ACE.Testing, Random
using ACE: evaluate, evaluate_d, SymmetricBasis, SimpleSparseBasis, PIBasis,
           PositionState
using ACEbase.Testing: fdtest
using Profile, ProfileSVG, TimerOutputs, BenchmarkTools

##

# construct the 1p-basis
maxdeg = 14
ord = 4
wL = 1.5
Bsel = SparseBasis(; maxorder = ord, p = 1, default_maxdeg = maxdeg,
                     weight = Dict(:n => 1.0, :l => wL))

B1p = ACE.Utils.RnYlm_1pbasis(maxdeg = maxdeg, maxL = ceil(Int, maxdeg/wL),
                              Bsel = Bsel)

# generate a configuration
nX = 10
Xs = rand(PositionState{Float64}, B1p.bases[1], nX)
cfg = ACEConfig(Xs)

φ = ACE.Invariant()
basis = SymmetricBasis(φ, B1p, Bsel)
@show length(basis)

BB = evaluate(basis, cfg)
c = rand(length(BB)) .- 0.5

standard = ACE.LinearACEModel(basis, c, evaluator = :standard)

##

@info("Time evaluate incl allocation")
(@btime evaluate($standard, $cfg))

##

@info("Time grad_config with and without allocation")
@info("This looks like a factor 4.5 of evaluate, so probably more we can do")
g = ACE.acquire_grad_config!(standard, cfg)
(@btime ACE.grad_config($standard, $cfg))
(@btime ACE.grad_config!($g, $standard, $cfg))

##

@info("grad_config vs rrule_evaluate")
g1 = ACE._rrule_evaluate(ACE._One(), standard, cfg)
g2 = ACE.grad_config(standard, cfg)
@show g1 ≈ g2
@btime ACE.grad_config($standard, $cfg)
@btime ACE._rrule_evaluate(ACE._One(), $standard, $cfg)

##

@info("Time grad_params with and without allocation")
@info("a little surprising we dont get closer to factor 1?")
g = ACE.acquire_grad_params!(standard, cfg)
(@btime ACE.grad_params($standard, $cfg))
(@btime ACE.grad_params!($g, $standard, $cfg))

##

@info("adjoint_EVAL_D - single property")
@info("This may be the single-most important function, and it looks fantastic!")
#create a random input emulating the pullback input
w = [ACE.DState(rr = randn(SVector{3, Float64})) for j in 1:length(cfg)]
@btime ACE.adjoint_EVAL_D($standard, $cfg, $w)

##

@info("Multi-property")

c_m = randn(SVector{2, Float64}, length(basis))
model2 = ACE.LinearACEModel(basis, c_m, evaluator = :standard)

@info(" - evaluate")
@btime evaluate($model2, $cfg)
@info(" - grad_params")
@btime ACE.grad_params($model2, $cfg)
@info(" - grad_config")
@btime ACE.grad_config($model2, $cfg)
w20 = randn(SVector{2, Float64})
@info(" - _rrule_evaluate")
@btime ACE._rrule_evaluate($w20, $model2, $cfg)

@info(" - adjoint_EVAL_D")
_w2() = SVector(ACE.DState(rr = randn(SVector{3, Float64})),
                ACE.DState(rr = randn(SVector{3, Float64})))
w2 = [ _w2() for j in 1:length(cfg)]
ACE.adjoint_EVAL_D(model2, cfg, w2)
@btime ACE.adjoint_EVAL_D($model2, $cfg, $w2)

##

# function runn(N, f, args...)
#    t = f(args...)
#    for n = 2:N
#      t = f(args...)
#    end
#    t
# end

# ##

# Profile.clear()
# @profile runn(10000, ACE.adjoint_EVAL_D, standard, cfg, w)
# Profile.print()

# ##

# ProfileSVG.view()

##

# # make sure to add suitable @timeit macros to the module
# @info("Detailed benchmarking of grad_config")
# reset_timer!()
# @timeit "grad_config" runn(10, ACE.grad_config, standard, cfg);
# print_timer()


module Flux

# Zero Flux Given

using Base: tail
using Statistics, Random, LinearAlgebra
using Zygote, MacroTools, Juno, Reexport
using MacroTools: @forward
@reexport using NNlib
using Zygote: Params, @adjoint, gradient, pullback, @nograd

export gradient

export Chain, Dense, Maxout, RNN, LSTM, GRU, SamePad, Conv, CrossCor, ConvTranspose,
       AdaptiveMaxPool, AdaptiveMeanPool, GlobalMaxPool, GlobalMeanPool, MaxPool,
       MeanPool, flatten, DepthwiseConv, Dropout, AlphaDropout, LayerNorm, BatchNorm,
       InstanceNorm, GroupNorm, SkipConnection, params, fmap, cpu, gpu, f32, f64,
       testmode!, trainmode!

include("optimise/Optimise.jl")
using .Optimise
using .Optimise: @epochs
using .Optimise: skip
export Descent, ADAM, Momentum, Nesterov, RMSProp,
  ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, OADAM,
  ADAMW, RADAM, AdaBelief, InvDecay, ExpDecay,
  WeightDecay, ClipValue, ClipNorm

using CUDA
const use_cuda = Ref(false)

include("utils.jl")
include("zeros.jl")
include("onehot.jl")
include("functor.jl")

include("layers/stateless.jl")
include("layers/basic.jl")
include("layers/conv.jl")
include("layers/recurrent.jl")
include("layers/normalise.jl")

include("data/Data.jl")

include("losses/Losses.jl")
using .Losses # TODO: stop importing Losses in Flux's namespace in v0.12

include("deprecations.jl")

include("cuda/cuda.jl")

function __init__()
  use_cuda[] = CUDA.functional() # Can be overridden after load with `Flux.use_cuda[] = false`
  if CUDA.functional()
    if !CUDA.has_cudnn()
      @warn "CUDA.jl found cuda, but did not find libcudnn. Some functionality will not be available."
    end
  end
end

end # module


# =============================================================================
# =
# = FlatZinc model.
# =
# =============================================================================

mutable struct VariableInfo
    index::MOI.VariableIndex
    name::String
    set::MOI.AbstractSet
end

mutable struct ConstraintInfo
    index::MOI.ConstraintIndex
    # FlatZinc does not allow constraint names.
    f::MOI.AbstractFunction
    s::MOI.AbstractSet
    output_as_part_of_variable::Bool
end

mutable struct Model <: MOI.AbstractOptimizer
    # A mapping from the MOI.VariableIndex to the variable object.
    # VariableInfo also stores some additional fields like the type of variable.
    variable_info::CleverDicts.CleverDict{MOI.VariableIndex, VariableInfo}

    # A mapping from the MOI.ConstraintIndex to the variable object.
    # ConstraintInfo also stores some additional fields like the type of
    # constraint. Deletion of constraints is not supported (having a vector
    # ensures that the order in which constraints are added is respected when
    # outputting the model, which makes testing easier).
    constraint_info::Vector{ConstraintInfo}

    # For FlatZinc-level sets/arrays, map the constraints that are defined on
    # sets/arrays to the assigned name of the set/array. This structure is
    # built when outputting the FZN file.
    sets_id::Dict{MOI.ConstraintIndex, String}
    arrs_id::Dict{MOI.ConstraintIndex, String}

    # Memorise the objective sense and the function separately.
    # The function can only be a single variable, as per FlatZinc limitations.
    objective_sense::MOI.OptimizationSense
    objective_function::Union{Nothing, MOI.VariableIndex}

    # Map names to MOI variable indices.
    name_to_var::Dict{String, MOI.VariableIndex}

    """
        Model()

    Create a new Model object.
    """
    function Model()
        model = new()
        model.variable_info =
            CleverDicts.CleverDict{MOI.VariableIndex, VariableInfo}()
        model.constraint_info = ConstraintInfo[]

        model.objective_sense = MOI.FEASIBILITY_SENSE
        model.objective_function = nothing

        model.sets_id = Dict{MOI.ConstraintIndex, String}()
        model.arrs_id = Dict{MOI.ConstraintIndex, String}()

        model.name_to_var = Dict{String, MOI.VariableIndex}()

        MOI.empty!(model)
        return model
    end
end

function Base.show(io::IO, ::Model)
    print(io, "A FlatZinc (fzn) model")
    return
end

function MOI.empty!(model::Model)
    empty!(model.variable_info)
    empty!(model.constraint_info)

    model.objective_sense = MOI.FEASIBILITY_SENSE
    model.objective_function = nothing

    model.sets_id = Dict{MOI.ConstraintIndex, String}()
    model.arrs_id = Dict{MOI.ConstraintIndex, String}()

    model.name_to_var = Dict{String, MOI.VariableIndex}()

    return
end

function MOI.is_empty(model::Model)
    !isempty(model.variable_info) && return false
    !isempty(model.constraint_info) && return false
    model.objective_sense != MOI.FEASIBILITY_SENSE && return false
    model.objective_function !== nothing && return false
    return true
end

# Set the objective.

function MOI.get(model::Model, ::MOI.ObjectiveFunction{MOI.VariableIndex})
    return model.objective_function
end

function MOI.set(
    model::Model,
    ::MOI.ObjectiveFunction{MOI.VariableIndex},
    f::MOI.VariableIndex,
)
    model.objective_function = f
    return nothing
end

function MOI.get(model::Model, ::MOI.ObjectiveSense)
    return model.objective_sense
end

function MOI.set(
    model::Model,
    ::MOI.ObjectiveSense,
    s::MOI.OptimizationSense,
)
    model.objective_sense = s
    return nothing
end

# Helpers.

function _create_variable(
    model::Model,
    set::Union{MOI.AbstractScalarSet, MOI.Reals},
)
    index = CleverDicts.add_item(
        model.variable_info,
        VariableInfo(MOI.VariableIndex(0), "", set),
    )
    model.variable_info[index].index = index
    return index
end

function _create_constraint(
    model::Model,
    f::F,
    set::S,
    as_part_of_variable::Bool,
) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
    index = MOI.ConstraintIndex{F, S}(length(model.constraint_info) + 1)
    push!(
        model.constraint_info,
        ConstraintInfo(index, f, set, as_part_of_variable),
    )
    return index
end

# Fallback for copying operations.
function MOI.copy_to(dest::Model, src::MOI.ModelLike)
    return MOIU.default_copy_to(dest, src)
end

function MOI.supports_incremental_interface(::Model)
    return true
end

# Names.
# No support for constraint names in fzn, hence no ConstraintName.

function MOI.get(model::Model, ::MOI.VariableName, v::MOI.VariableIndex)
    return model.variable_info[v].name
end

function MOI.get(model::Model, ::Type{MOI.VariableIndex}, n::String)::MOI.VariableIndex
    # TODO: test this.
    # TODO: not terribly efficient, and used in FZN for parsing a solution (finding a variable by its name). Rather build a dict to retrieve the indices faster.
    for i in keys(model.variable_info)
        if model.variable_info[i].name == n
            return i
        end
    end
    return nothing
end

function MOI.set(
    model::Model,
    ::MOI.VariableName,
    v::MOI.VariableIndex,
    name::AbstractString,
)
    model.variable_info[v].name = string(name)
    return
end

# Variables.

function MOI.supports_add_constrained_variable(
    ::Model,
    ::Type{F},
) where {
    F <: Union{
        MOI.EqualTo{Float64},
        MOI.LessThan{Float64},
        MOI.Interval{Float64},
        MOI.EqualTo{Int},
        MOI.LessThan{Int},
        MOI.Interval{Int},
        MOI.EqualTo{Bool},
        MOI.ZeroOne,
        MOI.Integer,
    },
}
    return true
end

function MOI.supports_add_constrained_variables(
    ::Model,
    ::Type{F},
) where {
    F <: Union{
        MOI.EqualTo{Float64},
        MOI.LessThan{Float64},
        MOI.Interval{Float64},
        MOI.EqualTo{Int},
        MOI.LessThan{Int},
        MOI.Interval{Int},
        MOI.EqualTo{Bool},
        MOI.ZeroOne,
        MOI.Integer,
    },
}
    return true
end

function MOI.add_variable(model::Model)
    return _create_variable(model, MOI.Reals(1))
end

function MOI.add_constrained_variables(
    model::Model,
    sets::AbstractVector{<:MOI.AbstractScalarSet},
)
    # TODO: memorise that these variables are part of the same call, so that
    # the generated FlatZinc is shorter (array of variables)? This would
    # require that all sets are identical, though.
    vidx = [_create_variable(model, sets[i]) for i in 1:length(sets)]
    cidx = [
        _create_constraint(model, vidx[i], sets[i], true) for i in 1:length(sets)
    ]
    return vidx, cidx
end

function MOI.add_constrained_variable(
    model::Model,
    set::MOI.AbstractScalarSet,
)
    # Unlike FZN, MOI does not assume that a variable >= 0.0 is any different
    # from a >= 0 (int).
    vidx = _create_variable(model, set)
    cidx = _create_constraint(model, vidx, set, true)
    return vidx, cidx
end

function MOI.is_valid(model::Model, v::MOI.VariableIndex)
    return haskey(model.variable_info, v)
end

function MOI.get(model::Model, ::MOI.NumberOfVariables)
    return length(model.variable_info)
end

function MOI.get(model::Model, ::MOI.ListOfVariableIndices)
    return collect(keys(model.variable_info))
end

# Constraints.

function MOI.is_valid(
    model::Model,
    c::MOI.ConstraintIndex{F, S},
) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
    info = get(model.constraint_info, c.value, nothing)
    return info !== nothing && typeof(info.s) == S
end

function MOI.add_constraint(
    model::Model,
    f::F,
    s::S,
) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
    index = MOI.ConstraintIndex{F, S}(length(model.constraint_info) + 1)
    push!(model.constraint_info, ConstraintInfo(index, f, s, false))
    return index
end

function MOI.get(
    model::Model,
    ::MOI.ConstraintFunction,
    c::MOI.ConstraintIndex{F, S},
) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
    return model.constraint_info[c.value].f
end

function MOI.get(
    model::Model,
    ::MOI.ConstraintSet,
    c::MOI.ConstraintIndex{F, S},
) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
    return model.constraint_info[c.value].s
end

function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VariableIndex},
    ::Type{MOI.LessThan{Int}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VariableIndex},
    ::Type{MOI.LessThan{Float64}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VariableIndex},
    ::Type{CP.Strictly{MOI.LessThan{Float64}}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VariableIndex},
    ::Type{CP.Domain{Int}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VariableIndex},
    ::Type{MOI.Interval{Float64}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VectorOfVariables},
    ::Type{CP.Element{Int}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VectorOfVariables},
    ::Type{CP.Element{Bool}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VectorOfVariables},
    ::Type{CP.Element{Float64}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VectorOfVariables},
    ::Type{CP.MaximumAmong},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.VectorOfVariables},
    ::Type{CP.MinimumAmong},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.ScalarAffineFunction{Int}},
    ::Type{MOI.EqualTo{Int}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.ScalarAffineFunction{Int}},
    ::Type{MOI.LessThan{Int}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.ScalarAffineFunction{Int}},
    ::Type{CP.DifferentFrom{Int}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.ScalarAffineFunction{Float64}},
    ::Type{MOI.EqualTo{Float64}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.ScalarAffineFunction{Float64}},
    ::Type{MOI.LessThan{Float64}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.ScalarAffineFunction{Float64}},
    ::Type{CP.Strictly{MOI.LessThan{Float64}}},
)
    return true
end
function MOI.supports_constraint(
    ::Model,
    ::Type{MOI.ScalarAffineFunction{Float64}},
    ::Type{CP.DifferentFrom{Float64}},
)
    return true
end

function MOI.get(
    model::Model,
    ::MOI.ListOfConstraintIndices{F, S},
) where {F, S}
    return [
        c.index for
        c in model.constraint_info if typeof(c.f) == F && typeof(c.s) == S
    ]
end

function MOI.get(model::Model, ::MOI.NumberOfConstraints{F, S}) where {F, S}
    count = 0
    for c in model.constraint_info
        if typeof(c.f) == F && typeof(c.s) == S
            count += 1
        end
    end
    return count
end

function MOI.get(model::Model, ::MOI.ListOfConstraintTypesPresent)
    types = Set{Tuple{Any, Any}}()
    for info in model.constraint_info
        push!(types, (typeof(info.f), typeof(info.s)))
    end
    return collect(types)
end


module PSSFSS

module GSMs
export GSM

mutable struct GSM{T1 <: AbstractMatrix, T2 <: AbstractMatrix}
    s11::T1
    s12::T2
    s21::T2
    s22::T1
end

end # module GSMs
##############################################
module Outputs
using ..GSMs: GSM
using JLD2
using StaticArrays
export runit

const SV2 = SVector{2,Float64}

SteerType = Union{NamedTuple{(:ψ₁, :ψ₂), Tuple{Float64, Float64}},
                  NamedTuple{(:θ, :ϕ), Tuple{Float64, Float64}}}

struct Result
    gsm::GSM
    steering::SteerType
    β⃗₀₀::SV2
    FGHz::Float64
    ϵᵣin::ComplexF64
    μᵣin::ComplexF64
    β₁in::SV2
    β₂in::SV2
    ϵᵣout::ComplexF64
    μᵣout::ComplexF64
    β₁out::SV2
    β₂out::SV2
end

function append_result_data(fname, gname, result)
    jldopen(fname, "a") do fid
        group = JLD2.Group(fid, gname)
        group["result"] = result
    end
    return
end

function read_result_file(fname::AbstractString)::Vector{Result}
    dat = load(fname) # a Dict
    ks = collect(keys(dat))
    sort!(ks, by = x -> parse(Int,split(x, '/')[1]))
    Result[dat[k] for k in ks]
end

function runit(n1=1)
    infile = "pssfss.res"
    outfile = "test.res"
    isfile(outfile) && rm(outfile)
    results = read_result_file(infile)
    for k = n1:length(results)
        println(k)
        gname = string(k)
        append_result_data(outfile, gname, results[k])
    end
end

end # module Outputs
##############################################

end # module PSSFSS


# Quick notes for tests. None is actually implemented. This file is not even supposed to work.

using FactCheck

facts("Module: plot") do
  context() do
    # a: no sc, i.e. only one scenario per reservoir.     => with sc, must throw an error.
    # d: no sc, check the second output argument is empty.
    a = getDataFrameForRiver([VesdreReservoir_adam2015_95_avg, VesdreReservoir_adam2015_95_low, VesdreReservoir_adam2015_95_high], 1:ruleLength, "Vesdre", multipleScenarios=:None)
    b = getDataFrameForRiver([sc, VesdreReservoir_adam2015_95_avg, VesdreReservoir_adam2015_95_low, VesdreReservoir_adam2015_95_high], 1:ruleLength, "Vesdre", multipleScenarios=:Merge)
    c = getDataFrameForRiver([sc, VesdreReservoir_adam2015_95_avg, VesdreReservoir_adam2015_95_low, VesdreReservoir_adam2015_95_high], 1:ruleLength, "Vesdre", multipleScenarios=:Split)
    d = getDataFrameForRiver([VesdreReservoir_adam2015_95_avg, VesdreReservoir_adam2015_95_low, VesdreReservoir_adam2015_95_high], 1:ruleLength, "Vesdre", multipleScenarios=:Split)
  end
end

facts("Module: read") do
  context() do
    b = makeScenarios(Vesdre_ts, sampleDuration=Week(1))
    c = makeScenarios(Vesdre_ts, sampleDuration=Day(1), resamplingRate=28) ---> throws()
  end
end

facts("Module: reservoir") do
  context("accessors: statistical") do
    splitIntoSets(Vesdre.scenarios, [1])
    splitIntoSets(Vesdre.scenarios, [.8, .2])
  end

  context("accessors") do
    totalInflow(VesdreReservoir, 1)
    totalInflow(VesdreReservoir, 1:22)
    totalInflow(VesdreReservoir, 1, asTimeSeries=true)
    totalInflow(VesdreReservoir, 1:22, asTimeSeries=true)
  end
end


module PDFEstimators

    using
        AverageShiftedHistograms,
        CUDA,
        ComputationalResources,
        DataFrames,
        DiffEqFlux,
        DiffEqSensitivity,
        DifferentialEquations,
        Distances,
        Distributions,
        Flux,
        ForwardDiff,
        GalacticFlux,
        GalacticOptim,
        GalacticOptimJL,
        KernelDensity,
        LineSearches,
        MLJBase,
        MLJFlux,
        MLJModelInterface,
        ModelingToolkit,
        Optim,
        Parameters,
        ReverseDiff,
        SciMLBase,
        ScientificTypes,
        Zygote,
        LinearAlgebra,
        Statistics

    include("core.jl")
    include("actual.jl")
    include("kde.jl")
    include("ash.jl")
    include("mvn.jl")
    include("ffjord.jl")
    include("utils.jl")

    MLJBase.metadata_pkg.(
        [ActualModel, KDEModel, FFJORDModel, MvnModel],
        package_name="PDFEstimators",
        package_uuid="4d826980-ec0a-4f65-a989-f1052e211ebd",
        package_url="https://github.com/prbzrg/PDFEstimators.jl",
        is_pure_julia=true,
        package_license="MIT",
        is_wrapper=false,
    )

    MLJBase.metadata_model(
        ActualModel,
        input_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        target_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        output_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        supports_weights=false,
        docstring="ActualModel",
        load_path="PDFEstimators.ActualModel",
    )
    MLJBase.metadata_model(
        KDEModel,
        input_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        target_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        output_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        supports_weights=false,
        docstring="KDEModel",
        load_path="PDFEstimators.KDEModel",
    )
    MLJBase.metadata_model(
        FFJORDModel,
        input_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        target_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        output_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        supports_weights=false,
        docstring="FFJORDModel",
        load_path="PDFEstimators.FFJORDModel",
    )
    MLJBase.metadata_model(
        MvnModel,
        input_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        target_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        output_scitype=Table{AbstractVector{ScientificTypes.Continuous}},
        supports_weights=false,
        docstring="MvnModel",
        load_path="PDFEstimators.MvnModel",
    )
end


#=

Draw radar map in resulting world frame (after solve)

=#

##

using Distributed
# addprocs(10)

using Colors
using Caesar
@everywhere using Caesar

using DistributedFactorGraphs: listDataEntrySequence

import Rotations as _Rot
using CoordinateTransformations

# using DistributedFactorGraphs.DocStringExtensions
using Dates
using JSON2
# using BSON
using Serialization
using FixedPointNumbers
using Manifolds
using StaticArrays
using TensorCast
using DataStructures

using ImageMagick, FileIO
using Images
using ImageDraw
using ImageView
using ImageTransformations
using ImageFeatures

using VideoIO

using RoMEPlotting
Gadfly.set_default_plot_size(35cm,20cm)

using Optim

## load the solved graph and reattach the datastores

dfg_datafolder = "/tmp/caesar/philos"

fg = loadDFG("$dfg_datafolder/results_5/x0_5_x280")

##

ds = FolderStore{Vector{UInt8}}(:radar, "$dfg_datafolder/data/radar")
addBlobStore!(fg, ds)

ds = FolderStore{Vector{UInt8}}(:gps_fix, "$dfg_datafolder/data/gps")
addBlobStore!(fg, ds)

ds = FolderStore{Vector{UInt8}}(:lidar, "$dfg_datafolder/data/lidar")
addBlobStore!(fg, ds)

ds = FolderStore{Vector{UInt8}}(:camera, "$dfg_datafolder/data/camera")
addBlobStore!(fg, ds)

## structure

# walk through all the poses in fg
#  pull out all world frame radar images
#  pull out all camera images at each (not all poses have world location solution)
#  draw side by side radar on left and camera on right (more camera than radar images)
#

##

M = SpecialEuclidean(2)
e0 = identity_element(M)

##

# get two radar sweep point clouds and rotate them to the world frame
function _fetchDataPointCloudWorld(fg::AbstractDFG, lb::Symbol)
  de,db = getData(fg, lb, :RADAR_SWEEP)
  # deserialize is not legacy safe, rather use BSON.@save(io, Caesar._PCL.PCLPointCloud2(pc))
  bPC = deserialize(PipeBuffer(db))
  wPb = exp(M, e0, hat(M, e0, getPPE(fg, lb).suggested))
  apply(M, wPb, bPC)
end

Base.clamp(img::AbstractMatrix{<:RGB},b::Real,t::Real) = map(px->RGB{N0f8}(clamp(px.r,b,t),clamp(px.g,b,t),clamp(px.b,b,t)), img)

function buildPoseImages( fg::AbstractDFG,
                          lb::Symbol;
                          # _X_ = 3000,
                          # _Y_ = 1200,
                          pattern::Regex=r"IMG_CENTER_",
                          T = RGB{N0f8},
                          color=Gray{Float64}(0.05),
                          scale_map=1.2,
                          ppose=zeros(3) )
  #
  global img_wPC, imgL

  M = SpecialEuclidean(2)
  e0 = identity_element(M)

  R_ = _Rot.RotMatrix(0.0) # 0.017*pi

  framestack = []

  # fetch the camera images
  ents = listDataEntrySequence(fg, lb, pattern, sortDFG)
  imgs_lb = fetchDataImage.(fg, lb, ents)
  rows,cols = size(imgs_lb[1])

  # expand the radar map
  if isSolved(fg, lb)
    wPC__ = _fetchDataPointCloudWorld(fg, lb);
    # add traj plotting here (draw points and lines on radar plot)

    nRw = _Rot.RotMatrix(140*pi/180.0)
    nTw = ProductRepr([0.0;0.0], nRw)
    nPC = apply(M, nTw, wPC__)
    __ppe = getPPE(fg, lb).suggested
    _ppe = Manifolds.compose(M, nTw, exp(M, e0, hat(M, e0, __ppe)))
    ppe = vee(M, e0, log(M, e0, _ppe))
    img_wPC_ = makeImage!(nPC, (-2500,1000),(-2000,1000); rows=trunc(Int, scale_map*rows), color, pose=ppe, ppose=ppose );
    ppose .= ppe
    if lb == :x0
      img_wPC = deepcopy(img_wPC_)
    else
      img_wPC .+= img_wPC_
    end
    imgG = (px->RGB(0,10*px,0)).(img_wPC_);

    imgL = RGB.(img_wPC) + imgG
  end

  imgs_R = map(img->warp(img, ImageTransformations.recenter(R_, center(img))), imgs_lb)

  frame_ = Matrix{T}(undef, size(imgs_R[1])...)

  # write mosaic frames to the video
  for frame in imgs_R
    ax, ay = axes(frame)
    I = ax.parent .+ ax.offset
    J = ay.parent .+ ay.offset
    for (i,i_) in enumerate(I), (j,j_) in enumerate(J)
      frame_[i,j] = frame[i_,j_]
    end
    imgF = mosaic([imgL, frame_]; nrow=1) #, npad=50)
    push!(framestack, imgF)
  end

  # typeof(framestack[1])
  framestack_ = Vector{Matrix{T}}(undef, length(framestack));
  for (i,fr) in enumerate(framestack)
    framestack_[i] = clamp(fr, 0,1);
  end

  return framestack_
end

##

lbls = sortDFG(ls(fg)) #[1:41]

# imgL = nothing    # Matrix{RGB{Float64}}()
global img_wPC = nothing
global imgL = nothing

ppose = zeros(3)
pattern = r"STAB_IMG_CENTER_"

framestack = buildPoseImages(fg, :x0; ppose, pattern );

##

encoder_options = (crf=23, preset="medium")
framerate=90
open_video_out("/tmp/caesar/philos/mosaic_stab.mp4", framestack[1], framerate=framerate, encoder_options=encoder_options) do writer

  for (i,lb) in enumerate(lbls)
    println("doing all frames for ", lb)

    framestack = buildPoseImages(fg, lb; ppose, pattern )

    for frame in framestack
      # frame_ .= T.(frame)
      write(writer, frame)
    end

  end # for

end # video

##

lb = :x48
des = listDataEntrySequence(fg, lb, r"IMG_CENTER", sortDFG)
imgs = fetchDataImage.(fg, lb, des)

##

#


# build factor graph from sandshark data
addprocs(7)
using Caesar
using RoME
using IncrementalInference, KernelDensityEstimate
using MAT
using JLD, HDF5
using Gadfly

vars = matread("/home/dehann/data/sandshark/cbf.mat");

len = length(vars["imu_gyr_x"]);

accs = zeros(len,3);
gyrs = zeros(len,3);

lbl = ["x";"y";"z"]
for i in 1:3 accs[:,i] = vars["imu_acc_$(lbl[i])"][:]; end
for i in 1:3 gyrs[:,i] = vars["imu_gyr_$(lbl[i])"][:]; end
imutime = vars["imu_time"];

nativeOdo = zeros(length(vars["nav_time"][:]), 5);
nativeOdo[:,1] = vars["nav_time"][:];
nativeOdo[:,2] = 0.68*vars["nav_speed"][:];
nativeOdo[:,3] = vars["nav_roll"][:]*pi/180.0;
nativeOdo[:,4] = vars["nav_pitch"][:]*pi/180.0;
nativeOdo[:,5] = (vars["nav_yaw"][:]-14.72)*pi/180.0; # assumed to be heading

nativeOdo[:,1] -= vars["nav_time"][1]
nativeOdo[:,1] *= 1e-9

for i in 1:size(nativeOdo,1)
  nativeOdo[i,5] = wrapRad( (2pi-nativeOdo[i,5])-3.0*pi/2.0 )
end

nativeOdo[2:end,5] = diff(nativeOdo[:,5])
nativeOdo[1,5] = 0
nativeOdo[nativeOdo[:,5] .> 2pi-0.1, 5] -= 2pi
nativeOdo[nativeOdo[:,5] .< -2pi+0.1, 5] += 2pi

nativeMeas = zeros(length(vars["CBF_MF_time"]),5);
nativeMeas[:,1] = (vars["CBF_MF_time"][:]-(vars["nav_time"][1]*1e-9));
# quality estimate of acoustic measurement
nativeMeas[:,2] = vars["MF_std"]
# elevation' and azimuth -- TODO remove roll, pitch from body relative azimuth
nativeMeas[:,3:4] = vars["CBF_max"][:,1:2]
# argmax range
for i in 1:length(vars["CBF_MF_time"])
  nativeMeas[i,5] =findmax(sum(abs(vars["MF_prob"][:,:,i]),2))[2]
end
# nativeMeas[:,5] = findmax(sum(abs(vars["MF_prob"]),2),1)[2][:]
nativeMeas[:,5] -= 400.0
nativeMeas[:,5] *= 1500.0/37500

@parallel for i in  1:905
  pls = []
  for j in 1:4
    push!(pls, plot(y=(abs(vars["MF_prob"][:,j,i])),Geom.line,Coord.Cartesian(ymax=0.2)));
  end
  push!(pls, plot(y=sum(abs(vars["MF_prob"][:,:,i]),2),Geom.line,Coord.Cartesian(ymax=0.2)));
  pl = vstack(pls...)
  draw(PNG("/home/dehann/Desktop/here/MF_prob_$(i).png",10cm,35cm),pl);
end

# integrate native measurements for dead reckoning
# Dxi_k1 is previous k-1 position in pose i frame, x_{k-1}=xk1
function sandsharkDeadReckoning2D(dt, x, speed, pitch, dyaw)
  Dxi_k1 = SE2(x)
  v = speed*cos(pitch)
  D = SE2([0.0;0.0;dyaw])*SE2(dt*[v;0.0;0.0])
  Dxi_k = Dxi_k1*D
  return vec(se2vee(Dxi_k))
end

function deltaOdom(xprev::Array{Float64,1}, xnow::Array{Float64,1})
  xk1 = SE2(xprev)
  xk = SE2(xnow)
  return se2vee((xk1 \ xk))
end

# odo has form odo[1:n,:] = [t, x, y, theta], with n rows of data in time
function advOdoByRules(odo::Array{Float64,2}, meas::Array{Float64,2};
                        distrule=10.0, timerule=30.0, yawrule=pi/6.0, trkfeats=false)
  # DBG = Dict{Int,Dict{Int64,Array{BallTreeDensity,1}}}()
  # DBG = Dict{Int,Dict{Int64,Feature}}()
  DBG = nothing
  cntr = 1

  # bfts0 = Array{Float64,2}() #lsrFeats[1].feats
  instSys = makeGenericInSituSys(zeros(3))
  instSys.T0 = odo[1,1] # in nanoseconds
  fdict = Dict{Int64,Array{Float64,1}}()
  # for f in instSys.trackers
  #   # if length(f[2].lastz) < 3
  #   #   error("HEY!! $(f[1])")
  #   # end
  #   fdict[f[2].id] = f[2].lastz
  # end
  # instSys.FeatAssc[instSys.poseid] = fdict
  for i in 1:size(odo,1)
    dt = (odo[i,1]-instSys.T0)
    bTbm = SE2(instSys.x)
    # instSys.x is the real time odometry estimate
    instSys.x = sandsharkDeadReckoning2D(dt, instSys.x, odo[i,2], odo[i,4],odo[i,5])
    # bTbp = SE2(instSys.x)

    trigmeas = false
    if cntr <= size(meas,1)
      if odo[i,1] >= meas[cntr,1]
        # TODO -- relax to include at least range on large azimuth
        if meas[cntr,2] < 5.0 &&  (meas[cntr,4] < 95.0 || 265.0 < meas[cntr,4])
          trigmeas = true
        end
        cntr+=1
      end
    end
  #   # subsample motion tracks for constructing the factor graph SLAM solution
    if poseTrigAndAdd!(instSys, (odo[i,1]), distrule, timerule, yawrule, auxtrig=trigmeas)
      fdict = Dict{Int64,Array{Float64,1}}()
      if trigmeas
        bearing = wrapRad(meas[cntr-1,4]*pi/180.0)
        fdict[1] = [meas[cntr-1,5];bearing;0.0]
      end
      instSys.FeatAssc[instSys.poseid] = fdict
    end
  #   # advance previous time point
    instSys.T0 = odo[i,1]
  end
  return instSys.dOdo, instSys.FeatAssc, DBG
end

XX = advOdoByRules(nativeOdo, nativeMeas);
XX[2][1] = Dict{Int64,Array{Float64,1}}()

progressExamplePlot(XX[1],XX[2])

MM = Dict{Int, Int}()

fg = emptyFactorGraph();
idx = appendFactorGraph!(fg, XX[1], XX[2], lcmode=:unimodal, MM=MM);

tree = prepBatchTree!(fg, ordering=:qr,drawpdf=true);
inferOverTree!(fg,tree, N=100);

isamdict, fgu = doISAMSolve(XX[1],XX[2], savejld=false, MM=MM, retfg=true)
isamdict[1] = zeros(3)
# @save "odoSandsharkISAM1.jld" isamdict fgu
@load "odoSandsharkISAM1.jld"

@save "sandsharkSolve01.jld" fg nativeOdo nativeMeas

pl = drawCompPosesLandm(fg, isamdict, fgu, lbls=false)
using Colors, Cairo
draw(PDF("ss01.pdf",20cm,40cm),pl)

function testSandsharkOdo(;iters::Int=10)
  odo = zeros(iters, 5)
  dt = 0.1
  odo[:,1] = collect(1:iters)*dt
  odo[:,2] = 1.0
  odo[4:end,5] = pi/4
  odo[8:end,5] = pi/2.0

  prevpose = zeros(3)
  prevx = zeros(3)
  X = Array{Float64,2}(iters,3)
  for i in 1:iters
    x = sandsharkDeadReckoning2D(dt, prevx, odo[i,2], odo[i,4], odo[i,5])
    X[i,:] = deepcopy(x)
    prevx = x
    # do pose trigger test
      # new pose
  end

  return X
end

# X = testSandsharkOdoReal(nativeOdo,iters=10000);
# plot(x=X[:,1],y=X[:,2], Geom.path())
function testSandsharkOdoReal(odo;iters::Int=100)

  Tprev = odo[1,1]

  prevx = zeros(3)
  X = Array{Float64,2}(iters,3)
  for i in 1:iters
    dt = (odo[i,1] - Tprev)
    x = sandsharkDeadReckoning2D(dt, prevx, odo[i,2], odo[i,4], odo[i,5])
    X[i,:] = deepcopy(x)
    prevx = x
    # do pose trigger test
      # new pose
    Tprev = odo[i,1]
  end

  return X
end


# AbstractFFTs interface

const DUALFORW = true
const DUALBACK = false
const INPLACE = true

# ---------------------------------------------------------------------------- #
# DualPlan type (for AbstractFFTs interface)
# ---------------------------------------------------------------------------- #
mutable struct DualPlan{T,forward,inplace} <: Plan{T}
    region
    msize
    pinv::ScaledPlan{T}
    DualPlan{T,forward,inplace}(region,msize) where {T,forward,inplace} = new{T,forward,inplace}(region,msize)
end

# ---------------------------------------------------------------------------- #
# AbstractFFTs interface functions
# ---------------------------------------------------------------------------- #
AbstractFFTs.plan_fft(  Z::Array{Complex{D}}, region=1:ndims(Z); kwargs...) where D <: Dual = DualPlan{Complex{D},DUALFORW,!INPLACE}(region,size(Z))
AbstractFFTs.plan_fft!( Z::Array{Complex{D}}, region=1:ndims(Z); kwargs...) where D <: Dual = DualPlan{Complex{D},DUALFORW,INPLACE}(region,size(Z))
AbstractFFTs.plan_bfft( Z::Array{Complex{D}}, region=1:ndims(Z); kwargs...) where D <: Dual = DualPlan{Complex{D},DUALBACK,!INPLACE}(region,size(Z))
AbstractFFTs.plan_bfft!(Z::Array{Complex{D}}, region=1:ndims(Z); kwargs...) where D <: Dual = DualPlan{Complex{D},DUALBACK,INPLACE}(region,size(Z))

AbstractFFTs.plan_inv(p::DualPlan{Complex{D},forward,inplace}) where {D<:Dual,forward,inplace} = ScaledPlan(DualPlan{Complex{D},!forward,inplace}(p.region,p.msize),
           normalization(basetype(D), p.msize, p.region))

LinearAlgebra.mul!(Y::Array{Complex{D}}, p::DualPlan{Complex{D},DUALFORW,inplace}, X::Array{Complex{D}}) where {D<:Dual,inplace} = (copy!(Y,X); dualfft!(Y,p.region); return Y)
LinearAlgebra.mul!(Y::Array{Complex{D}}, p::DualPlan{Complex{D},DUALBACK,inplace}, X::Array{Complex{D}}) where {D<:Dual,inplace} = (copy!(Y,X); dualbfft!(Y,p.region); return Y)

Base.:*(p::DualPlan{Complex{D},DUALFORW,!INPLACE}, X::Array{Complex{D}}) where D <: Dual = dualfft(X,p.region)
Base.:*(p::DualPlan{Complex{D},DUALFORW,INPLACE}, X::Array{Complex{D}}) where D <: Dual = (dualfft!(X,p.region); return X)
Base.:*(p::DualPlan{Complex{D},DUALBACK,!INPLACE}, X::Array{Complex{D}}) where D <: Dual = dualbfft(X,p.region)
Base.:*(p::DualPlan{Complex{D},DUALBACK,INPLACE}, X::Array{Complex{D}}) where D <: Dual = (dualbfft!(X,p.region); return X)


# Don't paste into REPL, use `include("prompt_times.jl")`

using REPL
using Printf
@eval REPL begin
     const LAST_TIME = Ref{Float64}(time())
     const PROMPT_CACHE  = Ref{Union{String, Nothing}}(nothing)
end

# Format the time printing here as you wish
function prettytime(t)
    if REPL.PROMPT_CACHE[] !== nothing
        return REPL.PROMPT_CACHE[]
    end
    if t < 1e2
        value, units = t, "ns"
    elseif t < 1e5
        value, units = t / 1e3, "μs"
    elseif t < 1e8
        value, units = t / 1e6, "ms"
    elseif t < 120e9
        value, units = t / 1e9, "s"
    elseif t < 1e9*60*90
	value, units = t / 60e9, "min"
    else
	value, units = t / 3600e9, "hr"
    end
    s = string(round(value, sigdigits=3), " ", units)
    REPL.PROMPT_CACHE[] = s
    return s
end

julia_repl = Base.active_repl.interface.modes[1]
old_f = julia_repl.on_done
julia_repl.on_done = (args...) -> (REPL.LAST_TIME[] = time(); REPL.PROMPT_CACHE[] = nothing; old_f(args...))
julia_repl.prompt = () -> string("julia [", prettytime((time() - REPL.LAST_TIME[])*1e9), "]> ")


#
println("Ed) DR rate coefficients for neon-like Ar: AverageSCA model.")
using JLD
#
if  true
    # Calculation of the Ne^+  2s^2 2p^6; 2s, 2p --> 3s, 3p, 3d, ... 5d electron capture and subsequent autoionization and stabilization
    setDefaults("print summary: open", "Ne-like-DR-rate.sum")
    setDefaults("method: continuum, Galerkin");            setDefaults("method: normalization, pure sine")
    grid = Radial.Grid(Radial.Grid(false), rnt = 1.0e-5, h = 5.0e-2, hp = 2.0e-2, rbox = 15.0)
    name = "DR rate coefficient for neon-like argon and for temperatures < 500 K"
    asfSettings = AsfSettings(AsfSettings(), maxIterationsScf = 36)
    wa   = Cascade.Computation(Cascade.Computation(); name=name, nuclearModel=Nuclear.Model(10.), grid=grid, asfSettings=asfSettings,
                               approach=Cascade.AverageSCA(),
                               scheme=Cascade.DrRateCoefficientScheme([ElecCapture(), Auger(), Radiative()], [E1], 2.5, 0.0, 1, 3, 2,
                                                                      [Shell("2s"), Shell("2p")]),
                               initialConfigs=[Configuration("1s^2 2s^2 2p^6")] )
    println(wa)
    wb = perform(wa; output=true)
    setDefaults("print summary: close", "")
    #
elseif false
    # Calculation of the Ne^+  2s^2 2p^5; photo-ionization cross sections
    setDefaults("print summary: open", "Ne-plus-photoionization.sum")

    name = "Photoionization of Si- "
    grid = Radial.Grid(Radial.Grid(false); rnt = 3.0e-6, h = 2.0e-2, hp = 3.0e-2, NoPoints=1110)
    wa   = Cascade.Computation(Cascade.Computation(); name=name, nuclearModel=Nuclear.Model(10.), grid=grid, approach=Cascade.AverageSCA(),
                               scheme=Cascade.PhotonIonizationScheme([Photo()], 1, [1.0, 3.0, 5.0]),
                               initialConfigs=[Configuration("1s^2 2s^2 2p^5")] )
    println(wa)
    wb = perform(wa; output=true)
    setDefaults("print summary: close", "")

else
    # Simulation of the Ne^+  2s^2 2p^5; photo-ionization cross sections
    setDefaults("print summary: open", "Ne-plus-photoabsorption-simulation.sum")

    data = [JLD.load("zzz-cascade-ionizing-computations-2020-07-10T19.jld"), JLD.load("zzz-cascade-excitation-computations-2020-07-10T19.jld")]
    ##x JLD2.@load "zzz-cascade-ionizing-computations-2020-06-23T09.jld"
    ##x resIon  = results
    ##x JLD2.@load "zzz-cascade-excitation-computations-2020-06-23T09.jld"
    ##x resExc = results
    ##x data = [resExc]

    name = "Simulation of photoabsorption for Ne^+"

    wc   = Cascade.Simulation(Cascade.Simulation(), name=name, property=Cascade.PhotoAbsorption(),
                              settings=Cascade.SimulationSettings(0., 0., 1., 4., 3., [(1, 1.0)]), computationData=data )
    println(wc)
    wd = perform(wc; output=true)
    setDefaults("print summary: close", "")

end


"""
    PhotometricFilm

The photometric film is a one-pixel detector, simply adding together all the intensity
coming from the scene to the camera.

TODO: This is really not in working condition, do not use it.

"""
mutable struct PhotometricFilm{S<:Spectrum} <: Film
    integral::S
    N::Int64
    area::Float64
    resX::Int64
    resY::Int64
end

"""
    add_sample!(F::PhotometricFilm, sample::Sample, L::Spectrum, isect::Union{Intersection,Nothing})

Add a ray sample to the `ImageFilm`. The spectral intensity `L` will be added to the
detector's single pixel. The ray intersection `isect` nor the camera sample `sample` are
not used, but they are part of the `Film` API.

TODO: This relies heavily on the addition of Spectrum types, which is a bit poorly defined
currently.

"""
function add_sample!(F::PhotometricFilm, sample::Sample, L::Spectrum, isect::Union{Intersection,Nothing})
    F.N += 1
    isblack(L) && return nothing
    F.integral += L
    return nothing
end

"""
    write_txt(F::PhotometricFilm, fname::String)

Write the value of the single pixel of the `PhotometricFilm` into a file named `fname`.

"""
function write_txt(F::PhotometricFilm, fname::String)
    S = F.area / F.N * F.integral
    f = open(fname, "w")
    write(f, string(S))
    return nothing
end