text:
using Statistics
using Printf
using Logging
using JLD2
using NCDatasets
using GeoData
using Oceanostics
using Oceananigans
using Oceananigans.Units

using Oceanostics.TKEBudgetTerms: TurbulentKineticEnergy, ZShearProduction

using LESbrary.Utils: SimulationProgressMessenger, fit_cubic, poly
using LESbrary.NearSurfaceTurbulenceModels: SurfaceEnhancedModelConstant
using LESbrary.TurbulenceStatistics: first_order_statistics,
                                     first_through_second_order,
                                     turbulent_kinetic_energy_budget,
                                     subfilter_momentum_fluxes,
                                     subfilter_tracer_fluxes,
                                     ViscousDissipation

try
    using CairoMakie
catch
    using GLMakie
finally
    @warn "Could not load either CairoMakie or GLMakie; animations are not available."
end

Logging.global_logger(OceananigansLogger())

@inline passive_tracer_forcing(x, y, z, t, p) = p.μ⁺ * exp(-(z - p.z₀)^2 / (2 * p.λ^2)) - p.μ⁻

# Code credit: https://discourse.julialang.org/t/collecting-all-output-from-shell-commands/15592
function execute(cmd::Cmd)
    out, err = Pipe(), Pipe()
    process = run(pipeline(ignorestatus(cmd), stdout=out, stderr=err))
    close(out.in)
    close(err.in)
    return (stdout = out |> read |> String, stderr = err |> read |> String, code = process.exitcode)
end

function three_layer_constant_fluxes_simulation(;
    name = "",
    size = (32, 32, 32),
    passive_tracers = true,
    extent = (512meters, 512meters, 256meters),
    architecture = CPU(),
    stop_time = 0.1hours,
    f = 1e-4,
    buoyancy_flux = 1e-8,
    momentum_flux = -1e-4,
    thermocline_type = "linear",
    surface_layer_depth = 48.0,
    thermocline_width = 24.0,
    surface_layer_buoyancy_gradient = 2e-6,
    thermocline_buoyancy_gradient = 1e-5,
    deep_buoyancy_gradient = 2e-6,
    surface_temperature = 20.0,
    pickup = false,
    jld2_output = true,
    netcdf_output = false,
    snapshot_time_interval = 2minutes,
    averages_time_interval = 3hours,
    averages_time_window = 15minutes,
    time_averaged_statistics = false,
    data_directory = joinpath(pwd(), "data"))

    Nx, Ny, Nz = size
    Lx, Ly, Lz = extent
    slice_depth = 8.0
    Qᵇ = buoyancy_flux
    Qᵘ = momentum_flux
    stop_hours = stop_time / hour

    ## Determine filepath prefix
    prefix = @sprintf("three_layer_constant_fluxes_%s_hr%d_Qu%.1e_Qb%.1e_f%.1e_Nh%d_Nz%d_",
                      thermocline_type, stop_hours, abs(Qᵘ), Qᵇ, f, Nx, Nz)

    data_directory = joinpath(data_directory, "data", prefix * name) # save data in /data/prefix

    @info "Mapping grid..."

    grid = RectilinearGrid(architecture,
                           size=(Nx, Ny, Nz),
                           x=(0, Lx), y=(0, Ly), z=(-Lz, 0),
                           halo=(3, 3, 3))

    # Buoyancy and boundary conditions

    @info "Enforcing boundary conditions..."

    buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(α=2e-4), constant_salinity=35.0)

    N²_surface_layer = surface_layer_buoyancy_gradient
    N²_thermocline   = thermocline_buoyancy_gradient
    N²_deep          = deep_buoyancy_gradient
    α = buoyancy.equation_of_state.α
    g = buoyancy.gravitational_acceleration

    Qᶿ = Qᵇ / (α * g)
    dθdz_surface_layer = N²_surface_layer / (α * g)
    dθdz_thermocline   = N²_thermocline   / (α * g)
    dθdz_deep          = N²_deep          / (α * g)

    θ_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᶿ),
                                    bottom = GradientBoundaryCondition(dθdz_deep))

    u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))

    # Tracer forcing

    @info "Forcing and sponging tracers..."

    # # Initial condition and sponge layer

    ## Fiddle with indices to get a correct discrete profile
    z = znodes(Center, grid)
    k_transition = searchsortedfirst(z, -surface_layer_depth)
    k_deep = searchsortedfirst(z, -(surface_layer_depth + thermocline_width))

    z_transition = z[k_transition]
    z_deep = z[k_deep]

    θ_surface = surface_temperature
    θ_transition = θ_surface + z_transition * dθdz_surface_layer
    θ_deep = θ_transition + (z_deep - z_transition) * dθdz_thermocline

    λ = 4.0
    μ⁺ = 1 / 6hour
    μ₀ = √(2π) * λ / grid.Lz * μ⁺ / 2
    μ∞ = √(2π) * λ / grid.Lz * μ⁺

    c₀_forcing = Forcing(passive_tracer_forcing, parameters=(z₀=  0.0, λ=λ, μ⁺=μ⁺, μ⁻=μ₀))
    c₁_forcing = Forcing(passive_tracer_forcing, parameters=(z₀=-48.0, λ=λ, μ⁺=μ⁺, μ⁻=μ∞))
    c₂_forcing = Forcing(passive_tracer_forcing, parameters=(z₀=-96.0, λ=λ, μ⁺=μ⁺, μ⁻=μ∞))

    # Sponge layer for u, v, w, and T
    gaussian_mask = GaussianMask{:z}(center=-grid.Lz, width=grid.Lz/10)

    u_sponge = v_sponge = w_sponge = Relaxation(rate=4/hour, mask=gaussian_mask)

    T_sponge = Relaxation(rate = 4/hour,
                          target = LinearTarget{:z}(intercept = θ_deep - z_deep*dθdz_deep, gradient = dθdz_deep),
                          mask = gaussian_mask)

    # # LES Model

    # Wall-aware AMD model constant which is 'enhanced' near the upper boundary.
    # Necessary to obtain a smooth temperature distribution.

    @info "Building the wall model..."

    Δz = grid.Δzᵃᵃᶜ
    Cᴬᴹᴰ = SurfaceEnhancedModelConstant(Δz, C₀ = 1/12, enhancement = 7, decay_scale = 4Δz)

    # # Instantiate Oceananigans.IncompressibleModel

    @info "Framing the model..."

    tracers = passive_tracers ? (:T, :c₀, :c₁, :c₂) : :T

    model = NonhydrostaticModel(; grid, buoyancy, tracers,
                                timestepper = :RungeKutta3,
                                advection = WENO5(),
                                coriolis = FPlane(f=f),
                                closure = AnisotropicMinimumDissipation(),
                                boundary_conditions = (T=θ_bcs, u=u_bcs),
                                forcing = (u=u_sponge, v=v_sponge, w=w_sponge, T=T_sponge,
                                           c₀=c₀_forcing, c₁=c₁_forcing, c₂=c₂_forcing))

    # # Set Initial condition

    @info "Setting initial conditions..."

    ## Noise with 8 m decay scale
    Ξ(z) = rand() * exp(z / 8)

    function thermocline_structure_function(thermocline_type, z_transition, θ_transition, z_deep,
                                            θ_deep, dθdz_surface_layer, dθdz_thermocline, dθdz_deep)

        if thermocline_type == "linear"
            return z -> θ_transition + dθdz_thermocline * (z - z_transition)

        elseif thermocline_type == "cubic"
            p1 = (z_transition, θ_transition)
            p2 = (z_deep, θ_deep)
            coeffs = fit_cubic(p1, p2, dθdz_surface_layer, dθdz_deep)
            return z -> poly(z, coeffs)

        else
            @error "Invalid thermocline type: $thermocline"
        end
    end

    θ_thermocline = thermocline_structure_function(thermocline_type, z_transition, θ_transition, z_deep, θ_deep,
                                                   dθdz_surface_layer, dθdz_thermocline, dθdz_deep)

    """
        initial_temperature(x, y, z)

    Returns a three-layer initial temperature distribution. The average temperature varies in z
    and is augmented by three-dimensional, surface-concentrated random noise.
    """
    function initial_temperature(x, y, z)
        noise = 1e-6 * Ξ(z) * dθdz_surface_layer * grid.Lz

        if z_transition < z <= 0
            return θ_surface + dθdz_surface_layer * z + noise
        elseif z_deep < z <= z_transition
            return θ_thermocline(z) + noise
        else
            return θ_deep + dθdz_deep * (z - z_deep) + noise
        end
    end

    set!(model, T = initial_temperature)

    # # Prepare the simulation

    @info "Conjuring the simulation..."

    simulation = Simulation(model; Δt=1.0, stop_time)

    # Adaptive time-stepping
    wizard = TimeStepWizard(cfl=0.8, max_change=1.1, min_Δt=0.01, max_Δt=30.0)
    simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))
    simulation.callbacks[:progress] = Callback(SimulationProgressMessenger(wizard), IterationInterval(10))

    # # Prepare Output

    @info "Strapping on checkpointer..."

    pickup = pickup
    force = pickup ? false : true

    simulation.output_writers[:checkpointer] =
        Checkpointer(model, schedule = TimeInterval(stop_time/3), prefix = prefix * "_checkpointer", dir = data_directory)

    @info "Squeezing out statistics..."

    # Prepare turbulence statistics
    k_xy_slice = searchsortedfirst(znodes(Face, grid), -slice_depth)

    b = BuoyancyField(model)
    p = Field(sum(model.pressures))

    ccc_scratch = Field{Center, Center, Center}(model.grid)
    ccf_scratch = Field{Center, Center, Face}(model.grid)
    fcf_scratch = Field{Face, Center, Face}(model.grid)
    cff_scratch = Field{Center, Face, Face}(model.grid)

    #primitive_statistics = first_through_second_order(model, b=b, p=p, w_scratch=ccf_scratch, c_scratch=ccc_scratch)
    primitive_statistics = first_order_statistics(model, b=b, p=p, w_scratch=ccf_scratch, c_scratch=ccc_scratch)

    subfilter_flux_statistics = merge(
        subfilter_momentum_fluxes(model, uz_scratch=ccf_scratch, vz_scratch=cff_scratch, c_scratch=ccc_scratch),
        subfilter_tracer_fluxes(model, w_scratch=ccf_scratch),
    )

    U = primitive_statistics[:u]
    V = primitive_statistics[:v]
    B = primitive_statistics[:b]

    e = Field(TurbulentKineticEnergy(model, U=U, V=V))
    shear_production = Field(ZShearProduction(model, U=U, V=V))
    dissipation = Field(ViscousDissipation(model))

    tke_budget_statistics = turbulent_kinetic_energy_budget(model,
                                                            b=b, p=p, U=U, V=V, e=e,
                                                            shear_production=shear_production,
                                                            dissipation=dissipation)

    dynamics_statistics = Dict(:Ri => Field(∂z(B) / (∂z(U)^2 + ∂z(V)^2)))
    fields_to_output = merge(model.velocities, model.tracers, (e=e, ϵ=dissipation))

    statistics_to_output = merge(primitive_statistics,
                                 subfilter_flux_statistics,
                                 tke_budget_statistics,
                                 dynamics_statistics)

    #statistics_to_output = merge(subfilter_flux_statistics, tke_budget_statistics)

    @info "Garnishing output writers..."

    global_attributes = (
        LESbrary_jl_commit_SHA1 = execute(`git rev-parse HEAD`).stdout |> strip,
        name = name,
        thermocline_type = thermocline_type,
        buoyancy_flux = Qᵇ,
        momentum_flux = Qᵘ,
        temperature_flux = Qᶿ,
        coriolis_parameter = f,
        thermal_expansion_coefficient = α,
        gravitational_acceleration = g,
        boundary_condition_θ_top = Qᶿ,
        boundary_condition_θ_bottom = dθdz_deep,
        boundary_condition_u_top = Qᵘ,
        boundary_condition_u_bottom = 0.0,
        surface_layer_depth = surface_layer_depth,
        thermocline_width = thermocline_width,
        N²_surface_layer = N²_surface_layer,
        N²_thermocline = N²_thermocline,
        N²_deep = N²_deep,
        dθdz_surface_layer = dθdz_surface_layer,
        dθdz_thermocline = dθdz_thermocline,
        dθdz_deep = dθdz_deep,
        θ_surface = θ_surface,
        θ_transition = θ_transition,
        θ_deep = θ_deep,
        z_transition = z_transition,
        z_deep = z_deep,
        k_transition = k_transition,
        k_deep = k_deep
    )

    function init_save_some_metadata!(file, model)
        for (name, value) in pairs(global_attributes)
            file["parameters/$(string(name))"] = value
        end
        return nothing
    end

    ## Add JLD2 output writers

    if jld2_output
        simulation.output_writers[Symbol(name, "_xy_jld2")] =
            JLD2OutputWriter(model, fields_to_output,
                             dir = data_directory,
                             prefix = name * "_xy_slice",
                             schedule = TimeInterval(snapshot_time_interval),
                             field_slicer = FieldSlicer(k=k_xy_slice, with_halos=true),
                             force = force,
                             init = init_save_some_metadata!)

        simulation.output_writers[Symbol(name, "_xz_jld2")] =
            JLD2OutputWriter(model, fields_to_output,
                             dir = data_directory,
                             prefix = name * "_xz_slice",
                             schedule = TimeInterval(snapshot_time_interval),
                             field_slicer = FieldSlicer(j=1, with_halos=true),
                             force = force,
                             init = init_save_some_metadata!)

        simulation.output_writers[Symbol(name, "_yz_jld2")] =
            JLD2OutputWriter(model, fields_to_output,
                             dir = data_directory,
                             prefix = name * "_yz_slice",
                             schedule = TimeInterval(snapshot_time_interval),
                             field_slicer = FieldSlicer(i=1, with_halos=true),
                             force = force,
                             init = init_save_some_metadata!)

        simulation.output_writers[Symbol(name, "_stats_jld2")] =
            JLD2OutputWriter(model, statistics_to_output,
                             dir = data_directory,
                             prefix = name * "_instantaneous_statistics",
                             schedule = TimeInterval(snapshot_time_interval),
                             field_slicer = nothing,
                             force = force,
                             init = init_save_some_metadata!)

        if time_averaged_statistics
            simulation.output_writers[Symbol(name, "_averaged_stats_jld2")] =
                JLD2OutputWriter(model, statistics_to_output,
                                 dir = data_directory,
                                 prefix = name * "_time_averaged_statistics",
                                 schedule = AveragedTimeInterval(averages_time_interval,
                                                                 window = averages_time_window),
                                 field_slicer = nothing,
                                 force = force,
                                 init = init_save_some_metadata!)
        end
    end

    if netcdf_output # Add NetCDF output writers
        statistics_to_output = Dict(string(k) => v for (k, v) in statistics_to_output)

        simulation.output_writers[:xy_nc] =
            NetCDFOutputWriter(model, fields_to_output,
                             mode = "c",
                         filepath = joinpath(data_directory, "xy_slice.nc"),
                         schedule = TimeInterval(snapshot_time_interval),
                     field_slicer = FieldSlicer(k=k_xy_slice),
                global_attributes = global_attributes)

        simulation.output_writers[:xz_nc] =
            NetCDFOutputWriter(model, fields_to_output,
                             mode = "c",
                         filepath = joinpath(data_directory, "xz_slice.nc"),
                         schedule = TimeInterval(snapshot_time_interval),
                     field_slicer = FieldSlicer(j=1),
                global_attributes = global_attributes)

        simulation.output_writers[:yz_nc] =
            NetCDFOutputWriter(model, fields_to_output,
                             mode = "c",
                         filepath = joinpath(data_directory, "yz_slice.nc"),
                         schedule = TimeInterval(snapshot_time_interval),
                     field_slicer = FieldSlicer(i=1),
                global_attributes = global_attributes)

        simulation.output_writers[:stats_nc] =
            NetCDFOutputWriter(model, statistics_to_output,
                             mode = "c",
                         filepath = joinpath(data_directory, "instantaneous_statistics.nc"),
                         schedule = TimeInterval(snapshot_time_interval),
                global_attributes = global_attributes)

        if time_averaged_statistics
            simulation.output_writers[:averaged_stats_nc] =
                NetCDFOutputWriter(model, statistics_to_output,
                                 mode = "c",
                             filepath = joinpath(data_directory, "time_averaged_statistics.nc"),
                             schedule = AveragedTimeInterval(averages_time_interval, window = averages_time_window),
                    global_attributes = global_attributes)
        end
    end

    return simulation
