using ArgParse
using Random
using StaticArrays
using NCDatasets
using Test
using DocStringExtensions
using LinearAlgebra
using DelimitedFiles
using Dierckx
using Printf

using ClimateMachine
using ClimateMachine.Atmos
using ClimateMachine.ArtifactWrappers
using ClimateMachine.Orientations
using ClimateMachine.ConfigTypes
using ClimateMachine.Diagnostics
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.GenericCallbacks
using ClimateMachine.ODESolvers
using ClimateMachine.Mesh.Filters
using ClimateMachine.TemperatureProfiles
using ClimateMachine.Thermodynamics
using ClimateMachine.TurbulenceClosures
using ClimateMachine.VariableTemplates
using ClimateMachine.BalanceLaws: vars_state, Prognostic

using CLIMAParameters
using CLIMAParameters.Planet
using CLIMAParameters.Atmos.Microphysics

struct LiquidParameterSet <: AbstractLiquidParameterSet end
struct IceParameterSet <: AbstractIceParameterSet end
struct RainParameterSet <: AbstractRainParameterSet end
struct SnowParameterSet <: AbstractSnowParameterSet end
struct MicropysicsParameterSet{L, I, R, S} <: AbstractMicrophysicsParameterSet
    liq::L
    ice::I
    rai::R
    sno::S
end
struct EarthParameterSet{M} <: AbstractEarthParameterSet
    microphys::M
end

const microphys = MicropysicsParameterSet(
    LiquidParameterSet(),
    IceParameterSet(),
    RainParameterSet(),
    SnowParameterSet(),
)
const param_set = EarthParameterSet(microphys)

"""
  Define initial conditions based on sounding data
"""
function init_squall_line!(problem, bl, state, aux, localgeo, t, args...)

    FT = eltype(state)

    spl_tinit, spl_qinit, spl_uinit, spl_vinit, spl_pinit = args[1]
    # interpolate data
    (x, y, z) = localgeo.coord
    data_t = FT(spl_tinit(z))
    data_q = FT(spl_qinit(z))
    data_u = FT(spl_uinit(z))
    data_v = FT(spl_vinit(z))
    data_p = FT(spl_pinit(z))
    u = data_u
    v = data_v
    w = FT(0)
    if z >= 14000
        data_q = FT(0)
    end
    θ_c = 3.0
    rx = 10000.0
    ry = 1500.0
    rz = 1500.0
    xc = 0.5 * (FT(0) + FT(0))
    yc = 0.5 * (FT(0) + FT(5000))
    zc = 2000.0
    cylinder_flg = 0.0
    r = sqrt(
        (x - xc)^2 / rx^2 +
        cylinder_flg * (y - yc)^2 / ry^2 +
        (z - zc)^2 / rz^2,
    )
    Δθ = 0.0
    if r <= 1.0
        Δθ = θ_c * (cospi(0.5 * r))^2
    end
    θ_liq = data_t + Δθ
    ts = PhaseNonEquil_pθq(param_set, data_p, θ_liq, PhasePartition(FT(data_q)))
    T = air_temperature(ts)
    ρ = air_density(ts)
    e_kin = FT(1 / 2) * FT((u^2 + v^2 + w^2))
    e_pot = gravitational_potential(bl.orientation, aux)
    E = ρ * total_energy(e_kin, e_pot, ts)

    state.ρ = ρ
    state.ρu = SVector(ρ * u, ρ * v, FT(0))
    state.ρe = E

    state.moisture.ρq_tot = ρ * data_q
    if bl.moisture isa NonEquilMoist
        state.moisture.ρq_liq = FT(0)
        state.moisture.ρq_ice = FT(0)
    end

    if bl.precipitation isa RainModel
        state.precipitation.ρq_rai = FT(0)
    end
    if bl.precipitation isa RainSnowModel
        state.precipitation.ρq_rai = FT(0)
        state.precipitation.ρq_sno = FT(0)
    end

    return nothing
end

"""
  Read the original squall sounding
"""
function read_sounding()

    # Artifact creation is not thread-safe
    #      https://github.com/JuliaLang/Pkg.jl/issues/1219
    # To avoid race conditions from multiple jobs running this
    # driver at the same time, we must store artifacts in a
    # separate folder.

    artifact_folder = mktempdir(@__DIR__; prefix = "artifacts_")

    soundings_dataset = ArtifactWrapper(
        joinpath(artifact_folder, "Artifacts.toml"),
        "soundings",
        ArtifactFile[ArtifactFile(
            url = "https://caltech.box.com/shared/static/rjnvt2dlw7etm1c7mmdfrkw5gnfds5lx.nc",
            filename = "sounding_gabersek.nc",
        ),],
    )
    data_folder = get_data_folder(soundings_dataset)
    fsounding = joinpath(data_folder, "sounding_gabersek.nc")
    sounding = Dataset(fsounding, "r")
    height = sounding["height [m]"][:]
    θ = sounding["theta [K]"][:]
    q_vap = sounding["qv [g kg⁻¹]"][:]
    u = sounding["u [m s⁻¹]"][:]
    v = sounding["v [m s⁻¹]"][:]
    p = sounding["pressure [Pa]"][:]
    close(sounding)
    return (height = height, θ = θ, q_vap = q_vap, u = u, v = v, p = p)
end

"""
  Get data from interpolated array onto vectors.
  Accepts data in 6 column format
"""
function spline_int()

    nt = read_sounding()

    # WARNING: Not all sounding data is formatted/scaled the same.
    spl_tinit = Spline1D(nt.height, nt.θ; k = 1)
    spl_qinit = Spline1D(nt.height, nt.q_vap .* 0.001; k = 1)
    spl_uinit = Spline1D(nt.height, nt.u; k = 1)
    spl_vinit = Spline1D(nt.height, nt.v; k = 1)
    spl_pinit = Spline1D(nt.height, nt.p; k = 1)
    return (
        t = spl_tinit,
        q = spl_qinit,
        u = spl_uinit,
        v = spl_vinit,
        p = spl_pinit,
    )
end

function config_squall_line(
    FT,
    N,
    resolution,
    xmax,
    ymax,
    zmax,
    xmin,
    ymin,
    moisture_model = "nonequilibrium",
    precipitation_model = "rainsnow",
)
    # Reference state
    nt = read_sounding()

    zinit = nt.height
    tinit = nt.θ
    qinit = nt.q_vap .* 0.001
    u_init = nt.u
    v_init = nt.v
    p_init = nt.p

    maxz = length(zinit)
    thinit = zeros(maxz)
    piinit = zeros(maxz)
    thinit[1] = tinit[1] / (1 + 0.61 * qinit[1])
    piinit[1] = 1
    for k in 2:maxz
        thinit[k] = tinit[k] / (1 + 0.61 * qinit[k])
        piinit[k] =
            piinit[k - 1] -
            9.81 / (1004 * 0.5 * (tinit[k] + tinit[k - 1])) *
            (zinit[k] - zinit[k - 1])
    end
    T_min = FT(thinit[maxz] * piinit[maxz])
    T_s = FT(thinit[1] * piinit[1])
    Γ_lapse = FT(9.81 / 1004)
    tvmax = T_s * (1 + 0.61 * qinit[1])
    deltatv = -(T_min - tvmax)
    tvmin = T_min * (1 + 0.61 * qinit[maxz])
    htv = 8000.0
    #T = DecayingTemperatureProfile(T_min, T_s, Γ_lapse)
    Tv = DecayingTemperatureProfile{FT}(param_set, tvmax, tvmin, htv)
    rel_hum = FT(0)
    ref_state = HydrostaticState(Tv, rel_hum)

    # Sponge
    c_sponge = FT(0.5)
    # Rayleigh damping
    u_relaxation = SVector(FT(0), FT(0), FT(0))
    zsponge = FT(15000.0)
    rayleigh_sponge =
        RayleighSponge(FT, zmax, zsponge, c_sponge, u_relaxation, 2)

    # Boundary conditions
    # SGS Filter constants
    C_smag = FT(0.18) # 0.21 for stable testing, 0.18 in practice
    C_drag = FT(0.0011)
    LHF = FT(50)
    SHF = FT(10)
    ics = init_squall_line!

    source = (Gravity(), rayleigh_sponge)

    # moisture model and its sources
    if moisture_model == "equilibrium"
        moisture = EquilMoist{FT}(; maxiter = 20, tolerance = FT(1))
    elseif moisture_model == "nonequilibrium"
        source = (source..., CreateClouds()...)
        moisture = NonEquilMoist()
    else
        @warn @sprintf(
            """
%s: unrecognized moisture_model in source terms, using the defaults""",
            moisture_model,
        )
        source = (source..., CreateClouds()...)
        moisture = NonEquilMoist()
    end

    # precipitation model and its sources
    if precipitation_model == "noprecipitation"
        precipitation = NoPrecipitation()
        source = (source..., RemovePrecipitation(true)...)
    elseif precipitation_model == "rain"
        source = (source..., WarmRain_1M()...)
        precipitation = RainModel()
    elseif precipitation_model == "rainsnow"
        source = (source..., RainSnow_1M()...)
        precipitation = RainSnowModel()
    else
        @warn @sprintf(
            """
%s: unrecognized precipitation_model in source terms, using the defaults""",
            precipitation_model,
        )
        source = (source..., RainSnow_1M()...)
        precipitation = RainSnowModel()
    end

    problem = AtmosProblem(
        boundaryconditions = (AtmosBC(), AtmosBC()),
        init_state_prognostic = ics,
    )

    model = AtmosModel{FT}(
        AtmosLESConfigType,
        param_set;
        problem = problem,
        ref_state = ref_state,
        moisture = moisture,
        precipitation = precipitation,
        turbulence = SmagorinskyLilly{FT}(C_smag),
        source = source,
    )

    ode_solver = ClimateMachine.IMEXSolverType()

    config = ClimateMachine.AtmosLESConfiguration(
        "Squall_line",
        N,
        resolution,
        xmax,
        ymax,
        zmax,
        param_set,
        init_squall_line!,
        xmin = xmin,
        ymin = ymin,
        solver_type = ode_solver,
        model = model,
        periodicity = (true, true, false),
        boundary = ((2, 2), (2, 2), (1, 2)),
    )
    return config
end

function config_diagnostics(driver_config, boundaries, resolution)

    interval = "3600ssecs"

    dgngrp_profiles = setup_atmos_default_diagnostics(
        AtmosLESConfigType(),
        interval,
        driver_config.name,
    )

    interpol = ClimateMachine.InterpolationConfiguration(
        driver_config,
        boundaries,
        resolution,
    )
    dgngrp_state = setup_dump_state_diagnostics(
        AtmosLESConfigType(),
        interval,
        driver_config.name,
        interpol = interpol,
    )
    dgngrp_aux = setup_dump_aux_diagnostics(
        AtmosLESConfigType(),
        interval,
        driver_config.name,
        interpol = interpol,
    )

    return ClimateMachine.DiagnosticsConfiguration([
        dgngrp_profiles,
        dgngrp_state,
        dgngrp_aux,
    ])
end

