include("required.jl")

## Conditional arguments
# create_distributions = true # should distributions be computed (optional, loaded otherwise)
# save_data = true # should data files be overwritten (optional)
# save_figures = true # should figures be overwritten (optional)

## Get & prepare data

# Define coordinates range
coords = (left=-145.0, right=-50.0, bottom=20.0, top=75.0)
copy_layer = SimpleSDMPredictor(WorldClim, BioClim, 1; coords...)

# Load data from CSV files
df = CSV.read(joinpath("data", "proc", "ebd_warblers_prep.csv"), DataFrame)

# Filter observations outside coordinates range
filter!(:longitude => >(coords.left), df)
filter!(:longitude => <(coords.right), df)
filter!(:latitude => >(coords.bottom), df)
filter!(:latitude => <(coords.top), df)

# Separate warblers species
warblers = groupby(df, :species)

# Reorder species by decreasing frequency
warblers = warblers |>
    x -> combine(nrow, x) |>
    x -> sortperm(x, :nrow, rev=true) |>
    x -> warblers[x]

# Extract species names
spenames = [w.species[1] for w in warblers]
specommon = [w.commonName[1] for w in warblers]

# Create index Dict for species names
speindex = indexmap(spenames)

## Get environmental data

# WorldClim data
wc_vars = SimpleSDMPredictor(WorldClim, BioClim, [1, 12]; resolution=10.0, coords...)

# Landcover data
lc_vars = SimpleSDMPredictor(Copernicus, LandCover, 1:10; resolution=10.0, coords...)

# Combine environmental data
env_vars = vcat(wc_vars, lc_vars)

## Get distribution for all species
# create_distributions = true # should distributions be computed (optional)
if (@isdefined create_distributions) && create_distributions == true
    @info "Raw distributions to be created"
    # Get raw distributions
    @time distributions = @showprogress [presence_absence(w, env_vars[1]) for w in warblers]
end

## Export distributions
# save_data = true # should data files be overwritten (optional)
if (@isdefined save_data) && save_data == true
    ## Export species names
    @info "Exporting species names to JLD2 file"
    @save joinpath("data", "jld2", "spenames.jld2") spenames specommon speindex

    ## Export glossary
    @info "Exporting species names to glossary"
    # Add species names in correct order to glossary
    include(joinpath("others", "data_glossary.jl"))

    ## Export layers

    # Export distribution layers
    geotiff(joinpath("data", "raster", "distributions_raw.tif"), distributions)

    # Export distributions for QC only
    coords_qc = (left=-80.0, right=-55.0, bottom=45.0, top=63.0)
    distributions_qc = [d[coords_qc] for d in distributions]
    geotiff(joinpath("data", "raster", "distributions_raw_qc.tif"), distributions_qc)

    ## Export to CSV as Y matrix
    @info "Exporting data to CSV as Y matrix (raw distributions data)"

    # Get Y matrix
    Y = Ymatrix(distributions)
    Yobs = Ymatrix(distributions; observed=true)
    inds_obs = _indsobs(Y)

    # Convert to dataframe
    spe_df = DataFrame(Yobs, :auto)
    rename!(spe_df, Symbol.("sp", 1:ncol(spe_df)))
    insertcols!(spe_df, 1, :site => inds_obs)

    # Export to CSV
    CSV.write(joinpath("data", "proc", "distributions_spe_full.csv"), spe_df; delim="\t")
else
    # Load data
    @info "Data imported from file (species names)"
    @load joinpath("data", "jld2", "spenames.jld2") spenames specommon speindex
    distributions = [
        geotiff(
            SimpleSDMPredictor, joinpath("data", "raster", "distributions_raw.tif"), i
        ) for i in eachindex(spenames)
    ]
end

## Count sites with presence per species
pres_counts = [
    length(filter(x -> !isnothing(x) && x > 0.0, species.grid)) for species in distributions
]
sort(pres_counts)

## Plot result

# Species 1
sp1 = "Setophaga_townsendi"
map_sp1 = plot_layer(
    distributions[speindex[sp1]];
    c=:BuPu,
    title="$(replace(sp1, "_" => " ")) distribution (raw)",
    colorbar=:none,
    dpi=200,
)

# Species 2
sp2 = "Setophaga_petechia"
map_sp2 = plot_layer(
    distributions[speindex[sp2]];
    c=:BuPu,
    title="$(replace(sp2, "_" => " ")) distribution (raw)",
    colorbar=:none,
    dpi=200,
)


@testset "Parsimonious flux balance analysis with StandardModel" begin
    model = test_toyModel()

    d = parsimonious_flux_balance_analysis_dict(
        model,
        Tulip.Optimizer;
        modifications = [
            change_constraint("EX_m1(e)", lb = -10.0),
            change_optimizer_attribute("IPM_IterationsLimit", 500),
        ],
        qp_modifications = [
            change_optimizer(OSQP.Optimizer),
            change_optimizer_attribute("polish", true),
            silence,
        ],
    )

    # The used optimizer doesn't really converge to the same answer everytime
    # here, we therefore tolerate a wide range of results.
    @test isapprox(d["biomass1"], 10.0, atol = QP_TEST_TOLERANCE)
end


"""
    SteppedLayeredModel <: SeisModel1D

A `SteppedLayeredModel` contains `n` layers each with constant properties.
"""
struct SteppedLayeredModel <: SeisModel1D
    "Earth radius in km"
    a :: Float64
    "Number of layers"
    n :: Int
    "Radii in km of top of each layer"
    r :: Vector{Float64}
    "Physical parameters"
    vp :: Vector{Float64}
    vs :: Vector{Float64}
    density :: Vector{Float64}
    "Optional anisotropic parameters"
    aniso :: Bool
    vph :: Vector{Float64}
    vpv :: Vector{Float64}
    vsh :: Vector{Float64}
    vsv :: Vector{Float64}
    eta :: Vector{Float64}
    "Optional attenuation parameters"
    attenuation :: Bool
    Qμ :: Vector{Float64}
    Qκ :: Vector{Float64}
end

"""
    SteppedLayeredModel(; kwargs...) -> m

Construct a `SteppedLayeredModel` by specifying keyword arguments.  The length of
arrays defines the number of layers and the presence of arguments determines
whether the model has density, anisotropy and attenuation.

Note that radii are given in km, velocities in km/s, densities in g/cm³
and frequency in Hz.

### Keyword arguments
#### Radii
- `r`: The upper radius of a set of constant layers in km.  The uppermost
  layer radius is taken to be the planetary radius.  This is the only
  mandatory argument.
#### Isotropic properties
These must be given each as a `Vector` of values for each layer.
Each parameter (`vp`, `density`, `Qμ`, etc.) must have the correct number
of layers.
- `vp`: P-wave velocity
- `vs`: S-wave velocity
- `density`: Material density
#### Anisotropic parameters
If the model is anisotropic and the radial anisotropy parameters are all
provided, then the Voigt average velocities are constructed automatically
unless `vp` and/or `vs` are also given separately.
- `vpv`: Vertical P-wave velocity
- `vph`: Horizontal P-wave velocity
- `vsv`: Vertical S-wave velocity
- `vsh`: Horizontally-polarised S-wave velocity
- `eta`: The ratio C₁₃/(Vph - 2Vsv), where cᵢⱼ is the Voigt elasticity
  matrix, the 1-direction is horizontal and the 3-direction is radial
#### Attenuation parameters
ASCII versions (`Qmu`, `Qkappa`) of these exist for ease of use as well as
the Unicode versions (`Qμ`, `Qκ`).  If both are supplied, an error is thrown.
- `Qμ` or `Qmu`: Shear modulus quality factor
- `Qκ` or `Qkappa`: Bulk modulus quality factor
"""
function SteppedLayeredModel(; r, vp=nothing, vs=nothing, density=nothing,
                             vph=nothing, vpv=nothing, vsh=nothing, vsv=nothing,
                             eta=nothing, Qμ=nothing, Qmu=nothing,
                             Qκ=nothing, Qkappa=nothing, fref=nothing)
    n = length(r)
    all(x->x>0, diff(r)) || throw(ArgumentError("radii must increase monotonically"))
    R = last(r)
    aniso = attenuation = false
    # Anisotropy
    if all(!isnothing, (vph, vpv, vsh, vsv, eta))
        aniso = true
        # Set isotropic velocities automatically if not provided
        if any(isnothing, (vp, vs))
            voigt_vels = voigt_velocities.(vpv, vsv, vph, vsh, eta)
            vp === nothing && (vp = first.(voigt_vels))
            vs === nothing && (vs = last.(voigt_vels))
        end
    end
    # Attenuation
    # Default to ASCII versions if given
    Qmu !== nothing && (Qμ = Qmu)
    Qkappa !== nothing && (Qκ = Qkappa)
    if all(!isnothing, (Qμ, Qκ))
        attenuation = true
    end
    # Convert to correct shape, and check for length
    vp = _set_coeffs_size(SteppedLayeredModel, n, vp)
    vs = _set_coeffs_size(SteppedLayeredModel, n, vs)
    density = _set_coeffs_size(SteppedLayeredModel, n, density)
    vph = _set_coeffs_size(SteppedLayeredModel, n, vph)
    vsh = _set_coeffs_size(SteppedLayeredModel, n, vsh)
    vpv = _set_coeffs_size(SteppedLayeredModel, n, vpv)
    vsv = _set_coeffs_size(SteppedLayeredModel, n, vsv)
    eta = _set_coeffs_size(SteppedLayeredModel, n, eta)
    Qμ = _set_coeffs_size(SteppedLayeredModel, n, Qμ)
    Qκ = _set_coeffs_size(SteppedLayeredModel, n, Qκ)
    SteppedLayeredModel(R, n, r, vp, vs, density, aniso, vph, vpv,
        vsh, vsv, eta, attenuation, Qμ, Qκ)
end

_set_coeffs_size(::Type{SteppedLayeredModel}, n, v::Nothing) = []
function _set_coeffs_size(::Type{SteppedLayeredModel}, n, v::Vector)
    length(v) == n || throw(ArgumentError("number of value must be $n"))
    v
end

for sym in model_variables_SeisModel1D
    @eval function ($sym)(m::SteppedLayeredModel, r::Real; depth::Bool=false)
        length(m.$sym) > 0 ||
            throw(ArgumentError("'$split(string($sym), ".")[end]' not defined for model"))
        depth && (r = radius(m, r))
        m.$(sym)[findlayer(m, r)]
    end
end

function evaluate(m::SteppedLayeredModel, field::Symbol, r; depth::Bool=false)
    length(getfield(m, field)) > 0 ||
        throw(ArgumentError("$field' not defined for model"))
    depth && (r = radius(m, r))
    getfield(m, field)[findlayer(m, r)]