end

function squeeze(A)
    singleton_dims = tuple((d for d in 1:ndims(A) if Base.size(A, d) == 1)...)
    return dropdims(A, dims=singleton_dims)
end

function three_layer_constant_fluxes_animate_xy_xz(;
                                                   name = "",
                                                   thermocline_type = "linear",
                                                   stop_time = 0.1hours,
                                                   momentum_flux = -1e-4,
                                                   buoyancy_flux = 1e-8,
                                                   f = 1e-4,
                                                   size = (32, 32, 32),
                                                   data_directory = nothing)

    if !isnothing(data_directory)
        prefix = @sprintf("three_layer_constant_fluxes_%s_hr%d_Qu%.1e_Qb%.1e_f%.1e_Nh%d_Nz%d_",
                          thermocline_type, stop_hours, abs(Qᵘ), Qᵇ, f, Nx, Nz)

        data_directory = joinpath(@__DIR__, "..", "data", prefix * name) # save data in /data/prefix
    end

    # # Load and plot turbulence statistics
    ds_xy = NCDstack(joinpath(data_directory, "xy_slice.nc"))
    ds_xz = NCDstack(joinpath(data_directory, "xz_slice.nc"))

    times = ds_xy[:time]
    Nt = length(times)

    xc = ds_xy[:xC].data
    xf = ds_xy[:xF].data
    yc = ds_xy[:yC].data
    yf = ds_xy[:yF].data
    zc = ds_xz[:zC].data
    zf = ds_xz[:zF].data

    fig = Figure(resolution=(3000, 750))

    u_max = max(maximum(abs, ds_xy[:u]), maximum(abs, ds_xz[:u]))
    v_max = max(maximum(abs, ds_xy[:v]), maximum(abs, ds_xz[:v]))
    w_max = max(maximum(abs, ds_xy[:w]), maximum(abs, ds_xz[:w]))
    U_max = max(u_max, v_max, w_max)
    U_lims = 0.5 .* (-U_max, +U_max)

    frame = Node(1)

    plot_title = @lift "LESbrary.jl three-layer constant fluxes $name: time = $(prettytime(times[$frame]))"

    u_xy = @lift ds_xy[:u][Ti=$frame].data |> squeeze
    v_xy = @lift ds_xy[:v][Ti=$frame].data |> squeeze
    w_xy = @lift ds_xy[:w][Ti=$frame].data |> squeeze
    T_xy = @lift ds_xy[:T][Ti=$frame].data |> squeeze
    c₀_xy = @lift ds_xy[:c₀][Ti=$frame].data |> squeeze
    c₁_xy = @lift ds_xy[:c₁][Ti=$frame].data |> squeeze
    c₂_xy = @lift ds_xy[:c₂][Ti=$frame].data |> squeeze
    e_xy = @lift ds_xy[:e][Ti=$frame].data |> squeeze .|> log10
    ϵ_xy = @lift ds_xy[:ϵ][Ti=$frame].data |> squeeze .|> log10

    u_xz = @lift ds_xz[:u][Ti=$frame].data |> squeeze
    v_xz = @lift ds_xz[:v][Ti=$frame].data |> squeeze
    w_xz = @lift ds_xz[:w][Ti=$frame].data |> squeeze
    T_xz = @lift ds_xz[:T][Ti=$frame].data |> squeeze
    c₀_xz = @lift ds_xz[:c₀][Ti=$frame].data |> squeeze
    c₁_xz = @lift ds_xz[:c₁][Ti=$frame].data |> squeeze
    c₂_xz = @lift ds_xz[:c₂][Ti=$frame].data |> squeeze
    e_xz = @lift ds_xz[:e][Ti=$frame].data |> squeeze .|> log10
    ϵ_xz = @lift ds_xz[:ϵ][Ti=$frame].data |> squeeze .|> log10

    ax_u_xy = fig[1, 1] = Axis(fig, title="u-velocity")
    hm_u_xy = heatmap!(ax_u_xy, xf, yc, u_xy, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_u_xy)

    ax_v_xy = fig[1, 2] = Axis(fig, title="v-velocity")
    hm_v_xy = heatmap!(ax_v_xy, xc, yf, v_xy, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_v_xy)

    ax_w_xy = fig[1, 3] = Axis(fig, title="w-velocity")
    hm_w_xy = heatmap!(ax_w_xy, xc, yc, w_xy, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_w_xy)

    ax_T_xy = fig[1, 4] = Axis(fig, title="temperature")
    hm_T_xy = heatmap!(ax_T_xy, xc, yc, T_xy, colormap=:thermal, colorrange=extrema(ds_xy[:T]))
    hidedecorations!(ax_T_xy)

    ax_c₀_xy = fig[1, 5] = Axis(fig, title="tracer 0")
    hm_c₀_xy = heatmap!(ax_c₀_xy, xc, yc, c₀_xy, colormap=:ice, colorrange=extrema(ds_xy[:c₀]))
    hidedecorations!(ax_c₀_xy)

    ax_c₁_xy = fig[1, 6] = Axis(fig, title="tracer 1")
    hm_c₁_xy = heatmap!(ax_c₁_xy, xc, yc, c₁_xy, colormap=:ice, colorrange=extrema(ds_xy[:c₁]))
    hidedecorations!(ax_c₁_xy)

    ax_c₂_xy = fig[1, 7] = Axis(fig, title="tracer 2")
    hm_c₂_xy = heatmap!(ax_c₂_xy, xc, yc, c₂_xy, colormap=:ice, colorrange=extrema(ds_xy[:c₂]))

    hidedecorations!(ax_c₂_xy)

    ax_e_xy = fig[1, 8] = Axis(fig, title="log TKE")
    hm_e_xy = heatmap!(ax_e_xy, xc, yc, e_xy, colormap=:deep, colorrange=(-5, -2))
    hidedecorations!(ax_e_xy)

    ax_ϵ_xy = fig[1, 9] = Axis(fig, title="log TKE dissipation")
    hm_ϵ_xy = heatmap!(ax_ϵ_xy, xc, yc, ϵ_xy, colormap=:dense, colorrange=(-10, -5))
    hidedecorations!(ax_ϵ_xy)

    ax_u_xz = fig[2, 1] = Axis(fig)
    hm_u_xz = heatmap!(ax_u_xz, xf, zc, u_xz, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_u_xz)

    ax_v_xz = fig[2, 2] = Axis(fig)
    hm_v_xz = heatmap!(ax_v_xz, xc, zc, v_xz, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_v_xz)

    ax_w_xz = fig[2, 3] = Axis(fig)
    hm_w_xz = heatmap!(ax_w_xz, xc, zf, w_xz, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_w_xz)

    ax_T_xz = fig[2, 4] = Axis(fig)
    hm_T_xz = heatmap!(ax_T_xz, xc, zc, T_xz, colormap=:thermal, colorrange=extrema(ds_xz[:T]))
    hidedecorations!(ax_T_xz)

    ax_c₀_xz = fig[2, 5] = Axis(fig)
    hm_c₀_xz = heatmap!(ax_c₀_xz, xc, zc, c₀_xz, colormap=:ice, colorrange=extrema(ds_xz[:c₀]))
    hidedecorations!(ax_c₀_xz)

    ax_c₁_xz = fig[2, 6] = Axis(fig)
    hm_c₁_xz = heatmap!(ax_c₁_xz, xc, zc, c₁_xz, colormap=:ice, colorrange=extrema(ds_xz[:c₁]))
    hidedecorations!(ax_c₁_xz)

    ax_c₂_xz = fig[2, 7] = Axis(fig)
    hm_c₂_xz = heatmap!(ax_c₂_xz, xc, zc, c₂_xz, colormap=:ice, colorrange=extrema(ds_xz[:c₂]))
    hidedecorations!(ax_c₂_xz)

    ax_e_xz = fig[2, 8] = Axis(fig)
    hm_e_xz = heatmap!(ax_e_xz, xc, zc, e_xz, colormap=:deep, colorrange=(-8, -2))
    hidedecorations!(ax_e_xz)

    ax_ϵ_xz = fig[2, 9] = Axis(fig)
    hm_ϵ_xz = heatmap!(ax_ϵ_xz, xc, zc, ϵ_xz, colormap=:dense, colorrange=(-15, -5))
    hidedecorations!(ax_ϵ_xz)

    supertitle = fig[0, :] = Label(fig, plot_title, textsize=30)

    filepath = joinpath(data_directory, "xy_xz_movie.mp4")
    record(fig, filepath, 1:Nt, framerate=15) do n
        @info "Animating xy/xz movie frame $n/$Nt..."
        frame[] = n
    end

    @info "Movie saved: $filepath"

    return filepath
end