function main()
    # TODO: this will move to the future namelist functionality
    squall_args = ArgParseSettings(autofix_names = true)
    add_arg_group!(squall_args, "SQUALL_LINE")
    @add_arg_table! squall_args begin
        "--moisture-model"
        help = "specify cloud condensate model"
        metavar = "equilibrium|nonequilibrium"
        arg_type = String
        default = "nonequilibrium"
        "--precipitation-model"
        help = "specify precipitation model"
        metavar = "noprecipitation|rain|rainsnow"
        arg_type = String
        default = "rainsnow"
        "--check-asserts"
        help = "should asserts be checked at the end of the simulation"
        metavar = "yes|no"
        arg_type = String
        default = "no"
    end

    cl_args =
        ClimateMachine.init(parse_clargs = true, custom_clargs = squall_args)
    moisture_model = cl_args["moisture_model"]
    precipitation_model = cl_args["precipitation_model"]
    check_asserts = cl_args["check_asserts"]

    FT = Float64

    # DG polynomial order
    N = 4

    # Domain resolution and size
    Δx = FT(250)
    Δy = FT(1000)
    Δv = FT(200)
    resolution = (Δx, Δy, Δv)

    xmax = FT(30000)
    ymax = FT(5000)
    zmax = FT(24000)
    xmin = FT(-30000)
    ymin = FT(0)
    # for diagnostics
    boundaries = [
        xmin ymin FT(0)
        xmax ymax zmax
    ]

    t0 = FT(0)
    timeend = FT(9000)
    spl_tinit, spl_qinit, spl_uinit, spl_vinit, spl_pinit = spline_int()
    Cmax = FT(0.4)

    # driver, solver and diagnostics configs
    driver_config = config_squall_line(
        FT,
        N,
        resolution,
        xmax,
        ymax,
        zmax,
        xmin,
        ymin,
        moisture_model,
        precipitation_model,
    )
    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        (spl_tinit, spl_qinit, spl_uinit, spl_vinit, spl_pinit);
        init_on_cpu = true,
        Courant_number = Cmax,
    )
    dgn_config = config_diagnostics(driver_config, boundaries, resolution)

    filter_vars = ("moisture.ρq_tot",)
    if moisture_model == "nonequilibrium"
        filter_vars = (filter_vars..., "moisture.ρq_liq", "moisture.ρq_ice")
    end
    if precipitation_model == "rain"
        filter_vars = (filter_vars..., "precipitation.ρq_rai")
    end
    if precipitation_model == "rainsnow"
        filter_vars =
            (filter_vars..., "precipitation.ρq_rai", "precipitation.ρq_rai")
    end

    cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(1) do
        Filters.apply!(
            solver_config.Q,
            filter_vars,
            solver_config.dg.grid,
            TMARFilter(),
        )
        nothing
    end

    result = ClimateMachine.invoke!(
        solver_config;
        diagnostics_config = dgn_config,
        user_callbacks = (cbtmarfilter,),
        check_euclidean_distance = true,
    )

    if check_asserts == "yes"

        m = driver_config.bl
        Q = solver_config.Q
        ρ_ind = varsindex(vars_state(m, Prognostic(), FT), :ρ)

        if moisture_model == "equilibrium"
            ρq_tot_ind =
                varsindex(vars_state(m, Prognostic(), FT), :moisture, :ρq_tot)

            min_q_tot = minimum(abs.(
                Array(Q[:, ρq_tot_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_tot = maximum(abs.(
                Array(Q[:, ρq_tot_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))

            @info(min_q_tot, max_q_tot)

            # test that moisture exists and is not NaN
            @test !isnan(max_q_tot)

            # test that there is reasonable amount of moisture
            @test abs(max_q_tot) > FT(1e-2)
        end
        if moisture_model == "nonequilibrium"
            ρq_tot_ind =
                varsindex(vars_state(m, Prognostic(), FT), :moisture, :ρq_tot)
            ρq_liq_ind =
                varsindex(vars_state(m, Prognostic(), FT), :moisture, :ρq_liq)
            ρq_ice_ind =
                varsindex(vars_state(m, Prognostic(), FT), :moisture, :ρq_ice)

            min_q_tot = minimum(abs.(
                Array(Q[:, ρq_tot_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_tot = maximum(abs.(
                Array(Q[:, ρq_tot_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            min_q_liq = minimum(abs.(
                Array(Q[:, ρq_liq_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_liq = maximum(abs.(
                Array(Q[:, ρq_liq_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            min_q_ice = minimum(abs.(
                Array(Q[:, ρq_ice_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_ice = maximum(abs.(
                Array(Q[:, ρq_ice_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            @info(min_q_tot, max_q_tot)
            @info(min_q_liq, max_q_liq)
            @info(min_q_ice, max_q_ice)

            # test that moisture exists and is not NaN
            @test !isnan(max_q_tot)
            @test !isnan(max_q_liq)
            @test !isnan(max_q_ice)

            # test that there is reasonable amount of moisture
            @test abs(max_q_tot) > FT(1e-2)
            @test abs(max_q_liq) > FT(1e-3)
            @test abs(max_q_ice) > FT(1e-3)
        end

        if precipitation_model == "rain"
            ρq_rai_ind = varsindex(
                vars_state(m, Prognostic(), FT),
                :precipitation,
                :ρq_rai,
            )

            min_q_rai = minimum(abs.(
                Array(Q[:, ρq_rai_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_rai = maximum(abs.(
                Array(Q[:, ρq_rai_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))

            @info(min_q_rai, max_q_rai)

            # test that rain variable exists and is not NaN
            @test !isnan(max_q_rai)

            # test that there is reasonable amount of rain water...
            @test abs(max_q_rai) > FT(1e-4)
        end
        if precipitation_model == "rainsnow"
            ρq_rai_ind = varsindex(
                vars_state(m, Prognostic(), FT),
                :precipitation,
                :ρq_rai,
            )
            ρq_sno_ind = varsindex(
                vars_state(m, Prognostic(), FT),
                :precipitation,
                :ρq_sno,
            )
            min_q_rai = minimum(abs.(
                Array(Q[:, ρq_rai_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_rai = maximum(abs.(
                Array(Q[:, ρq_rai_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            min_q_sno = minimum(abs.(
                Array(Q[:, ρq_sno_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))
            max_q_sno = maximum(abs.(
                Array(Q[:, ρq_sno_ind, :]) ./ Array(Q[:, ρ_ind, :]),
            ))

            @info(min_q_rai, max_q_rai)
            @info(min_q_sno, max_q_sno)

            # test that rain and snow variables exists and are not NaN
            @test !isnan(max_q_rai)
            @test !isnan(max_q_sno)

            # test that there is reasonable amount of precipitation
            @test abs(max_q_rai) > FT(1e-4)
            @test abs(max_q_sno) > FT(1e-6)
        end
    end
end

main()


isdefined(Base, :__precompile__) && __precompile__()

"""
Module for constructing self-contained power system objects.
"""
module PowerSystems

#################################################################################
# Exports

export System
export ConcreteSystem
export Bus
export LoadZones

export PowerSystemType
export Component
export Device
export Branch
export Injection
export Line
export MonitoredLine
export DCLine
export HVDCLine
export VSCDCLine
export Transformer2W
export TapTransformer
export PhaseShiftingTransformer

export Forecast
export Deterministic
export Scenarios
export Probabilistic

export Generator
export GenClasses

export HydroGen
export HydroFix
export HydroCurtailment
export HydroStorage
export TechHydro
export EconHydro

export RenewableGen
export TechRenewable
export EconRenewable
export RenewableFix
export RenewableCurtailment
export RenewableFullDispatch

export ThermalGen
export TechThermal
export EconThermal
export ThermalDispatch
export ThermalGenSeason

export ElectricLoad
export StaticLoad
export PowerLoad
export PowerLoadPF
export FixedAdmittance
export ControllableLoad
export InterruptibleLoad

export Storage
export GenericBattery

export Service
export Reserve
export ProportionalReserve
export StaticReserve
export Transfer

export parsestandardfiles
export parse_file
export ps_dict2ps_struct
export assign_ts_data
export read_data_files
export validate
export get_components
export get_mixed_components
export get_component_counts
export show_component_counts

#################################################################################
# Imports

import SparseArrays
import AxisArrays
import LinearAlgebra: LAPACK.getri!
import LinearAlgebra: LAPACK.getrf!
import LinearAlgebra: BLAS.gemm
import LinearAlgebra
import Dates
import TimeSeries
import DataFrames
import JSON
import CSV
import YAML

#################################################################################
# Includes

# supertype for all PowerSystems types
abstract type PowerSystemType end
abstract type Component <: PowerSystemType end
# supertype for "devices" (bus, line, etc.)
abstract type Device <: Component end
abstract type Injection <: Device end
# supertype for generation technologies (thermal, renewable, etc.)
abstract type TechnicalParams <: Component end

include("common.jl")

# Include utilities
include("utils/utils.jl")
include("utils/logging.jl")
include("utils/IO/base_checks.jl")

# PowerSystems models
include("models/topological_elements.jl")
include("models/forecasts.jl")
include("models/branches.jl")
#include("models/network.jl")

# Static types
include("models/generation.jl")
include("models/storage.jl")
include("models/loads.jl")
include("models/services.jl")

# Include Parsing files
include("parsers/pm_io.jl")
include("parsers/im_io.jl")
include("parsers/dict_to_struct.jl")
include("parsers/standardfiles_parser.jl")
include("parsers/cdm_parser.jl")
include("parsers/forecast_parser.jl")
include("parsers/pm2ps_parser.jl")

#Data Checks
include("utils/IO/system_checks.jl")
include("utils/IO/branchdata_checks.jl")

# Definitions of System
include("base.jl")
include("validation/powersystem.jl")

# Better printing
include("utils/print.jl")
include("utils/lodf_calculations.jl")

# Download test data
include("utils/data.jl")
import .UtilsData: TestData

end # module


# Copyright (c) 2015, Intel Corporation
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     * Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of Intel Corporation nor the names of its contributors
#       may be used to endorse or promote products derived from this software
#       without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""
Convert from an internal Latte tiled_loop to a normal Julia :for expression
"""
function unpack_tiled_loop(node::Expr)
    @assert node.head == :for
    range = node.args[1].args[2]
    node.args[1].args[2] = range.args[2]
    node
end

"""
Check if `node` is an internal Latte tiled_loop
"""
function is_tiled_loop(node)
    if !(isa(node, Expr) && node.head == :for)
        return false
    end
    range = node.args[1].args[2]
    isa(range, Expr) && range.head == :call && range.args[1] == :tiled_loop
end

"""
TODO: doc
"""
function get_tile_fusion_factor_forward(expr)
    return expr.args[3]
end

"""
TODO: doc
"""
function get_tile_fusion_factor_backward(expr)
    return expr.args[4]
end

"""
Replace indexing expressions with `tile_var` with the proper tiled expression
When `inputs` are not copied, the neuron transformer appends :NOTILE to the
index expression to force Latte not to tile the expression.
"""
function update_tile_var(node, tile_var, index, top_level, read)
    if isa(node, Expr) && node.head == :ref && contains(string(node.args[1]), "inputs")
        if node.args[end] != :NOTILE
            for i in 2:length(node.args)-1
                if node.args[i] == tile_var
                    node.args[i] = :(($(node.args[i]) - 1) % TILE_SIZE + 1)
                end
            end
            node.args[end] = :(_omp_get_thread_num() + 1)
        end
        return node
    end
    ASTWALK_RECURSE
end

"""
TODO: doc
"""
function inner_loop_tiler(node, cbdata, index, top_level, read)
    if isa(node, Expr) && node.head == :for && isa(get_loopvar(node), Symbol)
        loopvar = string(get_loopvar(node))
        if contains(loopvar, "_neuron_index_") && parse(Int, loopvar[end]) == 2
            loop_idx = parse(Int, loopvar[end])
            loop_range = node.args[1].args[2]
            loop_end = loop_range.args[2]
            unshift!(cbdata[:loop_lengths], loop_end)
            unshift!(cbdata[:loop_vars], symbol(loopvar, "_tile_idx"))
            loop_range.args[1] = :((($(symbol(loopvar, "_tile_idx")) - 1) * TILE_SIZE + 1))
            loop_range.args[2] = :(min($(symbol(loopvar, "_tile_idx")) * TILE_SIZE, $loop_end))
            node.args[2] = AstWalk(node.args[2], update_tile_var, get_loopvar(node))
            return node
        end
    end
    ASTWALK_RECURSE
end

global unique_id = 0

"""
TODO: doc
"""
function get_inner_loop_tiler(tile_fusion_factors)
    global unique_id += 1
    function tile_inner_loops(statement)
        if isa(statement, Expr) && statement.head == :for
            loopvar = string(get_loopvar(statement))
            if contains(loopvar, "_neuron_index_") && parse(Int, loopvar[end]) >= 4
                loop_lengths = []
                loop_vars = []
                cbdata = Dict(:loop_lengths => loop_lengths,
                              :TILE_SIZE => TILE_SIZE,
                              :loop_vars => loop_vars)
                statement = AstWalk(statement, inner_loop_tiler, cbdata)
                body = statement.args[2]
                num_tiles = []
                for len in loop_lengths
                    n_tiles = div(len, TILE_SIZE)
                    if len % TILE_SIZE != 0
                        n_tiles += 1
                    end
                    push!(num_tiles, n_tiles)
                end
                tile_loop = quote
                    $(body.args...)
                end
                for (loopvar, n_tiles, factor) in zip(loop_vars, num_tiles, tile_fusion_factors)
                    tile_loop = quote
                        for $loopvar = tiled_loop(1:$n_tiles, $(factor...), $unique_id)
                            $(tile_loop.args...)
                        end
                    end
                end
                statement.args[2] = tile_loop
            end
        end
        statement
    end
end

"""
TODO: doc
"""
function tile_size_inliner(node, cbdata, index, top_level, read)
    if isa(node, Expr)
        if node.head == :(*=) && node.args[1] == :TILE_SIZE
            cbdata[1] *= node.args[2]
            return ASTWALK_REMOVE
        elseif node.head == :(/=) && node.args[1] == :TILE_SIZE
            cbdata[1] /= node.args[2]
            return ASTWALK_REMOVE
        elseif is_tiled_loop(node)
            node = unpack_tiled_loop(node)
            node.args[2] = AstWalk(node.args[2], tile_size_inliner, cbdata)
            node.args[1] = AstWalk(node.args[1], tile_size_inliner, cbdata)
            return node
        elseif node.head == :ref && node.args[end] == :NOTILE
            pop!(node.args)
            return node
        elseif node.head in [:loophead, :parallel_loophead, :loopend, :parallel_loopend]
            return node
        end
    elseif isa(node, Symbol) && node == :TILE_SIZE
        return cbdata[1]
    end
    ASTWALK_RECURSE
end

function inline_tile_size(statements)
    map((x) -> AstWalk(x, tile_size_inliner, [TILE_SIZE]), statements)
end

"""
Tile loop variables use `_tile_idx` loopvars.
"""
function get_tile_loops(node, cbdata, index, top_level, read)
    if isa(node, Expr) && node.head == :for && isa(get_loopvar(node), Symbol)
        loopvar = string(get_loopvar(node))
        if contains(loopvar, "_tile_idx")
            push!(cbdata, node)
        end
    elseif isa(node, Expr) && node.head in [:loophead, :parallel_loophead, :loopend, :parallel_loopend]
        return node
    end
    ASTWALK_RECURSE
end


using DataStructures
using Statistics

"""
PS with PF
----------

Prioritized Sweeping with: ParticleFilterJump for the Tsas' estimation
and Perfect Integration for Rsas' estimation.

See learner/prioritizedsweepinglearner.jl for the core implementation and more info.

Note 1
------
With particle filtering we need to run multiple times with different seeds.
As final LL we report the MEAN across seeds.
For signals/results we keep track of the ones of all seeds
and in the end report the one with the MAXIMUM LL.
Note 2
------
get_all_puzzles(), get_all_episodes_in_puzzle(), get_data_by_puzzle_episode()
are in rl_data_toolsforalgos.jl
"""
function getLL_Prio_Sweep_TParticleFilterJump(data::Array{Int,2},
                                                gamma::Float64, temp::Float64,
                                                surpriseprobability::Float64,
                                                stochasticity::Float64,
                                                nr_upd_cycles::Union{Int, Float64},
                                                nparticles::Union{Int, Float64};
                                                particlefilterseeds = [1,2,3],
                                                init_R = 0.,
                                                do_debug_print = false
                                                )
    nr_upd_cycles = converttoInt(nr_upd_cycles)
    nparticles = converttoInt(nparticles)
    data_stats, nr_states, nr_actions, nr_samples = get_data_stats_and_dimensions(data)
    if data_stats.nr_subjects != 1
        error("Invalid argument: getLL_Prio_Sweep_TParticleFilter supports only one subject")
    end
    all_puzzles = get_all_puzzles(data)

    LL_seeds = zeros(Float64, length(particlefilterseeds))
    action_selection_probs_seeds = Array{Array{Float64, 1}, 1}(undef, length(particlefilterseeds))
    [action_selection_probs_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    # signals:
    Surprise_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [Surprise_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    γ_surprise_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [γ_surprise_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    QSA_seeds = Array{Any, 1}(undef, length(particlefilterseeds))
    Tsas_seeds = Array{Any, 1}(undef, length(particlefilterseeds))
    learner = []
    for i_seed in 1:length(particlefilterseeds)
        action_counter = 0
        for i_puzzle in all_puzzles
            # Puzzle = independent block of the experiment.
            # The code was made for general purposes, in this experiment we
            # have only ONE puzzle.
            learner = PrioritizedSweeping(nr_states = nr_states,
                                        nr_actions = nr_actions,
                                        γ = gamma, init_R = init_R,
                                        nr_upd_cycles = nr_upd_cycles,
                                        Testimatetype = TParticleFilterJump,
                                        nparticles = nparticles,
                                        surpriseprobability = surpriseprobability,
                                        stochasticity = stochasticity,
                                        seedlearner = particlefilterseeds[i_seed],
                                        do_debug_print = do_debug_print
                                        )
            all_epi_in_puzzle = get_all_episodes_in_puzzle(data, i_puzzle)
            for i_episode in all_epi_in_puzzle
                episode_data = get_data_by_puzzle_episode(data, i_puzzle, i_episode)
                path_length = size(episode_data)[1]  # = T
                if do_debug_print
                    println("i_puzzle:$i_puzzle | i_episode:$i_episode | path_length:$path_length")
                end
                for t in 1:path_length
                    action_counter += 1
                    state = episode_data[t, Int(c_state)]
                    action = episode_data[t, Int(c_action)]
                    reward = episode_data[t, Int(c_reward)]
                    next_state = episode_data[t, Int(c_nextState)]
                    done = reward>0 ? true : false
                    # println("t:$t | s:$state | a:$action | R:$reward | s':$next_state | done: $done")
                    if do_debug_print
                        println("t:$t | s:$state | a:$action | R:$reward | s':$next_state | done: $done")
                    end

                    p_actions = getSoftmaxProbs(learner.QSA[state,:], temp, [0., 0.])
                    action_selection_probs_seeds[i_seed][action_counter] = copy(p_actions[action])
                    LL_seeds[i_seed] += log(p_actions[action])
                    if do_debug_print
                        println("p(s,a)=$p_actions, action=$action")
                    end

                    # Prioritized Sweeping with Particle Filtering
                    update!(learner, state, action, next_state, reward, done)

                    Surprise_seeds[i_seed][action_counter] = copy(learner.Testimate.Sgm[1])
                    γ_surprise_seeds[i_seed][action_counter] = copy(learner.Testimate.γ_surprise[1])
                    QSA_seeds[i_seed] = copy(learner.QSA)
                    Tsas_seeds[i_seed]= deepcopy(learner.Testimate.Ps1a0s0)

                end # end t
            end # end i_episode
        end # end i_puzzle
    end # end i_seed

    LL = get_meanLL_acrossseeds(LL_seeds)

    action_selection_probs_cat = hcat(action_selection_probs_seeds...)
    action_selection_probs = mean(action_selection_probs_cat, dims=2)[:]
    trial_by_trial_p_sa = (action_selection_probs)

    # For signals: return the ones corresponding to the Max LL across seeds
    LL_max, seed_idx = findmax(LL_seeds)
    trial_by_trial_p_sa_max = (action_selection_probs_seeds[seed_idx])
    if do_debug_print
        println(finalQSA_record)
        println("action_selection_probs: $action_selection_probs")
        println("LL_mean: $LL_mean")
        println("LL_max: $LL_max")
        println("trial_by_trial_p_sa: $trial_by_trial_p_sa")
    end
    algo = "prio_sweep_particlefilter"
    param_v = [nr_upd_cycles, nparticles, gamma, temp, surpriseprobability, stochasticity]
    param_str = "nr_upd_cycles=$nr_upd_cycles, nparticles=$nparticles,
                gamma=$(round(gamma,digits=4)), temp=$(round(temp,digits=4)),
                surpriseprobability=$(round(surpriseprobability,digits=4)),
                stochasticity=$(round(stochasticity,digits=4))"
    results_dict = Dict{String,Any}("final_QSA" => QSA_seeds[seed_idx],
                                    "final_Tsas" => Tsas_seeds[seed_idx],
                                    "Surprise" => Surprise_seeds[seed_idx],
                                    "gamma_surprise" => γ_surprise_seeds[seed_idx],
                                    "nr_upd_cycles" => nr_upd_cycles,
                                    "nparticles" => nparticles,
                                    "particlefilterseeds" => particlefilterseeds,
                                    "trial_by_trial_p_sa_max" => trial_by_trial_p_sa_max)

    signals = Step_by_step_signals(data_stats.all_subject_ids[1],
    [], #RPE
    [], #SPE
    trial_by_trial_p_sa,
    LL,
    data_stats,
    algo,
    param_v,
    param_str,
    99, #biased_action_idx, dummy value
    results_dict)
    return LL, signals
end


dir = "../OCTHaGOn_benchmarks/gams/global/"
filenames = readdir(dir)
GAMS_DIR = dir

include(OCTHaGOn.PROJECT_ROOT * "/test/tools/gams.jl")

pd = DataFrame()
n_max = 100
pd = CSV.read(dir * "problem_stats.csv", DataFrame)
pd = pd[pd.all_bounded .>= 0.5, :] # Use only benchmarks that are bounded...
pd = pd[pd.n_vars .<= n_max, :]    # and have fewer than n_max variables.

function gs_and_time!(m::Union{JuMP.Model, GlobalModel}, n_tries = 10)
    t1 = time()
    if m isa JuMP.Model
        optimize!(m);
    else
        set_param(m, :sample_coeff, 400)
        ct = 0
        while ct < n_tries
            try
                globalsolve!(m)
                ct += 1e5
            catch
                ct += 1
            end
        end
        if ct == n_tries
            throw(OCTHaGOnException("Globalsolve failed for $(m.name)."))
        end
    end
    println("Model solution time: " * string(time()-t1) * " seconds.")
    return
end

# Benchmarking loop
gms = Dict()
results = DataFrame(:name => [], :oct_feas => [], :oct_cost => [], :oct_time => [])
for filename in pd.name
    @info "Trying $(filename)"
    gm = GAMS_to_GlobalModel(dir, filename * ".gms")
    set_optimizer(gm, SOLVER_SILENT)
    gms[gm.name] = gm
    t1 = time()
    try
        gs_and_time!(gm)
        push!(results, [gm.name,  is_feasible(gm), gm.cost[end], time() - t1])
    catch
        if !isempty(gm.cost)
            push!(results, [gm.name,  is_feasible(gm), gm.cost[end], time() - t1])
        else
            push!(results, [gm.name,  false, Inf, time() - t1])
        end
    end
end
results.name = [na[1:end-4] for na in results.name] # Trimming the .gms string

# Updating and rewriting data
orig_data = CSV.read(dir * "problem_stats.csv", DataFrame)
idxs = [idx for idx in 1:size(orig_data, 1) if orig_data.name[idx] in results.name]
orig_data[idxs,[:oct_feas, :oct_cost, :oct_time]] = results[:,[:oct_feas, :oct_cost, :oct_time]]
CSV.write(dir * "problem_stats.csv", orig_data)


module ICapeThermoPropertyRoutine

    export calcandgetlnphi!,calcsinglephaseprop,calctwophaseprop,checksinglephasepropspec!
    export checktwophasepropspec!,getsinglephaseproplist!,gettwophaseproplist!

    """
      This method is used to calculate the natural logarithm of the fugacity coefficients (and
      optionally their derivatives) in a single Phase mixture. The values of temperature, pressure
      and composition are specified in the argument list and the results are also returned through
      the argument list.
    """
    function calcandgetlnphi!(
        #= [in] =# phaseLabel::AbstractString,
        #= [in] =# temperature::Float64,
        #= [in] =# pressure::Float64,
        #= [in] =# moleNumbers::Vector{Float64},
        #= [in] =# fFlags::Int32,
        #= [out][in] =# lnPhi!::Vector{Float64},
        #= [out][in] =# lnPhiDT!::Vector{Float64},
        #= [out][in] =# lnPhiDP!::Vector{Float64},
        #= [out][in] =# lnPhiDn!::Vector{Float64})
		end

    """
      CalcSinglePhaseProp is used to calculate properties and property derivatives of a mixture in
      a single Phase at the current values of temperature, pressure and composition set in the
      Material Object. CalcSinglePhaseProp does not perform phase Equilibrium Calculations.
    """
    function calcsinglephaseprop(
        #= [in] =# props::Vector{AbstractString},
        #= [in] =# phaseLabel::AbstractString)
		end

    """
      CalcTwoPhaseProp is used to calculate mixture properties and property derivatives that depend on
      two Phases at the current values of temperature, pressure and composition set in the Material Object.
      It does not perform Equilibrium Calculations.
    """
    function calctwophaseprop(
        #= [in] =# props::Vector{AbstractString} ,
        #= [in] =# phaseLabels::Vector{AbstractString} )
		end

    """
      Checks whether it is possible to calculate a property with the CalcSinglePhaseProp method for a given Phase.
    """
    function checksinglephasepropspec!(
        #= [in] =# property::AbstractString,
        #= [in] =# phaseLabel::AbstractString,
        #= [retval][out] =# valid::Bool)
		end

    """
      Checks whether it is possible to calculate a property with the CalcTwoPhaseProp method for a given set of Phases.
    """
    function checktwophasepropspec!(
        #= [in] =# property::AbstractString,
        #= [in] =# phaseLabels::Vector{AbstractString},
        #= [retval][out] =# valid::Bool)
		end

    """
      Returns the list of supported non-constant single-phase Physical Properties.
    """
    function getsinglephaseproplist!(
        #= [retval][out] =# props::Vector{AbstractString})
		end

    """
      Returns the list of supported non-constant two-phase properties.
    """
    function gettwophaseproplist!(
        #= [retval][out] =# props::Vector{AbstractString})
		end

end


using Distributions
using FileIO
using GeometryBasics
using ImageMagick
using Images
using OffsetArrays
using Plots
using Random
using RecipesPipeline
using RDatasets
using SparseArrays
using StaticArrays
using Statistics
using StatsPlots
using Test
using TestImages

# makie_test_dir = joinpath(@__DIR__, "test_makie")
# mkpath(makie_test_dir)

@testset "RecipesPipeline.jl" begin
    # @testset "Makie integration" begin cd(makie_test_dir) do
    #     include("makie.jl")
    # end end
    @testset "Plots tests" begin
    for i in eachindex(Plots._examples)
        if i ∉ Plots._backend_skips[:gr]
            @test Plots.test_examples(:gr, i, disp=false) isa Plots.Plot
        end
    end
    end
end


struct Order_Report
    converged_K::Vector{Float64}
    converged_lh::Vector{Float64}
    failed_K::Vector{Float64}
    failed_lh::Vector{Float64}
end

struct Partition_Report
    partition_id::String
    naive_lh::Float64
    orddict::Dict{Integer,Order_Report}
    best_model::Tuple{Chain_ID,BHMM}
    best_repset::Vector{Chain_ID}
end

function Base.show(io::IO, report::Partition_Report)
    printstyled(io, "BACKGROUND HMM TRAINING REPORT\n", bold=true)
    printstyled(io, "Genome partition id: $(report.partition_id)\n", bold=true, color=:magenta)

    for (order, ordreport) in report.orddict
        if length(ordreport.failed_lh)>0
            ordplot=scatterplot(ordreport.converged_K, ordreport.converged_lh; name="Converged", title="Order $order HMMs", xlabel="K# States", ylabel="P(O|θ)", color=:green, xlim=(1,maximum(ordreport.converged_K)), ylim=(floor(min(minimum(ordreport.converged_lh),minimum(ordreport.failed_lh),report.naive_lh),sigdigits=2),ceil(max(maximum(ordreport.converged_lh),maximum(ordreport.failed_lh),report.naive_lh),sigdigits=2)))
            length(ordreport.failed_K)>=1 && scatterplot!(ordplot, ordreport.failed_K, ordreport.failed_lh; name="Unconverged",color=:red)
            lineplot!(ordplot, [report.naive_lh for i in 1:maximum(ordreport.converged_K)], name="Naive",color=:magenta)
            show(ordplot)
            println()
        else
            ordplot=scatterplot(ordreport.converged_K, ordreport.converged_lh; name="Converged", title="Order $order HMMs", xlabel="K# States", ylabel="P(O|θ)", color=:green, xlim=(1,maximum(ordreport.converged_K)), ylim=(floor(min(minimum(ordreport.converged_lh),report.naive_lh),sigdigits=2),ceil(max(maximum(ordreport.converged_lh),report.naive_lh),sigdigits=2)))
            lineplot!(ordplot, [report.naive_lh for i in 1:maximum(ordreport.converged_K)], name="Naive",color=:magenta)
            show(ordplot)
            println()
        end
    end
end

function report_partitions(chain_reports::Dict{Chain_ID,Chain_Report})
    partitions=Vector{String}()
    orders=Vector{Integer}()
    Ks=Vector{Integer}()
    reps=Vector{Integer}()
    for id in keys(chain_reports)
        !in(id.obs_id, partitions) && push!(partitions, id.obs_id)
        !in(id.order, orders) && push!(orders, id.order)
        !in(id.K, Ks) && push!(Ks, id.K)
        !in(id.replicate, reps) && push!(reps, id.replicate)
    end

    reports=Vector{Partition_Report}()
    for partition in partitions
        best_lh,best_rep=-Inf,Chain_ID("empty",1,0,1)
        naive_lh=1.
        orddict=Dict{Integer,Order_Report}()
        for order in orders
            conv_statevec=Vector{Float64}()
            fail_statevec=Vector{Float64}()
            conv_lh_vec=Vector{Float64}()
            fail_lh_vec=Vector{Float64}()
            for (id,report) in chain_reports
                if id.obs_id==partition && id.order==order
                    report.test_lh > best_lh && (best_lh=report.test_lh;best_rep=id)
                    naive_lh==1. && (naive_lh=chain_reports[id].naive_lh)
                    chain_reports[id].converged ? (push!(conv_statevec,id.K);
                    push!(conv_lh_vec,chain_reports[id].test_lh)) : (push!(fail_statevec,id.K);
                    push!(fail_lh_vec,chain_reports[id].test_lh))
                end
            end
            orddict[order]=Order_Report(conv_statevec,conv_lh_vec,fail_statevec,fail_lh_vec)
        end
        best_model=(best_rep,chain_reports[best_rep].final_hmm)
        best_repset=[Chain_ID(partition, best_rep.K, best_rep.order, rep) for rep in reps]
        push!(reports,Partition_Report(partition,naive_lh,orddict,best_model,best_repset))
    end
    return reports
end


module Atum

using Reexport
using StaticArrays
@reexport using Bennu

using KernelAbstractions
using KernelAbstractions.Extras: @unroll
using CUDAKernels
using UnPack

include("balancelaw.jl")
include("numericalfluxes.jl")
include("dgsem.jl")
include("esdgsem.jl")
include("odesolvers.jl")

include("balancelaws/advection.jl")
include("balancelaws/shallow_water.jl")
include("balancelaws/euler.jl")
include("balancelaws/euler_gravity.jl")

end # module


import Random.shuffle!
@testset "fsm_loopback_win.jl" begin
    DEFAULT_CAPACITY = 65000
    TIMEOUT_DFLT = 1000
    NREPS = 32
    MAX_PAYLOAD_SIZE = 1460
    @testset "loop segments back in a different order" begin
        for i in 0:NREPS
            rx_offset = WrappingInt32(rand(UInt32))
            conn = fsm_in_established(rx_offset - 1, rx_offset - 1; cap=DEFAULT_CAPACITY)
            send_ack!(conn, rx_offset, rx_offset, 65000)
            d = rand(UInt8(1):UInt8(127), DEFAULT_CAPACITY) |> pointer |> unsafe_string
            sendoff = 0
            segs = TCPSegment[]
            sendoff = 0
            while sendoff < length(d)
                len = min(length(d) - sendoff, rand(UInt) % 8192)
                len == 0 && continue
                write!(conn, @view(d[sendoff+1:sendoff+len]))
                tick!(conn, 1)
                @test bytes_in_flight(conn) == sendoff + len
                @test conn.segments_out |> !isempty
                while conn.segments_out |> !isempty
                    push!(segs, expect_seg(conn))
                end
                sendoff += len
            end
            seg_idx = collect(1:length(segs))
            shuffle!(seg_idx)
            acks = TCPSegment[]
            for idx in seg_idx
                send_seg!(conn, segs[idx])
                tick!(conn, 1)
                s = expect_one_seg(conn; ack=true)
                push!(acks, s)
                expect_no_seg(conn)
            end
            send_seg!(conn, acks[end])
            expect_no_seg(conn)
            expect_data!(conn, d)
        end
    end
end


const _Constant = JuMP._Constant
const _VariableAffExpr{C} = GenericAffExpr{C, VariableRef}
const _DecisionAffExpr{C} = GenericAffExpr{C, DecisionRef}
const _KnownAffExpr{C} = GenericAffExpr{C, KnownRef}
const _VAE = _VariableAffExpr{Float64}
const _DAE = _DecisionAffExpr{Float64}
const _KAE = _KnownAffExpr{Float64}

mutable struct DecisionAffExpr{C} <: JuMP.AbstractJuMPScalar
    variables::GenericAffExpr{C, VariableRef}
    decisions::GenericAffExpr{C, DecisionRef}
    knowns::GenericAffExpr{C, KnownRef}
end
const DAE = DecisionAffExpr{Float64}

DAE() = zero(DAE{Float64})

# Base overrides #
# ========================== #
Base.iszero(aff::DecisionAffExpr) =
    iszero(aff.variables) &&
    iszero(aff.decisions) &&
    iszero(aff.knowns)
Base.zero(::Type{DecisionAffExpr{C}}) where C =
    DecisionAffExpr{C}(
        zero(_VariableAffExpr{C}),
        zero(_DecisionAffExpr{C}),
        zero(_KnownAffExpr{C}))
Base.one(::Type{DecisionAffExpr{C}}) where C =
    DecisionAffExpr{C}(one(_VariableAffExpr{C}),
                       zero(_DecisionAffExpr{C}),
                       zero(_KnownAffExpr{C}))
Base.zero(aff::DecisionAffExpr) = zero(typeof(aff))
Base.one(aff::DecisionAffExpr) =  one(typeof(aff))
Base.copy(aff::DecisionAffExpr{C}) where C =
    DecisionAffExpr{C}(copy(aff.variables),
                       copy(aff.decisions),
                       copy(aff.knowns))
Base.broadcastable(aff::DecisionAffExpr) = Ref(aff)

Base.convert(::Type{DecisionAffExpr{C}}, v::VariableRef) where {C} =
    DecisionAffExpr{C}(_VariableAffExpr{C}(zero(C), v => one(C)),
                       zero(_DecisionAffExpr{C}),
                       zero(_KnownAffExpr{C}))
Base.convert(::Type{DecisionAffExpr{C}}, dv::DecisionRef) where C =
    DecisionAffExpr{C}(zero(_VariableAffExpr{C}),
                       _DecisionAffExpr{C}(zero(C), dv => one(C)),
                       zero(_KnownAffExpr{C}))
Base.convert(::Type{DecisionAffExpr{C}}, kv::KnownRef) where C =
    DecisionAffExpr{C}(zero(_VariableAffExpr{C}),
                       zero(_DecisionAffExpr{C}),
                       GenericAffExpr{C, KnownRef}(zero(C), kv => one(C)))
Base.convert(::Type{DecisionAffExpr{C}}, c::Number) where C =
    DecisionAffExpr{C}(_VariableAffExpr{C}(convert(C, r)),
                       zero(_DecisionAffExpr{C}),
                       zero(_KnownAffExpr{C}))
Base.convert(::Type{DecisionAffExpr{C}}, aff::_VariableAffExpr{C}) where C =
    DecisionAffExpr{C}(aff,
                       zero(_DecisionAffExpr{C}),
                       zero(_KnownAffExpr{C}))
Base.convert(::Type{DecisionAffExpr{C}}, aff::_DecisionAffExpr{C}) where C =
    DecisionAffExpr{C}(zero(_VariableAffExpr{C}),
                       aff,
                       zero(_KnownAffExpr{C}))
Base.convert(::Type{DecisionAffExpr{C}}, aff::_KnownAffExpr{C}) where C =
    DecisionAffExpr{C}(zero(_VariableAffExpr{C}),
                       zero(_DecisionAffExpr{C}),
                       aff)
function Base.convert(::Type{T}, aff::DecisionAffExpr{T}) where T
    if !isempty(aff.variables.terms) && !isempty(aff.decisions.terms) && !isempty(aff.knowns.terms)
        throw(InexactError(:convert, T, aff))
    end
    return convert(T, aff.variables.constant)
end

function Base.isequal(aff::DecisionAffExpr{C}, other::DecisionAffExpr{C}) where {C}
    return isequal(aff.variables, other.variables) &&
        isequal(aff.decisions, other.decisions) &&
        isequal(aff.knowns, other.knowns)
end

Base.hash(aff::DecisionAffExpr, h::UInt) =
    hash(aff.variables.constant,
         hash(aff.variables.terms, h),
         hash(aff.decisions.terms, h),
         hash(aff.knowns.terms, h))

# JuMP overrides #
# ========================== #
function JuMP.drop_zeros!(aff::DecisionAffExpr)
    JuMP.drop_zeros!(aff.variables)
    JuMP.drop_zeros!(aff.decisions)
    JuMP.drop_zeros!(aff.knowns)
    return nothing
end

function JuMP._affine_coefficient(f::DecisionAffExpr, variable::VariableRef)
    return JuMP._affine_coefficient(f.variables, variable)
end
function JuMP._affine_coefficient(f::DecisionAffExpr, decision::DecisionRef)
    return JuMP._affine_coefficient(f.decisions, decision)
end
function JuMP._affine_coefficient(f::DecisionAffExpr, known::KnownRef)
    return JuMP._affine_coefficient(f.knowns, known)
end
function JuMP._affine_coefficient(f::_VariableAffExpr{C}, decision::DecisionRef) where C
    return zero(C)
end
function JuMP._affine_coefficient(f::_VariableAffExpr{C}, known::KnownRef) where C
    return zero(C)
end

function JuMP.map_coefficients_inplace!(f::Function, aff::DecisionAffExpr)
    JuMP.map_coefficients_inplace!(f, aff.variables)
    for (coef, dvar) in linear_terms(aff.decisions)
        aff.decisions.terms[dvar] = f(coef)
    end
    for (coef, kvar) in linear_terms(aff.knowns)
        aff.knowns.terms[kvar] = f(coef)
    end
    return aff
end

function JuMP.map_coefficients(f::Function, aff::DecisionAffExpr)
    return JuMP.map_coefficients_inplace!(f, copy(aff))
end

function JuMP.value(aff::DecisionAffExpr, value::Function)
    return JuMP.value(aff.variables, value) +
        JuMP.value(aff.decisions, value) +
        JuMP.value(aff.knowns, value)
end

function JuMP.constant(aff::DecisionAffExpr)
    return aff.variables.constant
end

function SparseArrays.dropzeros(aff::DecisionAffExpr{C}) where C
    variables = SparseArrays.dropzeros(aff.variables)
    decisions = SparseArrays.dropzeros(aff.decisions)
    knowns = SparseArrays.dropzeros(aff.knowns)
    return DecisionAffExpr(variables, decisions, knowns)
end

function JuMP._assert_isfinite(aff::Union{_DecisionAffExpr, _KnownAffExpr})
    for (coef, dv) in linear_terms(aff)
        isfinite(coef) || error("Invalid coefficient $coef on decision $dv.")
    end
    return nothing
end

function JuMP._assert_isfinite(aff::DecisionAffExpr)
    JuMP._assert_isfinite(aff.variables)
    JuMP._assert_isfinite(aff.decisions)
    JuMP._assert_isfinite(aff.knowns)
    return nothing
end

function JuMP.check_belongs_to_model(aff::DecisionAffExpr, model::AbstractModel)
    JuMP.check_belongs_to_model(aff.variables, model)
    JuMP.check_belongs_to_model(aff.decisions, model)
    JuMP.check_belongs_to_model(aff.knowns, model)
    return nothing
end

function JuMP.function_string(mode, aff::DecisionAffExpr, show_constant=true)
    # Decisions
    ret = ""
    decision_terms = JuMP.function_string(mode, aff.decisions, false)
    if decision_terms != "0"
        ret = decision_terms
    end
    # Variables
    variable_terms = JuMP.function_string(mode, aff.variables, false)
    ret = _add_terms(ret, variable_terms)
    # Knowns
    known_terms = JuMP.function_string(mode, aff.knowns, false)
    ret = _add_terms(ret, known_terms)
    # Constant
    if !JuMP._is_zero_for_printing(aff.variables.constant) && show_constant
        ret = string(ret, JuMP._sign_string(aff.variables.constant),

                     JuMP._string_round(abs(aff.variables.constant)))
    end
    return ret
end

function _add_terms(ret::String, terms::String)
    if terms == ""
        return ret
    end
    if terms != "0"
        if ret == ""
            ret = terms
        else
            if terms[1] == '-'
                ret = ret * " - " * terms[2:end]
            else
                ret = ret * " + " * terms
            end
        end
    end
    return ret
end

# With one factor.
function JuMP.add_to_expression!(aff::DAE, other::Number)
    JuMP.add_to_expression!(aff.variables, other)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_var::VariableRef)
    JuMP.add_to_expression!(aff.variables, new_var)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_dvar::DecisionRef)
    JuMP.add_to_expression!(aff.decisions, new_dvar)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_kvar::KnownRef)
    JuMP.add_to_expression!(aff.knowns, new_kvar)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_vae::_VAE)
    JuMP.add_to_expression!(aff.variables, new_vae)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_dae::_DAE)
    JuMP.add_to_expression!(aff.decisions, new_dae)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_kae::_KAE)
    JuMP.add_to_expression!(aff.knowns, new_kae)
    return aff
end

function JuMP.add_to_expression!(lhs_aff::DAE, rhs_aff::DAE)
    JuMP.add_to_expression!(lhs_aff.variables, rhs_aff.variables)
    JuMP.add_to_expression!(lhs_aff.decisions, rhs_aff.decisions)
    JuMP.add_to_expression!(lhs_aff.knowns, rhs_aff.knowns)
    return lhs_aff
end

# With two factors.
function JuMP.add_to_expression!(aff::DAE, new_coef::_Constant, new_var::VariableRef)
    JuMP.add_to_expression!(aff.variables, new_coef, new_var)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_var::VariableRef, new_coef::_Constant)
    JuMP.add_to_expression!(aff.variables, new_coef, new_var)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_coef::_Constant, new_dvar::DecisionRef)
    JuMP.add_to_expression!(aff.decisions, new_coef, new_dvar)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_dvar::DecisionRef, new_coef::_Constant)
    JuMP.add_to_expression!(aff.decisions, new_coef, new_dvar)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_coef::_Constant, new_kvar::KnownRef)
    JuMP.add_to_expression!(aff.knowns, new_coef, new_kvar)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_kvar::KnownRef, new_coef::_Constant)
    JuMP.add_to_expression!(aff.knowns, new_coef, new_kvar)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_coef::_Constant, new_vae::_VAE)
    JuMP.add_to_expression!(aff.variables, new_coef, new_vae)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_vae::_VAE, new_coef::_Constant)
    JuMP.add_to_expression!(aff.variables, new_coef, new_vae)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_coef::_Constant, new_dae::_DAE)
    JuMP.add_to_expression!(aff.decisions, new_coef, new_dae)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_dae::_DAE, new_coef)
    JuMP.add_to_expression!(aff.decisions, new_coef, new_dae)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_coef::_Constant, new_kae::_KAE)
    JuMP.add_to_expression!(aff.knowns, new_coef, new_kae)
    return aff
end

function JuMP.add_to_expression!(aff::DAE, new_kae::_KAE, new_coef::_Constant)
    JuMP.add_to_expression!(aff.knowns, new_coef, new_kae)
    return aff
end

function JuMP.add_to_expression!(lhs_aff::DAE, new_coef::_Constant, rhs_aff::DAE)
    JuMP.add_to_expression!(lhs_aff, new_coef, rhs_aff.variables)
    JuMP.add_to_expression!(lhs_aff, new_coef, rhs_aff.decisions)
    JuMP.add_to_expression!(lhs_aff, new_coef, rhs_aff.knowns)
    return lhs_aff
end

function JuMP.add_to_expression!(lhs_aff::DAE, rhs_aff::DAE, new_coef::_Constant)
    JuMP.add_to_expression!(lhs_aff, rhs_aff.variables, new_coef)
    JuMP.add_to_expression!(lhs_aff, rhs_aff.decisions, new_coef)
    JuMP.add_to_expression!(lhs_aff, rhs_aff.knowns, new_coef)
    return lhs_aff
end


using Mill, JSON, Flux, JsonGrinder, Test, HierarchicalUtils, Setfield

using JsonGrinder: DictEntry, suggestextractor, schema
using Mill: reflectinmodel

@testset "testing pipeline mixing different types of arrays" begin
	j1 = JSON.parse("""{"a": 1, "b": "hello works", "c":{ "a":1 ,"b": "hello world"}}""")
	j2 = JSON.parse("""{"a": 2, "b": "hello world", "c":{ "a":2 ,"b": "hello"}}""")

	sch = schema([j1,j2])
	extractor = suggestextractor(sch)
	ds = map(s-> extractor(s), [j1,j2])
	dss = reduce(catobs, ds)

	m = reflectinmodel(dss, k -> Dense(k,10, relu));
	o = m(dss).data
	for i in 1:length(ds)
		@test o[:,i] ≈ m(ds[i]).data
	end
end

@testset "testing pipeline with simple arrays and missing values" begin
	j1 = JSON.parse("""{"a": 4, "b": {"a":[1,2,3],"b": 1},"c": { "a": {"a":[1,2,3],"b":[4,5,6]}}}""",inttype=Float64)
	j2 = JSON.parse("""{"a": 4, "c": {"a": {"a":[2,3],"b":[5,6]}}}""")
	j3 = JSON.parse("""{"a": 4, "b": {"a":[1,2,3], "b": 1}}""")
	j4 = JSON.parse("""{"a": 4, "b": {}}""")
	j5 = JSON.parse("""{"b": {}}""")
	j6 = JSON.parse("""{}""")

	sch = JsonGrinder.schema([j1,j2,j3])
	extractor = suggestextractor(sch)
	dss = map(extractor, [j1,j2,j3,j4,j5,j6])
	ds = reduce(catobs, dss)
	m = reflectinmodel(ds, k -> Dense(k,10, relu))
	o = m(ds).data

	for i in 1:length(dss)
		@test o[:,i] ≈ m(dss[i]).data
	end
end

# todo: dodělat test pro navrhnutý model a zkontrolovat, že je ok, i s missingy a jedním klíčem v dictu a se scalary
@testset "testing pipeline with arrays of dicts and missing values" begin
	j1 = JSON.parse("""{"a": [{"a":1},{"b":2}]}""")
	j2 = JSON.parse("""{"a": [{"a":1,"b":3},{"b":2,"a" : 1}]}""")
	j3 = JSON.parse("""{"a": [{"a":2,"b":3}]}""")
	j4 = JSON.parse("""{"a": []}""")
	j5 = JSON.parse("""{}""")

	sch = JsonGrinder.schema([j1,j2,j3])
	extractor = suggestextractor(sch)
	dss = map(s-> extractor(s), [j1,j2,j3,j4,j5])
	ds = reduce(catobs, dss)
	m = reflectinmodel(ds, k -> Dense(k,10, relu));
	o = m(ds).data
	for i in 1:length(dss)
		@test o[:,i] ≈ m(dss[i]).data
	end
end

@testset "testing pipeline with single-keyed dicts and merging scalars" begin
	j1 = JSON.parse("""{"a": [{"a":1},{"b":2,"c":"oh"}]}""")
	j2 = JSON.parse("""{"a": [{"a":1,"b":3,"c":"hi"},{"b":2,"a":1,"c":"Mark"}]}""")
	j3 = JSON.parse("""{"a": [{"a":2,"b":3}]}""")
	j4 = JSON.parse("""{"a": []}""")
	j5 = JSON.parse("""{}""")

	sch = JsonGrinder.schema([j1,j2,j3,j4, j5])
	ext = suggestextractor(sch)
	dss = ext.([j1,j2,j3,j4,j5])
	ds = reduce(catobs, dss)
	m = reflectinmodel(ds, k -> Dense(k, 10, relu))

	@test isequal(NodeIterator(m, ds) |> collect, [
		(m[""], ds[""]),
		(m["U"], ds["U"]),
		(m["k"], ds["k"]),
		(m["o"], ds["o"]),
		(m["s"], ds["s"]),
		(m["w"], ds["w"]),
	])
	@test NodeIterator(sch, ext) |> collect == [
		(sch[""], ext[""]),
		(sch["U"], ext["U"]),
		(sch["k"], ext["k"]),
		(sch["o"], ext["o"]),
		(sch["s"], ext["s"]),
		(sch["w"], ext["w"]),
	]
	@test NodeIterator(sch, m) |> collect == [
		(sch[""], m[""]),
		(sch["U"], m["U"]),
		(sch["k"], m["k"]),
		(sch["o"], m["o"]),
		(sch["s"], m["s"]),
		(sch["w"], m["w"]),
	]
	@test NodeIterator(ext, m) |> collect == [
		(ext[""], m[""]),
		(ext["U"], m["U"]),
		(ext["k"], m["k"]),
		(ext["o"], m["o"]),
		(ext["s"], m["s"]),
		(ext["w"], m["w"]),
	]
	@test isequal(NodeIterator(ext, ds) |> collect, [
		(ext[""], ds[""]),
		(ext["U"], ds["U"]),
		(ext["k"], ds["k"]),
		(ext["o"], ds["o"]),
		(ext["s"], ds["s"]),
		(ext["w"], ds["w"]),
	])
	@test buf_printtree(sch) == """
	[Dict] (updated = 5)
	  └── a: [List] (updated = 4)
	           └── [Dict] (updated = 5)
	                 ├── a: [Scalar - Int64], 2 unique values, updated = 4
	                 ├── b: [Scalar - Int64], 2 unique values, updated = 4
	                 └── c: [Scalar - String], 3 unique values, updated = 3"""
	@test buf_printtree(ext) == """
	Dict
	  └── a: Array of
	           └── Dict
	                 ├── a: Float32
	                 ├── b: Float32
	                 └── c: String"""
	@test buf_printtree(m) == """
	ProductModel … ↦ ArrayModel(identity)
	  └── a: BagModel … ↦ ⟨SegmentedMean(10)⟩ ↦ ArrayModel(Dense(11, 10, relu))
	           └── ProductModel … ↦ ArrayModel(Dense(12, 10, relu))
	                 ├── a: ArrayModel([pre_imputing]Dense(1, 1))
	                 ├── b: ArrayModel([pre_imputing]Dense(1, 1))
	                 └── c: ArrayModel([post_imputing]Dense(2053, 10, relu))"""
	@test buf_printtree(ds) == """
	ProductNode with 5 obs
	  └── a: BagNode with 5 obs
	           └── ProductNode with 5 obs
	                 ├── a: ArrayNode(1×5 Array with Union{Missing, Float32} elements) with 5 obs
	                 ├── b: ArrayNode(1×5 Array with Union{Missing, Float32} elements) with 5 obs
	                 └── c: ArrayNode(2053×5 NGramMatrix with Union{Missing, Int64} elements) with 5 obs"""

	@test m[""].m.m == identity

	@test nobs(dss[5]["s"].data) == 0
	@test nobs(dss[5]) == 1
end


module MadNLPHSL

import Libdl: dlopen, RTLD_DEEPBIND
import MadNLP: @kwdef, Logger, @debug, @warn, @error,
    AbstractOptions, AbstractLinearSolver, set_options!, SparseMatrixCSC, SubVector, StrideOneVector,
    SymbolicException,FactorizationException,SolveException,InertiaException,
    introduce, factorize!, solve!, improve!, is_inertia, inertia, findIJ, nnz,
    get_tril_to_full, transfer!

include(joinpath("..","deps","deps.jl"))

if @isdefined(libhsl)
    include("mc68.jl")
    include("ma27.jl")
    include("ma57.jl")
    include("ma77.jl")
    include("ma86.jl")
    include("ma97.jl")
    export MadNLPMa27, MadNLPMa57, MadNLPMa77, MadNLPMa86, MadNLPMa97
end

function __init__()
    check_deps()
    try
        @isdefined(libhsl) && dlopen(libhsl,RTLD_DEEPBIND)
    catch e
        println("HSL shared library cannot be loaded")
    end
end

end # module


# We wish to extend operations to identically named methods dispatched
# on `Machine`s. For example, we have from the model API
#
# `predict(model::M, fitresult, X) where M<:Supervised`
#
# but want also want to define
#
# 1. `predict(machine::Machine, X)` where `X` is concrete data
#
# and we would like the syntactic sugar (for `X` a node):
#
# 2. `predict(machine::Machine, X::Node) = node(predict, machine, X)`
#
# Finally, for a `model` that is `ProbabilisticComposite`,
# `DetermisiticComposite`, or `UnsupervisedComposite`, we want
#
# 3. `predict(model, fitresult, X) = fitresult.predict(X)`
#
# which makes sense because `fitresult` in those cases is a named
# tuple keyed on supported operations and with nodes as values.

## TODO: need to add checks on the arguments of
## predict(::Machine, ) and transform(::Machine, )

const OPERATIONS = (:predict, :predict_mean, :predict_mode, :predict_median, :predict_joint,
                    :transform, :inverse_transform)

for operation in OPERATIONS

    if operation != :inverse_transform

        ex = quote
            # 0. operations on machs, given empty data:
            function $(operation)(mach::Machine; rows=:)
                # Base.depwarn("`$($operation)(mach)` and "*
                #              "`$($operation)(mach, rows=...)` are "*
                #              "deprecated. Data or nodes "*
                #              "should be explictly specified, "*
                #              "as in `$($operation)(mach, X)`. ",
                #              Base.Core.Typeof($operation).name.mt.name)
                if isempty(mach.args) # deserialized machine with no data
                    throw(ArgumentError("Calling $($operation) on a "*
                                        "deserialized machine with no data "*
                                        "bound to it. "))
                end
                return ($operation)(mach, mach.args[1](rows=rows))
            end
        end
        eval(ex)

    end
end

for operation in (:inverse_transform,)
    ex = quote
        # 0. operations on machines, given empty data:
        $operation(mach::Machine; rows=:) =
            throw(ArgumentError("`$($operation)(mach)` and "*
                                "`$($operation)(mach, rows=...)` is "*
                                "not supported. Data or nodes "*
                                "must be explictly specified, "*
                                "as in `$($operation)(mach, X)`. "))
    end
    eval(ex)
end

_symbol(f) = Base.Core.Typeof(f).name.mt.name

for operation in OPERATIONS

    ex = quote
        # 1. operations on machines, given *concrete* data:
        function $operation(mach::Machine, Xraw)
            if mach.state > 0
                return $(operation)(mach.model, mach.fitresult,
                                    Xraw)
            else
                error("$mach has not been trained.")
            end
        end

        function $operation(mach::Machine{<:Static}, Xraw, Xraw_more...)
            isdefined(mach, :fitresult) || (mach.fitresult = nothing)
            return $(operation)(mach.model, mach.fitresult,
                                    Xraw, Xraw_more...)
        end

        # 2. operations on machines, given *dynamic* data (nodes):
        $operation(mach::Machine, X::AbstractNode) =
            node($(operation), mach, X)

        $operation(mach::Machine{<:Static}, X::AbstractNode, Xmore::AbstractNode...) =
            node($(operation), mach, X, Xmore...)
    end
    eval(ex)
end

## SURROGATE AND COMPOSITE MODELS

for operation in [:predict, :predict_joint, :transform, :inverse_transform]
    ex = quote
        $operation(model::Union{Composite,Surrogate}, fitresult,X) =
            fitresult.$operation(X)
    end
    eval(ex)
end

for (operation, fallback) in [(:predict_mode, :mode), (:predict_mean, :mean), (:predict_median, :median)]
    ex = quote
        function $(operation)(m::Union{ProbabilisticComposite,ProbabilisticSurrogate},
                                    fitresult,
                                    Xnew)
        if haskey(fitresult, $(QuoteNode(operation)))
                return fitresult.$(operation)(Xnew)
            end
            return $(fallback).(predict(m, fitresult, Xnew))
        end
    end
    eval(ex)
end


using ImageInpainting
using ColorTypes
using TestImages
using Plots; gr(size=(800,400),yflip=true,colorbar=false,ticks=false)
using VisualRegressionTests
using Test, Pkg

# list of maintainers
maintainers = ["juliohm"]

# environment settings
istravis = "TRAVIS" ∈ keys(ENV)
ismaintainer = "USER" ∈ keys(ENV) && ENV["USER"] ∈ maintainers
datadir = joinpath(@__DIR__,"data")

if ismaintainer
  Pkg.add("Gtk")
  using Gtk
end

# test images
blobs = testimage("blobs")
lighthouse = testimage("lighthouse")

# helper functions
function plot_crimisini_on_array(fname, img, inds, psize)
  fimg = Float64.(Gray.(img))
  mask = falses(size(fimg))
  mask[inds...] .= true
  fimg[mask] .= NaN
  fimg2 = inpaint(fimg, mask, Crimisini(psize))
  plt1 = heatmap(fimg, title="before inpainting")
  plt2 = heatmap(fimg2, title="after inpainting")
  plot(plt1, plt2)
  png(fname)
end

@testset "ImageInpainting.jl" begin
  @testset "Plain arrays" begin
    plot_blobs(fname) = plot_crimisini_on_array(fname, blobs, (50:150,50:150), (11,11))
    refimg = joinpath(datadir,"Blobs.png")
    @test test_images(VisualTest(plot_blobs, refimg), popup=!istravis) |> success

    plot_lighthouse(fname) = plot_crimisini_on_array(fname, lighthouse, (50:350,300:400), (30,30))
    refimg = joinpath(datadir,"LightHouse.png")
    @test test_images(VisualTest(plot_lighthouse, refimg), popup=!istravis) |> success
  end
end


"""
TODO: Description of the current population model.
"""
module population
importall common
using   core, _deme_, utility, species, breeder
export  Population, PopulationDefinition

"""
The base type used by all population definitions.
"""
abstract PopulationDefinition

"""
Used to hold the contents of a population, where a population is modelled
as a collection of demes (which may, or may not, be of different species).
"""
type Population
  demes::Vector{Deme}

  Population() = new([])
  Population(dl::Vector{Deme}) = new(dl)
end

"""
Produces the offspring for a given population at each generation.
"""
breed!(p::Population) = for deme in p.demes; breed!(deme); end

"""
Performs any necessary fitness scaling and processing on each of the
individuals within a given population.
"""
scale!(p::Population) = for deme in p.demes; scale!(deme); end
#scale!(d::Deme) = scale!(d.species.fitness, contents(d))

"""
TODO: IMPLEMENT SCALING.
"""
scale!(d::Deme) = true

# Load each of the population types.
include("population/simple.jl")
include("population/complex.jl")
end


@inline vzero(::Val{1}, ::Type{T}) where {T<:NativeTypes} = zero(T)
@inline vzero(::StaticInt{1}, ::Type{T}) where {T<:NativeTypes} = zero(T)
@generated function _vzero(::StaticInt{W}, ::Type{T}, ::StaticInt{RS}) where {W,T<:NativeTypes,RS}
    # isone(W) && return Expr(:block, Expr(:meta,:inline), Expr(:call, :zero, T))
    if W * sizeof(T) > RS
        d, r1 = divrem(sizeof(T) * W, RS)
        Wnew, r2 = divrem(W, d)
        (iszero(r1) & iszero(r2)) || throw(ArgumentError("If broadcasting to greater than 1 vector length, should make it an integer multiple of the number of vectors."))
        t = Expr(:tuple)
        for i ∈ 1:d
            push!(t.args, :v)
        end
        # return Expr(:block, Expr(:meta,:inline), :(v = vzero(StaticInt{$Wnew}(), $T)), :(VecUnroll{$(d-1),$Wnew,$T,Vec{$Wnew,$T}}($t)))
        return Expr(:block, Expr(:meta,:inline), :(v = _vzero(StaticInt{$Wnew}(), $T, StaticInt{$RS}())), :(VecUnroll($t)))
        # return Expr(:block, Expr(:meta,:inline), :(v = _vzero(StaticInt{$Wnew}(), $T, StaticInt{$RS}())), :(VecUnroll($t)::VecUnroll{$(d-1),$Wnew,$T,Vec{$Wnew,$T}}))
    end
    typ = LLVM_TYPES[T]
    instrs = "ret <$W x $typ> zeroinitializer"
    quote
        $(Expr(:meta,:inline))
        Vec($LLVMCALL($instrs, _Vec{$W,$T}, Tuple{}))
    end
end
@generated function _vundef(::StaticInt{W}, ::Type{T}) where {W,T<:NativeTypes}
    typ = LLVM_TYPES[T]
    if W == 1
        instrs = "ret $typ undef"
        quote
            $(Expr(:meta,:inline))
            $LLVMCALL($instrs, $T, Tuple{})
        end
    else
        instrs = "ret <$W x $typ> undef"
        quote
            $(Expr(:meta,:inline))
            Vec($LLVMCALL($instrs, _Vec{$W,$T}, Tuple{}))
        end
    end
end
@generated function _vbroadcast(::StaticInt{W}, s::_T, ::StaticInt{RS}) where {W,_T<:NativeTypes,RS}
    isone(W) && return :s
    if _T <: Integer && sizeof(_T) * W > RS
        intbytes = max(4, RS ÷ W)
        T = integer_of_bytes(intbytes)
        if _T <: Unsigned
            T = unsigned(T)
        end
        # ssym = :(s % $T)
        ssym = :(convert($T, s))
    elseif sizeof(_T) * W > RS
        d, r1 = divrem(sizeof(_T) * W, RS)
        Wnew, r2 = divrem(W, d)
        (iszero(r1) & iszero(r2)) || throw(ArgumentError("If broadcasting to greater than 1 vector length, should make it an integer multiple of the number of vectors."))
        t = Expr(:tuple)
        for i ∈ 1:d
            push!(t.args, :v)
        end
        return Expr(:block, Expr(:meta,:inline), :(v = vbroadcast(StaticInt{$Wnew}(), s)), :(VecUnroll($t)))
    else
        T = _T
        ssym = :s
    end
    typ = LLVM_TYPES[T]
    vtyp = vtype(W, typ)
    instrs = """
        %ie = insertelement $vtyp undef, $typ %0, i32 0
        %v = shufflevector $vtyp %ie, $vtyp undef, <$W x i32> zeroinitializer
        ret $vtyp %v
    """
    quote
        $(Expr(:meta,:inline))
        Vec($LLVMCALL($instrs, _Vec{$W,$T}, Tuple{$T}, $ssym))
    end
end
@inline vzero(::Union{Val{W},StaticInt{W}}, ::Type{T}) where {W,T} = _vzero(StaticInt{W}(), T, register_size(T))
@inline vbroadcast(::Union{Val{W},StaticInt{W}}, s::T) where {W,T} = _vbroadcast(StaticInt{W}(), s, register_size(T))
@inline function _vbroadcast(::StaticInt{W}, vu::VecUnroll{N, 1, T, T}, ::StaticInt{RS}) where {W,N,T,RS}
    VecUnroll(fmap(_vbroadcast, StaticInt{W}(), data(vu), StaticInt{RS}()))
end

@generated function vbroadcast(::Union{Val{W},StaticInt{W}}, ptr::Ptr{T}) where {W, T}
    isone(W) && return Expr(:block, Expr(:meta, :inline), :(vload(ptr)))
    typ = LLVM_TYPES[T]
    ptyp = JULIAPOINTERTYPE
    vtyp = "<$W x $typ>"
    alignment = Base.datatype_alignment(T)
    instrs = """
        %ptr = inttoptr $ptyp %0 to $typ*
        %res = load $typ, $typ* %ptr, align $alignment
        %ie = insertelement $vtyp undef, $typ %res, i32 0
        %v = shufflevector $vtyp %ie, $vtyp undef, <$W x i32> zeroinitializer
        ret $vtyp %v
    """
    quote
        $(Expr(:meta,:inline))
        Vec($LLVMCALL( $instrs, _Vec{$W,$T}, Tuple{Ptr{$T}}, ptr ))
    end
end

@inline vbroadcast(::Union{Val{W},StaticInt{W}}, v::AbstractSIMDVector{W}) where {W} = v

@generated function vbroadcast(::Union{Val{W},StaticInt{W}}, v::V) where {W,L,T,V<:AbstractSIMDVector{L,T}}
    N, r = divrem(L, W)
    @assert iszero(r)
    V = if T === Bit
        :(Mask{$W,$(mask_type_symbol(W))})
    else
        :(Vec{$W,$T})
    end
    Expr(:block, Expr(:meta,:inline), :(vconvert(VecUnroll{$(N-1),$W,$T,$V}, v)))
end

@inline Vec{W,T}(v::Vec{W,T}) where {W,T} = v

@inline Base.zero(::Type{Vec{W,T}}) where {W,T} = _vzero(StaticInt{W}(), T, StaticInt{W}() * static_sizeof(T))
@inline Base.zero(::Vec{W,T}) where {W,T} = zero(Vec{W,T})
@inline Base.one(::Vec{W,T}) where {W,T} = vbroadcast(Val{W}(), one(T))

@inline Base.one(::Type{Vec{W,T}}) where {W,T} = vbroadcast(Val{W}(), one(T))
@inline Base.oneunit(::Type{Vec{W,T}}) where {W,T} = vbroadcast(Val{W}(), one(T))
@inline vzero(::Type{T}) where {T<:Number} = zero(T)
@inline vzero() = vzero(pick_vector_width(Float64), Float64)

@inline Vec{W,T}(s::Real) where {W,T} = vbroadcast(Val{W}(), T(s))
@inline Vec{W}(s::T) where {W,T<:NativeTypes} = vbroadcast(Val{W}(), s)
@inline Vec(s::T) where {T<:NativeTypes} = vbroadcast(pick_vector_width(T), s)

@generated function _vzero(::Type{VecUnroll{N,W,T,V}}, ::StaticInt{RS}) where {N,W,T,V,RS}
    t = Expr(:tuple); foreach(_ -> push!(t.args, :(_vzero(StaticInt{$W}(), $T, StaticInt{$RS}()))), 0:N)
    Expr(:block, Expr(:meta, :inline), :(VecUnroll($t)))
end
@inline Base.zero(::Type{VecUnroll{N,W,T,V}}) where {N,W,T,V} = _vzero(VecUnroll{N,W,T,V}, register_size())
@inline Base.zero(::VecUnroll{N,W,T,V}) where {N,W,T,V} = zero(VecUnroll{N,W,T,V})

@inline Base.one(::Type{VecUnroll{N,W,T,V}}) where {N,W,T,V} = VecUnroll{N}(one(V))

@generated function VecUnroll{N,W,T,V}(x::S) where {N,W,T,V<:AbstractSIMDVector{W,T},S<:Real}
    t = Expr(:tuple)
    for n ∈ 0:N
        push!(t.args, :(convert($V, x)))
    end
    Expr(:block, Expr(:meta,:inline), :(VecUnroll($t)))
end
@generated function VecUnroll{N,1,T,T}(x::S) where {N,T<:NativeTypes,S<:Real}
    t = Expr(:tuple)
    for n ∈ 0:N
        push!(t.args, :(convert($T, x)))
    end
    Expr(:block, Expr(:meta,:inline), :(VecUnroll($t)))
end
@inline VecUnroll{N,W,T}(x::NativeTypesV) where {N,W,T} = VecUnroll{N,W,T,Vec{W,T}}(x)
@inline VecUnroll{N}(x::V) where {N,W,T,V <: AbstractSIMDVector{W,T}} = VecUnroll{N,W,T,V}(x)
@inline VecUnroll{N}(x::T) where {N,T<:NativeTypes} = VecUnroll{N,1,T,T}(x)

@generated function zero_vecunroll(::StaticInt{N}, ::StaticInt{W}, ::Type{T}, ::StaticInt{RS}) where {N,W,T,RS}
    Expr(:block, Expr(:meta, :inline), :(_vzero(VecUnroll{$(N-1),$W,$T,Vec{$W,$T}}, StaticInt{$RS}())))
end
@inline zero_init(::Type{T}, ::StaticInt{1}, ::StaticInt{0}, ::StaticInt{RS}) where {T,RS} = zero(T)
@inline zero_init(::Type{T}, ::StaticInt{W}, ::StaticInt{0}, ::StaticInt{RS}) where {W,T,RS} = _vzero(StaticInt{W}(), T, StaticInt{RS}())
@inline zero_init(::Type{T}, ::StaticInt{W}, ::StaticInt{U}, ::StaticInt{RS}) where {W,U,T,RS} = _vzero(VecUnroll{U,W,T,Vec{W,T}}, StaticInt{RS}())
@inline zero_init(::Type{T}, ::Tuple{StaticInt{W},StaticInt{U}}, ::StaticInt{RS}) where {W,U,T,RS} = zero_init(T, StaticInt{W}(), StaticInt{U}(), StaticInt{RS})
@generated function vbroadcast_vecunroll(::StaticInt{N}, ::StaticInt{W}, s::T, ::StaticInt{RS}) where {N,W,T,RS}
    q = Expr(:block, Expr(:meta,:inline), :(v = _vbroadcast(StaticInt{$W}(), s, StaticInt{$RS}())))
    t = Expr(:tuple)
    for n ∈ 1:N
        push!(t.args, :v)
    end
    push!(q.args, :(VecUnroll($t)))
    q
end


mutable struct LatLonAccelDirection
    a_lat::Float64
    a_lon::Float64
    direction::Int
end
Base.show(io::IO, a::LatLonAccelDirection) = @printf(io, "LatLonAccelDirection(%6.3f, %6.3f, %d)", a.a_lat, a.a_lon, a.direction)
Base.length(::Type{LatLonAccelDirection}) = 3
Base.convert(::Type{LatLonAccelDirection}, v::Vector{Float64}) = LatLonAccel(v[1], v[2],convert(Int,v[3]))
function Base.copyto!(v::Vector{Float64}, a::LatLonAccelDirection)
    v[1] = a.a_lat
    v[2] = a.a_lon
    v[3] = convert(Float64,a.direction)
    v
end
function AutomotiveDrivingModels.propagate(veh::Entity{VehicleState, D, Int}, action::LatLonAccelDirection,  roadway::Roadway, Δt::Float64) where {D<:Union{VehicleDef, BicycleModel}}
    previousInd = veh.state.posF.roadind
    a_lat = action.a_lat
    a_lon = action.a_lon

     v = veh.state.v
     ϕ = veh.state.posF.ϕ
    ds = v*cos(ϕ)
     t = veh.state.posF.t
    dt = v*sin(ϕ)

    ΔT = Δt
    ΔT² = ΔT*ΔT
    Δs = ds*ΔT + 0.5*a_lon*ΔT²
    Δt = dt*ΔT + 0.5*a_lat*ΔT²

    ds₂ = ds + a_lon*ΔT
    dt₂ = dt + a_lat*ΔT
    speed₂ = sqrt(dt₂*dt₂ + ds₂*ds₂)
    v₂ = sign(ds₂)*speed₂
    ϕ₂ = atan(dt₂, ds₂) + (v₂ < 0.0)*π # handle negative speeds
    roadind = move_along_with_direction(veh.state.posF.roadind, roadway, Δs, direction = action.direction)
    if t+Δt <= roadway[roadind.tag].width/2.0
        posF = Frenet(roadind, roadway, t=t+Δt, ϕ = ϕ₂)
        return VehicleState(posF, roadway, v₂)
    else
        footpoint = roadway[roadind]
        posG = convert(VecE2, footpoint.pos) + polar(t + Δt, footpoint.pos.θ + π/2)
        posG = VecSE2(posG.x, posG.y, footpoint.pos.θ + ϕ₂)

        state = VehicleState(posG, roadway, v₂)
        projections = in_lanes(posG, roadway)
        if length(roadway[previousInd.tag].exits)>0
            for projection in projections
                dir = clamp(action.direction,1,length(roadway[previousInd.tag].exits))
                if roadway[previousInd.tag].exits[dir].target.tag.segment == projection.tag.segment && !isempty(roadway[previousInd.tag].exits)
                    posF = Frenet(projection, roadway)
                    state = VehicleState(posG, posF, v₂)
                    #println(posF)
                    return state
                end
            end
        end

        for projection in projections
            dir = clamp(action.direction,1,length(roadway[previousInd.tag].exits))
            if  previousInd.tag.segment == projection.tag.segment
                posF = Frenet(projection, roadway)
                state = VehicleState(posG, posF, v₂)
                #println(posF)
                return state
            end
        end
        return state
    end
end

function Base.get(::Type{LatLonAccelDirection}, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    accel_lat = get(ACCFT, rec, roadway, vehicle_index, pastframe)
    accel_lon = get(ACCFS, rec, roadway, vehicle_index, pastframe)
    LatLonAccelDirection(accel_lat, accel_lon,1)
end
function pull_action!(::Type{LatLonAccelDirection}, a::Vector{Float64}, rec::SceneRecord, roadway::Roadway, vehicle_index::Int, pastframe::Int=0)
    a[1] = get(ACCFT, rec, roadway, vehicle_index, pastframe)
    a[2] = get(ACCFS, rec, roadway, vehicle_index, pastframe)
    a
end

mutable struct AccelSteeringDirection
    a::Float64 # accel [m/s²]
    δ::Float64 # steering angle [rad]
    direction::Int64 # turning direction
end
Base.show(io::IO, a::AccelSteeringDirection) = @printf(io, "AccelSteeringDirection(%6.3f,%6.3f,%d)", a.a, a.δ,a.direction)
Base.length(::Type{AccelSteeringDirection}) = 3
Base.convert(::Type{AccelSteeringDirection}, v::Vector{Float64}) = AccelSteeringDirection(v[1], v[2], convert(Int,v[3]))
function Base.copyto!(v::Vector{Float64}, a::AccelSteeringDirection)
    v[1] = a.a
    v[2] = a.δ
    v[3] = convert(Float64,a.direction)
    v
end
#function propagate(veh::Entity{VehicleState, BicycleModel, Int}, action::AccelSteeringAngle, roadway::Roadway, Δt::Float64)
function AutomotiveDrivingModels.propagate(veh::Entity{VehicleState, D, Int}, action::AccelSteeringDirection, roadway::Roadway, Δt::Float64) where {D<:Union{VehicleDef, BicycleModel}}
    previousInd = veh.state.posF.roadind

    L = veh.def.a + veh.def.b
    l = -veh.def.b

    a = action.a # accel [m/s²]
    δ = action.δ # steering wheel angle [rad]

    x = veh.state.posG.x
    y = veh.state.posG.y
    θ = veh.state.posG.θ
    v = veh.state.v

    s = v*Δt + a*Δt*Δt/2 # distance covered
    v′ = v + a*Δt

    posG = nothing
    if abs(δ) < 0.01 # just drive straight
        posG = veh.state.posG + polar(s, θ)
    else # drive in circle

        R = L/tan(δ) # turn radius

        β = s/R
        xc = x - R*sin(θ) + l*cos(θ)
        yc = y + R*cos(θ) + l*sin(θ)

        θ′ = mod(θ+β, 2π)
        x′ = xc + R*sin(θ+β) - l*cos(θ′)
        y′ = yc - R*cos(θ+β) - l*sin(θ′)

        posG = VecSE2(x′, y′, θ′)
    end

    #VehicleState(posG, roadway, v′)
    state = VehicleState(posG, roadway, v′)
    projections = in_lanes(posG, roadway)
    if length(roadway[previousInd.tag].exits)>0
        for projection in projections
            dir = clamp(action.direction,1,length(roadway[previousInd.tag].exits))
            if roadway[previousInd.tag].exits[dir].target.tag.segment == projection.tag.segment && !isempty(roadway[previousInd.tag].exits)
                posF = Frenet(projection, roadway)
                state = VehicleState(posG, posF, v′)
                #println(posF)
                return state
            end
        end
    end

    for projection in projections
        dir = clamp(action.direction,1,length(roadway[previousInd.tag].exits))
        if  previousInd.tag.segment == projection.tag.segment
            posF = Frenet(projection, roadway)
            state = VehicleState(posG, posF, v′)
            #println(posF)
            return state
        end
    end
    return state
end

mutable struct NextState
    x::Float64
    y::Float64
    theta::Float64
    speed::Float64
    direction::Int
end

function NextState()
    return NextState(0.0,0.0,0.0,0.0,1)
end

function AutomotiveDrivingModels.propagate(veh::Entity{VehicleState, D, Int}, action::NextState,  roadway::Roadway, Δt::Float64) where {D<:Union{VehicleDef, BicycleModel}}
    previousInd = veh.state.posF.roadind

    posG = VecSE2(action.x, action.y, action.theta)
    v′ = action.speed

    state = VehicleState(posG, roadway, v′)
    projections = in_lanes(posG, roadway)
    if length(roadway[previousInd.tag].exits)>0
        for projection in projections
            dir = clamp(action.direction,1,length(roadway[previousInd.tag].exits))
            if roadway[previousInd.tag].exits[dir].target.tag.segment == projection.tag.segment && !isempty(roadway[previousInd.tag].exits)
                posF = Frenet(projection, roadway)
                state = VehicleState(posG, posF, v′)
                #println(posF)
                return state
            end
        end
    end

    for projection in projections
        dir = clamp(action.direction,1,length(roadway[previousInd.tag].exits))
        if  previousInd.tag.segment == projection.tag.segment
            posF = Frenet(projection, roadway)
            state = VehicleState(posG, posF, v′)
            #println(posF)
            return state
        end
    end
    return state
end


using BHAPtfem, DataFrames, GR

ProjDir = dirname(@__FILE__)
ProjName = split(ProjDir, "/")[end]

#bha = BHA(ProjName, ProjDir)
#bha.ratio = 1.7
#bha.medium = :lightmud

segs2 = [
# Element type,  Material,    Length,     OD,         ID,        fc
  :bit,           :steel ,       0.00,    2.75,      12.25,     0.0,
  :collar ,       :monel ,      25.00,    2.75,       7.75,     0.0,
  :stabilizer,    :steel,        0.00,    2.75,      12.125,    0.0,
  :collar,        :steel,       30.00,    2.75,       7.75,     0.0,
  :stabilizer,    :steel,        0.00,    2.75,      12.125,    0.0,
  :collar,        :steel,       30.00,    2.75,       7.75,     0.0,
  :stabilizer,    :steel,        0.00,    2.75,      12.125,    0.0,
  :collar,        :steel,       90.00,    2.75,       7.75,     0.0,
  :stabilizer,    :steel,        0.00,    2.75,      12.125,    0.0,
  :pipe,          :steel,      100.00,    2.75,       7.75,     0.0
];

segs1 = reshape(segs2, (6, :));
syms = Array{Symbol, 2}(reshape(segs1[1:2, :], (2, :)))
segs = Array{Float64, 2}(reshape(segs1[3:6, :], (4, :)))

df = DataFrame([Symbol, Symbol, Float64], [:eltype, :material, :length], 0)
append!(df, DataFrame(eltype=:bit, material=:steel, length=0.0))

using GR
x = 0:pi/10:2*pi
y1 = sin.(x)
y2 = sin.(x .- 0.25)
y3 = sin.(x .- 0.5)

plot(x,y1,"g",x,y2,"--",x,y2,"g--o",x,y3,"-",x,y3,"c*")
savefig("$(ProjDir)/gr_plot_test.pdf")

df


using DataFrames
using Gadfly

dirnames = ["barrier_ciris", "new_paper_output", "1_stepfac", "10_stepfac"]

dirlabels = ["Base model","Ea=300K", "Δz=1 θ (θ ≈ 104)", "Δz=10 θ (θ ≈ 103)"]

# Input parameters
path = "/scratch/cns7ae/ciris/"
abundanceFile = "abundance.wsv"
plotname = "f3.pdf" #"abundance_plot.pdf"
thick = 1.0e-5
edge  = 3.5e-6
denom = 2*thick*edge*edge*1.0e20
iceSize = (170^3)/3
xmax    = 1.0e18
xmin    = 1.0e10
ymin    = (10^(-5.0)) #0.5
species = "O3" # all, O2, O, O3

plottheme = Gadfly.Theme(
    line_width=2pt,
    minor_label_font_size=13pt,
    minor_label_color=colorant"black",
    major_label_font_size=15pt,
    major_label_color=colorant"black",
    key_label_color=colorant"black",
    key_title_color=colorant"black",
    key_title_font_size=12pt,
    key_label_font_size=11pt,
    grid_color=colorant"black")

println("Initializing dataframes")
ab_tot_df = DataFrame()
O3ab_df   = DataFrame()
Oab_df    = DataFrame()
O2ab_df   = DataFrame()

for i = 1:size(dirnames,1)
    println("Reading input file in $(dirnames[i])")

    modelOutput = readtable(
        path*dirnames[i]*"/"*abundanceFile,
        names = [:Fluence, :Time, :Geminacy, :O2, :O, :O3, :Oex,:O2ex,:O3ex,:Atoms],
        eltypes = [Float64,Float64, Float64, Int64,Int64,Int64,Int64,Int64,Int64,Int64],
        header = false
    )

    println("Populating dataframes")
    Oab_df = vcat(Oab_df,
                  DataFrame(
                      Fluence=modelOutput[:Fluence],
                      Abundance=modelOutput[:O]/denom,
                      Species="O",
                      Type=dirlabels[i]
                  )
                  )

    O3ab_df = vcat(
        O3ab_df,
        DataFrame(
            Fluence=modelOutput[:Fluence],
            Abundance=modelOutput[:O3]/denom,
            Species="O3",
            Type=dirlabels[i]
        )
    )

    O2ab_df = vcat(
        O2ab_df,
        DataFrame(
            Fluence=modelOutput[:Fluence],
            Abundance=modelOutput[:O2]/denom,
            Species="O2",
            Type=dirlabels[i]
        )
    )

    ab_tot_df = vcat(ab_tot_df,
                     DataFrame(
                         Fluence=modelOutput[:Fluence],
                         Abundance=modelOutput[:O2]/denom,
                         Species="O2",
                         Type=dirlabels[i]
                     ),
                     DataFrame(
                         Fluence=modelOutput[:Fluence],
                         Abundance=modelOutput[:O]/denom,
                         Species="O",
                         Type=dirlabels[i]
                     ),
                     DataFrame(
                         Fluence=modelOutput[:Fluence],
                         Abundance=modelOutput[:O3]/denom,
                         Species="O3",
                         Type=dirlabels[i]
                     )
                     )

end

expFluence = [8.e11,5.e12,1.1e13,4.e13,1.15e14,2.e14,7.e14,2.1e15,7.5e15]
expAbundance = [0.5,1.5,2.15,2.9,3.4,3.8,4.1,3.9,4.17]
O3exp = DataFrame(
    Fluence=expFluence,
    Abundance=expAbundance,
    Species="O3",
    Type="Experiment"
)

#O3total = vcat(O3total,O3ab_df,O3exp)

# species_ab_af = DataFrame()
# species_ab_df = vcat(species_ab_df,O3ab_df, O2ab_df, Oab_df)
# Select the dataframe to plot
plot_df = DataFrame()
if species == "all"
    plot_df  = vcat(plot_df,ab_tot_df)
elseif species == "O2"
    plot_df = vcat(plot_df,O2ab_df)
elseif species == "O"
    plot_df = vcat(plot_df,Oab_df)
elseif species == "O3"
    plot_df = vcat(plot_df,O3ab_df)
end

O3exp[:Abundance] = O3exp[:Abundance].*1.0e20
O3exp[:Abundance] = O3exp[:Abundance]./6.02214e23
plot_df[:Abundance] = plot_df[:Abundance].*1.0e20
plot_df[:Abundance] = plot_df[:Abundance]./6.02214e23

println("Generating plots")
o3 = plot(
    O3exp,
    x="Fluence",
    y="Abundance",
    color="Type",
    #          Geom.smooth,
    Geom.point,
    # Guide.title("Species vs. Fluence"),
    Guide.xlabel("Fluence (ions cm-2)"),
    Guide.ylabel("[O3] (mol cm-3)"),
    Scale.x_log10(minvalue=xmin,maxvalue=xmax),
    Scale.y_log10
)

plot_df = plot_df[xmax.>plot_df[:Fluence].>xmin,:]
plot_df = plot_df[plot_df[:Abundance].>ymin,:]

push!(o3,layer(
    plot_df,
    x="Fluence",
    y="Abundance",
    color="Type",
    Geom.line))
push!(o3,plottheme)

o2 = plot(
    O2ab_df,
    x="Fluence",
    y="Abundance",
    color="Type",
    Geom.line,
    # Guide.title("Species vs. Fluence"),
    Guide.xlabel("Fluence (ions cm-2)"),
    Guide.ylabel("Abundance (cm-3) * 1020"),
    Scale.x_log10(minvalue=xmin,maxvalue=xmax),
    Scale.y_log10
)
push!(o2,plottheme)

plot_df = vcat(plot_df,O2ab_df)

println("Now exporting plot")
draw(PDF(path*"f3.pdf",6inch,4inch),o3)
#draw(PDF(path*plotname, 6inch, 5inch), p1)
println("Ending script!!")
#quit()


#usegnum.jl creates a @bind_operation macro which binds unum and gnum operations
#to their respective operations depending on whether or not the :usegnum option
#is set.

doc"""
  `@bind_operation` uses julia's multiple dispatch capability to bind standard
  binary numerical operations to standard "word" functions.  For example, the
  macro:

    @bind_operation + add

  will generate functions mapping all relevant + combinations to the add_create
  function or the add! functions (which release gnums), if :use_gnum is specified,
  otherwise, mapping + to the simple "add" function
"""
macro bind_operation(op, fsymbol)
  gsymbol = Symbol(fsymbol, "!")
  gcreate = Symbol(fsymbol, "_create")
  if options[:usegnum]
    esc(quote
      @universal $op(a::Unum,   b::Unum)   = $gcreate(a, b)
      @universal $op(a::Unum,   b::Ubound) = $gcreate(a, b)
      @universal $op(a::Ubound, b::Unum)   = $gcreate(a, b)
      @universal $op(a::Ubound, b::Ubound) = $gcreate(a, b)

      @universal $op(a::Unum,   b::Gnum)   = $gsymbol(a, b)
      @universal $op(a::Ubound, b::Gnum)   = $gsymbol(a, b)
      @universal $op(a::Gnum,   b::Unum)   = $gsymbol(a, b)
      @universal $op(a::Gnum,   b::Ubound) = $gsymbol(a, b)
      @universal $op(a::Gnum,   b::Gnum)   = $gsymbol(a, b)
    end)
  else
    esc(quote
      @universal $op(a::Unum, b::Unum)     = $fsymbol(a, b)
      @universal $op(a::Unum, b::Ubound)   = $fsymbol(a, b)
      @universal $op(a::Ubound, b::Unum)   = $fsymbol(a, b)
      @universal $op(a::Ubound, b::Ubound) = $fsymbol(a, b)
    end)
  end
end


export
    fovray,
    fovtrg,
    frame,
    frinfo,
    frmnam,
    ftncls,
    furnsh

"""
    fovray(inst, raydir, rframe, abcorr, observer, et)

Determine if a specified ray is within the field-of-view (FOV) of a specified instrument
at a given time.

### Arguments ###

- `inst`: Name or ID code string of the instrument
- `raydir`: Ray's direction vector
- `rframe`: Body-fixed, body-centered frame for target body
- `abcorr`: Aberration correction flag
- `observer`: Name or ID code string of the observer
- `et`: Time of the observation (seconds past J2000)

### Output ###

Returns `true` if the ray is visible.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/fovray_c.html)
"""
function fovray(inst, raydir, rframe, abcorr, observer, et)
    @checkdims 3 raydir
    visible = Ref{SpiceBoolean}()
    ccall((:fovray_c, libcspice), Cvoid,
          (Cstring, Ref{SpiceDouble}, Cstring, Cstring, Cstring, Ref{SpiceDouble}, Ref{SpiceBoolean}),
          inst, raydir, rframe, abcorr, observer, et, visible)
    handleerror()
    Bool(visible[])
end

"""
    fovtrg(inst, target, tshape, tframe, abcorr, obsrvr, et)

Determine if a specified ephemeris object is within the field-of-view (FOV) of a specified
instrument at a given time.

### Arguments ###

- `inst`: Name or ID code string of the instrument.
- `target`: Name or ID code string of the target.
- `tshape`: Type of shape model used for the target.
- `tframe`: Body-fixed, body-centered frame for target body.
- `abcorr`: Aberration correction flag.
- `obsrvr`: Name or ID code string of the observer.
- `et`: Time of the observation (seconds past J2000).

### Output ###

Returns `true` if the object is visible.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/fovtrg_c.html)
"""
function fovtrg(inst, target, tshape, tframe, abcorr, obsrvr, et)
    visible = Ref{SpiceBoolean}()
    ccall((:fovtrg_c, libcspice), Cvoid,
          (Cstring, Cstring, Cstring, Cstring, Cstring, Cstring, Ref{SpiceDouble}, Ref{SpiceBoolean}),
          inst, target, tshape, tframe, abcorr, obsrvr, et, visible)
    handleerror()
    Bool(visible[])
end

"""
    frame(x)

Given a vector `x`, this routine builds a right handed orthonormal frame `x`, `y`, `z`
where the output `x` is parallel to the input `x`.

### Arguments ###

- `x`: Input vector

### Output ###

- `x`: Unit vector parallel to `x` on output
- `y`: Unit vector in the plane orthogonal to `x`
- `z`: Unit vector given by `x × y`

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/frame_c.html)
"""
function frame(x)
    @checkdims 3 x
    x = copy(x)
    y = Array{SpiceDouble}(undef, 3)
    z = Array{SpiceDouble}(undef, 3)
    ccall((:frame_c, libcspice), Cvoid,
          (Ref{SpiceDouble}, Ref{SpiceDouble}, Ref{SpiceDouble}),
          x, y, z)
    x, y, z
end

"""
    frinfo(frcode)

Retrieve the minimal attributes associated with a frame needed for converting transformations
to and from it.

### Arguments ###

- `frcode`: The id code for a reference frame

### Output ###

- `cent`: The center of the frame
- `frclss`: The class (type) of the frame
- `clssid`: The idcode for the frame within its class

Returns `nothing` if no frame with id `frcode` could be found.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/frinfo_c.html)
"""
function frinfo(frcode)
    cent = Ref{SpiceInt}()
    frclss = Ref{SpiceInt}()
    clssid = Ref{SpiceInt}()
    found = Ref{SpiceBoolean}()
    ccall((:frinfo_c, libcspice), Cvoid,
          (SpiceInt, Ref{SpiceInt}, Ref{SpiceInt}, Ref{SpiceInt}, Ref{SpiceBoolean}),
          frcode, cent, frclss, clssid, found)

    Bool(found[]) || return nothing
    Int(cent[]), Int(frclss[]), Int(clssid[])
end

"""
    frmnam(frcode)

Retrieve the name of a reference frame associated with an id code.

### Arguments ###

- `frcode`: The id code for a reference frame

### Output ###

Returns the name associated with the reference frame.

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/frmnam_c.html)
"""
function frmnam(frcode)
    lenout = 33
    frname = Array{SpiceChar}(undef, lenout)
    ccall((:frmnam_c, libcspice), Cvoid, (SpiceInt, SpiceInt, Ref{SpiceChar}),
          frcode, lenout, frname)
    chararray_to_string(frname)
end

@deprecate ftncls close

"""
    furnsh(kernels...)

Load one or more SPICE kernels into a program.

### Arguments ###

- `kernels`: Path(s) of SPICE kernels to load

### References ###

- [NAIF Documentation](https://naif.jpl.nasa.gov/pub/naif/toolkit_docs/C/cspice/furnsh_c.html)
"""
function furnsh(kernels...)
    for kernel in kernels
        ccall((:furnsh_c, libcspice), Cvoid, (Cstring,), kernel)
        handleerror()
    end
end