end

function mass(m::SteppedLayeredModel, r; depth::Bool=false)
    hasdensity(m) || throw(ArgumentError("model does not contain density"))
    depth && (r = radius(m, r))
    l = findlayer(m, r)
    r *= 1.e3
    M = 0.
    for i in 1:l
        R0 = i == 1 ? 0. : m.r[i-1]*1.e3
        R = i < l ? m.r[i]*1.e3 : r
        M += 1.e3*m.density[i]*(R^3 - R0^3)
    end
    4/3*π*M
end


"""
    simulate_detector(sim_config_filename)

AbstractString -> <detector simulation>

Simulate detector according to given simulation configuration.
What type <detector simulation> is depends on the given method of simulation.
"""
function simulate_detector(sim_config::PropDict)
    det_meta = propdict(sim_config.detector_metadata)
    env = Environment(sim_config)
    ps_simulator = PSSimulator(sim_config)

    simulate_detector(det_meta, env, sim_config.pss.cached_name, ps_simulator)
end

function simulate_detector(sim_config_filename::AbstractString)
    simulate_detector(propdict(sim_config_filename))
end

function simulate_detector(det_metadata::AbstractString, sim_config_filename::AbstractString)
    # merge detector info with the rest
    sim_config = load_config(det_metadata, sim_config_filename)
    simulate_detector(sim_config)
end

"""
    simulate_detector(det_meta, env, config_name, ::SSDSimulator)

PropDict, Environment, AbstractString -> SSD.Simulation

Look up cached SSD simulation .h5f file, and if does not exist,
    simulate the detector according to geometry as given in LEGEND metadata <det_meta>
    and environment settings given in <env> using SolidStateDetectors.
The simulation will be cached based on <config_name>.

"""
function simulate_detector(det_meta::PropDict, env::Environment, cached_name::AbstractString, ssd_sim::SSDSimulator)
    det_h5 = joinpath("cache", cached_name*"_ssd.h5f")
    if isfile(det_h5)
        @info("Reading SSD simulation from cached file $det_h5")
        simulation = SolidStateDetectors.ssd_read(det_h5, Simulation)
    else
        @info("Simulating $cached_name with SSD from scratch for given settings")
        ssd_conf = ssd_config(det_meta, env, ssd_sim)
        simulation = simulate_ssd(ssd_conf)

        if !ispath(dirname(det_h5)) mkpath(dirname(det_h5)) end
        SolidStateDetectors.ssd_write(det_h5, simulation)
        @info("-> Saved cached simulation to $det_h5")
    end

    simulation
end

"""
    simulate_detector(det_meta, env, config_name, ps_simulator)

PropDict, Environment, AbstractString, Siggen -> SigGenSetup

Construct a fieldgen/siggen configuration file
    according to geometry as given in LEGEND metadata <det_meta>
    and environment settings given in <env>.
Look up fieldgen generated electric field and weighting potential files
    based on given <config_name>, and if not found, call fieldgen.
"""
function simulate_detector(det_meta::PropDict, env::Environment, cached_name::AbstractString, ps_simulator::SiggenSimulator)
    @info "Constructing Fieldgen/Siggen configuration file"
    # returns the name of the resulting siggen config file
    # and the name of (already or to be) generated weighting potential file
    siggen_config_name, fieldgen_wp_name = siggen_config(det_meta, env, ps_simulator, cached_name)
    fieldgen_wp_name = joinpath("cache", fieldgen_wp_name)

    # if the WP file with such a name exists...
    if isfile(fieldgen_wp_name)
        #...do nothing, siggen will later read the files based on the generated conifg
        @info "Reading cached fields from $fieldgen_wp_name"
    else
        #...otherwise call fieldgen
        @info "_|~|_|~|_|~|_ Fieldgen simulation"
        fieldgen(siggen_config_name)
        @info "_|~|_|~|_|~|_ Fieldgen simulation complete"
    end

    # a SigGenSetup object
    SigGenSetup(siggen_config_name)
end

# ------------------------------------------------------------------- SSD

"""
    simulate_ssd(ssd_config)

Dict -> SSD.Simulation

Simulate detector based on given SSD config
"""
function simulate_ssd(ssd_config::Dict)

    @info "_|~|_|~|_|~|_ SSD simulation"

    simulation = Simulation(SolidStateDetector{Float32}(ssd_config))

    println("...electric potential")
    calculate_electric_potential!( simulation,
                               max_refinements = 4)

    println("...electric field")
    calculate_electric_field!(simulation, n_points_in_φ = 72)

    println("...drift field")
    calculate_drift_fields!(simulation)

    println("...weighting potential")
    for contact in simulation.detector.contacts
        calculate_weighting_potential!(simulation, contact.id, max_refinements = 4, n_points_in_φ = 2, verbose = false)
    end

    @info "_|~|_|~|_|~|_ SSD detector simulation complete"

    simulation
end

"""
    ssd_config(meta, env)

PropDict, Environment -> Dict

Construct SSD type config based on
    geometry as given in LEGEND metadata <meta>
    and environment settings given in <env>.
"""
function ssd_config(meta::PropDict, env::Environment, ssd_sim::SSDSimulator)
    ## some parameters that will be used multiple times later
    # or need a calculation
    cylinder_height = meta.geometry.height_in_mm
    cylinder_radius = meta.geometry.radius_in_mm
    borehole_height = cylinder_height - meta.geometry.well.gap_in_mm
    borehole_r_bottom = meta.geometry.well.radius_in_mm
    borehole_r_top = borehole_r_bottom + borehole_height * tan(meta.geometry.taper.top.inner.angle_in_deg / 180. * π)
    cone_r_top = cylinder_radius - meta.geometry.taper.top.outer.height_in_mm * tan(meta.geometry.taper.top.outer.angle_in_deg / 180. * π)
    cone_height = meta.geometry.taper.top.outer.height_in_mm

    ## grid
    # kind of initialize so that it exists out of the loop?
    grid = Dict()

    if ssd_sim.coord == "cylindrical"
        grid = Dict(
            "axes" => Dict(
                "r" => Dict("to" => cylinder_radius*1.2, "boundaries" => "inf"), # leave some margin
                "phi" => Dict("from" => 0, "to" => ssd_sim.comp == "3D" ? 360 : 0, "boundaries" => "periodic"),
                "z" => Dict(
                    "from" => -0.2*cylinder_height,
                    "to" => cylinder_height*1.2, # leave some margin
                    "boundaries" => Dict("left" => "inf", "right" => "inf")
                )
            )
        )
    else # cartesian
        grid = Dict(
            "axes" => Dict(
                "x" => Dict("from" => -cylinder_radius*1.2, "to" => cylinder_radius*1.2, "boundaries" => "inf"),
                "y" => Dict("from" => -cylinder_radius*1.2, "to" => cylinder_radius*1.2, "boundaries" => "inf"),
                "z" => Dict("from" => -0.2*cylinder_height, "to" => cylinder_height*1.2, "boundaries" => "inf")
            )
        )
    end

    grid = merge(grid, Dict("coordinates" => ssd_sim.coord))

    ## charge density model
    # kind of initialize so that it exists out of the loop?
    charge_density_model = Dict()

    if ssd_sim.coord == "cylindrical"
        charge_density_model = Dict(
            "r" => Dict("init" => 0, "gradient" => 0),
            "phi" => Dict("init" => 0.0, "gradient" => 0.0),
            "z" => Dict("init" => -1e7, "gradient" => 1e5)
        )
    else # cartesian
        charge_density_model = Dict(
            "x" => Dict("init" => 0.0, "gradient" => 0.000),
            "y" => Dict("init" => 0.0, "gradient" => 0.0),
            "z" => Dict("init" => 1e7, "gradient" => 5e4)
        )
    end

    # for now we only use linear charge density model for either cartesian or cylindrical coordinates
    # charge_density_model["name"] = "linear"
    charge_density_model = merge(charge_density_model, Dict("name" => "linear"))

    ## geometry
    dct = Dict(
        "name" => meta.det_name,
        "medium" => env.medium,
        "units" => Dict(
            "length" => "mm",
            "angle" => "deg",
            "potential" => "V",
            "temperature" => "K"
        ),
        "grid" => grid,
        "objects" => [
            # crystal
            Dict(
                "type" => "semiconductor",
                "material" => "HPGe",
                "bulk_type" => "p",
                "temperature" => env.crystal_t,
                "charge_density_model" => charge_density_model,
                "geometry" => Dict(
                    "type" => "difference",
                    "parts" => [
                        # crystal
                        Dict(
                            "type" => "difference",
                            "parts" => [
                                # cylinder
                                Dict(
                                    "name" => "Initial Cylinder",
                                    "type" => "tube",
                                    "r" => Dict(
                                        "from" => 0.0,
                                        "to" => cylinder_radius,
                                    ),
                                    "phi" => Dict("from" => 0.0, "to" => 360.0),
                                    "h" => cylinder_height,
                                    "translate" => Dict("z" => 0.0)
                                )#, # cylinder
                                # upper cone
                                Dict(
                                    "name" => "Upper cone",
                                    "type" => "cone",
                                    "r" => Dict(
                                        "bottom" => Dict("from" => cylinder_radius, "to" => cylinder_radius+1),
                                        "top" => Dict("from" => cone_r_top, "to" => cylinder_radius+1)
                                    ),
                                    "phi" => Dict("from" => 0.0, "to" => 360.0),
                                    "h" => meta.geometry.taper.top.outer.height_in_mm,
                                    "translate" => Dict("z" => cylinder_height - cone_height)
                                ) # upper cone
                            ] # crystal parts
                        ), # crystal
                        # borehole
                        Dict(
                            "name" => "borehole",
                            "type" => "cone",
                            "r" => Dict(
                                "bottom" => Dict(
                                    "from" => 0.0,
                                    "to" => borehole_r_bottom
                                ),
                                "top" => Dict(
                                    "from" => 0.0,
                                    "to" => borehole_r_top
                                )
                            ),
                            "phi" => Dict("from" => 0.0, "to" => 360.0),
                            "h" => borehole_height,
                            "translate" => Dict("z" => meta.geometry.well.gap_in_mm)
                        ), # borehole
                        # groove
                        Dict(
                            "name" => "ditch",
                            "type" => "tube",
                            "r" => Dict(
                                "from" => meta.geometry.groove.outer_radius_in_mm - meta.geometry.groove.width_in_mm,
                                "to" => meta.geometry.groove.outer_radius_in_mm
                            ),
                            "phi" => Dict("from" => 0.0, "to" => 360.0),
                            "h" => meta.geometry.groove.depth_in_mm
                        ) # groove
                    ] # parts
                ) # geometry
            ), # crystal
            # p+ contact
            Dict(
                "type" => "contact",
                "material" => "HPGe",
                "channel" => 1,
                "potential" => 0.0,
                "geometry" => Dict(
                    "type" => "tube",
                    "r" => Dict(
                        "from" => 0.0,
                        "to" => meta.geometry.contact.radius_in_mm
                    ),
                    "phi" => Dict("from" => 0.0, "to" => 360.0),
                    "h" => meta.geometry.contact.depth_in_mm
                )
            ), # p+ contact
            # n+ contact
            Dict(
                "type" => "contact",
                "material" => "HPGe",
                "channel" => 2,
                "potential" => env.op_voltage == 0 ? meta.characterization.manufacturer.op_voltage_in_V : env.op_voltage,
                "geometry" => Dict(
                    "type" => "union",
                    "parts" => [Dict() for i in 1:6] # to be filled later
                )
            ) # n+ contact
        ] # objects
    )

    ## Construct n+ surface

    # bottom lid
    dct["objects"][3]["geometry"]["parts"][1] = Dict(
        "name" => "bottom lid",
        "type" => "tube",
        "r" => Dict(
            "from" => meta.geometry.groove.outer_radius_in_mm,
            "to" => cylinder_radius # cylinder radius
        ),
        "phi" => Dict("from" => 0.0, "to" => 360.0),
        "h" => 0
    )

    # side lining
    dct["objects"][3]["geometry"]["parts"][2] = Dict(
        "name" => "side lining",
        "type" => "tube",
        "r" => Dict(
            "from" => cylinder_radius,
            "to" => cylinder_radius
        ),
        "phi" => Dict("from" => 0.0, "to" => 360.0),
        "h" => cylinder_height - cone_height
    )

    # upper cone
    dct["objects"][3]["geometry"]["parts"][3] = Dict(
        "name" => "Upper cone",
        "type"=> "cone",
        "r" => Dict(
            "bottom" => Dict("from" => cylinder_radius, "to" => cylinder_radius),
            "top" => Dict("from" => cone_r_top, "to" => cone_r_top)
        ),
        "phi" => Dict("from" => 0.0, "to" => 360.0),
        "h" => cone_height,
        "translate" => Dict("z" => cylinder_height - cone_height)
    )

    # top lid
    dct["objects"][3]["geometry"]["parts"][4] = Dict(
        "name" => "top lid",
        "type" => "tube",
        "r" => Dict(
            "from" => borehole_r_top,
            "to" => cone_r_top
        ),
        "phi" => Dict("from" => 0.0, "to" => 360.0),
        "h" => 0.0,
        "translate" => Dict("z" => cylinder_height)
    )

    # borehole lining
    dct["objects"][3]["geometry"]["parts"][5] = Dict(
        "name" => "borehole linning",
        "type" => "cone",
        "r" => Dict(
            "bottom" => Dict(
                "from" => borehole_r_bottom,
                "to" => borehole_r_bottom
            ),
            "top" => Dict(
                "from" => borehole_r_top,
                "to" => borehole_r_top
            )
        ),
        "phi" => Dict("from" => 0.0, "to" => 360.0),
        "h" => borehole_height,
        "translate" => Dict("z" => meta.geometry.well.gap_in_mm)
    )

    # borehole bottom lid
    dct["objects"][3]["geometry"]["parts"][6] = Dict(
        "name" => "borehole bottom lid",
        "type" => "tube",
        "r" => Dict(
            "from" => 0.0,
            "to" => borehole_r_bottom
        ),
        "phi" => Dict("from" => 0.0, "to" => 360.0),
        "h" => 0.0,
        "translate" => Dict("z" => meta.geometry.well.gap_in_mm)
    )

    dct