function three_layer_constant_fluxes_animate_statistics(;
                                                        name = "",
                                                        thermocline_type = "linear",
                                                        stop_hours = 0.1,
                                                        momentum_flux = -1e-4,
                                                        buoyancy_flux = 1e-8,
                                                        f = 1e-4,
                                                        size = (32, 32, 32),
                                                        data_directory = nothing)

    if !isnothing(data_directory)
        prefix = @sprintf("three_layer_constant_fluxes_%s_hr%d_Qu%.1e_Qb%.1e_f%.1e_Nh%d_Nz%d_",
                          thermocline_type, stop_hours, abs(Qᵘ), Qᵇ, f, Nx, Nz)

        data_directory = joinpath(@__DIR__, "..", "data", prefix * name) # save data in /data/prefix
    end

    ds = NCDstack(joinpath(data_directory, "instantaneous_statistics.nc"))

    times = ds[:time]
    Nt = length(times)

    zc = ds[:zC].data
    zf = ds[:zF].data

    fig = Figure(resolution=(3000, 2000))

    frame = Node(1)

    plot_title = @lift "LESbrary.jl three-layer constant fluxes $name: time = $(prettytime(times[$frame]))"

    u = @lift ds[:u][Ti=$frame]
    v = @lift ds[:v][Ti=$frame]
    T = @lift ds[:T][Ti=$frame]
    c₀ = @lift ds[:c₀][Ti=$frame]
    c₁ = @lift ds[:c₁][Ti=$frame]
    c₂ = @lift ds[:c₂][Ti=$frame]
    uu = @lift ds[:uu][Ti=$frame]
    vv = @lift ds[:vv][Ti=$frame]
    ww = @lift ds[:ww][Ti=$frame]
    uv = @lift ds[:uv][Ti=$frame]
    wu = @lift ds[:wu][Ti=$frame]
    wv = @lift ds[:wv][Ti=$frame]
    uT = @lift ds[:uT][Ti=$frame]
    vT = @lift ds[:vT][Ti=$frame]
    wT = @lift ds[:wT][Ti=$frame]
    Ri = @lift ds[:Ri][Ti=$frame]
    νₑ_∂z_u = @lift ds[:νₑ_∂z_u][Ti=$frame]
    νₑ_∂z_v = @lift ds[:νₑ_∂z_v][Ti=$frame]
    νₑ_∂z_w = @lift ds[:νₑ_∂z_w][Ti=$frame]
    κₑ_∂z_T = @lift ds[:κₑ_∂z_T][Ti=$frame]
    κₑ_∂z_c₀ = @lift ds[:κₑ_∂z_c₀][Ti=$frame]
    κₑ_∂z_c₁ = @lift ds[:κₑ_∂z_c₁][Ti=$frame]
    κₑ_∂z_c₂ = @lift ds[:κₑ_∂z_c₂][Ti=$frame]
    e = @lift ds[:e][Ti=$frame]
    tke_dissipation = @lift ds[:tke_dissipation][Ti=$frame]
    tke_advective_flux = @lift ds[:tke_advective_flux][Ti=$frame]
    tke_buoyancy_flux = @lift ds[:tke_buoyancy_flux][Ti=$frame]
    tke_pressure_flux = @lift ds[:tke_pressure_flux][Ti=$frame]
    tke_shear_production = @lift ds[:tke_shear_production][Ti=$frame]

    colors = ["dodgerblue2", "crimson", "forestgreen"]

    ax_uv = fig[1, 1] = Axis(fig, xlabel="m/s", ylabel="z (m)")
    line_u = lines!(ax_uv, u, zc, label="U", linewidth=3, color=colors[1])
    line_v = lines!(ax_uv, v, zc, label="V", linewidth=3, color=colors[2])
    axislegend(ax_uv, position=:rb, framevisible=false)
    xlims!(ax_uv, extrema([extrema(ds[:u])..., extrema(ds[:v])...]))
    ylims!(ax_uv, extrema(zf))

    ax_T = fig[1, 2] = Axis(fig, xlabel="°C")
    line_T = lines!(ax_T, T, zc, linewidth=3, color=colors[1])
    xlims!(ax_T, extrema(ds[:T]))
    ylims!(ax_T, extrema(zf))
    hideydecorations!(ax_T, grid=false)

    ax_c = fig[1, 3] = Axis(fig, xlabel="c")
    line_c₀ = lines!(ax_c, c₀, zc, linewidth=3, label="c₀", color=colors[1])
    line_c₁ = lines!(ax_c, c₁, zc, linewidth=3, label="c₁", color=colors[2])
    line_c₂ = lines!(ax_c, c₂, zc, linewidth=3, label="c₂", color=colors[3])
    axislegend(ax_c, position=:rb, framevisible=false)
    xlims!(ax_c, extrema([extrema(ds[:c₀])..., extrema(ds[:c₁])..., extrema(ds[:c₂])...]))
    ylims!(ax_c, extrema(zf))
    hideydecorations!(ax_c, grid=false)

    ax_ke = fig[1, 4] = Axis(fig, xlabel="m²/s²")
    line_uu = lines!(ax_ke, uu, zc, linewidth=3, label="uu", color=colors[1])
    line_vv = lines!(ax_ke, vv, zc, linewidth=3, label="vv", color=colors[2])
    line_ww = lines!(ax_ke, ww, zf, linewidth=3, label="ww", color=colors[3])
    axislegend(ax_ke, position=:rb, framevisible=false)
    xlims!(ax_ke, extrema([extrema(ds[:c₀])..., extrema(ds[:c₁])..., extrema(ds[:c₂])...]))
    ylims!(ax_ke, extrema(zf))
    hideydecorations!(ax_ke, grid=false)

    ax_UΣ = fig[1, 5] = Axis(fig, xlabel="m²/s²")
    line_uv = lines!(ax_UΣ, uv, zc, linewidth=3, label="uv", color=colors[1])
    line_wu = lines!(ax_UΣ, wu, zf, linewidth=3, label="wu", color=colors[2])
    line_wv = lines!(ax_UΣ, wv, zf, linewidth=3, label="wv", color=colors[3])
    axislegend(ax_UΣ, position=:rb, framevisible=false)
    xlims!(ax_UΣ, extrema([extrema(ds[:uv])..., extrema(ds[:wu])..., extrema(ds[:wv])...]))
    ylims!(ax_UΣ, extrema(zf))
    hideydecorations!(ax_UΣ, grid=false)

    ax_UT = fig[1, 6] = Axis(fig, xlabel="m·K/s")
    line_uT = lines!(ax_UT, uT, zc, linewidth=3, label="uT", color=colors[1])
    line_vT = lines!(ax_UT, vT, zc, linewidth=3, label="vT", color=colors[2])
    line_wT = lines!(ax_UT, wT, zf, linewidth=3, label="wT", color=colors[3])
    axislegend(ax_UT, position=:rb, framevisible=false)
    xlims!(ax_UT, extrema([extrema(ds[:uT])..., extrema(ds[:vT])..., extrema(ds[:wT])...]))
    ylims!(ax_UT, extrema(zf))
    hideydecorations!(ax_UT, grid=false)

    ax_νₑ = fig[1, 7] = Axis(fig, xlabel="m²/s²")
    line_νₑ_∂z_u = lines!(ax_νₑ, νₑ_∂z_u, zf, linewidth=3, label="νₑ ∂z(u)", color=colors[1])
    line_νₑ_∂z_v = lines!(ax_νₑ, νₑ_∂z_v, zf, linewidth=3, label="νₑ ∂z(v)", color=colors[2])
    line_νₑ_∂z_w = lines!(ax_νₑ, νₑ_∂z_w, zc, linewidth=3, label="νₑ ∂z(w)", color=colors[3])
    axislegend(ax_νₑ, position=:rb, framevisible=false)
    xlims!(ax_νₑ, extrema([extrema(ds[:νₑ_∂z_u])..., extrema(ds[:νₑ_∂z_v])..., extrema(ds[:νₑ_∂z_w])...]))
    ylims!(ax_νₑ, extrema(zf))
    hideydecorations!(ax_νₑ, grid=false)

    ax_κₑ∂zT = fig[1, 8] = Axis(fig, xlabel="m⋅K/s")
    line_κₑ∂zT = lines!(ax_κₑ∂zT, κₑ_∂z_T, zf, label="κₑ ∂z(T)", linewidth=3, color=colors[1])
    axislegend(ax_κₑ∂zT, position=:rb, framevisible=false)
    xlims!(ax_κₑ∂zT, extrema(ds[:κₑ_∂z_T]))
    ylims!(ax_κₑ∂zT, extrema(zf))
    hideydecorations!(ax_κₑ∂zT, grid=false)

    ax_κₑ∂zC = fig[2, 1] = Axis(fig, xlabel="m²/s²")
    line_κₑ_∂z_c₀ = lines!(ax_κₑ∂zC, κₑ_∂z_c₀, zf, linewidth=3, label="κₑ ∂z(c₀)", color=colors[1])
    line_κₑ_∂z_c₁ = lines!(ax_κₑ∂zC, κₑ_∂z_c₂, zf, linewidth=3, label="κₑ ∂z(c₁)", color=colors[2])
    line_κₑ_∂z_c₂ = lines!(ax_κₑ∂zC, κₑ_∂z_c₂, zf, linewidth=3, label="κₑ ∂z(c₂)", color=colors[3])
    axislegend(ax_κₑ∂zC, position=:rb, framevisible=false)
    xlims!(ax_κₑ∂zC, extrema([extrema(ds[:κₑ_∂z_c₀])..., extrema(ds[:κₑ_∂z_c₂])..., extrema(ds[:κₑ_∂z_c₂])...]))
    ylims!(ax_κₑ∂zC, extrema(zf))

    ax_Ri = fig[2, 2] = Axis(fig, xlabel="Ri")
    line_Ri = lines!(ax_Ri, Ri, zf, linewidth=3, color=colors[1])
    xlims!(ax_Ri, (-0.5, 4))
    ylims!(ax_Ri, extrema(zf))
    hideydecorations!(ax_Ri, grid=false)

    ax_e = fig[2, 3] = Axis(fig, xlabel="TKE")
    line_e = lines!(ax_e, e, zc, linewidth=3, color=colors[1])
    xlims!(ax_e, extrema(ds[:e]))
    ylims!(ax_e, extrema(zf))
    hideydecorations!(ax_e, grid=false)

    ax_ϵ = fig[2, 4] = Axis(fig, xlabel="TKE dissipation")
    line_ϵ = lines!(ax_ϵ, tke_dissipation, zc, linewidth=3, color=colors[1])
    xlims!(ax_ϵ, extrema(ds[:tke_dissipation]))
    ylims!(ax_ϵ, extrema(zf))
    hideydecorations!(ax_ϵ, grid=false)

    ax_tke_a = fig[2, 5] = Axis(fig, xlabel="TKE advective flux")
    line_tke_a = lines!(ax_tke_a, tke_advective_flux, zf, linewidth=3, color=colors[1])
    xlims!(ax_tke_a, extrema(ds[:tke_advective_flux]))
    ylims!(ax_tke_a, extrema(zf))
    hideydecorations!(ax_tke_a, grid=false)

    ax_tke_b = fig[2, 6] = Axis(fig, xlabel="TKE buoyancy flux")
    line_tke_b = lines!(ax_tke_b, tke_buoyancy_flux, zc, linewidth=3, color=colors[1])
    xlims!(ax_tke_b, extrema(ds[:tke_buoyancy_flux]))
    ylims!(ax_tke_b, extrema(zf))
    hideydecorations!(ax_tke_b, grid=false)

    ax_tke_p = fig[2, 7] = Axis(fig, xlabel="TKE pressure flux")
    line_tke_p = lines!(ax_tke_p, tke_pressure_flux, zf, linewidth=3, color=colors[1])
    xlims!(ax_tke_p, extrema(ds[:tke_pressure_flux]))
    ylims!(ax_tke_p, extrema(zf))
    hideydecorations!(ax_tke_p, grid=false)

    ax_tke_sp = fig[2, 8] = Axis(fig, xlabel="TKE shear production")
    line_tke_sp = lines!(ax_tke_sp, tke_shear_production, zc, linewidth=3, color=colors[1])
    xlims!(ax_tke_sp, extrema(ds[:tke_shear_production]))
    ylims!(ax_tke_sp, extrema(zf))
    hideydecorations!(ax_tke_sp, grid=false)

    supertitle = fig[0, :] = Label(fig, plot_title, textsize=30)

    filepath = joinpath(data_directory, "instantaneous_statistics.mp4")
    record(fig, filepath, 1:Nt, framerate=15) do n
        @info "Animating instantaneous statistics movie frame $n/$Nt..."
        frame[] = n
    end

    @info "Movie saved: $filepath"

    return filepath