end

"""
    ssd_config(det_meta)

AbstractString -> Dict

Construct SSD type config based on geometry as given in LEGEND metadata <det_meta>.
Set arbitrary envorinment variables and simulation settings (used only for geometry visualization)
"""
function ssd_config(det_meta::AbstractString)
    ssd_config(propdict(det_meta), Environment(), SSDSimulator())
end

"""
    siggen_config(meta, env, siggen_sim, config_name)

PropDict, Environment, SiggenSimulator, String -> String, String

Construct Siggen/Fieldgen config based on
    geometry as given in LEGEND metadata <meta>,
    environment settings given in <env>,
    and fieldgen settings contained in <siggen_sim>.

The resulting siggen config will be cached based on given <config_name>
The function returns the path to the cached siggen config file, and
    the relative path to the fieldgen WP file (to check if already exists later)
"""
function siggen_config(meta::PropDict, env::Environment, siggen_sim::SiggenSimulator, cached_name::AbstractString)
    # function pss_config(meta::PropDict, env::Environment, siggen_sim::SiggenSimulatorss, config_name::AbstractString)
    # construct siggen geometry based on LEGEND metadata JSON + extra environment settings
    println("...detector geometry")
    siggen_geom = meta2siggen(meta, env)

    # construct lines for future siggen config file
    detector_lines = Vector{String}()
    push!(detector_lines, "# detector related parameters")
    for (param, value) in siggen_geom
        push!(detector_lines, "$param   $value")
    end

    println("...fieldgen configuration")
    # create filenames for future/existing fieldgen output
    fieldgen_wp_name = joinpath("fieldgen", cached_name*"_fieldgen_WP.dat")
    fieldgen_names = Dict(
        "drift_name" => joinpath("..", siggen_sim.drift_vel),
        # "drift_name" => joinpath("..", "data", "drift_vel_tcorr.tab"),
        "field_name" => joinpath("fieldgen", cached_name*"_fieldgen_EV.dat"),
        "wp_name" => fieldgen_wp_name
    )

    # add corresponding lines to existing detector geometry lines
    push!(detector_lines, "")
    push!(detector_lines, "# fieldgen filenames")
    for (param, value) in fieldgen_names
        push!(detector_lines, "$param   $value")
    end

    # read lines from user input for fieldgen
    fieldgen_lines = readlines(open(siggen_sim.fieldgen_config))
    fieldgen_lines = replace.(fieldgen_lines, "\t" => "   ")

    # unite constructed detector geometry and fieldgen settings
    total_lines = vcat(detector_lines, [""], fieldgen_lines)

    # write a siggen/fieldgen config file
    sig_conf_name = joinpath("cache", cached_name*"_siggen_config.txt")
    writedlm(sig_conf_name, total_lines)
    println("...fieldgen/siggen config written to $sig_conf_name")

    # return resulting fieldgen/siggen config name
    # (currently kinda redundant but that's because we have no idea what's gonna happen later)
    sig_conf_name, fieldgen_wp_name
end

"""
    meta2siggen(meta, env)

PropDict, Environment -> Dict

Construct siggen type config in a form of Dict based on
    geometry as given in LEGEND metadata <meta>
    and environment settings given in <env>.
"""
function meta2siggen(meta::PropDict, env::Environment)
    Dict(
    # crystal
        # z length: crystal height
        "xtal_length" => meta.geometry.height_in_mm,
        # radius: crystal radius
        "xtal_radius" => meta.geometry.radius_in_mm,
    # contact
        # point contact length
        "pc_length" => meta.geometry.contact.depth_in_mm,
        # point contact radius
        "pc_radius" => meta.geometry.contact.radius_in_mm,
    # groove
        # wrap-around radius. Set to zero for ORTEC: groove outer radius
        "wrap_around_radius" => meta.geometry.groove.outer_radius_in_mm,
        # depth of ditch next to wrap-around for BEGes. Set to zero for ORTEC: groove height
        "ditch_depth" => meta.geometry.groove.depth_in_mm,
        # width of ditch next to wrap-around for BEGes. Set to zero for ORTEC: groove inner radius
        "ditch_thickness" => meta.geometry.groove.width_in_mm,
    # borehole
        # length of hole, for inverted-coax style
        "hole_length" => meta.geometry.height_in_mm - meta.geometry.well.gap_in_mm,
        # radius of hole, for inverted-coax style
        "hole_radius" => meta.geometry.well.radius_in_mm,
    # tapering

    # comment from David Radford:

    # The bottom taper is always at 45 degrees, but the top outer taper is not.
    # It's width/angle is defined using either the outer_taper_width parameter or the taper_angle parameter.
    # Likewise, the inner taper width/angle is defined using either the inner_taper_width parameter or the taper_angle parameter.
    # The inner_taper_length can be anything from zero to the hole_length
    # TODO: check if bottom taper angle is 45 deg, if not issue a warning for siggen case

        # size of 45-degree taper at bottom of ORTEC-type crystal (for r=z)
        "bottom_taper_length" => meta.geometry.taper.bottom.outer.height_in_mm,
        # z-length of outside taper for inverted-coax style
        "outer_taper_length" => meta.geometry.taper.top.outer.height_in_mm,
        # z-length of inside (hole) taper for inverted-coax style
        "inner_taper_length" => meta.geometry.taper.top.inner.height_in_mm,
        # taper angle in degrees, for inner or outer taper
        "taper_angle" => meta.geometry.taper.top.outer.angle_in_deg,
        # "taper_angle" => meta.geometry.taper.top.inner.angle_in_deg,
        # depth of full-charge-collection boundary for Li contact (not currently used)
        "Li_thickness" => meta.geometry.dl_thickness_in_mm,

    # configuration for mjd_fieldgen (calculates electric fields & weighing potentials)
        # detector bias for fieldgen, in Volts
        "xtal_HV" => env.op_voltage == 0 ? meta.characterization.manufacturer.op_voltage_in_V : env.op_voltage,

    # configuration for signal calculation
        # crystal temperature in Kelvin
        "xtal_temp" => env.crystal_t
    )
end


################################################################################
#  Copyright 2020, Marta Vanin, Tom Van Acker                                  #
################################################################################
# PowerModelsDistributionStateEstimation.jl                                    #
# An extention package of PowerModels(Distribution).jl for Static Power System #
# State Estimation.                                                            #
################################################################################
# Contributor: Sander Claeys (@sanderclaeys)                                   #
################################################################################

## Electricity North West Data Parser
# constants
const ENWL_ROOT_PATH    = joinpath(BASE_DIR,"test/data/enwl")
const ENWL_PROFILE_PATH = "$ENWL_ROOT_PATH/profiles"
const ENWL_NETWORK_PATH = "$ENWL_ROOT_PATH/networks"
const ENWL_PROFILES     = Dict(
    ("summer","load") =>"Summer_Load_Profiles.csv",
    ("summer","pv")   =>"Summer_PV_Profiles.csv",
    ("winter","ehp")  =>"Winter_EHP_Profiles.csv",
    ("winter","ev")   =>"Winter_EV_Profiles.csv",
    ("winter","load") =>"Winter_Load_Profiles.csv",
    ("winter","uchp") =>"Winter_uCHP_Profiles.csv",
)
const ENWL_PROFILE_SIGN = Dict(
    ("summer","load") =>  1,
    ("summer","pv")   => -1,
    ("winter","ehp")  =>  1,
    ("winter","ev")   =>  1,
    ("winter","load") =>  1,
    ("winter","uchp") => -1,
)

# functions
get_profiles(ssn, dev) =
    convert(Matrix, _CSV.read("$ENWL_PROFILE_PATH/$(ENWL_PROFILES[(ssn, dev)])"))
function get_enwl_dss_path(network::Int, feeder::Int)
    return "$ENWL_NETWORK_PATH/network_$network/Feeder_$feeder/Master.dss"
end

"""
    insert_profiles!(data, season, devices, pfs; t=missing, useactual=true)

This function adds the load profiles to the parsed ENWL ENGINEERING data `data`.
# Arguments
- data: ENGINEERING data model of the feeder
- season: `"summer"`, `"winter"`
- devices: `"load"`, `"pv"`, `"ev"`, `"ehp"`, `"uchp"`
- pfs: power factor of the devices
- t: time-step

# Example
```julia-repl
julia> data = parse_file(get_enwl_dss_path(10, 1))
julia> insert_profiles!(data, "summer", ["load", "pv"], [0.95, 0.90], t = 144)
```
"""
function insert_profiles!(data, season, devices, pfs; t=missing, useactual=true)
    N = length(data["load"])
    mult = 1E3/data["settings"]["power_scale_factor"]
    sd_tot = zeros(287, N)
    for i in 1:length(devices)
        device = devices[i]; pf = pfs[i]
        linds  = ((1:N).-1).%100 .+1
        sign   = ENWL_PROFILE_SIGN[(season, device)]
        pd     = mult*sign*get_profiles(season, device)[:,linds]
        sd_tot += (1 + im*sqrt(1-pf^2))*pd
    end

    if ismissing(t)
        t = argmax(abs.(reshape(sum(real.(sd_tot), dims=2), 287)))
    end

    for i in 1:N
        load = data["load"]["load$i"]
        dims = length(load["pd_nom"])
        if useactual
            load["pd_nom"] = fill(real(sd_tot[t,i])/dims, dims)
            load["qd_nom"] = fill(imag(sd_tot[t,i])/dims, dims)
        else
            load["pd_nom"] *= real(sd_tot[t,i])
            load["qd_nom"] *= imag(sd_tot[t,i])
        end
    end
end
"""
    rm_enwl_transformer!(data_eng)

This function removes the transformer from a parsed ENWL `ENGINEERING` data file.
"""
function rm_enwl_transformer!(data_eng)
    if haskey(data_eng, "transformer")
        line1 = data_eng["line"]["line1"]
        trans = data_eng["transformer"]["tr1"]
        vprim_scale = trans["vm_nom"][2]/trans["vm_nom"][1]

        vsource = data_eng["voltage_source"]["source"]

        vsource["vm"] *= vprim_scale
        vsource["rs"] *= vprim_scale^2
        vsource["xs"] *= vprim_scale^2
        vsource["bus"] = "1"

        delete!(data_eng, "transformer")
        delete!(data_eng["bus"], "sourcebus")

        vbases_default = data_eng["settings"]["vbases_default"]
        vbases_default["1"] = vbases_default["sourcebus"]*vprim_scale
        delete!(vbases_default, "sourcebus")
    end
end
"""
    reduce_enwl_lines_eng!(data_eng)

This function removes all trailing lines from a parsed ENWL `ENGINEERING` data
file.
"""
function reduce_enwl_lines_eng!(data_eng)
    rm_trailing_lines_eng!(data_eng)
    join_lines_eng!(data_eng)
end
function rm_trailing_lines_eng!(data_eng)

    buses_exclude = []
    for comp_type in ["load", "shunt", "generator", "voltage_source"]
        if haskey(data_eng, comp_type)
            buses_exclude = union(buses_exclude, [comp["bus"] for (_, comp) in data_eng[comp_type]])
        end
    end
    if haskey(data_eng, "transformer")
        buses_exclude = union(buses_exclude, hcat([tr["bus"] for (_, tr) in data_eng["transformer"]]...))
    end

    line_has_shunt = Dict()
    bus_lines = Dict(k=>[] for k in keys(data_eng["bus"]))
    for (id, line) in data_eng["line"]
        lc = data_eng["linecode"][line["linecode"]]
        line_has_shunt[id] = !all(iszero(lc[k]) for k in ["b_fr", "b_to", "g_fr", "g_to"])
        push!(bus_lines[line["f_bus"]], id)
        push!(bus_lines[line["t_bus"]], id)
    end

    eligible_buses = [bus_id for (bus_id, line_ids) in bus_lines if length(line_ids)==1 && !(bus_id in buses_exclude) && !line_has_shunt[line_ids[1]]]

    while !isempty(eligible_buses)
        for bus_id in eligible_buses
            # this trailing bus has one associated line
            line_id = bus_lines[bus_id][1]
            line = data_eng["line"][line_id]

            delete!(data_eng["line"], line_id)
            delete!(data_eng["bus"],  bus_id)

            other_end_bus = line["f_bus"]==bus_id ? line["t_bus"] : line["f_bus"]
            bus_lines[other_end_bus] = setdiff(bus_lines[other_end_bus], [line_id])
            delete!(bus_lines,  bus_id)
        end

        eligible_buses = [bus_id for (bus_id, line_ids) in bus_lines if length(line_ids)==1 && !(bus_id in buses_exclude) && !line_has_shunt[line_ids[1]]]
    end
end
function _line_reverse_eng!(line)
    prop_pairs = [("f_bus", "t_bus")]

    for (x,y) in prop_pairs
        tmp = line[x]
        line[x] = line[y]
        line[y] = tmp
    end
end
function join_lines_eng!(data_eng)
    # a bus is eligible for reduction if it only appears in exactly two lines
    buses_all = collect(keys(data_eng["bus"]))
    buses_exclude = []

    # start by excluding all buses that appear in components other than lines
    for comp_type in ["load", "shunt", "generator", "voltage_source"]
        if haskey(data_eng, comp_type)
            buses_exclude = union(buses_exclude, [comp["bus"] for (_, comp) in data_eng[comp_type]])
        end
    end

    # per bus, list all inbound or outbound lines
    bus_lines = Dict(bus=>[] for bus in buses_all)
    for (id, line) in data_eng["line"]
        push!(bus_lines[line["f_bus"]], id)
        push!(bus_lines[line["t_bus"]], id)
    end

    # exclude all buses that do not have exactly two lines connected to it
    buses_exclude = union(buses_exclude, [bus for (bus, lines) in bus_lines if length(lines)!=2])

    # now loop over remaining buses
    candidates = setdiff(buses_all, buses_exclude)
    for bus in candidates
        line1_id, line2_id = bus_lines[bus]
        line1 = data_eng["line"][line1_id]
        line2 = data_eng["line"][line2_id]

        # reverse lines if needed to get the order
        # (x)--fr-line1-to--(bus)--to-line2-fr--(x)
        if line1["f_bus"]==bus
            _line_reverse_eng!(line1)
        end
        if line2["f_bus"]==bus
            _line_reverse_eng!(line2)
        end

        reducable = true
        reducable = reducable && line1["linecode"]==line2["linecode"]
        reducable = reducable && all(line1["t_connections"].==line2["t_connections"])
        if reducable

            line1["length"] += line2["length"]
            line1["t_bus"] = line2["f_bus"]
            line1["t_connections"] = line2["f_connections"]

            delete!(data_eng["line"], line2_id)
            delete!(data_eng["bus"], bus)
            for x in candidates
                if line2_id in bus_lines[x]
                    bus_lines[x] = [setdiff(bus_lines[x], [line2_id])..., line1_id]
                end
            end
        end
    end

    return data_eng
end
"""
    reduce_enwl_lines_math!(data_math)

This function removes all trailing lines from a parsed ENWL `MATHEMATICAL` data
file.
"""
function reduce_enwl_lines_math!(data_math)

    # a bus is eligible for reduction if it only appears in exactly two lines
    buses_all = collect(keys(data_math["bus"]))
    buses_exclude = []

    # start by excluding all buses that appear in components other than lines
    for comp_type in ["load", "shunt", "gen"]
        if haskey(data_math, comp_type)
            buses_exclude = union(buses_exclude, [string(comp["$(comp_type)_bus"]) for (_, comp) in data_math[comp_type]])
        end
    end

    # per bus, list all inbound or outbound lines
    bus_lines = Dict(bus=>[] for bus in buses_all)
    for (id, line) in data_math["branch"]
        push!(bus_lines[string(line["f_bus"])], id)
        push!(bus_lines[string(line["t_bus"])], id)
    end

    # exclude all buses that do not have exactly two lines connected to it
    buses_exclude = union(buses_exclude, [bus for (bus, lines) in bus_lines if length(lines)!=2])

    # now loop over remaining buses
    candidates = setdiff(buses_all, buses_exclude)
    for bus in candidates
        line1_id, line2_id = bus_lines[bus]
        line1 = data_math["branch"][line1_id]
        line2 = data_math["branch"][line2_id]

        # reverse lines if needed to get the order
        # (x)--fr-line1-to--(bus)--to-line2-fr--(x)
        if line1["f_bus"]==bus
            _line_reverse!(line1)
        end
        if line2["f_bus"]==bus
            _line_reverse!(line2)
        end

        reducable = true
        reducable = reducable && iszero(line1["g_to"]) && iszero(line1["b_to"])
        reducable = reducable && iszero(line2["g_fr"]) && iszero(line2["b_fr"])
        if reducable
            line1["br_r"] = line1["br_r"] .+ line2["br_r"]
            line1["br_x"] = line1["br_x"] .+ line2["br_x"]
            line1["g_to"] = line2["g_fr"]
            line1["b_to"] = line2["b_fr"]
            line1["t_bus"] = line2["f_bus"]

            delete!(data_math["branch"], line2_id)
            for x in candidates
                if line2_id in bus_lines[x]
                    bus_lines[x] = [setdiff(bus_lines[x], [line2_id])..., line1_id]
                end
            end
        end
    end

    return data_math