end

max_stars_repo_name: ali-ramadhan/LESbrary

max_issues_repo_name: ali-ramadhan/LESbrary

max_forks_repo_name: ali-ramadhan/LESbrary


text:
function get_parameter_upper_bound(tree)
    ranges = Dict('R'=>1.0e9,'C'=>10,'L'=>5,'P'=>[1.0e9,1],'W'=>1.0e9,'+'=>0,'-'=>0)
    return [ranges[node.Type] for node in tree]
end

function get_parameter_upper_bound(readablecircuit::String)
    elements = foldl(replace,["["=>"","]"=>"","-"=>"",","=>""],init = denumber_circuit(readablecircuit))
    ranges = Dict('R'=>1.0e9,'C'=>0.01,'L'=>5,'P'=>[1.0e9,1],'W'=>1.0e9,'+'=>0,'-'=>0)
    return flatten([ranges[e] for e in elements])
end

function func_and_params_for_optim(tree)
    circuit,circuit_parameters,param_inds = tree_to_circuit_with_inds(tree)
    circuitfunc = circuitfunction(circuit)
    upperbounds = get_parameter_upper_bound(tree)[param_inds]
    return circuitfunc, flatten(circuit_parameters) , flatten(upperbounds) , param_inds
end

function optimizeparameters(objective,initial_parameters,upper)
    lower = zeros(length(initial_parameters))
    inner_optimizer = NelderMead()
    results = optimize(objective, lower, upper, initial_parameters, Fminbox(inner_optimizer), Optim.Options(time_limit = 20.0))
    return results.minimizer,results.minimum
end

"""
    parameteroptimisation(circuit::String,measurements::Array{Complex{Float64},1},frequencies::Array{Float64,1})

Fit the parameters of a given equivalent circuit to measurement values, using the Nelder-Mead simplex algorithm.

The inputs are a circuit (e.g. "R1-[C2,R3]-P4"), an array of complex-valued impedance measurements and their corresponding frequencies.
The output is NamedTuple of the circuit's components with their corresponding parameter values.

# Example
```julia

julia> using EquivalentCircuits, Random

julia> Random.seed!(25);

julia> measurements = [5919.9 - 15.7, 5918.1 - 67.5im, 5887.1 - 285.7im, 5428.9 - 997.1im, 3871.8 - 978.9im,
3442.9 - 315.4im, 3405.5 - 242.5im, 3249.6 - 742.0im, 1779.4 - 1698.9im,  208.2 - 777.6im, 65.9 - 392.5im];

julia> frequencies = [0.10, 0.43, 1.83, 7.85, 33.60, 143.84, 615.85,  2636.65, 11288.38, 48329.30, 100000.00];

julia> parameteroptimisation("R1-[C2,R3-[C4,R5]]",measurements,frequencies)
(R1 = 19.953805651358255, C2 = 3.999778355811269e-9, R3 = 3400.0089192843684, C4 = 3.999911415903211e-6, R5 = 2495.2493215522577)
"""
function parameteroptimisation(circuit::String,measurements,frequencies)
    #   generate initial parameters.
        elements = foldl(replace,["["=>"","]"=>"","-"=>"",","=>""],init = denumber_circuit(circuit))
        initial_parameters = flatten(karva_parameters(elements))
    #   get circuitfunction and objective.
        circfunc = circuitfunction(circuit)
        objective = objectivefunction(circfunc,measurements,frequencies)
    #   get bounds.
        lower = zeros(length(initial_parameters))
        upper = get_parameter_upper_bound(circuit)
    #   optimize.
        inner_optimizer = NelderMead()
        results = optimize(objective, lower, upper, initial_parameters, Fminbox(inner_optimizer), Optim.Options(time_limit = 20.0))
        return parametertuple(circuit,results.minimizer)
end

"""
    parameteroptimisation(circuit::String,filepath::String)

Fit the parameters of a given equivalent circuit to measurement values, using the Nelder-Mead simplex algorithm.

The inputs are a circuit (e.g. "R1-[C2,R3]-P4") and a filepath to a CSV file containing the three following columns:
the real part of the impedance, the imaginary part of the impedance, and the frequencies corresponding to the measurements.
The output is NamedTuple of the circuit's components with their corresponding parameter values.
"""

function parameteroptimisation(circuit::String,data::String)
        meansurement_file = readdlm(data,',')
        # convert the measurement data into usable format.
        reals = meansurement_file[:,1]
        imags = meansurement_file[:,2]
        frequencies = meansurement_file[:,3]
        measurements = reals + imags*im
    #   generate initial parameters.
        elements = foldl(replace,["["=>"","]"=>"","-"=>"",","=>""],init = denumber_circuit(circuit))
        initial_parameters = flatten(karva_parameters(elements))
    #   get circuitfunction and objective.
        circfunc = circuitfunction(circuit)
        objective = objectivefunction(circfunc,measurements,frequencies)
    #   get bounds.
        lower = zeros(length(initial_parameters))
        upper = get_parameter_upper_bound(circuit)
    #   optimize.
        inner_optimizer = NelderMead()
        results = optimize(objective, lower, upper, initial_parameters, Fminbox(inner_optimizer), Optim.Options(time_limit = 20.0))
        return parametertuple(circuit,results.minimizer)
end

function deflatten_parameters(parameters,tree,param_inds)
    correct_value_lengths = length.(get_tree_parameters(tree)[param_inds])
    correct_length = length(correct_value_lengths)
    deflattened_parameters = Array{Any}(undef,length(correct_value_lengths))
    flat_index_counter = 1
    for (e,v) in enumerate(correct_value_lengths)
        if v == 1
            deflattened_parameters[e] = parameters[flat_index_counter]
            flat_index_counter += 1
        else
            deflattened_parameters[e] = [parameters[flat_index_counter],parameters[flat_index_counter+1]]
            flat_index_counter += 2
        end
    end
    return deflattened_parameters
end

function deflatten_parameters(parameters,circuit)
    elements = foldl(replace,["["=>"","]"=>"","-"=>"",","=>""],init = denumber_circuit(circuit))
    correct_value_lengths = [e=='P' ? 2 : 1 for e in elements]
    correct_length = length(correct_value_lengths)
    deflattened_parameters = Array{Any}(undef,length(correct_value_lengths))
    flat_index_counter = 1
    for (e,v) in enumerate(correct_value_lengths)
        if v == 1
            deflattened_parameters[e] = parameters[flat_index_counter]
            flat_index_counter += 1
        else
            deflattened_parameters[e] = [parameters[flat_index_counter],parameters[flat_index_counter+1]]
            flat_index_counter += 2
        end
    end
    return deflattened_parameters
end

ext: jl


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.4
#   kernelspec:
#     display_name: Julia 1.2.0
#     language: julia
#     name: julia-1.2
# ---

# # 2019-nCoV outbreak analysis
#
# ## Background
# A novel coronavirus (2019-nCoV) outbreak has been continuing in China, while a few imported cases are observed in the neighbouring countries. Although the spread across China suggests a certain amount of human-to-human transmisison, there has not been any observed in those countries which saw travel cases. This suggests that not all symptomatic cases lead to a secondary transmission, which was also the case in the past SARS/MERS outbreaks. Furthermore, even if any subclinical cases had been imported into these countries undetected, at least such cases did not contributed to secondary transmissions from which a detectable case originates.
#

# ## Model
# Let $J_t$ and $I_t$ be the incidence of imported and local cases of nCoV detected in a country outside China.
# Let $q_t$ be the (potentially time-dependent) detection probability of nCoV infection, which is understood as the combination of both symptomatic ratio and proper reporting, and $s_\tau$ be the distribution of the serial interval (i.e. time between the linked infection pairs) of nCoV.
#
# Assuming that the risk of importing an nCoV case follows an exponential hazard function, reflecting an ongoing outbreak in China as of February 2020, the overall (including those undetected) number of imported cases $j_t$ is given as
# $$
# j_t\sim \mathrm{Pois}(\lambda_t), \\
# h_t=h_0\exp(r t)
# $$
#
# We model the offspring distribution of nCoV by a negative binomial distribution $\mathrm{NB}\left(k,\frac{k}{R_0+k}\right)$, where $R$ is the mean (i.e. the basic reproduction number) and $k$ is the overdispersion parameter. The reproductive property of the negative binomial distriubution assures that the number of secondary infections caused by the overall infector at time $t$ follows $\mathrm{NB}\left(k(i_t+j_t),\frac{k}{R_0+k}\right)$, and we assume these offsprings are distributed on the timeline $t$ according to the serial interval distribution $s_\tau$.
#
# Since sampling from a negative binomial distribution is identical to sequentially sampling from Gamma and Poisson distributions, we can construct the renewal process as
# $$
# i_t \sim \mathrm{Pois}\left(\sum_{\tau=1}^\infty \lambda_{t-\tau},s_\tau \right),\\
# \lambda_t\sim\mathrm{Gamma}\left(k(i_{t-\tau}+j_{t-\tau}),\frac{R_0}{k}\right).
# $$
# Here, $\lambda_t$ is the total force of infection caused by $i_t$: infectious individuals with onset $t$.
#
# The observation of cases is assumed to follow the binomial sampling:
# $$
# I_t\sim \mathrm{Binom}(i_t,q_t), \\
# J_t\sim \mathrm{Binom}(j_t,q_t)
# $$
#
# We assumed the probability of detection $q_t$ might have been lower before the epidemic became widely recognised. Because of reporting delays, the most recent data might have also been underreported. We model $q_t$ as
# $$
# q_t=Q(t)\delta(t), \\
# \delta(t)=1-\exp\left(\frac{1}{d}(T-t)\right),
# $$
# where $Q(t)$ is the baseline reporting probability at time $t$ and $\delta(t)$ is the factor reflecting the delay. We assume the distribution of reporting delays follow an exponential distribution with mean $d$.

# ## Statistical analysis
# We assumed the observed imported/local cases $I_t,J_t$ and the serial interval distribution $S_t$ are given. Of the unkown variables, $i_t$, $j_t$ and $\lambda_t$ are integrated out by the particle-Metropolis-Hastings algorithm and the remaining variables $R_0, k, h_0,r,q,d$ are inferred.
#

# ## Data
# Reported cases outsie China by onset dates were extracted from the WHO situation reports (Situation Report 14 as of 4/2/2020, accessed on the same day). We considered cases with a travel history in china as imported cases, and cases labelled as "locally acquired" as local cases. We did not included those labelled as "Under investigation" in this analysis.

# Show incidence data
using DataFrames, Dates, PyPlot
# Imported and local cases outside China
# Source (accessed 4/2/2020): https://www.who.int/docs/default-source/coronaviruse/situation-reports/20200204-sitrep-15-ncov.pdf
dates=Date("2019-12-31"):Day(1):Date("2020-2-2")
china_hubei  =[1,0,0,2,0,1,0,0,0,1,1,0,0,1,2,0,3,3,1,3,3,7,3,9,6,8,5,4,3,3,4,1,0,0]
china_unknown=[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0]
localcases   =[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,1,1,2,0,0,0,1,1,0,0,1,2,0,0,0,0]
u_inv        =[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,0,1,0,0,1,0,1,0]
casedata=DataFrame([dates,china_hubei,china_unknown,localcases,u_inv], [:date,:china_hubei,:china_unknown,:loc,:underinv])
barplots=PyPlot.bar.(Ref(1:length(casedata.date)),[Vector(casedata[:,c]) for c in 2:5],0.8,[cumsum(hcat(zeros(Int,length(casedata.date)),Matrix(casedata[:,2:5])),dims=2)[:,c] for c in 1:4])
PyPlot.xticks(1:length(casedata.date),dates,rotation="vertical")
PyPlot.legend(getindex.(barplots,1),["China (Hubei)","China(Location unknown)","Locally acquired","Under investigation"]);

# Data preparation
const initdate=Date("2019-9-13")
const newyeardate=Date("2020-1-1")
# Possible (earliest bound of) the estimted start of the outbreak
# Source: http://virological.org/t/preliminary-phylogenetic-analysis-of-11-ncov2019-genomes-2020-01-19/329
const timelines=initdate:Day(1):Date("2020-2-3")
const tlen=length(timelines)
imported=zeros(Int,tlen)
loc=zeros(Int,tlen)
datesid=findall(x->x in dates,timelines)
imported[datesid].=casedata.china_hubei.+casedata.china_unknown
loc[datesid].=casedata.loc
observed=(imported=imported,loc=loc);

# ## Setups

# +
# Packages
using Mamba, Distributions, LinearAlgebra
# Distributions
NBmu(mu,k)=NegativeBinomial(k,mu/(mu+k))
Gmusd(mu,sd)=Gamma(mu^2/sd^2, sd^2/mu)
module DSI
using Distributions
struct DiscreteSerialInterval{T<:NamedTuple,F<:AbstractFloat, D<:Distribution}
    params::T; dist::D; distvec::Vector{F}
end
end
function DSIconstruct(distconstructor,params,tlen)
    dist=distconstructor(params...)
    distvec=cdf.(dist,0:tlen) |> diff
    DSI.DiscreteSerialInterval(params,dist,distvec)
end
SIparams=(μ=7.0,σ=3.0)
SIdist=Gmusd(SIparams...)
serialint=DSIconstruct(Gamma,NamedTuple{(:α,:θ)}(params(SIdist)),tlen)

# utils
function sysresample(weights)
    size=length(weights)
    weights=weights/sum(weights) # normalise
    cumweights=cumsum(weights)
    randomiser=floor.(Int,cumweights.*size.+rand())
    freqs=randomiser.-[0; randomiser[1:end-1]]
    return(reduce(vcat,fill.(1:size,freqs)))
end
# -

# ## Simulation functions

using Distributions, StatsFuns, StatsBase, SpecialFunctions, Parameters
function smc!(infness_pts::NamedTuple{(:imported,:loc),T} where T,cases_pts::NamedTuple{(:imported,:loc),T} where T,hazard_pts::NamedTuple{(:imported,:loc),T} where T,nbparm,detectprob,gtimevec,observed)
    nsample=size(infness_pts.imported,2)
    tlen=length(observed.loc)
    lweights=zeros(nsample)
    llkh=0.0
    α,p=nbparm
    θ=(1-p)/p
    for t in 1:length(observed.loc)
        for tag in keys(cases_pts)
            # draw cases
            @views cases_pts[tag][t,:].=observed[tag][t].+rand.(Poisson.((1-detectprob[t]).*hazard_pts[tag][t,:]))
            #draw gamma: total offsprings
            nonzerocase=cases_pts[tag][t,:].!=0 # to avoid Gamma(0,θ)
            if sum(nonzerocase)!=0 @views infness_pts[tag][t,nonzerocase].=rand.(Gamma.(α.*cases_pts[tag][t,nonzerocase],θ)) end
            #distribute infness on timeline
            hazard_pts.loc[t+1:end,:] .+= (@view infness_pts[tag][t:t,:]).*gtimevec[1:tlen-t]
            @views lweights.+=logpdf.(Poisson.(detectprob[t].*hazard_pts[tag][t,:]),observed[tag][t])
        end
        if all(lweights.≤-Inf) llkh=-Inf;break end
        if 2logsumexp(lweights)-logsumexp(2 .*lweights)< log(nsample)-log(2) || t==tlen
            # resample
            llkh+=logsumexp(lweights)-log(nsample)
            lweights.-=maximum(lweights)
            newid=sysresample(exp.(lweights))
            for tag in keys(cases_pts)
                @views infness_pts[tag][1:t,:].=infness_pts[tag][1:t,newid]
                @views cases_pts[tag][1:t,:].=cases_pts[tag][1:t,newid]
            end
            @views hazard_pts.loc[1:t,:].=hazard_pts.loc[1:t,newid]
            lweights.=0.0
        end
    end
    return(llkh)
end
function llkhpf!(paths,nsample,branchdist,gtimedist,observed,detectprob,tlen)
    particles=(infness_pts=(imported=zeros(tlen,nsample), loc=zeros(tlen,nsample)),
    hazard_pts=(imported=paths.hazard.imported,loc=zeros(tlen,nsample)),
    cases_pts=(imported=zeros(Int,tlen,nsample),loc=zeros(Int,tlen,nsample)))
    @unpack infness_pts,cases_pts,hazard_pts = particles
    @unpack infness,cases,hazard=paths
    gtimevec=diff(cdf.(gtimedist,0:tlen))
    counts=0
    ll=0.0
    while true
        infness_pts.imported.=0.0
        infness_pts.loc.=0.0
        hazard_pts.loc.=0.0
        ll=smc!(infness_pts,cases_pts,hazard_pts,params(branchdist),detectprob,gtimevec,observed)
        counts+=1
        if ll>-Inf break end
        if counts>100 print("infness could not be sampled in 100 SMC iterations") end
    end
    # sample one particle
    sampleid=sample(1:nsample)
    @views infness.imported.=infness_pts.imported[:,sampleid]
    @views infness.loc.=infness_pts.loc[:,sampleid]
    @views cases.imported.=cases_pts.imported[:,sampleid]
    @views cases.loc.=cases_pts.loc[:,sampleid]
    @views hazard.loc.=hazard_pts.loc[:,sampleid]
    ll
end
function importhazard(param)
    @unpack h0,r=param
    @. h0*exp(r*((1:tlen)-Dates.value(newyeardate-initdate)-1))
end
function detectprob(param)
    @unpack q,delayrate=param
    @. q*(1.0-exp(-delayrate*(tlen-(1:tlen))))
end

# +
# test simulation run
nb=NBmu(2,0.5)
gt=Gmusd(7,1)
R0=2
k=0.5
h0=10
r=0.05
qt=fill(0.1,tlen)
nsample=1000

paths=(hazard=(imported=importhazard((h0=h0,r=r)),loc=zeros(tlen)),
        cases=(imported=zeros(Int,tlen),loc=zeros(Int,tlen)),
        infness=(imported=observed.imported.+0.0,loc=observed.loc.+0.0))

@time lls=llkhpf!(paths,nsample,nb,gt,observed,qt,tlen)
#@time lls=infnessgibbs!(infness,hazard.loc,500,nb,gt,tlen,cases);
#@time casescondsampler!(cases,infness,hazard,params(nb),observed,q);
#@time llnbdist(params(nb),infness,cases)
#@time lldetectprob(qt,cases,observed)
#@time llimporthazard((h0,r),cases)
#using RCall;@rimport base as R
# -

# ## MCMC sampling

# +
# unknown variables: λ, i, j, R₀, k, q
parms=Dict{Symbol,Any}(
    :h₀=>0.1,
    :r=>0.1,
    :R₀=>1.0,
    :k=>0.5,
    :nlogq=>0.5,
    :delayrate=>0.1
)
priors=Dict{Symbol,Any}()
for parname in keys(parms)
    priors[parname]=Stochastic(()->Uniform(0,5))
end
#priors[:nlogq]=Stochastic(1,()->Uniform(0,5))
dotter=[100]
inputs=Dict{Symbol,Any}(
    :observed=>observed,
    :paths=>paths,
    :zerotrick=>0.0,
    :invtemp=>1.0,
    :smcsize=>100,
    :counter=>([0],dotter)
)

inits=merge(parms,inputs)
#inits=copy(parms)
inits=[inits]

model=Model(
    j=Logical(1,(paths,k)-> paths.cases.imported),
    i=Logical(1,(paths,k)-> paths.cases.loc),
    ll_smc=Logical((R₀,k,h₀,r,nlogq,delayrate,observed,paths,smcsize)->begin
        paths.hazard.imported.=importhazard((h0=h₀,r=r))
        nb=NBmu(R₀,k)
        qt=detectprob((q=exp(-nlogq),delayrate=delayrate))
        ll=llkhpf!(paths,smcsize,nb,serialint.dist,observed,qt,tlen)
        ll
    end),
    lltotal=Logical((ll_smc,invtemp)->ll_smc*invtemp
        , false),
    zerotrick=Stochastic((lltotal,i,j)->Poisson(-lltotal),false),
    count=Logical((counter,k)->0);
    priors...
)
countup=Sampler([:count],
    (count,counter)->begin
        counter[1].+=1
        if counter[1][1]==counter[2][1]
            counter[1].=0
            print(".")
        end
        count+=1
    end
)

setsamplers!(model,[AMM(collect(keys(parms)),Matrix{Float64}(I,fill(length(keys(parms)),2)...).*0.05),countup]);

# +
mcmclen=100000
burn=div(mcmclen,5)
dotter.=div(mcmclen,100)

chain = mcmc(model, inputs, inits, mcmclen, burnin=burn, thin=div(mcmclen,500), chains=1, verbose=true)
# -

showparam=[:R₀,:k,:nlogq,:delayrate,:h₀,:r]
@show Mamba.draw(Mamba.plot(chain[:,showparam,:]))
# Visualise
i=chain[:,:i,:].value
im=median(i,dims=1)[1,:,1]
j=chain[:,:j,:].value
jm=median(j,dims=1)[1,:,1]
q=median(hcat([detectprob((q=exp(-chain[t,:nlogq,1].value[1]),delayrate=chain[t,:delayrate,1].value[1])) for t in chain.range]...),dims=2)
inferredvars=PyPlot.plot.([jm,im,q.*100])
PyPlot.legend(getindex.(inferredvars,1),["jₜ: imported cases","iₜ: local cases","detection probability (%)"]);

@show fittodata=PyPlot.plot.([observed.imported,jm.*q[:,1],observed.loc,im.*q[:,1]])
PyPlot.legend(getindex.(fittodata,1),["Jₜ","E(Jₜ)","Iₜ","E(Iₜ)"]);
using RCall;@rimport graphics as rg;
rg.pairs(chain[:,[:nlogq,:R₀,:k,:h₀,:r],:].value[:,:,1])

# Save chain
using Serialization
serialize(open("pmcmc_chain100000", "w"), chain)


text:
let # preventing global scope
    # ======================================================
    # Testing continuous-time homogeneous system (no input)
    # ======================================================
    A = sparse([1, 1, 2, 3, 4], [1, 2, 2, 4, 3], [1., 2., 3., 4., 5.], 4, 4)
    X0 = BallInf(zeros(4), 0.1)
    cont_sys_homog = IVP(LCS(A), X0)

    # Check if the input is constant
    @test inputdim(cont_sys_homog) == 0
    # Check if the input is empty
    #@test isa(next_set(cont_sys_homog.U), ZeroSet) # there is no input set
    # Check data fields
    @test cont_sys_homog.s.A == A
    @test cont_sys_homog.x0.center == zeros(4) && cont_sys_homog.x0.radius == 0.1

    # ===================================================
    # Testing continuous-time system with constant input
    # ===================================================
    U = ConstantInput(Ball2(ones(4), 0.5))
    cont_sys = IVP(CLCCS(A, Matrix(1.0I, 4, 4), nothing, U), X0)

    # check initial state
    @test cont_sys.x0.center ≈ zeros(4) && cont_sys.x0.radius ≈ 0.1

    # recover input
    inputs = next_set(cont_sys.s.U) #collect(nextinput(cont_sys.s.U))[1]

    @test inputs.center ≈ ones(4) && inputs.radius ≈ 0.5

    # ========================================================
    # Testing continuous-time system with time-varying input
    # ========================================================
    U = VaryingInput([Ball2(0.01*i*ones(4), i*0.2) for i in 1:3])
    cont_sys = cont_sys = IVP(CLCCS(A, Matrix(1.0I, 4, 4), nothing, U), X0)
    for (i, inputs) in enumerate(cont_sys.s.U)
        if i == 1
            @test inputs.center ≈ 0.01*ones(4)
            @test inputs.radius ≈ 0.2
        elseif i == 2
            @test inputs.center ≈ 0.02*ones(4)
            @test inputs.radius ≈ 0.4
        else
            @test inputs.center ≈ 0.03*ones(4)
            @test inputs.radius ≈ 0.6
        end
    end

    # =========================================
    # Testing discrete-time homogeneous system
    # =========================================
    δ = 0.01
    discr_sys_homog = IVP(LDS(A), X0)

    # Check if the input is constant
    #@test isa(discr_sys_homog.s.U, ConstantInput) # no field U
    # Check if the input is empty
    #@test isa(next_set(discr_sys_homog.s.U), ZeroSet)
    # Check data fields
    @test discr_sys_homog.s.A == A
    @test discr_sys_homog.x0.center == zeros(4) && discr_sys_homog.x0.radius == 0.1

    # =================================================
    # Testing discrete-time system with constant input
    # =================================================
    U = ConstantInput(Ball2(ones(4), 0.5))
    discr_sys = IVP(CLCDS(A, Matrix(1.0I, 4, 4), nothing, U), X0)

    # check initial state
    @test discr_sys.x0.center ≈ zeros(4) && discr_sys.x0.radius ≈ 0.1

    # recover input
    inputs = collect(nextinput(discr_sys.s.U))[1]

    @test inputs.center ≈ ones(4) && inputs.radius ≈ 0.5

    # =====================================================
    # Testing discrete-time system with time-varying input
    # =====================================================
    U = [Ball2(0.01*i*ones(4), i*0.2) for i in 1:3]
    discr_sys = IVP(CLCDS(A, Matrix(1.0I, 4, 4), nothing, U), X0)
    Uset = inputset(discr_sys.s)
    for (i, inputs) in enumerate(cont_sys.s.U)
        if i == 1
            @test inputs.center ≈ 0.01*ones(4)
            @test inputs.radius ≈ 0.2
        elseif i == 2
            @test inputs.center ≈ 0.02*ones(4)
            @test inputs.radius ≈ 0.4
        else
            @test inputs.center ≈ 0.03*ones(4)
            @test inputs.radius ≈ 0.6
        end
    end