end
function _line_reverse!(line)
    prop_pairs = [("f_bus", "t_bus"), ("g_fr", "g_to"), ("b_fr","b_to")]

    for (x,y) in prop_pairs
        tmp = line[x]
        line[x] = line[y]
        line[y] = tmp
    end
end


abstract type AbstractRALF end

# RALF1
mutable struct RALF1{LC} <: AbstractRALF
    f::Function
    f0::Function
    cache::LC
end

get_cache_type(ral::RALF1{LC}) where {LC} = LC

function RALF1(f::Function, x1::C1, cache::AbstractArray{<: Number};
               chunksize::Int = ForwardDiff.pickchunksize(length(x1))) where {C1 <: AbstractArray{<: Number}, N}
    if applicable(f, cache, x1)
        fnew = function _f_ip(cache::LCN, x1::C1N) where {LCN <: DiffCache, C1N <: AbstractArray{<: Number}}
            target_cache = get_tmp(cache, x1)
            if size(target_cache) != size(cache.du)
                target_cache = reshape(target_cache, size(cache.du))
            end
            f(target_cache, x1)
            return target_cache
        end
    else
        fnew = function _f_oop(cache::LCN, x1::C1N) where {LCN <: DiffCache, C1N <: AbstractArray{<: Number}}
            target_cache = get_tmp(cache, x1)
            if size(target_cache) != size(cache.du)
                target_cache = reshape(target_cache, size(cache.du))
            end
            target_cache .= f(x1)
            return target_cache
        end
    end
    return RALF1(fnew, f, dualcache(cache, Val{chunksize}))
end

function RALF1(fin::LC) where {LC <: AbstractArray{<: Number}}
    f(cache::LCN, x1::C1N) where {LCN <: AbstractMatrix{ <: Number}, C1N <: AbstractArray{<: Number}} = cache
    return RALF1{LC}(f, x -> fin, fin)
end

function (ralf::RALF1)(x1::C1) where {C1 <: AbstractArray{<: Number}}
    return ralf.f(ralf.cache, x1)
end

# RALF2
mutable struct RALF2{LC} <: AbstractRALF
    f::Function
    f0::Function
    cache::LC
end

get_cache_type(ral::RALF2{LC}) where {LC} = LC

function RALF2(f::Function, x1::C1, x2::C2, cache::AbstractArray{<: Number}, chunksizes::NTuple{Nc, Int} =
               (ForwardDiff.pickchunksize(min(length(x1), length(x2))), )) where {C1 <: AbstractArray{<: Number}, C2 <: AbstractArray{<: Number}, N, Nc}
    if applicable(f, cache, x1, x2)
        if length(chunksizes) == 1 # Figure out which type of DiffCache is needed
            diffcache = dualcache(cache, Val{chunksizes[1]})
            fnew      = function _f_ip1(cache::LCN, x1::C1N, x2::C2N, select::Tuple{Int, Int}) where {LCN <: DiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                f(target_cache, x1, x2)
                return target_cache
            end
        elseif length(chunksizes) == 2
            diffcache = twodualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]})
            fnew      = function _f_ip2(cache::LCN, x1::C1N, x2::C2N, select::Tuple{Int, Int}) where {LCN <: TwoDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                f(target_cache, x1, x2)
                return target_cache
            end
        elseif length(chunksizes) == 3
            diffcache = threedualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]}, Val{chunksizes[3]})
            fnew      = function _f_ip3(cache::LCN, x1::C1N, x2::C2N, select::Tuple{Int, Int}) where {LCN <: ThreeDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                f(target_cache, x1, x2)
                return target_cache
            end
        else
            throw(MethodError("The length of the sixth input argument, chunksizes, must be 1, 2, or 3."))
        end
    else
        if length(chunksizes) == 1 # Figure out which type of DiffCache is needed
            diffcache = dualcache(cache, Val{chunksizes[1]})
            fnew      = function _f_oop1(cache::LCN, x1::C1N, x2::C2N, select::Tuple{Int, Int}) where {LCN <: DiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                target_cache .= f(x1, x2)
                return target_cache
            end
        elseif length(chunksizes) == 2
            diffcache = twodualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]})
            fnew      = function _f_oop2(cache::LCN, x1::C1N, x2::C2N, select::Tuple{Int, Int}) where {LCN <: TwoDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                target_cache .= f(x1, x2)
                return target_cache
            end
        elseif length(chunksizes) == 3
            diffcache = threedualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]}, Val{chunksizes[3]})
            fnew      = function _f_oop3(cache::LCN, x1::C1N, x2::C2N, select::Tuple{Int, Int}) where {LCN <: ThreeDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                target_cache .= f(x1, x2)
                return target_cache
            end
        else
            throw(MethodError("The length of the sixth input argument, chunksizes, must be 1, 2, or 3."))
        end
    end
    return RALF2(fnew, f, diffcache)
end

function RALF2(fin::LC) where {LC <: AbstractArray{<: Number}}
    f(cache::LCN, x1::C1N, x2::C2N, select::Tuple{Int, Int}) where {LCN <: AbstractArray{<: Number}, C1N <: AbstractArray{<: Number}, C2N <: AbstractArray{<: Number}} = cache
    return RALF2{LC}(f, x -> fin, fin)
end

# Using the default of (1, 1) for `select` is important. This way, we can use autodiff
# during the homotopy algorithm without requiring additional arguments when calling `update!`
# to ensure the correct cache is used.
function (ralf::RALF2)(x1::C1, x2::C2, select::Tuple{Int, Int} = (1, 1)) where {C1 <: AbstractArray{<: Number}, C2 <: AbstractArray{<: Number}}
    return ralf.f(ralf.cache, x1, x2, select)
end

# RALF3
mutable struct RALF3{LC} <: AbstractRALF
    f::Function
    f0::Function
    cache::LC
end

get_cache_type(ral::RALF3{LC}) where {LC} = LC

function RALF3(f::Function, x1::C1, x2::C2, x3::C3, cache::AbstractArray{<: Number},
               chunksizes::NTuple{Nc, Int} =
               (ForwardDiff.pickchunksize(min(length(x1),
                                              length(x2), length(x3))), )) where {C1 <: AbstractArray{<: Number}, C2 <: AbstractArray{<: Number},
                                                                                C3 <: AbstractArray{<: Number}, N, Nc}
    if applicable(f, cache, x1, x2, x3)
        if length(chunksizes) == 1 # Figure out which type of DiffCache is needed
            diffcache = dualcache(cache, Val{chunksizes[1]})
            fnew      = function _f_ip1(cache::LCN, x1::C1N, x2::C2N,
                                        x3::C3N, select::Tuple{Int, Int}) where {LCN <: DiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, x3, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                f(target_cache, x1, x2, x3)
                return target_cache
            end
        elseif length(chunksizes) == 2
            diffcache = twodualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]})
            fnew      = function _f_ip2(cache::LCN, x1::C1N, x2::C2N,
                                        x3::C3N, select::Tuple{Int, Int}) where {LCN <: TwoDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, x3, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                f(target_cache, x1, x2, x3)
                return target_cache
            end
        elseif length(chunksizes) == 3
            diffcache = threedualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]}, Val{chunksizes[3]})
            fnew      = function _f_ip3(cache::LCN, x1::C1N, x2::C2N,
                                        x3::C3N, select::Tuple{Int, Int}) where {LCN <: ThreeDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, x3, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                f(target_cache, x1, x2, x3)
                return target_cache
            end
        else
            throw(MethodError("The length of the seventh input argument, chunksizes, must be 1, 2, or 3."))
        end
    else
        if length(chunksizes) == 1 # Figure out which type of DiffCache is needed
            diffcache = dualcache(cache, Val{chunksizes[1]})
            fnew      = function _f_oop1(cache::LCN, x1::C1N, x2::C2N,
                                         x3::C3N, select::Tuple{Int, Int}) where {LCN <: DiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, x3, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                target_cache .= f(x1, x2, x3)
                return target_cache
            end
        elseif length(chunksizes) == 2
            diffcache = twodualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]})
            fnew      = function _f_oop2(cache::LCN, x1::C1N, x2::C2N,
                                         x3::C3N, select::Tuple{Int, Int}) where {LCN <: TwoDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, x3, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                target_cache .= f(x1, x2, x3)
                return target_cache
            end
        elseif length(chunksizes) == 3
            diffcache = threedualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]}, Val{chunksizes[3]})
            fnew      = function _f_oop3(cache::LCN, x1::C1N, x2::C2N,
                                         x3::C3N, select::Tuple{Int, Int}) where {LCN <: ThreeDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, x3, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                target_cache .= f(x1, x2, x3)
                return target_cache
            end
        else
            throw(MethodError("The length of the seventh input argument, chunksizes, must be 1, 2, or 3."))
        end
    end
    return RALF3(fnew, f, diffcache)
end

function RALF3(fin::LC) where {LC <: AbstractArray{<: Number}}
    f(cache::LCN, x1::C1N, x2::C2N, x3::C3N, select::Tuple{Int, Int}) where {LCN <: AbstractArray{<: Number}, C1N <: AbstractArray{<: Number}, C2N <: AbstractArray{<: Number}, C3N <: AbstractArray{<: Number}} = cache
    return RALF3{LC}(f, x -> fin, fin)
end

# Using the default of (1, 1) for `select` is important. This way, we can use autodiff
# during the homotopy algorithm without requiring additional arguments when calling `update!`
# to ensure the correct cache is used.
function (ralf::RALF3)(x1::C1, x2::C2, x3::C3,
                       select::Tuple{Int, Int} = (1, 1)) where {C1 <: AbstractArray{<: Number}, C2 <: AbstractArray{<: Number},
                                                                C3 <: AbstractArray{<: Number}}
    return ralf.f(ralf.cache, x1, x2, x3, select)
end

# RALF4
mutable struct RALF4{LC} <: AbstractRALF
    f::Function
    f0::Function
    cache::LC
end

get_cache_type(ral::RALF4{LC}) where {LC} = LC

function RALF4(f::Function, x1::C1, x2::C2, x3::C3, x4::C4, cache::AbstractArray{<: Number},
               chunksizes::NTuple{Nc, Int} =
               (ForwardDiff.pickchunksize(min(length(x1),
                                              length(x2), legnth(x3), length(x4))), )) where {C1 <: AbstractArray{<: Number}, C2 <: AbstractArray{<: Number},
                                                                                C3 <: AbstractArray{<: Number}, C4 <: AbstractArray{<: Number}, N, Nc}

    if applicable(f, cache, x1, x2, x3, x4)
        if length(chunksizes) == 1 # Figure out which type of DiffCache is needed
            diffcache = dualcache(cache, Val{chunksizes[1]})
            fnew      = function _f_ip1(cache::LCN, x1::C1N, x2::C2N,
                                        x3::C3N, x4::C4N, select::Tuple{Int, Int}) where {LCN <: DiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number},
                                                                                C4N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, x3, x4, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                f(target_cache, x1, x2, x3, x4)
                return target_cache
            end
        elseif length(chunksizes) == 2
            diffcache = twodualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]})
            fnew      = function _f_ip2(cache::LCN, x1::C1N, x2::C2N,
                                        x3::C3N, x4::C4N, select::Tuple{Int, Int}) where {LCN <: TwoDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number},
                                                                                C4N <: AbstractArray{<: Number}}
                f(get_tmp(cache, x1, x2, x3, x4, select), x1, x2, x3, x4)
                return get_tmp(cache, x1, x2, x3, x4, select)
            end
        elseif length(chunksizes) == 3
            diffcache = threedualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]}, Val{chunksizes[3]})
            fnew      = function _f_ip3(cache::LCN, x1::C1N, x2::C2N,
                                        x3::C3N, x4::C4N, select::Tuple{Int, Int}) where {LCN <: ThreeDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number},
                                                                                C4N <: AbstractArray{<: Number}}
                f(get_tmp(cache, x1, x2, x3, x4, select), x1, x2, x3, x4)
                return get_tmp(cache, x1, x2, x3, x4, select)
            end
        else
            throw(MethodError("The length of the eighth input argument, chunksizes, must be 1, 2, or 3."))
        end
    else
        if length(chunksizes) == 1 # Figure out which type of DiffCache is needed
            diffcache = dualcache(cache, Val{chunksizes[1]})
            fnew      = function _f_oop1(cache::LCN, x1::C1N, x2::C2N,
                                         x3::C3N, x4::C4N, select::Tuple{Int, Int}) where {LCN <: DiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number},
                                                                                C4N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, x3, x4, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                target_cache .= f(x1, x2, x3, x4)
                return target_cache
            end
        elseif length(chunksizes) == 2
            diffcache = twodualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]})
            fnew      = function _f_oop2(cache::LCN, x1::C1N, x2::C2N,
                                         x3::C3N, x4::C4N, select::Tuple{Int, Int}) where {LCN <: TwoDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number},
                                                                                C4N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, x3, x4, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                target_cache .= f(x1, x2, x3, x4)
                return target_cache
            end
        elseif length(chunksizes) == 3
            diffcache = threedualcache(cache, Val{chunksizes[1]}, Val{chunksizes[2]}, Val{chunksizes[3]})
            fnew      = function _f_oop3(cache::LCN, x1::C1N, x2::C2N,
                                         x3::C3N, x4::C4N, select::Tuple{Int, Int}) where {LCN <: ThreeDiffCache,
                                                                                C1N <: AbstractArray{<: Number},
                                                                                C2N <: AbstractArray{<: Number},
                                                                                C3N <: AbstractArray{<: Number},
                                                                                C4N <: AbstractArray{<: Number}}
                target_cache = get_tmp(cache, x1, x2, x3, x4, select)
                if size(target_cache) != size(cache.du)
                    target_cache = reshape(target_cache, size(cache.du))
                end
                target_cache .= f(x1, x2, x3, x4)
                return target_cache
            end
        else
            throw(MethodError("The length of the eighth input argument, chunksizes, must be 1, 2, or 3."))
        end
    end
    return RALF4(fnew, f, diffcache)