end


text:
##===----------------------------------------------------------------------===##
##                                     YEPPP.JL                               ##
##                                                                            ##
##===----------------------------------------------------------------------===##
##                                                                            ##
##  This file provides access to Yeppp and handles checking for the existence ##
##  of the Yeppp library during import.                                       ##
##                                                                            ##
##===----------------------------------------------------------------------===##
module Yeppp
import Vectorize: functions, addfunction

using Libdl

# Cross-version compatibility
OS = Sys.KERNEL

# Attempt to locate Yeppp library before we proceed with compiling file
libyeppp_ = Libdl.find_library(["libyeppp"])
if libyeppp_ != "" # using system installed yeppp
    const global libyeppp = libyeppp_
else # using Vectorize.jl provided yppp
    currdir = @__FILE__
    bindir = currdir[1:end-12]*"deps/src/yeppp/binaries/"
    if OS == :Darwin
        @eval const global libyeppp = bindir*"macosx/x86_64/libyeppp.dylib"
    elseif OS == :Linux
        @eval const global libyeppp = bindir*"linux/x86_64/libyeppp.so"
    elseif OS == :NT
        @eval const global libyeppp = bindir*"windows/amd64/yeppp.dll"
    end
end

"""
`__init__()` Cvoid

This function initializes the Yeppp library whenever the module is imported.
This allows us to precompile the rest of the module without asking the user
to explictly initialize Yeppp! on load.
"""
function __init__()
    # Yeppp Initialization
    if isfile(libyeppp)
        status = ccall(("yepLibrary_Init", libyeppp), Cint,
                        (), )
        status != 0 && error("Error initializing Yeppp library (error: ", status, ")")
    end

    return true
end

## This defines the argument and return types for every yepCore function; this are as follow:
##
## (First Argument Type, Second Argument Type, Return Type
##
const yepppcore = ((Int8, Int8, Int8), (UInt8, UInt8, UInt16),  (Int16, Int16, Int16),
                   (UInt16, UInt16, UInt32), (Int32, Int32, Int32),
                   (UInt32, UInt32, UInt64), (Int64, Int64, Int64),
                   (Float32, Float32, Float32), (Float64, Float64, Float64))

## This defines the mapping between argument types and Yeppp function name specifications
const identifier = Dict(Int8 => "V8s", Int16 => "V16s", Int32 => "V32s", Int64 => "V64s",
                        UInt8 => "V8u", UInt16 => "V16u", UInt32 => "V32u", UInt64 => "V64u",
                        Float32 => "V32f", Float64 => "V64f")

#### Yeppp Arithmetic taking two arguments and returning vector
for (f, fname, name) in ((:add, "Add", "addition"),  (:sub, "Subtract", "subtraction"),
                   (:mul, "Multiply", "multiplication"), (:max, "Max", "maximum"),
                   (:min, "Min", "minimum"))
    for (argtype1, argtype2, returntype) in yepppcore
        # generate Yeppp function name
        yepppname = string("yepCore_$(fname)_", identifier[argtype1], identifier[argtype2],
                           "_", identifier[returntype])
        addfunction(functions, (f, (Array{Float64}, Array{Float64})), "Vectorize.Yeppp.$f")
        @eval begin
            @doc """
            `$($f)(X::Array{$($argtype1)}, Y::Array{$($argtype2)})`
            Implements element-wise **$($name)** over two **Array{$($argtype1)}**. Allocates
            memory to store result. *Returns:* **Array{$($returntype)}**
            """
            function ($f)(X::Array{$argtype1}, Y::Array{$argtype2})
                len = length(X)
                out = Array{$returntype}(undef, len)
                ccall(($(yepppname, libyeppp)), Cint,
                          (Ptr{$argtype1}, Ptr{$argtype2}, Ptr{$returntype},  Clonglong),
                          X, Y, out, len)
                return out
            end
        end
    end
end

#### YepppMath functions returning vectors
for (f, fname) in [(:sin, "Sin"),  (:cos, "Cos"),  (:tan, "Tan"), (:log, "Log"), (:exp, "Exp")]
    yepppname = string("yepMath_$(fname)_V64f_V64f")
    name = string(f)
    f! = Symbol("$(f)!")
    # register functions for build
    addfunction(functions, (f, (Array{Float64},)), "Vectorize.Yeppp.$f")
    addfunction(functions, (f!, (Array{Float64}, Array{Float64})), "Vectorize.Yeppp.$(f!)")
    @eval begin
         @doc """
        `$($f)(X::Array{Float64})`
        Implements element-wise **$($name)** over a **Array{Float64}**. Allocates
        memory to store result. *Returns:* **Array{Float64}**
        """
        function ($f)(X::Array{Float64}) # regular syntax
            out = Array{Float64}(undef, length(X))
            return ($f!)(out, X)
        end
        @doc """
        `$($f!)(result::Array{Float64}, X::Array{Float64})`
        Implements element-wise **$($name)** over a **Array{Float64}** and overwrites
        the result vector with computed value. *Returns:* **Array{Float64}** `result`
        """
        function ($f!)(out::Array{Float64}, X::Array{Float64}) # in place syntax
            len = length(X)
            ccall(($(yepppname, libyeppp)), Cint,
                  (Ptr{Float64}, Ptr{Float64},  Clonglong),
                  X, out, len)
            return out
        end
    end
end

## YepppCore - functions that return scalars
for (T, Tscalar) in ((Float32, "S32f"), (Float64, "S64f"))
    for (f, fname, name) in [(:sum, "Sum", "sum"), (:sumsqr, "SumSquares", "sum-of-squares")]
        yepppname = string("yepCore_$(fname)_", identifier[T], "_", Tscalar)
        addfunction(functions, (f, (Array{T},)), "Vectorize.Yeppp.$f")
        @eval begin
            @doc """
            `$($f)(X::Array{$($T)})`
            Computes the **$($name)** of a **Array{$($T)}**. Allocates
            memory to store result. *Returns:* **Array{$($T)}**
            """
            function ($f)(X::Array{$T})
                len = length(X)
                out = Array{$T}(undef, 1)
                ccall(($(yepppname, libyeppp)), Cint,
                      (Ptr{$T}, Ptr{$T},  Clonglong),
                      X, out, len)
                return out[1]
            end
        end
    end
end

end # End Module


text:
"""
    conv_bn(kernelsize, inplanes, outplanes, activation = relu;
            rev = false,
            stride = 1, pad = 0, dilation = 1, groups = 1, [bias, weight, init],
            initβ = Flux.zeros32, initγ = Flux.ones32, ϵ = 1f-5, momentum = 1f-1)

Create a convolution + batch normalization pair with ReLU activation.

# Arguments
- `kernelsize`: size of the convolution kernel (tuple)
- `inplanes`: number of input feature maps
- `outplanes`: number of output feature maps
- `activation`: the activation function for the final layer
- `rev`: set to `true` to place the batch norm before the convolution
- `stride`: stride of the convolution kernel
- `pad`: padding of the convolution kernel
- `dilation`: dilation of the convolution kernel
- `groups`: groups for the convolution kernel
- `bias`, `weight`, `init`: initialization for the convolution kernel (see [`Flux.Conv`](#))
- `initβ`, `initγ`: initialization for the batch norm (see [`Flux.BatchNorm`](#))
- `ϵ`, `momentum`: batch norm parameters (see [`Flux.BatchNorm`](#))
"""
function conv_bn(kernelsize, inplanes, outplanes, activation = relu;
                 rev = false,
                 initβ = Flux.zeros32, initγ = Flux.ones32, ϵ = 1f-5, momentum = 1f-1,
                 kwargs...)
  layers = []

  if rev
    activations = (conv = activation, bn = identity)
    bnplanes = inplanes
  else
    activations = (conv = identity, bn = activation)
    bnplanes = outplanes
  end

  push!(layers, Conv(kernelsize, Int(inplanes) => Int(outplanes), activations.conv; kwargs...))
  push!(layers, BatchNorm(Int(bnplanes), activations.bn;
                          initβ = initβ, initγ = initγ, ϵ = ϵ, momentum = momentum))

  return rev ? reverse(layers) : layers
end

"""
    cat_channels(x, y)

Concatenate `x` and `y` along the channel dimension (third dimension).
Equivalent to `cat(x, y; dims=3)`.
Convenient binary reduction operator for use with `Parallel`.
"""
cat_channels(x, y) = cat(x, y; dims = 3)

"""
    weights(model)

Load the pre-trained weights for `model` using the stored artifacts.
"""
function weights(model)
  try
    path = joinpath(@artifact_str(model), "$model.bson")
    return BSON.load(path, @__MODULE__)[:weights]
  catch e
    throw(ArgumentError("No pre-trained weights available for $model."))
  end
end

"""
    loadpretrain!(model, name)

Load the pre-trained weight artifacts matching `<name>.bson` into `model`.
"""
loadpretrain!(model, name) = Flux.loadparams!(model, weights(name))

function _maybe_big_show(io, model)
  if isdefined(Flux, :_big_show)
    if isnothing(get(io, :typeinfo, nothing)) # e.g. top level in REPL
      Flux._big_show(io, model)
    else
      show(io, model)
    end
  else
    show(io, model)
  end
end


text:
# Automatically generated using Clang.jl wrap_c, version 0.0.0

export
    OBJC_NEW_PROPERTIES,
    LIBAVDEVICE_VERSION_MAJOR,
    LIBAVDEVICE_VERSION_MINOR,
    LIBAVDEVICE_VERSION_MICRO,
    LIBAVDEVICE_BUILD,
    AVDeviceRect,
    AVAppToDevMessageType,
    AV_APP_TO_DEV_NONE,
    AV_APP_TO_DEV_WINDOW_SIZE,
    AV_APP_TO_DEV_WINDOW_REPAINT,
    AV_APP_TO_DEV_PAUSE,
    AV_APP_TO_DEV_PLAY,
    AV_APP_TO_DEV_TOGGLE_PAUSE,
    AV_APP_TO_DEV_SET_VOLUME,
    AV_APP_TO_DEV_MUTE,
    AV_APP_TO_DEV_UNMUTE,
    AV_APP_TO_DEV_TOGGLE_MUTE,
    AV_APP_TO_DEV_GET_VOLUME,
    AV_APP_TO_DEV_GET_MUTE,
    AVDevToAppMessageType,
    AV_DEV_TO_APP_NONE,
    AV_DEV_TO_APP_CREATE_WINDOW_BUFFER,
    AV_DEV_TO_APP_PREPARE_WINDOW_BUFFER,
    AV_DEV_TO_APP_DISPLAY_WINDOW_BUFFER,
    AV_DEV_TO_APP_DESTROY_WINDOW_BUFFER,
    AV_DEV_TO_APP_BUFFER_OVERFLOW,
    AV_DEV_TO_APP_BUFFER_UNDERFLOW,
    AV_DEV_TO_APP_BUFFER_READABLE,
    AV_DEV_TO_APP_BUFFER_WRITABLE,
    AV_DEV_TO_APP_MUTE_STATE_CHANGED,
    AV_DEV_TO_APP_VOLUME_LEVEL_CHANGED,
    AVDeviceCapabilitiesQuery,
    AVDeviceInfo,
    AVDeviceInfoList

const OBJC_NEW_PROPERTIES  =  1
const LIBAVDEVICE_VERSION_MAJOR  =  57
const LIBAVDEVICE_VERSION_MINOR  =  0
const LIBAVDEVICE_VERSION_MICRO  =  101

# Skipping MacroDefinition: LIBAVDEVICE_VERSION_INT AV_VERSION_INT ( LIBAVDEVICE_VERSION_MAJOR , LIBAVDEVICE_VERSION_MINOR , LIBAVDEVICE_VERSION_MICRO )
# Skipping MacroDefinition: LIBAVDEVICE_VERSION AV_VERSION ( LIBAVDEVICE_VERSION_MAJOR , LIBAVDEVICE_VERSION_MINOR , LIBAVDEVICE_VERSION_MICRO )

# const LIBAVDEVICE_BUILD  =  LIBAVDEVICE_VERSION_INT

# Skipping MacroDefinition: LIBAVDEVICE_IDENT "Lavd" AV_STRINGIFY ( LIBAVDEVICE_VERSION )

immutable AVDeviceRect
    x::Cint
    y::Cint
    width::Cint
    height::Cint
end

# begin enum AVAppToDevMessageType
const AVAppToDevMessageType = UInt32
const AV_APP_TO_DEV_NONE  =  (UInt32)(1313820229)
const AV_APP_TO_DEV_WINDOW_SIZE  =  (UInt32)(1195724621)
const AV_APP_TO_DEV_WINDOW_REPAINT  =  (UInt32)(1380274241)
const AV_APP_TO_DEV_PAUSE  =  (UInt32)(1346458912)
const AV_APP_TO_DEV_PLAY  =  (UInt32)(1347174745)
const AV_APP_TO_DEV_TOGGLE_PAUSE  =  (UInt32)(1346458964)
const AV_APP_TO_DEV_SET_VOLUME  =  (UInt32)(1398165324)
const AV_APP_TO_DEV_MUTE  =  (UInt32)(541939028)
const AV_APP_TO_DEV_UNMUTE  =  (UInt32)(1431131476)
const AV_APP_TO_DEV_TOGGLE_MUTE  =  (UInt32)(1414354260)
const AV_APP_TO_DEV_GET_VOLUME  =  (UInt32)(1196838732)
const AV_APP_TO_DEV_GET_MUTE  =  (UInt32)(1196250452)
# end enum AVAppToDevMessageType

# begin enum AVDevToAppMessageType
const AVDevToAppMessageType = UInt32
const AV_DEV_TO_APP_NONE  =  (UInt32)(1313820229)
const AV_DEV_TO_APP_CREATE_WINDOW_BUFFER  =  (UInt32)(1111708229)
const AV_DEV_TO_APP_PREPARE_WINDOW_BUFFER  =  (UInt32)(1112560197)
const AV_DEV_TO_APP_DISPLAY_WINDOW_BUFFER  =  (UInt32)(1111771475)
const AV_DEV_TO_APP_DESTROY_WINDOW_BUFFER  =  (UInt32)(1111770451)
const AV_DEV_TO_APP_BUFFER_OVERFLOW  =  (UInt32)(1112491596)
const AV_DEV_TO_APP_BUFFER_UNDERFLOW  =  (UInt32)(1112884812)
const AV_DEV_TO_APP_BUFFER_READABLE  =  (UInt32)(1112687648)
const AV_DEV_TO_APP_BUFFER_WRITABLE  =  (UInt32)(1113018912)
const AV_DEV_TO_APP_MUTE_STATE_CHANGED  =  (UInt32)(1129141588)
const AV_DEV_TO_APP_VOLUME_LEVEL_CHANGED  =  (UInt32)(1129729868)
# end enum AVDevToAppMessageType

immutable AVDeviceCapabilitiesQuery
    av_class::Ptr{AVClass}
    device_context::Ptr{AVFormatContext}
    codec::AVCodecID
    sample_format::AVSampleFormat
    pixel_format::AVPixelFormat
    sample_rate::Cint
    channels::Cint
    channel_layout::Int64
    window_width::Cint
    window_height::Cint
    frame_width::Cint
    frame_height::Cint
    fps::AVRational
end

immutable AVDeviceInfo
    device_name::Cstring
    device_description::Cstring
end

immutable AVDeviceInfoList
    devices::Ptr{Ptr{AVDeviceInfo}}
    nb_devices::Cint
    default_device::Cint
end

max_stars_repo_name: asafasban/AAVideoIO

max_issues_repo_name: asafasban/AAVideoIO

max_forks_repo_name: asafasban/AAVideoIO


text:
# Usage: run.jl SOLVERNAME TIMELIMIT DATAFOLDER INSTANCENAME

using ConicBenchmarkUtilities
using Pajarito
using ConicNonlinearBridge

# Print with full precision
function vec_to_string(v)
    terms = String[]
    for x in v
        b = IOBuffer()
        show(b, x)
        push!(terms, String(take!(b)))
    end
    return string("[",join(terms,","),"]")
end

function solveprint(instance, solver)
    # Convert from cbf to our conic format
    (c, A, b, con_cones, var_cones, vartypes, sense, objoffset) = cbftompb(instance)
    if sense == :Max
        scale!(c, -1)
    end

    # Load, solve, and time
    m = MathProgBase.ConicModel(solver)
    timeall = time()
    MathProgBase.loadproblem!(m, c, A, b, con_cones, var_cones)
    if !all(t->t==:Cont, vartypes)
        MathProgBase.setvartype!(m, vartypes)
    end
    MathProgBase.optimize!(m)
    timeall = time() - timeall
    timesolver = MathProgBase.getsolvetime(m)
    status = MathProgBase.status(m)

    x = []
    try
        x = MathProgBase.getsolution(m)
    end

    objval = NaN
    objbound = NaN
    nodecount = NaN
    try
        objval = MathProgBase.getobjval(m)
        if sense == :Max
            objval = -objval
        end
        objval += objoffset
    end
    try
        objbound = MathProgBase.getobjbound(m)
        if sense == :Max
            objbound = -objbound
        end
        objbound += objoffset
    end
    try
        nodecount = MathProgBase.getnodecount(m)
    end

    println("#STATUS# $status")
    println("#OBJVAL# $objval")
    println("#OBJBOUND# $objbound")
    println("#NODECOUNT# $nodecount")
    println("#TIMESOLVER# $timesolver")
    println("#TIMEALL# $timeall")
    println("#SOLUTION# $(vec_to_string(x))")
end

solvermap = Dict(
    # BONMIN
    "BONMIN_OA" =>
    (["AmplNLWriter"], quote ConicNLPWrapper(
    nlp_solver=BonminNLSolver(["bonmin.algorithm=\"B-OA\"", "bonmin.time_limit=$tlim", "halt_on_ampl_error=yes", "bonmin.allowable_fraction_gap=$rgap", "bonmin.oa_log_level=1", "bonmin.bb_log_level=1"]),
    disaggregate_soc=false,
    ) end),

    "BONMIN_OADIS" =>
    (["AmplNLWriter"], quote ConicNLPWrapper(
    nlp_solver=BonminNLSolver(["bonmin.algorithm=\"B-OA\"", "bonmin.time_limit=$tlim", "halt_on_ampl_error=yes", "bonmin.allowable_fraction_gap=$rgap", "bonmin.oa_log_level=1", "bonmin.bb_log_level=1"]),
    disaggregate_soc=true,
    ) end),

    "BONMIN_BB" =>
    (["AmplNLWriter"], quote ConicNLPWrapper(
    nlp_solver=BonminNLSolver(["bonmin.algorithm=\"B-BB\"", "bonmin.time_limit=$tlim", "halt_on_ampl_error=yes", "bonmin.allowable_fraction_gap=$rgap", "bonmin.oa_log_level=1", "bonmin.bb_log_level=1"]),
    disaggregate_soc=false,
    ) end),

    # SCIP
    "SCIP_MISOCP" =>
    (["SCIP"], quote SCIPSolver(
    "limits/gap", rgap, "numerics/feastol", tol_feas, "limits/time", tlim,
    ) end),

    # CPLEX
    "CPLEX_MISOCP" =>
    (["CPLEX"], quote CplexSolver(
    CPX_PARAM_THREADS=1, CPX_PARAM_TILIM=tlim, CPX_PARAM_SCRIND=1, CPX_PARAM_EPINT=tol_int, CPX_PARAM_EPRHS=tol_feas, CPX_PARAM_EPGAP=rgap,
    ) end),

    # Paj CBC
    "PAJ_CBC_ECOS" =>
    (["Cbc","ECOS"], quote PajaritoSolver(
    mip_solver=CbcSolver(logLevel=0, integerTolerance=tol_int, primalTolerance=tol_feas, ratioGap=tol_gap, check_warmstart=false),
    cont_solver=ECOSSolver(verbose=false, reltol=1e-10, feastol=1e-10, reltol_inacc=1e-5, feastol_inacc=1e-8),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    ) end),

    # Paj GLPK
    "PAJ_GLPK_ECOS" =>
    (["GLPKMathProgInterface","ECOS"], quote PajaritoSolver(
    mip_solver=GLPKSolverMIP(msg_lev=GLPK.MSG_OFF, tol_int=tol_int, tol_bnd=tol_feas, mip_gap=tol_gap, presolve=true),
    cont_solver=ECOSSolver(verbose=false, reltol=1e-10, feastol=1e-10, reltol_inacc=1e-5, feastol_inacc=1e-8),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    ) end),

    # Paj CPLEX MOSEK
    "PAJ_CPLEX_MOSEK" =>
    (["CPLEX","Mosek"], quote PajaritoSolver(
    mip_solver=CplexSolver(CPX_PARAM_THREADS=1, CPX_PARAM_SCRIND=0, CPX_PARAM_EPINT=tol_int, CPX_PARAM_EPRHS=tol_feas, CPX_PARAM_EPGAP=tol_gap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=1, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    ) end),

    "PAJ_CPLEX_MOSEK_msd" =>
    (["CPLEX","Mosek"], quote PajaritoSolver(
    mip_solver=CplexSolver(CPX_PARAM_THREADS=1, CPX_PARAM_SCRIND=1, CPX_PARAM_EPINT=tol_int, CPX_PARAM_EPRHS=tol_feas, CPX_PARAM_EPGAP=rgap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=1, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    mip_solver_drives=true,
    ) end),

    # Paj Gurobi sep only
    "PAJ_Gurobi_sep" =>
    (["Gurobi"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=0, Threads=8, IntFeasTol=tol_int, FeasibilityTol=tol_feas, MIPGap=tol_gap),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    prim_cuts_only=true, solve_relax=false, solve_subp=false,
    ) end),

    "PAJ_Gurobi_msd_sep" =>
    (["Gurobi"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=1, Threads=8, IntFeasTol=tol_int, FeasibilityTol=tol_feas, MIPGap=rgap),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    mip_solver_drives=true,
    prim_cuts_only=true, solve_relax=false, solve_subp=false,
    ) end),

    # Paj Gurobi MOSEK
    "PAJ_Gurobi_MOSEK" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=0, Threads=8, IntFeasTol=tol_int, FeasibilityTol=tol_feas, MIPGap=tol_gap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    ) end),

    "PAJ_Gurobi_MOSEK_msd" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=1, Threads=8, IntFeasTol=tol_int, FeasibilityTol=tol_feas, MIPGap=rgap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    mip_solver_drives=true,
    ) end),

    # Paj Gurobi MOSEK no init
    "PAJ_Gurobi_MOSEK_noinit" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=0, Threads=8, IntFeasTol=tol_int, FeasibilityTol=tol_feas, MIPGap=tol_gap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    init_soc_one=false, init_soc_inf=false, init_exp=false, init_sdp_lin=false, init_sdp_soc=false,
    ) end),

    "PAJ_Gurobi_MOSEK_msd_noinit" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=1, Threads=8, IntFeasTol=tol_int, FeasibilityTol=tol_feas, MIPGap=rgap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    mip_solver_drives=true,
    init_soc_one=false, init_soc_inf=false, init_exp=false, init_sdp_lin=false, init_sdp_soc=false,
    ) end),

    # Paj Gurobi MOSEK subp only, no init
    "PAJ_Gurobi_MOSEK_subponly_noinit" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=0, Threads=8, IntFeasTol=tol_int, FeasibilityTol=tol_feas, MIPGap=tol_gap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    prim_cuts_assist=false,
    init_soc_one=false, init_soc_inf=false, init_exp=false, init_sdp_lin=false, init_sdp_soc=false,
    ) end),

    "PAJ_Gurobi_MOSEK_msd_subponly_noinit" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=1, Threads=8, IntFeasTol=tol_int, FeasibilityTol=tol_feas, MIPGap=rgap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    mip_solver_drives=true,
    prim_cuts_assist=false,
    init_soc_one=false, init_soc_inf=false, init_exp=false, init_sdp_lin=false, init_sdp_soc=false,
    ) end),

    # Paj Gurobi MOSEK no disagg, subp only, no init
    "PAJ_Gurobi_MOSEK_nodisagg_subponly_noinit" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=0, Threads=8, IntFeasTol=tol_int, FeasibilityTol=tol_feas, MIPGap=tol_gap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    prim_cuts_assist=false,
    init_soc_one=false, init_soc_inf=false, init_exp=false, init_sdp_lin=false, init_sdp_soc=false,
    all_disagg=false, soc_disagg=false, sdp_eig=false
    ) end),

    "PAJ_Gurobi_MOSEK_msd_nodisagg_subponly_noinit" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=1, Threads=8, IntFeasTol=tol_int, FeasibilityTol=tol_feas, MIPGap=rgap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=tol_feas,
    mip_solver_drives=true,
    prim_cuts_assist=false,
    init_soc_one=false, init_soc_inf=false, init_exp=false, init_sdp_lin=false, init_sdp_soc=false,
    all_disagg=false, soc_disagg=false, sdp_eig=false
    ) end),

    # Paj Gurobi MOSEK scale, subp only, no init
    # NOTE feas tol is 1e-6
    "PAJ_Gurobi_MOSEK_scale_subponly_noinit" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=0, Threads=8, IntFeasTol=tol_int, FeasibilityTol=1e-6, MIPGap=tol_gap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=1e-6,
    prim_cuts_assist=false,
    init_soc_one=false, init_soc_inf=false, init_exp=false, init_sdp_lin=false, init_sdp_soc=false,
    ) end),

    "PAJ_Gurobi_MOSEK_msd_scale_subponly_noinit" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=1, Threads=8, IntFeasTol=tol_int, FeasibilityTol=1e-6, MIPGap=rgap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=1e-6,
    mip_solver_drives=true,
    prim_cuts_assist=false,
    init_soc_one=false, init_soc_inf=false, init_exp=false, init_sdp_lin=false, init_sdp_soc=false,
    ) end),

    # Paj Gurobi MOSEK noscale, subp only, no init
    # NOTE feas tol is 1e-6
    "PAJ_Gurobi_MOSEK_noscale_subponly_noinit" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=0, Threads=8, IntFeasTol=tol_int, FeasibilityTol=1e-6, MIPGap=tol_gap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=1e-6,
    scale_subp_cuts=false,
    prim_cuts_assist=false,
    init_soc_one=false, init_soc_inf=false, init_exp=false, init_sdp_lin=false, init_sdp_soc=false,
    ) end),

    "PAJ_Gurobi_MOSEK_msd_noscale_subponly_noinit" =>
    (["Gurobi","Mosek"], quote PajaritoSolver(
    mip_solver=Gurobi.GurobiSolver(OutputFlag=1, Threads=8, IntFeasTol=tol_int, FeasibilityTol=1e-6, MIPGap=rgap),
    cont_solver=MosekSolver(LOG=0, NUM_THREADS=8, MSK_DPAR_INTPNT_CO_TOL_REL_GAP=1e-10, MSK_DPAR_INTPNT_CO_TOL_PFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_DFEAS=1e-10, MSK_DPAR_INTPNT_CO_TOL_NEAR_REL=1e3),
    log_level=logl, timeout=tlim, rel_gap=rgap, prim_cut_feas_tol=1e-6,
    mip_solver_drives=true,
    scale_subp_cuts=false,
    prim_cuts_assist=false,
    init_soc_one=false, init_soc_inf=false, init_exp=false, init_sdp_lin=false, init_sdp_soc=false,
    ) end),
)