end

function RALF4(fin::LC) where {LC <: AbstractArray{<: Number}}
    f(cache::LCN, x1::C1N, x2::C2N, x3::C3N, x4::C4N, select::Tuple{Int, Int}) where {LCN <: AbstractArray{<: Number}, C1N <: AbstractArray{<: Number}, C2N <: AbstractArray{<: Number}, C3N <: AbstractArray{<: Number}, C4N <: AbstractArray{<: Number}} = cache
    return RALF4{LC}(f, x -> fin, fin)
end

# Using the default of (1, 1) for `select` is important. This way, we can use autodiff
# during the homotopy algorithm without requiring additional arguments when calling `update!`
# to ensure the correct cache is used.
function (ralf::RALF4)(x1::C1, x2::C2, x3::C3,
                       x4::C4, select::Tuple{Int, Int} = (1, 1)) where {C1 <: AbstractArray{<: Number}, C2 <: AbstractArray{<: Number},
                                                                        C3 <: AbstractArray{<: Number}, C4 <: AbstractArray{<: Number}}
    return ralf.f(ralf.cache, x1, x2, x3, x4, select)
end


"""
    FEMMDeforLinearModule

Module for operations on interiors of domains to construct system matrices and
system vectors for linear deformation models.
"""
module FEMMDeforLinearModule

using FinEtools.FTypesModule: FInt, FFlt, FCplxFlt, FFltVec, FIntVec, FFltMat, FIntMat, FMat, FVec, FDataDict
import FinEtools.FENodeSetModule: FENodeSet
import FinEtools.FESetModule: FESet, manifdim
import FinEtools.IntegDataModule: IntegData
import FinEtools.FEMMDeforLinearBaseModule: FEMMDeforLinearAbstract
import FinEtools.DeforModelRedModule: DeforModelRed, DeforModelRed2DAxisymm
import FinEtools.MatDeforModule: MatDefor
import FinEtools.CSysModule: CSys

"""
    FEMMDeforLinear{S<:FESet, F<:Function, P<:PropertyDeformationLinear}

Class for linear deformation finite element modeling machine.
"""
mutable struct FEMMDeforLinear{MR<:DeforModelRed,  S<:FESet, F<:Function, M<:MatDefor} <: FEMMDeforLinearAbstract
    mr::Type{MR}
    integdata::IntegData{S, F} # geometry data
    mcsys::CSys # updater of the material orientation matrix
    material::M # material object
end

function FEMMDeforLinear(mr::Type{MR}, integdata::IntegData{S, F}, material::M) where {MR<:DeforModelRed, S<:FESet, F<:Function, M<:MatDefor}
    # @show mr
    # @show material.mr
    @assert mr === material.mr "Model reduction is mismatched"
    @assert (integdata.axisymmetric) || (mr != DeforModelRed2DAxisymm) "Axially symmetric requires axisymmetric to be true"
    return FEMMDeforLinear(mr, integdata, CSys(manifdim(integdata.fes)), material)
end

function FEMMDeforLinear(mr::Type{MR}, integdata::IntegData{S, F}, material::M, print) where {MR<:DeforModelRed, S<:FESet, F<:Function, M<:MatDefor}
    @show mr
    @show material.mr
    @assert mr === material.mr "Model reduction is mismatched"
    @assert (integdata.axisymmetric) || (mr != DeforModelRed2DAxisymm) "Axially symmetric requires axisymmetric to be true"
    return FEMMDeforLinear(mr, integdata, CSys(manifdim(integdata.fes)), material)
end

end


@with_kw mutable struct PedestrianAvoidanceSystem <: DriverModel{LatLonAccel}
    timestep::Float64 = 0
    t_current::Float64 = 0
    tick::Int64 = 0

    ego_vehicle::Vehicle = Vehicle(VehicleState(VecSE2(0., 0., 0.), 0.), VehicleDef(), 1)
    a::LatLonAccel = LatLonAccel(0.0, 0)

    sensor_observations::Vector{Vehicle} = []

    update_tick_high_level_planner::Int64 = 4
    update_tick_emergency_braking_system::Int64 = 1

    pedestrian_pomdp_frenet::PedestrianAvoidancePOMDPFrenet{SingleOCFUpdater} = PedestrianAvoidancePOMDPFrenet()
    emergency_braking_system::EmergencyBrakingSystem.EmergencySystem = EmergencyBrakingSystem.EmergencySystem()

    b_dict::Dict{Int64, SingleOCFBelief} = Dict{Int64, SingleOCFBelief}()

    prediction_obstacle::Array{Float64} = []

    brake_request::Bool = false
    collision_rate::Float64 = 0.0
    risk::Float64 = 0.0
    ttc::Float64 = 0.0

end

function AutomotiveDrivingModels.observe!(model::PedestrianAvoidanceSystem, scene::Scene, roadway::Roadway, egoid::Int)

    ego = scene[findfirst(egoid, scene)]
    model.ego_vehicle = ego

    #println("model.t_current: ", model.t_current)

    ################ Emergency Braking System #############################################
    #println("-> emergency braking system")
    observe!(model.emergency_braking_system, scene, roadway, egoid)
    model.prediction_obstacle = model.emergency_braking_system.prediction_obstacle
    model.brake_request = model.emergency_braking_system.brake_request
    model.ttc = model.emergency_braking_system.ttc
    model.collision_rate = model.emergency_braking_system.collision_rate

    ################ High Level Planner ###################################################
    if (  model.tick % model.update_tick_high_level_planner == 0 )
        #println("--------------------------POMDP high level planner----------------------- t: ", model.t_current)
        observe!(model.pedestrian_pomdp_frenet, scene, roadway, egoid)
        # println("Action high-level-planner: ", model.pedestrian_pomdp_frenet.a)
        model.b_dict = model.pedestrian_pomdp_frenet.b_dict
    end

    ################ Combination Emergency System + High Level Planner ####################
    # lateral movement only when no strong decceleration
    if ( model.emergency_braking_system.a.a_lon < -2. )
        a_lat = 0.0
    else
        a_lat = model.pedestrian_pomdp_frenet.a.a_lat
    end

    # stronger decceleration is used
    if ( (model.emergency_braking_system.a.a_lon < 0 && ego.state.v > 10 / 3.6) || ( model.emergency_braking_system.a.a_lon < -4. ) )
       # a_lon = min(model.pedestrian_pomdp_frenet.a.a_lon, model.emergency_braking_system.a.a_lon)
        a_lon = model.emergency_braking_system.a.a_lon
    else
       a_lon = model.pedestrian_pomdp_frenet.a.a_lon
    end

    model.a = LatLonAccel(a_lat, a_lon)
    #println("Action combined: ", model.a)

    model.tick += 1
    model.t_current = model.t_current + model.timestep
end

function AutomotiveDrivingModels.get_name(model::PedestrianAvoidanceSystem)
    return "Pedestrian Avoidance System (POMDP + Emergency System)"
end

AutomotiveDrivingModels.rand(rng::AbstractRNG, model::PedestrianAvoidanceSystem) = model.a


include("../imports.jl")