function getsolver(solvername, tlim, logl, rgap)
    packages, solvername = solvermap[solvername]
    for p in packages
        eval(parse("using $p"))
    end
    return eval(solvername)
end

@assert length(ARGS) >= 3

# Save process ID for runmeta.jl
open("mypid", "w") do fd
    print(fd, getpid())
end

# General options
logl = 3
rgap = 1e-5

# Pajarito MIP solver options
tol_int = 1e-9
tol_feas = 1e-8
tol_gap = 0.

solvername = ARGS[1]
tlim = parse(Float64, ARGS[2])
datafolder = ARGS[3]

# Force solver to compile on a small instance for the solver to avoid measuring compilation time, keep quiet
TT = STDOUT
solver = getsolver(solvername, 60., 3, rgap)
open("/dev/null", "w") do fd
    redirect_stdout(fd)
    instance = readcbfdata(joinpath(datafolder, "compile.cbf.gz"))
    solveprint(instance, solver)
end
redirect_stdout(TT)

instancename = ARGS[4]

# Interpret instance data as cbf
instance = readcbfdata(joinpath(datafolder, instancename))
solver = getsolver(solvername, tlim, logl, rgap)

println("#SOLVERNAME# $solvername")
println("#SOLVER# $solver")
println("#INSTANCE# $instancename")
println("#TIMELIMIT# $tlim")

# Attempt to solve, print solve info
timesolve = time()
try
    solveprint(instance, solver)
catch e
    println("#ERROR# $e")
end

max_stars_repo_name: chriscoey/PajaritoSupplement

max_issues_repo_name: chriscoey/PajaritoSupplement

max_forks_repo_name: chriscoey/PajaritoSupplement


text:
import Base: values, find

# when ############################

when(ta::TimeArray, period::Function, t::Int) =
    ta[find(period.(ta.timestamp) .== t)]
when(ta::TimeArray, period::Function, t::String) =
    ta[find(period.(ta.timestamp) .== t)]

# from, to ######################

from(ta::TimeArray{T, N, D}, d::D) where {T, N, D} =
    length(ta) == 0 ? ta :
        d < ta.timestamp[1] ? ta :
        d > ta.timestamp[end] ? ta[1:0] :
        ta[searchsortedfirst(ta.timestamp, d):end]

to(ta::TimeArray{T, N, D}, d::D) where {T, N, D} =
    length(ta) == 0 ? ta :
        d < ta.timestamp[1] ? ta[1:0] :
        d > ta.timestamp[end] ? ta :
        ta[1:searchsortedlast(ta.timestamp, d)]

###### find ##################

find(ta::TimeArray{Bool, 1}) = find(ta.values)

###### findwhen #################

findwhen(ta::TimeArray{Bool, 1}) = ta.timestamp[find(ta.values)]

###### head, tail ###########

function head(ta::TimeArray, n::Int=1)
    ncol          = length(ta.colnames)
    new_timestamp = ta.timestamp[1:n]
    new_values    = ta.values[1:n, 1:ncol]
    TimeArray(new_timestamp, new_values, ta.colnames, ta.meta)
end

function tail(ta::TimeArray, n::Int=1)
    ncol          = length(ta.colnames)
    tail_start = length(ta)-n+1
    new_timestamp = ta.timestamp[tail_start:end]
    new_values    = ta.values[tail_start:end, 1:ncol]
    TimeArray(new_timestamp, new_values, ta.colnames, ta.meta)
end

###### element wrapers ###########

timestamp(ta::TimeArray) = ta.timestamp
values(ta::TimeArray)    = ta.values
colnames(ta::TimeArray)  = ta.colnames
meta(ta::TimeArray)      = ta.meta


text:
module DianoiaML
    using LoopVectorization, LinearAlgebra
    BLAS.set_num_threads(1)

    include("./activation_function/relu.jl")
    using .ReLU
    export ReLU
    include("./activation_function/sigmoid.jl")
    using .Sigmoid
    export Sigmoid
    include("./activation_function/softmax.jl")
    using .Softmax
    export Softmax
    using .Softmax_CEL
    export Softmax_CEL
    include("./activation_function/tanh.jl")
    using .tanH
    export tanH
    include("./activation_function/none.jl")
    using .None
    export None

    include("./layer/dense.jl")
    using .DenseM:Dense
    export Dense
    include("./layer/conv2d.jl")
    using .Conv2DM:Conv2D
    export Conv2D
    include("./layer/maxpooling2d.jl")
    using .MaxPooling2DM:MaxPooling2D
    export MaxPooling2D
    include("./layer/upsampling2d.jl")
    using .UpSampling2DM:UpSampling2D
    export UpSampling2D
    include("./layer/flatten.jl")
    using .FlattenM:Flatten
    export Flatten
    include("./layer/constructive.jl")
    using .ConstructiveM:Constructive
    export Constructive
    include("./layer/dropout.jl")
    using .DropoutM:Dropout
    export Dropout
    include("./layer/tconv2d.jl")
    using .TConv2DM:TConv2D
    export TConv2D
    include("./layer/residual.jl")
    using .ResidualM:Residual
    export Residual

    include("./network/sequential.jl")
    using .sequential:Sequential, Hidden_Output_Layer
    export Sequential, Hidden_Output_Layer
    include("./network/gan.jl")
    using .gan:GAN
    export GAN
    include("./network/resnet.jl")
    using .resnet:ResNet
    export ResNet

    include("./loss_function/categorical_cross_entropy_loss.jl")
    using .Categorical_Cross_Entropy_Loss
    export Categorical_Cross_Entropy_Loss
    include("./loss_function/binary_cross_entropy_loss.jl")
    using .Binary_Cross_Entropy_Loss
    export Binary_Cross_Entropy_Loss
    include("./loss_function/quadratic_loss.jl")
    using .Quadratic_Loss
    export Quadratic_Loss
    include("./loss_function/mean_squared_error.jl")
    using .Mean_Squared_Error
    export Mean_Squared_Error
    include("./loss_function/absolute_loss.jl")
    using .Absolute_Loss
    export Absolute_Loss

    include("./monitor/absolute.jl")
    using .Absolute
    export Absolute
    include("./monitor/classification.jl")
    using .Classification
    export Classification

    include("./optimizer/minibatch_gd.jl")
    using .Minibatch_GD
    export Minibatch_GD
    include("./optimizer/adam.jl")
    using .Adam
    export Adam
    include("./optimizer/adabelief.jl")
    using .AdaBelief
    export AdaBelief
    include("./optimizer/sgd.jl")
    using .SGD
    export SGD
    include("./optimizer/ga.jl")
    using .GA
    export GA

    include("./tools/model_management.jl")
    export save_Sequential, load_Sequential
    export save_GAN, load_GAN
    include("./tools/oneHot.jl")
    export oneHot
end


text:
for N in [Float64, Rational{Int}, Float32]
    # random universe
    rand(Universe)

    U = Universe{N}(2)
    B = BallInf(ones(N, 2), N(1))

    # universe is an absorbing element for
    # - convex hull
    @test CH(B, U) == CH(U, B) == CH(U, U) == U
    cha = ConvexHullArray([B, N(2) * B, N(3) * B])
    @test CH(cha, U) == CH(U, cha) == U
    # - union
    @test B ∪ U == U ∪ B == U ∪ U == U
    # - Minkowski sum
#     @test B ⊕ U == U ⊕ B == U ⊕ U == U  # TODO problematic
#     msa = MinkowskiSumArray([B, N(2) * B, N(3) * B])
#     @test msa ⊕ U == U ⊕ msa == U

    # universe is a neutral element for
    # - intersection
    @test B ∩ U == U ∩ B == B
    @test U ∩ U == U
    ia = IntersectionArray([B, N(2) * B, N(3) * B])
    @test ia ∩ U == U ∩ ia == ia

    # dim
    @test dim(U) == 2

    # support function and support vector
    @test ρ(N[0, 1], U) == N(Inf)
    @test ρ(N[0, 0], U) == zero(N)
    @test σ(N[0, 1], U) == N[0, Inf]
    @test σ(N[-1, 2], U) == N[-Inf, Inf]

    # boundedness
    @test !isbounded(U)

    # membership
    @test N[0, 0] ∈ U
    @test_throws AssertionError N[0] ∈ U

    # emptiness
    @test !isempty(U)

    # universality
    @test isuniversal(U)

    # an_element
    @test an_element(U) ∈ U

    # norm/radius/diameter functions
    @test_throws ErrorException norm(U)
    @test_throws ErrorException radius(U)
    @test_throws ErrorException diameter(U)

    # translation
    @test translate(U, N[1, 2]) == U

    # concrete intersection
    @test intersection(B, U) == intersection(U, B) == B
    @test intersection(U, U) == U

    # intersection emptiness
    res, w = isdisjoint(B, U, true)
    @test !isdisjoint(B, U) && !res && w ∈ B && w ∈ U
    res, w = isdisjoint(U, B, true)
    @test !isdisjoint(U, B) && !res && w ∈ B && w ∈ U
    res, w = isdisjoint(U, U, true)
    @test !isdisjoint(U, U) && !res && w ∈ U
    E = EmptySet{N}()
    res, w = isdisjoint(E, U, true)
    @test isdisjoint(E, U) && res && w == N[]
    res, w = isdisjoint(U, E, true)
    @test isdisjoint(U, E) && res && w == N[]

    # subset
    res, w = ⊆(B, U, true)
    @test B ⊆ U && res && w == N[]
    # res, w = ⊆(U, B, true)  # currently not supported
    # @test !(U ⊆ B) && !res && w ∉ B
    @test !(U ⊆ B)
    res, w = ⊆(U, U, true)
    @test U ⊆ U && res && w == N[]
end

# default Float64 constructor
@test Universe(2) == Universe{Float64}(2)