@testset "TabularClassificationSingle" begin
    df = DataFrame(A = 1:4, B = ["M", "F", "F", "M"], C = ["P", "F", "P", "F"])
    td = TableDataset(df)

    method = TabularClassificationSingle(["P", "F"], td; catcols=(:B,), contcols=(:A,), )
    testencoding(method.encodings, method.blocks)
    DLPipelines.checkmethod_core(method)
    @test_nowarn methodlossfn(method)
    @test_nowarn methodmodel(method)

    @testset "`encodeinput`" begin
        row = mockblock(method.blocks[1])

        xtrain = encodeinput(method, Training(), row)
        @test length(xtrain[1]) == length(method.blocks[1].catcols)
        @test length(xtrain[2]) == length(method.blocks[1].contcols)

        @test eltype(xtrain[1]) <: Number
    end

    @testset "`encodetarget`" begin
        category = "P"
        y = encodetarget(method, Training(), category)
        @test y ≈ [1, 0]
    end

    @test_nowarn TabularClassificationSingle(["P", "F"], td)
    @test TabularClassificationSingle(["P", "F"], td).blocks[1].catcols == (:B, :C)

    FastAI.test_method_show(method, ShowText(Base.DevNull()))
end


# Unit tests for implementations/nbodysimulator/nbodysimulator_result.jl

@testset "nbodysimulator_result.jl" begin
    particles = [
        AtomsBase.Atom(:Ar, (@SVector [7, 7, 7])u"bohr", 6e-5(@SVector randn(3))u"bohr * hartree / ħ_au"; meta=:data),
        AtomsBase.Atom(:Ar, (@SVector [7, 7, 21])u"bohr", 6e-5(@SVector randn(3))u"bohr * hartree / ħ_au"; hello="world"),
        AtomsBase.Atom(:Ar, (@SVector [7, 21, 7])u"bohr", 6e-5(@SVector randn(3))u"bohr * hartree / ħ_au"),
        AtomsBase.Atom(:Ar, (@SVector [7, 21, 21])u"bohr", 6e-5(@SVector randn(3))u"bohr * hartree / ħ_au"),
        AtomsBase.Atom(:Ar, (@SVector [21, 7, 7])u"bohr", 6e-5(@SVector randn(3))u"bohr * hartree / ħ_au"),
        AtomsBase.Atom(:Ar, (@SVector [21, 7, 21])u"bohr", 6e-5(@SVector randn(3))u"bohr * hartree / ħ_au"),
        AtomsBase.Atom(:Ar, (@SVector [21, 21, 7])u"bohr", 6e-5(@SVector randn(3))u"bohr * hartree / ħ_au"),
        AtomsBase.Atom(:Ar, (@SVector [21, 21, 21])u"bohr", 6e-5(@SVector randn(3))u"bohr * hartree / ħ_au")
    ]
    box = [[28.0, 0.0, 0.0], [0.0, 28.0, 0.0], [0.0, 0.0, 28.0]]u"bohr"
    system = periodic_system(particles, box)

    simulator1 = NBSimulator(400, 10, t₀=1000, thermostat=NBodySimulator.AndersenThermostat(2e-4, 2e-4))
    simulator2 = NBSimulator(400, 10, t₀=1000)

    potential = InteratomicPotentials.LennardJones(1.657e-21u"J", 0.34u"nm", 0.765u"nm", [:Ar])

    result = simulate(system, simulator1, potential)
    result2 = simulate(system, simulator2, potential)

    @test get_time_range(result) == (1000:400:5000)u"ħ_au / hartree"
    @test get_num_bodies(result) == 8
    @test get_bounding_box(result) == box
    @test get_boundary_conditions(result) == [Periodic(), Periodic(), Periodic()]

    @test reference_temperature(result) == 2e-4u"hartree / k_au"
    @test ismissing(reference_temperature(result2))

    @test get_positions(result) isa AbstractVector{<:StaticVector{3,<:Unitful.Length}}
    @test get_velocities(result) isa AbstractVector{<:StaticVector{3,<:Unitful.Velocity}}
    @test get_particles(result) isa AbstractVector{<:AtomsBase.Atom}

    @test all(get_positions(result, t) isa AbstractVector{<:StaticVector{3,<:Unitful.Length}} for t ∈ 1:11)
    @test all(get_velocities(result, t) isa AbstractVector{<:StaticVector{3,<:Unitful.Velocity}} for t ∈ 1:11)
    @test all(get_particles(result, t) isa AbstractVector{<:AtomsBase.Atom} for t ∈ 1:11)

    @test all(all(all(0.0u"bohr" ≤ c < 28.0u"bohr" for c ∈ p) for p ∈ get_positions(result, t)) for t ∈ 1:11)
    @test all(get_particles(result, t)[1].meta == :data && get_particles(result, t)[2].hello == "world" for t ∈ 1:11)

    @test Atomistic.temperature(result) isa Unitful.Temperature
    @test Atomistic.kinetic_energy(result) isa Unitful.Energy
    @test Atomistic.potential_energy(result) isa Unitful.Energy

    @test all(isapprox(Atomistic.total_energy(result2, t), Atomistic.total_energy(result2), rtol=0.1) for t ∈ 1:11)
end


export twri_objective, TWRIOptions

# TWRI options
mutable struct TWRIOptions
    grad_corr::Bool
    comp_alpha::Bool
    weight_fun
    eps
    params
    Invq::String
end

"""
    TWRIOptions
        grad_corr::Bool
        comp_alpha::Bool
        weight_fun
        eps
        params::Symbol
        Invq::String

Options structure for TWRI.

`grad_corr`: Whether to add the gradient correction J'(m0, q)*∇_y

`comp_alpha`: Whether to compute optimal alpha (alpha=1 if not)

`weight_fun`: Whether to apply focusing/weighting function to F(m0)'*y and its norm

`eps`: Epsilon (noise level) value (default=0)

`Invq`: How to compute F'Y, either as full field or as a rank 1 approximation `w(t)*q(x)` using the source wavelet for w

`param`: Which gradient to compute. Choices are `nothing` (objective only), `:m`, `:y` or `:all`

Constructor
==========

All arguments are optional keyword arguments with the following default values:

TWRIOptions(;grad_corr=false, comp_alpha=true, weight_fun=nothing, eps=0, params=:m)
"""

TWRIOptions(;grad_corr=false, comp_alpha=true,
            weight_fun=nothing, eps=0, params=:m, Invq="standard")=
            TWRIOptions(grad_corr, comp_alpha, weight_fun, eps, params, Invq)

function subsample(opt::TWRIOptions, srcnum::Int)
    eloc = length(opt.eps) == 1 ? opt.eps : opt.eps[srcnum]
    return TWRIOptions(opt.grad_corr, opt.comp_alpha, opt.weight_fun, eloc, opt.params, opt.Invq)
end

twri_objective(model_full::Model, source::judiVector, dObs::judiVector, y::Union{judiVector, Nothing}) =
    twri_objective(model_full, source, dObs, y, Options(), TWRIOptions())
twri_objective(model_full::Model, source::judiVector, dObs::judiVector, y::Union{judiVector, Nothing}, opt::Options) =
    twri_objective(model_full, source, dObs, y, opt, TWRIOptions())
twri_objective(model_full::Model, source::judiVector, dObs::judiVector, y::Union{judiVector, Nothing}, opt::TWRIOptions) =
    twri_objective(model_full, source, dObs, y, Options(), opt)
twri_objective(model_full::Model, source::judiVector, dObs::judiVector, y::Union{judiVector, Nothing}, twri_opt::TWRIOptions, opt::Options) =
    twri_objective(model_full, source, dObs, y, opt, twri_opt)

function twri_objective(model_full::Model, source::judiVector, dObs::judiVector, y::Union{judiVector, Nothing},
                        options::Options, optionswri::TWRIOptions)
    # Load full geometry for out-of-core geometry containers
    dObs.geometry = Geometry(dObs.geometry)
    source.geometry = Geometry(source.geometry)

    # Limit model to area with sources/receivers
    if options.limit_m == true
        model = deepcopy(model_full)
        model = limit_model_to_receiver_area(source.geometry, dObs.geometry, model, options.buffer_size)
    else
        model = model_full
    end

    # Set up Python model structure
    modelPy = devito_model(model, options)
    dtComp = get_dt(model; dt=options.dt_comp)

    # Extrapolate input data to computational grid
    qIn = time_resample(source.data[1],source.geometry,dtComp)[1]
    obsd = typeof(dObs.data[1]) == SegyIO.SeisCon ? convert(Array{Float32,2}, dObs.data[1][1].data) : dObs.data[1]
    dObserved = time_resample(obsd, dObs.geometry, dtComp)[1]

    isnothing(y) ? Y = nothing : Y = time_resample(y.data[1], y.geometry, dtComp)[1]

    # Set up coordinates
    src_coords = setup_grid(source.geometry, model.n)  # shifts source coordinates by origin
    rec_coords = setup_grid(dObs.geometry, model.n)    # shifts rec coordinates by origin

    ~isempty(options.frequencies) ? freqs = options.frequencies : freqs = nothing
    ~isempty(options.frequencies) ? (wfilt, freqs) =  filter_w(qIn, dtComp, freqs) : wfilt = nothing
    obj, gradm, grady = pycall(ac."wri_func", PyObject,
                               modelPy, src_coords, qIn, rec_coords, dObserved, Y,
                               t_sub=options.subsampling_factor, space_order=options.space_order,
                               grad=optionswri.params, grad_corr=optionswri.grad_corr, eps=optionswri.eps,
                               alpha_op=optionswri.comp_alpha, w_fun=optionswri.weight_fun,
                               freq_list=freqs, wfilt=wfilt)

    if (optionswri.params in [:m, :all])
        gradm = remove_padding(gradm, modelPy.padsizes; true_adjoint=options.sum_padding)
        options.limit_m==true && (gradm = extend_gradient(model_full, model, gradm))
    end
    if ~isnothing(grady)
        grady = time_resample(grady, dtComp, dObs.geometry)
        grady = judiVector(dObs.geometry, grady)
    end

    return filter_out(obj, gradm, grady)
end

function filter_w(qIn, dt, freqs)
    ff = FFTW.fftfreq(length(qIn), 1/dt)
    df = ff[2] - ff[1]
    inds = [findmin(abs.(ff.-f))[2] for f in freqs]
    DFT = joDFT(length(qIn); DDT=Float32)
    R = joRestriction(size(DFT,1), inds; RDT=Complex{Float32}, DDT=Complex{Float32})
    qfilt = DFT'*R'*R*DFT*qIn
    return qfilt, freqs
end

filter_out(obj, ::Nothing, ::Nothing) = obj
filter_out(obj, m, ::Nothing) = obj, m
filter_out(obj, ::Nothing, y) = obj, y
filter_out(obj, m, y) = obj, m, y


using Statistics
using CUDA.CUDNN:
    cudnnReduceTensor,
    cudnnReduceTensor!,
    cudnnGetReductionIndicesSize,
    cudnnGetReductionWorkspaceSize,
    cudnnReduceTensorDescriptor,
        cudnnReduceTensorDescriptor_t,
        cudnnCreateReduceTensorDescriptor,
        cudnnSetReduceTensorDescriptor,
        cudnnGetReduceTensorDescriptor,
        cudnnDestroyReduceTensorDescriptor,
    cudnnReduceTensorOp_t,
        CUDNN_REDUCE_TENSOR_ADD,          # 0,
        CUDNN_REDUCE_TENSOR_MUL,          # 1,
        CUDNN_REDUCE_TENSOR_MIN,          # 2,
        CUDNN_REDUCE_TENSOR_MAX,          # 3,
        CUDNN_REDUCE_TENSOR_AMAX,         # 4,
        CUDNN_REDUCE_TENSOR_AVG,          # 5,
        CUDNN_REDUCE_TENSOR_NORM1,        # 6,
        CUDNN_REDUCE_TENSOR_NORM2,        # 7,
        CUDNN_REDUCE_TENSOR_MUL_NO_ZEROS, # 8,
    cudnnNanPropagation_t,
        CUDNN_NOT_PROPAGATE_NAN, # 0
        CUDNN_PROPAGATE_NAN,     # 1
    cudnnReduceTensorIndices,
    cudnnReduceTensorIndices_t,
        CUDNN_REDUCE_TENSOR_NO_INDICES,        # 0,
        CUDNN_REDUCE_TENSOR_FLATTENED_INDICES, # 1,
    cudnnIndicesType,
    cudnnIndicesType_t,
        CUDNN_32BIT_INDICES, # 0,
        CUDNN_64BIT_INDICES, # 1,
        CUDNN_16BIT_INDICES, # 2,
        CUDNN_8BIT_INDICES,  # 3,
    cudnnDataType

@testset "cudnn/reduce" begin
    @test cudnnReduceTensorDescriptor(C_NULL) isa cudnnReduceTensorDescriptor
    @test Base.unsafe_convert(Ptr, cudnnReduceTensorDescriptor(C_NULL)) isa Ptr
    @test cudnnReduceTensorDescriptor(CUDNN_REDUCE_TENSOR_ADD,cudnnDataType(Float32),CUDNN_NOT_PROPAGATE_NAN,CUDNN_REDUCE_TENSOR_NO_INDICES,CUDNN_32BIT_INDICES) isa cudnnReduceTensorDescriptor

    (ax,ay) = randn(Float32,10,10), randn(Float32,10,1)
    (cx,cy) = CuArray.((ax,ay))

    function reducetensortest(;
        op::cudnnReduceTensorOp_t = CUDNN_REDUCE_TENSOR_ADD,
        compType::DataType = (eltype(ax) <: Float64 ? Float64 : Float32),
        nanOpt::cudnnNanPropagation_t = CUDNN_NOT_PROPAGATE_NAN,
        indices::Union{Vector{<:Unsigned},Nothing} = nothing,
        d::cudnnReduceTensorDescriptor = cudnnReduceTensorDescriptor(op, cudnnDataType(compType), nanOpt, cudnnReduceTensorIndices(op, indices), cudnnIndicesType(indices)),
        alpha::Real = 1,
        beta::Real = 0,
    )
        f0 = (op === CUDNN_REDUCE_TENSOR_ADD          ? sum(ax, dims=2) :
              op === CUDNN_REDUCE_TENSOR_MUL          ? prod(ax, dims=2) :
              op === CUDNN_REDUCE_TENSOR_MIN          ? minimum(ax, dims=2) :
              op === CUDNN_REDUCE_TENSOR_MAX          ? maximum(ax, dims=2) :
              op === CUDNN_REDUCE_TENSOR_AMAX         ? maximum(abs, ax, dims=2) :
              op === CUDNN_REDUCE_TENSOR_AVG          ? mean(ax, dims=2) :
              op === CUDNN_REDUCE_TENSOR_NORM1        ? sum(abs, ax, dims=2) :
              op === CUDNN_REDUCE_TENSOR_NORM2        ? sqrt.(sum(abs2, ax, dims=2)) :
              op === CUDNN_REDUCE_TENSOR_MUL_NO_ZEROS ? (ax1=copy(ax);ax1[ax.==0].=1;prod(ax1,dims=2)) :
              error("Unknown reducetensor"))
        f1 = alpha * f0
        f2 = f1 + beta * ay
        dims = size(ay)
        @test f1 ≈ cudnnReduceTensor(cx; dims, op, compType, nanOpt, indices, alpha) |> Array
        @test f1 ≈ cudnnReduceTensor(cx, d; dims, indices, alpha) |> Array
        @test f2 ≈ cudnnReduceTensor!(copy(cy), cx; op, compType, nanOpt, indices, alpha, beta) |> Array
        @test f2 ≈ cudnnReduceTensor!(copy(cy), cx, d; indices, alpha, beta) |> Array
    end

    reducetensortest()
    reducetensortest(op = CUDNN_REDUCE_TENSOR_MUL)
    reducetensortest(op = CUDNN_REDUCE_TENSOR_MIN)
    reducetensortest(op = CUDNN_REDUCE_TENSOR_MAX)
    reducetensortest(op = CUDNN_REDUCE_TENSOR_AMAX)
    reducetensortest(op = CUDNN_REDUCE_TENSOR_AVG)
    reducetensortest(op = CUDNN_REDUCE_TENSOR_NORM1)
    reducetensortest(op = CUDNN_REDUCE_TENSOR_NORM2)
    reducetensortest(op = CUDNN_REDUCE_TENSOR_MUL_NO_ZEROS)
    reducetensortest(nanOpt = CUDNN_PROPAGATE_NAN)
    reducetensortest(alpha = 2)
    reducetensortest(beta = 2)
end


### Intelligent pedestrian model that stops if a crosswalk is busy

@with_kw mutable struct IntelligentPedestrian <: DriverModel{ConstantSpeedDawdling}
    a::ConstantSpeedDawdling = ConstantSpeedDawdling(0., 0.)
    v_noise::Float64 = 0.
    dt::Float64 = 0.1
    motion::DriverModel = ConstantPedestrian(v_noise = v_noise, dt = dt)
    ttc_threshold::Float64 = 4.5
    crosswalk::Lane = Lane()
    conflict_lanes::Vector{Lane{Float64}} = Lane{Float64}[]
    dtol::Float64 = 0.1
    # states
    wait::Bool = true
    crossing::Bool = false
    done::Bool = false
end

AutomotiveDrivingModels.get_name(model::IntelligentPedestrian) = "IntelligentPedestrian"
Base.rand(model::IntelligentPedestrian) = model.a

AutomotiveDrivingModels.reset_hidden_state!(model::IntelligentPedestrian) = model

function AutomotiveDrivingModels.observe!(model::IntelligentPedestrian,
                                        scene::EntityFrame{VehicleState, VehicleDef, Int64},
                                        roadway::Roadway,
                                        egoid::Int)
    ped = scene[findfirst(egoid, scene)]
    observe!(model.motion, scene, roadway, egoid)
    model.crossing = update_crossing(ped, model.crosswalk, model.conflict_lanes, roadway)
    model.done = update_done(ped, model.crossing, model.crosswalk)
    if !model.crossing
        model.wait = update_wait(ped, scene, model, roadway)
    else
        model.wait = false
    end
    if !model.wait || model.done# if engaged, keep going
        model.a = model.motion.a
    else
        if model.wait
            model.a = ConstantSpeedDawdling(0., 0.)
        else
            model.a = model.motion.a
        end
    end
end

function update_crossing(ped::Vehicle, crosswalk::Lane{T}, conflict_lanes::Vector{Lane{T}}, roadway::Roadway) where T
    # check if the pedestrian is in the conflict zone
    for lane in conflict_lanes
        ped_f = Frenet(proj(ped.state.posG, lane, roadway, move_along_curves=false), roadway)
        if abs(ped_f.t) < lane.width/2 && get_lane(roadway, ped).tag == crosswalk.tag
            # println(" position on lane : ", ped_f)
            return true
        end
    end
    return false
end

function update_done(ped::Vehicle, crossing::Bool, crosswalk::Lane)
    # check if the pedestrian is done crossing
    return !crossing && direction_from_center(ped, crosswalk) .< 0.
end

function get_distance_to_crosswalk(veh::Vehicle, crosswalk::Lane, roadway::Roadway)
    lane = get_lane(roadway, veh)
    cw_length = get_end(crosswalk)
    cw_center = posg(Frenet(crosswalk, cw_length/2), roadway)
    collision_point = VecSE2(cw_center.x, veh.state.posG.y)
    cw_proj = Frenet(collision_point, lane, roadway)
    Δs = cw_proj.s - veh.state.posF.s
    return Δs
end

function update_wait(ped::Vehicle,
                    scene::Scene,
                    model::IntelligentPedestrian,
                    roadway::Roadway)
    # check if the pedestrian should wait or not
    for veh in scene
        if veh.def.class == AgentClass.PEDESTRIAN
            continue
        end
        # check if vehicle is on the crosswalk
        posG = veh.state.posG
        cw_posF = Frenet(proj(posG, model.crosswalk, roadway, move_along_curves=false), roadway)
        lane = get_lane(roadway, veh)
        # if lane ∈ model.conflict_lanes
        # println("vehicle position on cw $(cw_posF)")
        # println("pedestrian position $(ped.state.posF)")
        # println(" difference: $(cw_posF.s - ped.state.posF.s)")
        # t_car = cw_posF.t
        # println( abs(cw_posF.t) < model.crosswalk.width/2 + veh.def.length/2 - model.dtol)
        # println( (cw_posF.s - ped.state.posF.s)*cos(ped.state.posF.ϕ) > 0.)
        if abs(cw_posF.t) < model.crosswalk.width/2 + veh.def.length/2 - model.dtol &&
                                (cw_posF.s - ped.state.posF.s)*cos(ped.state.posF.ϕ) > 0.
            # println("vehicle $(veh.id) is on the crosswalk")
            # vehicle is on the crosswalk
            # check if pedestrian is passed it or not
            return true
        end
        if lane ∈ model.conflict_lanes
            # check the ttc
            Δs = get_distance_to_crosswalk(veh, model.crosswalk, roadway)
            ttc = Δs/veh.state.v
            if 0. < ttc < model.ttc_threshold
                # println("vehicle $(veh.id) ttc: $ttc, Δs: $Δs")
                return true
            end
        end
    end
    return false
end