module TestTrees

using Phylo
using DataFrames
using JuliaDB

using Test
using IterableTables: getiterator

species = ["Dog", "Cat", "Human"]
ntips = 10
df = DataFrame(species = species, count = [10, 20, 3])
observations = ["Dog", "Cat", "Dog", "Dog"]
jdb = table((species = observations, count = 1:4))

@testset "NamedBinaryTree()" begin
    ntn = NamedBinaryTree(ntips)
    @test length(nodefilter(isroot, ntn)) == ntips
    @test length(nodefilter(isleaf, ntn)) == ntips
    @test length(nodefilter(isinternal, ntn)) == 0
    @test validate!(ntn)
    nt = NamedBinaryTree(species)
    @test_throws Exception noderoute(nt, "Dog", "Cat")
    @test_throws Exception branchroute(nt, "Dog", "Human")
    @test validate!(nt)
    n = getnodename(nt, createnode!(nt))
    b1 = createbranch!(nt, n, "Dog", 2.0)
    b2 = createbranch!(nt, n, "Cat", 2.0)
    @test_throws Exception createbranch!(nt, n, "Human", 2.0)
    @test validate!(nt)
    r = getnodename(nt, createnode!(nt))
    @test_throws Exception createbranch!(nt, r, "Potato", 2.0)
    b3 = createbranch!(nt, r, "Human", 5.0)
    b4 = createbranch!(nt, r, n, 3.0)
    @test maximum(distances(nt)) ≈ 10.0
    @test validate!(nt)
    @test noderoute(nt, "Human", "Dog") == ["Human", r, n, "Dog"]
    @test branchroute(nt, "Human", "Dog") == [b3, b4, b1]
    nb = BinaryTree{OneRoot}(DataFrame(name=["Human", "Cat"]))
    @test_throws Exception BinaryTree(nb)
    nb = BinaryTree(DataFrame(name=["Human", "Cat"]))
    @test_throws Exception BinaryTree(nb)
    nb = BinaryTree(DataFrame(name=["Human"]))
    @test_nowarn BinaryTree(nb)
    nb = BinaryTree{ManyRoots}(DataFrame(name=["Human", "Cat"]))
    @test getleafinfo(nb) === getleafinfo(BinaryTree(nb))
    @test getleafinfo(nb) !== getleafinfo(BinaryTree(nb; copyinfo = true))
    @test_nowarn createnode!(nb, "Dog")
    @test !validate!(nb)
    np = PolytomousTree{OneRoot}(DataFrame(name=["Human", "Cat"]))
    @test_throws Exception PolytomousTree(np)
    np = PolytomousTree(DataFrame(name=["Human", "Cat"]))
    @test_throws Exception BinaryTree(np)
    np = PolytomousTree(DataFrame(name=["Human"]))
    @test_nowarn PolytomousTree(np)
    np = PolytomousTree{ManyRoots}(DataFrame(name=["Human", "Cat"]))
    @test getleafinfo(np) !== getleafinfo(PolytomousTree(np; copyinfo = true))
    @test_nowarn createnode!(np, "Dog")
    @test !validate!(np)
    @test_throws Exception PolytomousTree(np)
    @test getleafnames(nb) ⊆ getleafnames(np)
    @test getleafinfo(BinaryTree(df)) == df
    @test getleafinfo(PolytomousTree(df)) == df
    @test branchdatatype(typeof(np)) ≡ Nothing
    @test nodedatatype(typeof(np)) ≡ Dict{String, Any}
    @test leafinfotype(typeof(np)) ≡ DataFrame
end

@testset "BinaryTree()" begin
    btn = BinaryTree{OneRoot, DataFrame, Vector{Float64}}(ntips)
    @test length(nodefilter(isroot, btn)) == ntips
    @test length(nodefilter(isleaf, btn)) == ntips
    @test length(nodefilter(isinternal, btn)) == 0

    tmp = BinaryTree{ManyRoots}(DataFrame(names=species))
    @test validate!(tmp)
    tmp = BinaryTree(DataFrame(names=species))
    @test !validate!(tmp)
    names = createnodes!(tmp, 2)
    @test_nowarn createbranch!(tmp, names[1], species[1])
    @test_nowarn createbranch!(tmp, names[1], species[2])
    @test_nowarn createbranch!(tmp, names[2], species[3])
    @test_nowarn createbranch!(tmp, names[2], names[1])
    @test Set(getleafnames(tmp)) == Set(species)
    @test validate!(tmp)
    nt = BinaryTree{ManyRoots, DataFrame, Vector{Float64}}(species)
    @test validate!(nt)
    n = getnodename(nt, createnode!(nt))
    @test Set(getleafnames(nt)) == Set(species) ∪ Set([n])
    createbranch!(nt, n, "Dog", 2.0)
    createbranch!(nt, n, "Cat", 2.0)
    @test_throws Exception createbranch!(nt, n, "Human", 2.0)
    @test validate!(nt)
    r = getnodename(nt, createnode!(nt))
    @test_throws Exception createbranch!(nt, r, "Potato", 2.0)
    createbranch!(nt, r, "Human", 4.0)
    createbranch!(nt, r, n, 2.0)
    @test validate!(nt)
    setnodedata!(nt, "Dog", [1.0])
    @test getnodedata(nt, "Dog")[1] ≈ 1.0
    @test length(getnodedata(nt, "Cat")) == 0
    nt2 = BinaryTree(nt)
    @test length(getnodedata(nt2, "Dog")) == 0
    @test Set(getleafnames(nt2)) == Set(species)
    @test validate!(nt)
    nt3 = BinaryTree(nt)
    nt4 = BinaryTree(nt3)
    a=IOBuffer()
    b=IOBuffer()
    show(a, nt3)
    show(b, nt4)
    @test all(a.data .== b.data)

    tdf = BinaryTree(df)
    @test nleaves(tdf) == 3
    tj = BinaryTree(jdb)
    @test nleaves(tj) == 2

    lij = getleafinfo(tj, "Dog")
    @test sum(map(line -> line.count, lij)) == 8

    lij = getleafinfo(tj, "Cat")
    @test sum(map(line -> line.count, lij)) == 2

    lidf = getleafinfo(tdf, "Dog")
    @test sum(map(line -> line.count, lidf)) == 10

    @test branchdatatype(typeof(tdf)) ≡ Nothing
    @test nodedatatype(typeof(tdf)) ≡ Dict{String, Any}
    @test leafinfotype(typeof(tdf)) ≡ DataFrame
    @test branchdatatype(typeof(tj)) ≡ Nothing
    @test nodedatatype(typeof(tj)) ≡ Dict{String, Any}
    @test leafinfotype(typeof(tj)) <: IndexedTable
end

end


# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

include("tree_mesh.jl")
include("structured_mesh.jl")
include("surface_interpolant.jl")
include("unstructured_mesh.jl")
include("face_interpolant.jl")
include("transfinite_mappings_3d.jl")
include("p4est_mesh.jl")
include("mesh_io.jl")
include("dgmulti_meshes.jl")

end # @muladd


# Julia wrapper for header: /usr/include/scip/pub_bandit_ucb.h
# Automatically generated using Clang.jl wrap_c

function SCIPcreateBanditUcb(scip, ucb, priorities, alpha, nactions, initseed)
    ccall((:SCIPcreateBanditUcb, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{Ptr{SCIP_BANDIT}}, Ptr{Cdouble}, Cdouble, Cint, UInt32), scip, ucb, priorities, alpha, nactions, initseed)
end

function SCIPgetConfidenceBoundUcb(ucb, action)
    ccall((:SCIPgetConfidenceBoundUcb, libscip), Cdouble, (Ptr{SCIP_BANDIT}, Cint), ucb, action)
end

function SCIPgetStartPermutationUcb(ucb)
    ccall((:SCIPgetStartPermutationUcb, libscip), Ptr{Cint}, (Ptr{SCIP_BANDIT},), ucb)
end


"""
This module provides utilities to build neural networks with Flux,
along with a library of standard architectures.
"""
module FluxNets

export SimpleNet, SimpleNetHP, ResNet, ResNetHP

using ..Network
using Base: @kwdef
import ..GameInterface
import ..Util

using CUDAapi

# Import CuArrays only if CUDA is installed
if has_cuda()
  try
    using CuArrays
    @show CuArrays.usage_limit[]
    @eval const CUARRAYS_IMPORTED = true
  catch ex
    @warn(
      "CUDA is installed, but CuArrays.jl fails to load.",
      exception=(ex,catch_backtrace()))
    @eval const CUARRAYS_IMPORTED = false
  end
else
  @eval const CUARRAYS_IMPORTED = false
end

import Flux

if CUARRAYS_IMPORTED
  @eval begin
    CuArrays.allowscalar(false)
    on_gpu(::Type{<:Array}) = false
    on_gpu(::Type{<:CuArray}) = true
    on_gpu(x) = on_gpu(typeof(x))
  end
else
  @eval begin
    on_gpu(x) = false
  end
end

using Flux: relu, softmax
using Flux: Chain, Dense, Conv, BatchNorm, SkipConnection

#####
##### Flux Networks
#####

abstract type FluxNetwork{Game} <: AbstractNetwork{Game} end

function Base.copy(nn::Net) where Net <: FluxNetwork
  #new = Net(Network.hyperparams(nn))
  #Flux.loadparams!(new, Flux.params(nn))
  #return new
  return Base.deepcopy(nn)
end

Network.to_cpu(nn::FluxNetwork) = Flux.cpu(nn)

function Network.to_gpu(nn::FluxNetwork)
  CUARRAYS_IMPORTED && CuArrays.allowscalar(false)
  return Flux.gpu(nn)
end

# TODO: dirty hack to cope with the fact that Flux 0.10 removed Flux.testmode!
const TEST_MODE = IdDict()

function Network.set_test_mode!(nn::FluxNetwork, mode)
  TEST_MODE[nn] = mode
end

Network.convert_input(nn::FluxNetwork, x) =
  Network.on_gpu(nn) ? Flux.gpu(x) : x

Network.convert_output(nn::FluxNetwork, x) = Flux.cpu(x)

Network.params(nn::FluxNetwork) = Flux.params(nn)

function Network.train!(nn::FluxNetwork, loss, data, lr)
  optimizer = Flux.ADAM(lr)
  Flux.train!(loss, Flux.params(nn), data, optimizer)
end

regularized_child_leaves(l) = []
regularized_child_leaves(l::Flux.Dense) = [l.W]
regularized_child_leaves(l::Flux.Conv) = [l.weight]

# Reimplementation of what used to be Flux.prefor, does not visit leafs
function foreach_flux_node(f::Function, x, seen = IdDict())
  Flux.isleaf(x) && return
  haskey(seen, x) && return
  seen[x] = true
  f(x)
  for child in Flux.trainable(x)
    foreach_flux_node(f, child, seen)
  end
end

function Network.regularized_params(net::FluxNetwork)
  ps = Flux.Params()
  foreach_flux_node(net) do p
    for r in regularized_child_leaves(p)
      any(x -> x === r, ps) || push!(ps, r)
    end
  end
  return ps
end

function Network.gc(::FluxNetwork)
  CUARRAYS_IMPORTED || return
  GC.gc(true)
  CuArrays.reclaim()
end

#####
##### Common functions between two-head neural networks
#####

abstract type TwoHeadNetwork{G} <: FluxNetwork{G} end

function Network.forward(nn::TwoHeadNetwork, board)
  # TODO: eliminate this horror when Flux catches up
  # @eval Flux.istraining() = $(!get(TEST_MODE, nn, false))
  c = nn.common(board)
  v = nn.vhead(c)
  p = nn.phead(c)
  return (p, v)
end

# Flux.@functor does not work do to Network being parametric
function Flux.functor(nn::Net) where Net <: TwoHeadNetwork
  children = (nn.common, nn.vhead, nn.phead)
  constructor = cs -> Net(nn.hyper, cs...)
  return (children, constructor)
end

Network.hyperparams(nn::TwoHeadNetwork) = nn.hyper

Network.on_gpu(nn::TwoHeadNetwork) = on_gpu(nn.vhead[end].b)

#####
##### Utilities for the networks library
#####

linearize(x) = reshape(x, :, size(x)[end])

#####
##### Include networks library
#####

include("architectures/simplenet.jl")
include("architectures/resnet.jl")

end


export Graph, graph
export predict
export forward
export backward
export update!
export zerograds!
export clear

"""
    mutable struct Graph
# Field
`backward`: store backward operations
"""
mutable struct Graph
    # 存储反向传播操作
    backward::Vector
    function Graph()
        new(Vector(undef,0))
    end
end

import Base.size
import Base.length
import Base.getindex
import Base.deepcopy

Base.size(g::Graph)           = size(g.backward)
Base.size(g::Graph, dim::Int) = size(g.backward, dim)
Base.length(g::Graph)         = length(g.backward)
Base.getindex(g::Graph, k...) =  g.backward[k...]

# -- 全局图变量，存储所有反向运算的中间
# -- 变量，但是在参数更新后需要将其置空
global graph = Graph()

function clear()
    graph.backward = [];
end

function backward()
    for i = length(graph):-1:1
        graph[i]()
    end
    clear()
end


testOnAllContexts("GA") do ctx
    println("\tsolves a simple problem")
    testSimpleProblem(
        comparer ->
        begin
            ga = GA(ctx, 60, comparer, dna -> dna)
            @test ga.dnaSize == 60
            @test ga.populationSize == 100
            @test ga.mutationChance == 0.02f0
            @test ga.mutationStrength == 0.05f0
            @test ga.selectionStdDev == 0.25f0
            @test ga.keepElitesRate == 0.05f0
            ga
        end)

    println("\tsolves a problem with synthesizer")
    testSimpleProblem(
        ctx,
        (comparer, synth) ->
        begin
            ga = GA(
                ctx,
                dnaSize!(synth),
                comparer,
                dna -> decode(synth, dna),
                mutationChance = 0.05f0,
                selectionStdDev = 0.3f0)
            @test ga.dnaSize == dnaSize!(synth)
            @test ga.populationSize == 100
            @test ga.mutationChance == 0.05f0
            @test ga.mutationStrength == 0.05f0
            @test ga.selectionStdDev == 0.3f0
            @test ga.keepElitesRate == 0.05f0
            ga
        end)
end


function assign_scenarios_to_procs(scen_tree::ScenarioTree
                                   )::Dict{Int,Set{ScenarioID}}
    sp_map = Dict{Int, Set{ScenarioID}}()

    nprocs = nworkers()
    wrks = workers()
    sp_map = Dict(w => Set{ScenarioID}() for w in wrks)

    for (k,s) in enumerate(scenarios(scen_tree))
        w = wrks[(k-1) % nprocs + 1]
        push!(sp_map[w], s)
    end

    return sp_map
end

function _initialize_subproblems(sp_map::Dict{Int,Set{ScenarioID}},
                                 wi::WorkerInf,
                                 scen_tree::ScenarioTree,
                                 constructor::Function,
                                 constructor_args::Tuple,
                                 r::AbstractPenaltyParameter,
                                 warm_start::Bool,
                                 subproblem_callbacks::Vector{SubproblemCallback};
                                 kwargs...
                                 )

    # Send initialization commands
    @sync for (wid, scenarios) in pairs(sp_map)
        @async _send_message(wi,
                             wid,
                             Initialize(constructor,
                                        constructor_args,
                                        (;kwargs...),
                                        typeof(r),
                                        scenarios,
                                        scen_tree,
                                        warm_start,
                                        subproblem_callbacks)
                             )
    end

    # Wait for and process initialization replies
    var_maps = Dict{ScenarioID,Dict{VariableID,VariableInfo}}()
    remaining_maps = copy(scenarios(scen_tree))

    while !isempty(remaining_maps)

        msg = _retrieve_message_type(wi, VariableMap)
        var_maps[msg.scen] = msg.var_info

        delete!(remaining_maps, msg.scen)

    end

    return var_maps
end

function _initialize_lb_subproblems(sp_map::Dict{Int,Set{ScenarioID}},
                                    wi::WorkerInf,
                                    scen_tree::ScenarioTree,
                                    constructor::Function,
                                    constructor_args::Tuple,
                                    warm_start::Bool;
                                    kwargs...
                                    )::Nothing

    @sync for (wid, scenarios) in pairs(sp_map)
        @async _send_message(wi,
                             wid,
                             InitializeLowerBound(constructor,
                                                  constructor_args,
                                                  (;kwargs...),
                                                  scenarios,
                                                  scen_tree,
                                                  warm_start)
                             )
    end

    return
end

function _set_initial_values(phd::PHData,
                             wi::WorkerInf,
                             )::Float64

    # Wait for and process mapping replies
    remaining_init = copy(scenarios(phd.scenario_tree))

    while !isempty(remaining_init)

        msg = _retrieve_message_type(wi, ReportBranch)

        _verify_report(msg)
        _update_values(phd, msg)

        delete!(remaining_init, msg.scen)

    end

    # Start computation of initial PH values -- W values are computed at
    # the end of `_set_penalty_parameter`.
    xhat_res_sq = compute_and_save_xhat(phd)

    return xhat_res_sq
end

function _set_penalty_parameter(phd::PHData,
                                wi::WorkerInf,
                                )::Float64

    if is_subproblem_dependent(typeof(phd.r))

        # Wait for and process penalty parameter messages
        remaining_maps = copy(scenarios(phd.scenario_tree))

        while !isempty(remaining_maps)

            msg = _retrieve_message_type(wi, PenaltyInfo)

            process_penalty_subproblem(phd.r, phd, msg.scen, msg.penalty)

            delete!(remaining_maps, msg.scen)

        end

    end

    if is_initial_value_dependent(typeof(phd.r))

        # Compute penalty parameter values
        process_penalty_initial_value(phd.r, phd)

    end

    if is_variable_dependent(typeof(phd.r))

        # Convert dictionary to subproblem format
        pen_map = Dict{ScenarioID,Dict{VariableID,Float64}}(
            s => Dict{VariableID,Float64}() for s in scenarios(phd)
        )

        for (xhid,penalty) in pairs(penalty_map(phd.r))
            for vid in convert_to_variable_ids(phd, xhid)
                pen_map[scenario(vid)][vid] = penalty
            end
        end

    else

        pen_map = Dict{ScenarioID,Float64}(
            s => get_penalty_value(phd.r) for s in scenarios(phd)
        )

    end

    # Send penalty parameter values to workers
    @sync for (scid, sinfo) in pairs(phd.scenario_map)
        @async begin
            wid = sinfo.pid
            penalty = pen_map[scid]
            _send_message(wi, wid, PenaltyInfo(scid, penalty))
        end
    end

    # Complete computation of initial PH values -- Xhat values are computed at
    # the end of `_set_initial_values`.
    x_res_sq = compute_and_save_w(phd)

    return x_res_sq
end

function initialize(scen_tree::ScenarioTree,
                    model_constructor::Function,
                    r::AbstractPenaltyParameter,
                    user_sp_map::Dict{Int,Set{ScenarioID}},
                    warm_start::Bool,
                    timo::TimerOutputs.TimerOutput,
                    report::Int,
                    lower_bound::Int,
                    subproblem_callbacks::Vector{SubproblemCallback},
                    constructor_args::Tuple,;
                    kwargs...
                    )::Tuple{PHData,WorkerInf}

    # Assign scenarios to processes
    scen_proc_map = isempty(user_sp_map) ? assign_scenarios_to_procs(scen_tree) : user_sp_map
    scen_per_worker = maximum(length.(collect(values(scen_proc_map))))
    n_scenarios = length(scenarios(scen_tree))

    # Start worker loops
    if report > 0
        println("...launching workers...")
    end
    worker_inf = @timeit(timo,
                         "Launch Workers",
                         _launch_workers(2*scen_per_worker, n_scenarios)
                         )

    # Initialize subproblems
    if report > 0
        println("...initializing subproblems...")
        flush(stdout)
    end

    var_map = @timeit(timo,
                      "Initialize Subproblems",
                      _initialize_subproblems(scen_proc_map,
                                              worker_inf,
                                              scen_tree,
                                              model_constructor,
                                              constructor_args,
                                              r,
                                              warm_start,
                                              subproblem_callbacks;
                                              kwargs...)
                      )

    if lower_bound > 0

        if report > 0
            println("...initializing lower-bound subproblems...")
            flush(stdout)
        end

        @timeit(timo,
                "Initialze Lower Bound Subproblems",
                _initialize_lb_subproblems(scen_proc_map,
                                           worker_inf,
                                           scen_tree,
                                           model_constructor,
                                           constructor_args,
                                           warm_start;
                                           kwargs...)
                )

    end

    # Construct master ph object
    phd = @timeit(timo,
                  "Other",
                  PHData(r,
                         scen_tree,
                         scen_proc_map,
                         var_map,
                         timo)
                  )

    # Initial values
    xhat_res_sq = @timeit(timo,
                          "Initial Values",
                          _set_initial_values(phd, worker_inf)
                          )

    # Check for penalty parameter update
    x_res_sq = @timeit(timo,
                       "Penalty Parameter",
                       _set_penalty_parameter(phd,
                                              worker_inf)
                       )

    # Save residual
    _save_residual(phd, -1, xhat_res_sq, x_res_sq, 0.0, 0.0)

    return (phd, worker_inf)
end

function build_extensive_form(tree::ScenarioTree,
                              model_constructor::Function,
                              constructor_args::Tuple,
                              sub_type::Type{S};
                              kwargs...
                              )::JuMP.Model where {S <: AbstractSubproblem}

    model = JuMP.Model()
    JuMP.set_objective_sense(model, MOI.MIN_SENSE)
    JuMP.set_objective_function(model, zero(JuMP.QuadExpr))

    # Below for mapping subproblem variables onto existing extensive form variables
    node_var_map = ef_node_dict_constructor(sub_type)

    for s in scenarios(tree)

        smod = model_constructor(s, constructor_args...; kwargs...)

        snode_var_map = ef_copy_model(model, smod, s, tree, node_var_map)

        @assert(isempty(intersect(keys(node_var_map), keys(snode_var_map))))
        merge!(node_var_map, snode_var_map)
    end

    return model
end


# using ..CovidSurvey
#include("../data.jl")

const regions = ["Sweden","France","Hungary","UK","USA","Italy","Germany","Denmark"]
const nregions = length(regions)
const populations = fill(1e10, nregions)

function regions!(data)
    data["regions"] = regions
    data["num_regions"] = nregions
    data["populations"] = populations
    return nothing
end

function observables!(data)
    epilong = data["data"] = @_ projectdir("data/epidata/", "epidata_longformat.csv") |>
        CSV.read( __, DataFrame ) |>
        preprocess

    data["cases"] = @_ DataFrames.select(epilong, :date, :cases, :region) |>
        unstack(__, :region, :cases) |>
        disallowmissing(__)

    data["hospit"] = nothing

    data["deaths"] = @_ DataFrames.select(epilong, :date, :deaths, :region) |>
        unstack(__, :region, :deaths) |>
        disallowmissing(__)

    return nothing
end

function preprocess(df::DataFrame)
    startdate = Date("2020-01-05") #Date("2020-01-05") #avoid missing value at very early date
    enddate = Date("2021-04-01")
    @_ df |>
        filter(:date => >=(startdate), __) |>
        filter(:date => <=(enddate), __) |>
        disallowmissing
end

# max_vacc_rate = 0.05
# min_case_rate = 20
# epidemic = CSV.read(projectdir("data/epidata/epidata_region=$(region).csv"), DataFrame)
# is = findfirst(>(22), epidemic.cases_per_100000)
# ds = epidemic.date[is]
# de1 = epidemic.date[is:end][findfirst(<(10), epidemic.cases_per_100000[is:end])]
# de2 = epidemic.date[findfirst(>(0.05), epidemic.fraction_vaccinated)]
# de  = min(de1, de2)

# ============================================================================
# observation model defaults
# ============================================================================
@kwdef struct CaseInit <: CovidSurvey.ObservationInit
    obs_start        ::String  = ""
    obs_stop         ::String  = ""
    delay_dist                 = NegativeBinomial2
    delay_length     ::Int64   = 40
    delay_dispersion ::Float64 = 5.41
    populations      ::Vector{Int64} = International.populations
end

@kwdef struct DeathInit <: CovidSurvey.ObservationInit
    obs_start        ::String  = ""
    obs_stop         ::String  = ""
    delay_dist                 = NegativeBinomial2
    delay_length     ::Int64   = 40
    delay_dispersion ::Float64 = 14.26
    populations      ::Vector{Int64} = International.populations
end

ObservParams!(data, o::CaseInit)   = data["casemodel"]   = CovidSurvey.ObservParams(data, o)
ObservParams!(data, o::DeathInit)  = data["deathmodel"]  = CovidSurvey.ObservParams(data, o)

# ============================================================================

function lockdown!(data, lockdowns, include_early_dynamic)
    @unpack dates_turing, num_regions = data
    include_early_dynamic && (data["lockdown_indices"] = fill(0, num_regions); return)
    li = Vector{Int64}()
    for (d, lockdown) in zip(dates_turing, lockdowns)
        i = findfirst(==(Date(lockdown)), d)
        push!(li, i)
    end
    data["lockdown_indices"] = li
    return nothing
end

# ===============================================================================
#                           main function
# ===============================================================================
function load_data(;
    observationsend  = nothing,
    predictors       = nothing,
    rwstep           = 7,
    epidemicstart    = Date("2020-06-01"),
    epidemic_threshold = 30,
    numimpute        = 6,
    casemodel         = CaseInit(),
    deathmodel        = DeathInit(),
    link              = KLogistic(3.),
    invlink           = KLogit(3.),
    lockdowns         = nothing,
    predictive        = false,
    include_early_dynamic = true,
    covariates_kwargs = Dict(
        :semiparametric => false,
        :fname => projectdir( homedir(), "data/covidsurvey/smoothed_contacts.csv" ),
        :shift => 0,
        :startdate => nothing,
        :enddate => nothing,
        :aggregation => 1,
        :mobility => nothing,
        :standardize => false
    )
)

    data = Dict{String, Any}()
    regions!(data)
    observables!(data)
    CovidSurvey.consistent!(data; startdates = epidemicstart)
    CovidSurvey.turingformat!(data, observationsend, predictive)
    lockdown!(data, lockdowns, include_early_dynamic)
    # cases_start_date!(data, cases_start)
    CovidSurvey.covariates!(data, predictors; covariates_kwargs...)
    CovidSurvey.randomwalk!(data, rwstep, covariates_kwargs)
    ObservParams!(data, casemodel)
    ObservParams!(data, deathmodel)

    # iar_idx = get_iar_idx(num_tot, num_observations, cases_start_idx, iar_step)
    # num_iar_steps = length(unique(iar_idx))

    turing_data = (
        num_impute                  = numimpute,
        num_total_days              = data["num_tot"],
        dates                       = data["dates_turing"],
        cases                       = data["cases_turing"],
        casemodel                   = data["casemodel"],
        deaths                      = data["deaths_turing"],
        deathmodel                  = data["deathmodel"],
        epidemic_start              = epidemicstart,
        serial_interval             = CovidSurvey.serialinterval(15),#padzeros(serialinterval(30), 0, num_tot), #data_usa.turing_data.serial_intervals[1:50],
        num_si                      = 15,
        lockdown_indices            = data["lockdown_indices"],
        # cases_start_indices         = data["cases_start_idx"],
        # num_case_observations       = data["num_case_obs"],
        num_rt_steps                = data["num_rt_steps"],
        rt_step_indices             = data["rt_step_index"],
        num_regions                 = length(regions),
        num_observations            = data["num_observations"],
        link                        = link,
        invlink                     = invlink,
        num_covariates              = data["num_covariates"],
        covariates                  = data["covariates"],
        covariates_starts           = data["covariates_starts"],
        semiparametric              = covariates_kwargs[:semiparametric],
        include_early_dynamic       = include_early_dynamic,
        rwscale                     = sqrt(rwstep),
        populations                 = populations,
        init_infected               = first.(data["cases_turing"])
        # num_i2h                     = 40,
        # ϕ_i2h                       = 5.41,
        # num_i2d                     = 60,
        # ϕ_i2d                       = 14.26,
        # num_i2c                     = 40,
        # ϕ_i2c                       = 5.41,
    )

    dates = [ collect(d:Day(1):d+Day(n-1)) for (d, n) in zip(data["startdates"], data["num_tot"])]
    startdates = [d[data["epidemicstart"]] for d in data["dates_turing"]]

    CovidSurvey.Data(
        turing_data,
        data["cases"],
        data["deaths"],
        data["hospit"],
        data["num_observations"], #num_observations
        data["num_tot"],
        dates,
        populations,
        predictors,
        nothing,
        regions,
        nothing,
        data["observationsend"],
        nothing,
        startdates
    )
end


"""
$(DocStringExtensions.TYPEDSIGNATURES)

Description

```jldoctest
julia> 1 + 1
2
```
"""
function graph_to_kmers(g)
    kmers = [g.vprops[v][:kmer] for v in Graphs.vertices(g)]
    return kmers
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

Description

```jldoctest
julia> 1 + 1
2
```
"""
function graph_to_edge_sequences(g)
    edges = Set{BioSequences.BigDNAMer{g.gprops[:k]+1}}()
    for edge in Graphs.edges(g)
        src_kmer = g.vprops[edge.src][:kmer]
        dst_kmer = g.vprops[edge.dst][:kmer]
        for orientation in g.eprops[edge][:orientations]
            if orientation.source_orientation
                oriented_src_kmer = src_kmer
            else
                oriented_src_kmer = BioSequences.reverse_complement(src_kmer)
            end
            if orientation.destination_orientation
                oriented_dst_kmer = dst_kmer
            else
                oriented_dst_kmer = BioSequences.reverse_complement(dst_kmer)
            end
            for i in 1:g.gprops[:k]-1
                @assert oriented_src_kmer[i+1] == oriented_dst_kmer[i]
            end
            edge_mer = BioSequences.BigDNAMer((nuc for nuc in oriented_src_kmer)..., last(oriented_dst_kmer))
            push!(edges, BioSequences.canonical(edge_mer))
        end
    end
    return edges
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

Description

```jldoctest
julia> 1 + 1
2
```
"""
function kmer_graph_distances(g1, g2)
    g1_kmers = Set(graph_to_kmers(g1))
    g1_edges = graph_to_edge_sequences(g1)

    g2_kmers = Set(graph_to_kmers(g2))
    g2_edges = graph_to_edge_sequences(g2)

    kmer_distance = 1 - LSHFunctions.jaccard(g1_kmers, g2_kmers)
    edge_distance = 1 - LSHFunctions.jaccard(g1_edges, g2_edges)

    result = (
        kmer_distance = kmer_distance,
        edge_distance = edge_distance
    )
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

A short description of the function

```jldoctest
julia> 1 + 1
2
```
"""
function set_metadata!(kmer_graph, vertex::V, key, value) where V <: Integer
    if MetaGraphs.has_prop(kmer_graph, vertex, key)
        push!(kmer_graph.vprops[vertex][key], value)
    else
        MetaGraphs.set_prop!(kmer_graph, vertex, key, Set([value]))
    end
    return true
end

function set_metadata!(kmer_graph, edge::E, key, value) where E <: Graphs.Edge
    if MetaGraphs.has_prop(kmer_graph, edge, key)
        current_value = MetaGraphs.get_prop(kmer_graph, edge, key)
        updated_value = push!(current_value, value)
        MetaGraphs.set_prop!(kmer_graph, edge, key, updated_value)
    else
        MetaGraphs.set_prop!(kmer_graph, edge, key, Set([value]))
    end
    return true
end


#!/usr/bin/env julia
using FASTX
using BioAlignments

# Read sequence dictionary for HLA
seqFile = "data/S2_callseqs.fa"
name2seq = Dict()
open(FASTA.Reader,seqFile) do reader
	for record in reader
		name2seq[FASTA.identifier(record)]=FASTA.sequence(record)
	end
end

target = "HLA:HLA00602" # >HLA:HLA00602 DQA1*01:02:01:01 6484 bp
chrStart=32628178 # DQA1
chrEnd=32637531
hlaseq = name2seq[target]

# Read chromosome 6
chr6File= "data/chr6_hg38.fa"
chr6record = FASTA.Record()
reader = open(FASTA.Reader,chr6File)
read!(reader,chr6record)
close(reader)

# Pull out region of chromosome6
chr6seq = FASTA.sequence(chr6record)
hlaref = chr6seq[chrStart:chrEnd]

# Align that region to big HLA sequence
scoremodel = AffineGapScoreModel(EDNAFULL, gap_open=-5, gap_extend=-1);
res = pairalign(LocalAlignment(),hlaref,hlaseq,scoremodel)

print(res)

# HLA-DQA1
# ./cutandalign.jl  2.58s user 0.86s system 101% cpu 3.398 total


# `_vscalef` for architectures without `vscalef`.
# magic rounding constant: 1.5*2^52 Adding, then subtracting it from a float rounds it to an Int.
MAGIC_ROUND_CONST(::Type{Float64}) = 6.755399441055744e15
MAGIC_ROUND_CONST(::Type{Float32}) = 1.2582912f7
@inline function vscalef(x::Union{T,AbstractSIMD{<:Any,T}}, y::Union{T,AbstractSIMD{<:Any,T}}, ::False) where {T<:Union{Float32,Float64}}
    _vscalef(x, floor(y))
end
@inline signif_bits(::Type{Float32}) = 0x00000017 # 23
@inline signif_bits(::Type{Float64}) = 0x0000000000000034 # 52
@inline function _vscalef(x::Union{T,AbstractSIMD{<:Any,T}}, y::Union{T,AbstractSIMD{<:Any,T}}) where {T<:Union{Float32,Float64}}
    N = reinterpret(Base.uinttype(T), y + MAGIC_ROUND_CONST(T))
    k = N# >>> 0x00000008

    small_part = reinterpret(Base.uinttype(T), x)
    twopk = (k % Base.uinttype(T)) << signif_bits(T)
    reinterpret(T, twopk + small_part)
end
@inline vscalef(m::AbstractMask, v1::AbstractSIMD, v2::AbstractSIMD, v3::AbstractSIMD, ::False) = vifelse(m, vscalef(v1, v2, False()), v3)
@inline vscalef(v1::T,v2::T) where {T<:AbstractSIMD} = vscalef(v1,v2,has_feature(Val(:x86_64_avx512f)))
@inline vscalef(m::AbstractMask, v1::T, v2::T, v3::T) where {T<:AbstractSIMD} = vscalef(m,v1,v2,v3,has_feature(Val(:x86_64_avx512f)))
@inline vscalef(v1::T,v2::T) where {T<:Union{Float32,Float64}} = vscalef(v1,v2,False())
@inline vscalef(b::Bool,v1::T,v2::T,v3::T) where {T<:Union{Float32,Float64}} = b ? vscalef(v1,v2,False()) : zero(T)
@inline vscalef(v1,v2) = ((v3,v4) = promote(v1,v2); vscalef(v3,v4))
@generated function vscalef(v1::Vec{W,T}, v2::Vec{W,T}, ::True) where {W,T<:Union{Float32,Float64}}
    bits = 8W*sizeof(T)
    bits ∈ (128,256,512) || throw(ArgumentError("Vectors are $bits bits, but only 128, 256, and 512 bits are supported."))
    ltyp = LLVM_TYPES[T]
    vtyp = "<$W x $ltyp>"
    dors = T === Float64 ? 'd' : 's'
    instr = "$vtyp @llvm.x86.avx512.mask.scalef.p$(dors).$bits"
    mtyp = W == 16 ? "i16" : "i8"
    if bits == 512
        decl = "declare $instr($vtyp, $vtyp, $vtyp, $mtyp, i32)"
        # instrs = "%res = call $instr($vtyp %0, $vtyp %1, $vtyp undef, $mtyp -1, i32 11)\nret $vtyp %res"
        instrs = "%res = call $instr($vtyp %0, $vtyp %1, $vtyp undef, $mtyp -1, i32 8)\nret $vtyp %res"
    else
        decl = "declare $instr($vtyp, $vtyp, $vtyp, $mtyp)"
        instrs = "%res = call $instr($vtyp %0, $vtyp %1, $vtyp undef, $mtyp -1)\nret $vtyp %res"
    end
    arg_syms = [:(data(v1)), :(data(v2))]
    llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T},_Vec{$W,$T}}), vtyp, fill(vtyp, 2), arg_syms)
end
@generated function vscalef(m::AbstractMask{W}, v1::Vec{W,T}, v2::Vec{W,T}, v3::Vec{W,T}, ::True) where {W,T<:Union{Float32,Float64}}
    bits = 8W*sizeof(T)
    bits ∈ (128,256,512) || throw(ArgumentError("Vectors are $bits bits, but only 128, 256, and 512 bits are supported."))
    ltyp = LLVM_TYPES[T]
    vtyp = "<$W x $ltyp>"
    dors = T === Float64 ? 'd' : 's'
    mtyp = W == 16 ? "i16" : "i8"
    mtypj = W == 16 ? :UInt16 : :UInt8
    instr = "$vtyp @llvm.x86.avx512.mask.scalef.p$(dors).$bits"
    if bits == 512
        decl = "declare $instr($vtyp, $vtyp, $vtyp, $mtyp, i32)"
        instrs = "%res = call $instr($vtyp %0, $vtyp %1, $vtyp %2, $mtyp %3, i32 11)\nret $vtyp %res"
        instrs = "%res = call $instr($vtyp %0, $vtyp %1, $vtyp %2, $mtyp %3, i32 8)\nret $vtyp %res"
    else
        decl = "declare $instr($vtyp, $vtyp, $vtyp, $mtyp)"
        instrs = "%res = call $instr($vtyp %0, $vtyp %1, $vtyp %2, $mtyp %3)\nret $vtyp %res"
    end
    arg_syms = [:(data(v1)), :(data(v2)), :(data(v3)), :(data(m))]
    llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$mtypj}), vtyp, [vtyp,vtyp,vtyp,mtyp], arg_syms)
end
@generated function vsreduce(v::Vec{W,T}, ::Val{M}) where {W,T<:Union{Float32,Float64}, M}
    bits = 8W*sizeof(T)
    bits ∈ (128,256,512) || throw(ArgumentError("Vectors are $bits bits, but only 128, 256, and 512 bits are supported."))
    M isa Integer || throw(ArgumentError("M must be an integer, but received $M of type $(typeof(M))."))
    ltyp = LLVM_TYPES[T]
    vtyp = "<$W x $ltyp>"
    dors = T === Float64 ? 'd' : 's'
    instr = "$vtyp @llvm.x86.avx512.mask.reduce.p$(dors).$bits"
    mtyp = W == 16 ? "i16" : "i8"
    if bits == 512
        decl = "declare $instr($vtyp, i32, $vtyp, $mtyp, i32)"
        instrs = "%res = call $instr($vtyp %0, i32 $M, $vtyp undef, $mtyp -1, i32 8)\nret $vtyp %res"
    else
        decl = "declare $instr($vtyp, i32, $vtyp, $mtyp)"
        instrs = "%res = call $instr($vtyp %0, i32 $M, $vtyp undef, $mtyp -1)\nret $vtyp %res"
    end
    arg_syms = [:(data(v))]
    llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T}}), vtyp, [vtyp], arg_syms)
end

@generated function vpermi2pd(c::Vec{8,UInt64}, v1::Vec{8,Float64}, v2::Vec{8,Float64}) #where {W,T<:Union{Float32,Float64}, M}
    W = 8; T = Float64
    bits = 8W*sizeof(T)
    # bits ∈ (128,256,512) || throw(ArgumentError("Vectors are $bits bits, but only 128, 256, and 512 bits are supported."))
    ityp = "i$(8sizeof(T))"
    vityp = "<$W x $ityp>"
    ltyp = LLVM_TYPES[T]
    vtyp = "<$W x $ltyp>"
    dors = T === Float64 ? 'd' : 's'
    instr = "$vtyp @llvm.x86.avx512.vpermi2var.p$(dors).$bits"
    decl = "declare $instr($vtyp, $vityp, $vtyp)"
    instrs = "%res = call $instr($vtyp %0, $vityp %1, $vtyp %2)\nret $vtyp %res"
    arg_syms = [:(data(v1)), :(data(c)), :(data(v2))]
    jityp = T === Float64 ? :UInt64 : :UInt32
    llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T},_Vec{$W,$jityp},_Vec{$W,$T}}), vtyp, [vtyp,vityp,vtyp], arg_syms)
end

@inline vscalef(v1::VecUnroll, v2::VecUnroll) = VecUnroll(fmap(vscalef, getfield(v1, :data), getfield(v2, :data)))
@inline vscalef(m::VecUnroll, v1::VecUnroll, v2::VecUnroll, v3::VecUnroll) = VecUnroll(fmap(vscalef, getfield(m, :data), getfield(v1, :data), getfield(v2, :data), getfield(v3, :data)))
@inline vsreduce(v::VecUnroll, ::Val{M}) where {M} = VecUnroll(fmap(vsreduce, getfield(v, :data), Val{M}()))
@inline vpermi2pd(v1::VecUnroll, v2::VecUnroll, v3::VecUnroll) = VecUnroll(fmap(vpermi2pd, getfield(v1, :data), getfield(v2, :data), getfield(v3, :data)))
@inline vpermi2pd(v1::VecUnroll, v2::Vec, v3::Vec) = VecUnroll(fmap(vpermi2pd, getfield(v1, :data), v2, v3))

# magic rounding constant: 1.5*2^52 Adding, then subtracting it from a float rounds it to an Int.

# min and max arguments by base and type
MAX_EXP(::Val{2},  ::Type{Float64}) =  1024.0                   # log2 2^1023*(2-2^-52)
MIN_EXP(::Val{2},  ::Type{Float64}) = -1022.0                   # log2(big(2)^-1023*(2-2^-52))
MAX_EXP(::Val{2},  ::Type{Float32}) =  128f0                    # log2 2^127*(2-2^-52)
MIN_EXP(::Val{2},  ::Type{Float32}) = -126f0                    # log2 2^-1075
MAX_EXP(::Val{ℯ},  ::Type{Float64}) =  709.782712893383996732   # log 2^1023*(2-2^-52)
MIN_EXP(::Val{ℯ},  ::Type{Float64}) = -708.396418532264106335  # log 2^-1075
MAX_EXP(::Val{ℯ},  ::Type{Float32}) =  88.72283905206835f0      # log 2^127 *(2-2^-23)
MIN_EXP(::Val{ℯ},  ::Type{Float32}) = -87.3365448101577555f0#-103.97207708f0           # log 2^-150
MAX_EXP(::Val{10}, ::Type{Float64}) =  308.25471555991675       # log10 2^1023*(2-2^-52)
MIN_EXP(::Val{10}, ::Type{Float64}) = -307.65260000       # log10 2^-1075
MAX_EXP(::Val{10}, ::Type{Float32}) =  38.531839419103626f0     # log10 2^127 *(2-2^-23)
MIN_EXP(::Val{10}, ::Type{Float32}) = -37.9297794795476f0      # log10 2^-127 *(2-2^-23)

# 256/log(base, 2) (For Float64 reductions)
LogBo256INV(::Val{2}, ::Type{Float64})    = 256.
LogBo256INV(::Val{ℯ}, ::Type{Float64})    = 369.3299304675746
LogBo256INV(::Val{10}, ::Type{Float64})   = 850.4135922911647
LogBo16INV(::Val{2}, ::Type{Float64})    = 16.
LogBo16INV(::Val{ℯ}, ::Type{Float64})    = 23.083120654223414
LogBo16INV(::Val{10}, ::Type{Float64})   = 53.150849518197795
# -log(base, 2)/256 in upper and lower bits
LogBo256U(::Val{2}, ::Type{Float64})      = -0.00390625
LogBo256U(::Val{ℯ}, ::Type{Float64})      = -0.0027076061740622863
LogBo256U(::Val{10}, ::Type{Float64})     = -0.0011758984205624266
LogBo256L(base::Val{2}, ::Type{Float64})  =  0.
LogBo256L(base::Val{ℯ}, ::Type{Float64})  = -9.058776616587108e-20
LogBo256L(base::Val{10}, ::Type{Float64}) = 1.0952062999160822e-20

LogBo16U(::Val{2}, ::Type{Float64})      = -0.0625
LogBo16U(::Val{ℯ}, ::Type{Float64})      = -0.04332169878499658183857700759113603550471875839751595338254250059333710137310597
LogBo16U(::Val{10}, ::Type{Float64})     = -0.01881437472899882470085868092028081417301186759138178383190171632044426182965149
LogBo16L(base::Val{2}, ::Type{Float64})  = Zero()
LogBo16L(base::Val{ℯ}, ::Type{Float64})  = -1.4494042586539372e-18
LogBo16L(base::Val{10}, ::Type{Float64}) = 1.7523300798657315e-19

# 1/log(base, 2) (For Float32 reductions)
LogBINV(::Val{2}, ::Type{Float32})    =  1f0
LogBINV(::Val{ℯ}, ::Type{Float32})    =  1.442695f0
LogBINV(::Val{10}, ::Type{Float32})   =  3.321928f0
# -log(base, 2) in upper and lower bits
LogBU(::Val{2}, ::Type{Float32})      = -1f0
LogBU(::Val{ℯ}, ::Type{Float32})      = -0.6931472f0
LogBU(::Val{10}, ::Type{Float32})     = -0.30103f0
LogBL(base::Val{2}, ::Type{Float32})  =  0f0
LogBL(base::Val{ℯ}, ::Type{Float32})  =  1.9046542f-9
LogBL(base::Val{10}, ::Type{Float32}) =  1.4320989f-8

const FloatType64 = Union{Float64,AbstractSIMD{<:Any,Float64}}
const FloatType32 = Union{Float32,AbstractSIMD{<:Any,Float32}}
# Range reduced kernels
@inline function expm1b_kernel(::Val{2}, x::FloatType64)
    # c6 = 0.6931471807284470571335252997834339128744539291358546258980326560263434831636494
    # c5 = 0.2402265119815758621794630410361025063296309075509484445159065872903725193960909
    # c4 = 0.05550410353447979823044149277158612685395896745775325243210607075766620053156177
    # c3 = 0.009618027253668450057706478612143223628979891379942570690446533010539871321541621
    # c2 = 0.001333392256353875413926876917141786686018234585146983223440727245459444740967253
    # c1 = 0.0001546929114168849728971327603158937595919966441732209337930866845915899223829891
    # c0 = 1.520192159457321441849564286267892034534060236471603225598783028117591315796835e-05
    # c5 = 0.6931472067096466099497350107329038640311532915014328403152023514862215769240471
    # c4 = 0.2402265150505520926831534797602254284855354178135282005410994764797061783074115
    # c3 = 0.05550327215766594452554739168596479012109775178907146059734348050730310383358696
    # c2 = 0.00961799451416147891836707565892019722415069604010430702590438969636015825337285
    # c1 = 0.001340043166700788064581996335332499076913713747844545061461125917537819216640726
    # c0 = 0.0001547802227945780278842074081393459334029013917349238368250485753166922523500416
    # x * muladd(muladd(muladd(muladd(muladd(muladd(muladd(muladd(c0,x,c1),x,c2),x,c3),x,c4),x,c5),x,c6),x,c7),x,c8)
    # x * muladd(muladd(muladd(muladd(muladd(muladd(c0,x,c1),x,c2),x,c3),x,c4),x,c5),x,c6)
    # x * muladd(muladd(muladd(muladd(muladd(c0,x,c1),x,c2),x,c3),x,c4),x,c5)
  x * muladd(muladd(muladd(0.009618130135925114, x, 0.055504115022757844), x, 0.2402265069590989), x, 0.6931471805599393)
end
@inline function expm1b_kernel(::Val{ℯ}, x::FloatType64)
  x * muladd(muladd(muladd(0.04166666762124105, x, 0.1666666704849642), x, 0.49999999999999983), x, 0.9999999999999998)
end
@inline function expm1b_kernel(::Val{10}, x::FloatType64)
  x * muladd(muladd(muladd(muladd(0.5393833837413015, x, 1.1712561359457612), x, 2.0346785922926713), x, 2.6509490552382577), x, 2.302585092994046)
end
@inline function expb_kernel(::Val{2}, x::FloatType32)
  muladd(muladd(muladd(muladd(muladd(muladd(muladd(1.5316464f-5, x, 0.00015478022f0), x, 0.0013400431f0), x, 0.009617995f0), x, 0.05550327f0), x, 0.24022652f0), x, 0.6931472f0), x, 1.0f0)
end
@inline function expb_kernel(::Val{ℯ}, x::FloatType32)
  muladd(muladd(muladd(muladd(muladd(muladd(muladd(0.00019924171f0, x, 0.0013956056f0), x, 0.008375129f0), x, 0.041666083f0), x, 0.16666415f0), x, 0.5f0), x, 1.0f0), x, 1.0f0)
end
@inline function expb_kernel(::Val{10}, x::FloatType32)
  muladd(muladd(muladd(muladd(muladd(muladd(muladd(0.06837386f0, x, 0.20799689f0), x, 0.54208815f0), x, 1.1712388f0), x, 2.034648f0), x, 2.6509492f0), x, 2.3025851f0), x, 1.0f0)
end

const J_TABLE= Float64[2.0^(big(j-1)/256) for j in 1:256];

@inline fast_fma(a, b, c, ::True) = fma(a, b, c)
@inline function fast_fma(a, b, c, ::False)
    d = dadd(dmul(Double(a),Double(b),False()),Double(c))
    add_ieee(d.hi, d.lo)
end

@static if (Sys.ARCH === :x86_64) | (Sys.ARCH === :i686)
const TABLE_EXP_64_0= Vec(ntuple(j -> Core.VecElement(Float64(2.0^(big(j-1)/16))), Val(8)))
const TABLE_EXP_64_1= Vec(ntuple(j -> Core.VecElement(Float64(2.0^(big(j+7)/16))), Val(8)))

@inline target_trunc(v, ::VectorizationBase.True) = v
@inline target_trunc(v, ::VectorizationBase.False) = v % UInt32
@inline target_trunc(v) = target_trunc(v, VectorizationBase.has_feature(Val(:x86_64_avx512dq)))

# @inline function vexp2_v1(x::AbstractSIMD{8,Float64})
#     x16 = x
#     # x16 = 16x
#     r =  vsreduce(x16, Val(4))
#     m = x16 - r
#     mfrac = m
#     inds = (reinterpret(UInt64, mfrac) >> 0x000000000000002d) & 0x000000000000000f
#     # @show r m mfrac reinterpret(UInt64, m) reinterpret(UInt64, mfrac)
#     # js = vpermi2pd(inds, TABLE_EXP_64_0, TABLE_EXP_64_1)
#     # @show m mfrac r
#     small_part = expm1b_kernel(Val(2), r) + 1.0
#     # js = 1.0
#     # small_part = vfmadd(js, expm1b_kernel(Val(2), r), js)
#     vscalef(small_part, mfrac)
# end

@inline function expm1b_kernel_16(::Val{2}, x)
    c5 = 0.6931471805599457533351827593325319924753473772859614915719486459595837313933663
    c4 = 0.2402265069591009431089489060897837825648676480621950809556237945205562267112511
    c3 = 0.05550410865663929372911461843767669974316894963735870580154522796380984673567634
    c2 = 0.00961812910613182376367867689426991348318504321185094738294343470767871628697879
    c1 = 0.001333378157683735211078403326604752238340853209789619792858391909299167771871147
    c0 = 0.0001540378851029625623114826398060979330719673345637296642237670082377277446583639

    x * vmuladd_fast(vmuladd_fast(vmuladd_fast(vmuladd_fast(vmuladd_fast(c0,x,c1),x,c2),x,c3),x,c4),x,c5)
end
@inline function expm1b_kernel_16(::Val{ℯ}, x)
    c5 = 1.000000000000000640438225946852982701258391604480638275588427888399057119915227
    c4 = 0.5000000000000004803287542332217715766116071510522583538834274474935075866898906
    c3 = 0.1666666666420970554527207281834431226735741940161719645985080160507073865025253
    c2 = 0.04166666666018301921665935823120024420659933010915524936059730214858712998209511
    c1 = 0.008333472974984405879148046293292753978131598285368755170712233472964279745777974
    c0 = 0.001388912162496018623851591184688043066476532389093553363939521815388727152058955
    x * vmuladd_fast(vmuladd_fast(vmuladd_fast(vmuladd_fast(vmuladd_fast(c0,x,c1),x,c2),x,c3),x,c4),x,c5)
end
@inline function expm1b_kernel_16(::Val{10}, x)
    c5 = 2.302585092994047158681503503460480873860999793973827515869365761962430703319599
    c4 = 2.650949055239201551934947671858339219424413336755573194223955910456821842421188
    c3 = 2.034678591993528625083054018110717593492391962926107057414972651492780932119609
    c2 = 1.171255148730010832148986815739840479496404930758611797062000621828706728750671
    c1 = 0.5393919676343165962473794696403862709895176169091768718271240564783020989456929
    c0 = 0.2069993173257113377172910397724414085027323868592924170210484263853229417141011
    x * vmuladd_fast(vmuladd_fast(vmuladd_fast(vmuladd_fast(vmuladd_fast(c0,x,c1),x,c2),x,c3),x,c4),x,c5)
end

# @inline function expm1b_kernel_6(::Val{2}, x)
#     c6 = 0.6931471805599453094157857128856867906777808839530204388709017060018142940776171
#     c5 = 0.2402265069591008469523412357777710806331204676211882505524452303635329805655358
#     c4 = 0.05550410866482161698117481145997657675888637052903928922860089114193615876343007
#     c3 = 0.009618129106525681209446910869436511674612664593938506025921861852732097374219765
#     c2 = 0.001333355814485111028814491629997422287837080138262113640339762815224264878390602
#     c1 = 0.0001540375624549508734984308915404920724081959405205139643458055193131059150331185
#     c0 = 1.525295744513115862409484203574886089068776710280414076942740530839031792795809e-05
#     x * vmuladd_fast(vmuladd_fast(vmuladd_fast(vmuladd_fast(vmuladd_fast(vmuladd_fast(c0,x,c1),x,c2),x,c3),x,c4),x,c5),x,c6)
# end
# # using Remez
# # N,D,E,X = ratfn_minimax(x -> (exp2(x) - big"1")/x, [big(nextfloat(-0.03125)),big(0.03125)], 4, 0); @show(E); N
# @inline function expm1b_kernel_4(::Val{2}, x)
#     c4 = 0.6931471805599461972549081995383434692316977327912755704234013405443109498729026
#     c3 = 0.2402265069131940842333497738928958607054740795078596615709864445611497846077303
#     c2 = 0.05550410865300270379171299778517151376504051870524903806295523325435530249981495
#     c1 = 0.009618317140648284298097106744730186251913149278152053357630395863210686828434175
#     c0 = 0.001333381881551676348461495248002642715422207072457864472267417920610122672570108
#     x * vmuladd_fast(vmuladd_fast(vmuladd_fast(vmuladd_fast(c0,x,c1),x,c2),x,c3),x,c4)
# end

# @inline function vexp2_v4(x::AbstractSIMD{W,Float64}) where {W}
#     # r - VectorizationBase.vroundscale(r, Val(16*(4)))
#     r = VectorizationBase.vsreduce(x,Val(0))
#     rscale = VectorizationBase.vroundscale(r, Val(64))
#     rs = r - rscale
#     inds = convert(UInt, vsreduce(rscale, Val(1))*16.0)
#     expr = expm1b_kernel_5(Val(2), rs)
#     N_float = x - rs
#     # @show inds rs N_float

#     js = vpermi2pd(inds, TABLE_EXP_64_0, TABLE_EXP_64_1)
#     small_part = vfmadd(js, expr, js)
#     res = vscalef(small_part, N_float)
#     return res

# end

####################################################################################################
################################## Non-AVX512 implementation #######################################
####################################################################################################

# With AVX512, we use a tiny look-up table of just 16 numbers for `Float64`,
# because we can perform the lookup using `vpermi2pd`, which is much faster than gather.
# To compensate, we need a larger polynomial.

# Because of the larger polynomial, this implementation works better on systems with 2 FMA units.

@inline function vexp2(x::AbstractSIMD{8,Float64}, ::True)
    r = vsreduce(16.0x, Val(0)) * 0.0625
    N_float = x - r
    expr = expm1b_kernel_16(Val(2), r)
    inds = convert(UInt64, vsreduce(N_float, Val(1))*16.0)
    # inds = ((trunc(Int64, 16.0*N_float)%UInt64)) & 0x000000000000000f
    js = vpermi2pd(inds, TABLE_EXP_64_0, TABLE_EXP_64_1)
    small_part = vfmadd(js, expr, js)
    res = vscalef(small_part, N_float)
    return res
end
# @inline function vexp2_v2(x::AbstractSIMD{8,Float64}, ::True)#, ::Val{N}) where {N}
#     r1 = vsreduce(x, Val(0))
#     m = x - r1
#     r = vfmsub(vsreduce(r1 * 16.0, Val(1)), 0.0625, 0.5)
#     j = r1 - r
#     js = vpermi2pd(convert(UInt, j), TABLE_EXP_64_0, TABLE_EXP_64_1)

#     expr = expm1b_kernel_5(Val(2), r) # 2^r - 1

#     small_part = vfmadd(js, expr, js)
#     # @show r m j
#     vscalef(small_part, m)
# end
# @inline function vexp_v3(x::AbstractSIMD{8,Float64}, ::True)#, ::Val{N}) where {N}
#     xl2e = mul_ieee(1.4426950408889634, x)
#     r1 = vsreduce(xl2e, Val(0))
#     m = xl2e - r1
#     r = vfmsub(vsreduce(r1 * 16.0, Val(1)), 0.0625, 0.5)
#     j = r1 - r
#     js = vpermi2pd(convert(UInt, j), TABLE_EXP_64_0, TABLE_EXP_64_1)
#     rs = vfnmadd(0.6931471805599453094172321214581765680755001343602552541206800094933936219696955, m+j, x)
#     expr = expm1b_kernel_5(Val(ℯ), r) # 2^r - 1

#     small_part = vfmadd(js, expr, js)
#     @show r m j
#     vscalef(small_part, m)
# end

# _log2(::Val{ℯ}) = 1.4426950408889634
# invlog2hi(::Val{ℯ}) = 0.6931471805599453094172321214581765680755001343602552541206800094933936219696955
# invlog2lo(::Val{ℯ}) = -2.319046813846299615494855463875478650412068000949339362196969553467383712860567e-17
# _log2(::Val{10}) = 3.321928094887362347870319429489390175864831393024580612054756395815934776608624
# invlog2hi(::Val{10}) = 0.3010299956639811952137388947244930267681898814621085413104274611271081892744238
# invlog2lo(::Val{10}) = 2.803728127785170339013117338996875833689572538872891810725576172209659522828247e-18

# Requires two more floating point μops, but 8 less loading μops than the default version.
# This thus microbenchmarks a little worse, but the theory is that using less cache than the
# 256 Float64 * 8 bytes/Float64 = 2 KiB table may improve real world performance / reduce
# random latency.
@inline function vexp_avx512(x::AbstractSIMD{8,Float64}, ::Val{B}) where {B}
    N_float = round(x*LogBo16INV(Val(B), Float64))
    r = muladd(N_float, LogBo16U(Val(B), Float64), x)
    r = muladd(N_float, LogBo16L(Val(B), Float64), r)
    inds = ((trunc(Int64, N_float)%UInt64)) & 0x000000000000000f
    expr = expm1b_kernel_16(Val(B), r)
    js = vpermi2pd(inds, TABLE_EXP_64_0, TABLE_EXP_64_1)
    small_part = vfmadd(js, expr, js)
    res = vscalef(small_part, 0.0625*N_float)
    return res
end

@inline function vexp_avx512(x::AbstractSIMD{W,Float64}, ::Val{B}) where {W,B}
    N_float = muladd(x, LogBo256INV(Val{B}(), Float64), MAGIC_ROUND_CONST(Float64))
    N = target_trunc(reinterpret(UInt64, N_float))
    N_float = N_float - MAGIC_ROUND_CONST(Float64)
    r = fma(N_float, LogBo256U(Val{B}(), Float64), x)
    r = fma(N_float, LogBo256L(Val{B}(), Float64), r)
    # @show (N & 0x000000ff) % Int
    # @show N N & 0x000000ff
    js = vload(VectorizationBase.zero_offsets(stridedpointer(J_TABLE)), (N & 0x000000ff,))
    # k = N >>> 0x00000008
    # small_part = reinterpret(UInt64, vfmadd(js, expm1b_kernel(Val{B}(), r), js))
    small_part = vfmadd(js, expm1b_kernel(Val{B}(), r), js)
    # return reinterpret(Float64, small_part), r, k, N_float, js
    res = vscalef(small_part, 0.00390625*N_float)
    # twopk = (k % UInt64) << 0x0000000000000034
    # res = reinterpret(Float64, twopk + small_part)
    return res
end
@inline function vexp_avx512(x::Union{Float32,AbstractSIMD{<:Any,Float32}}, ::Val{B}) where {B}
    N_float = vfmadd(x, LogBINV(Val{B}(), Float32), MAGIC_ROUND_CONST(Float32))
    N = reinterpret(UInt32, N_float)
    N_float = (N_float - MAGIC_ROUND_CONST(Float32))

    r = fast_fma(N_float, LogBU(Val{B}(), Float32), x, fma_fast())
    r = fast_fma(N_float, LogBL(Val{B}(), Float32), r, fma_fast())

    small_part = expb_kernel(Val{B}(), r)
    res = vscalef(small_part, N_float)
    # twopk = N << 0x00000017
    # res = reinterpret(Float32, twopk + small_part)
    return res
end

else# if !((Sys.ARCH === :x86_64) | (Sys.ARCH === :i686))
  const target_trunc = identity
end

# @inline function vexp_avx512(vu::VecUnroll{1,8,Float64,Vec{8,Float64}}, ::Val{B}) where {B}
#   x, y = data(vu)
#   N_float₁ = round(x*LogBo16INV(Val(B), Float64))
#   N_float₂ = muladd(y, LogBo256INV(Val{B}(), Float64), MAGIC_ROUND_CONST(Float64))
#   r₁ = muladd(N_float₁, LogBo16U(Val(B), Float64), x)
#   N₂ = target_trunc(reinterpret(UInt64, N_float₂))
#   r₁ = muladd(N_float₁, LogBo16L(Val(B), Float64), r₁)
#   js₂ = vload(VectorizationBase.zero_offsets(stridedpointer(J_TABLE)), (N₂ & 0x000000ff,))
#   N_float₂ = N_float₂ - MAGIC_ROUND_CONST(Float64)
#   inds₁ = ((trunc(Int64, N_float₁)%UInt64)) & 0x000000000000000f
#   r₂ = fma(N_float₂, LogBo256U(Val{B}(), Float64), y)
#   expr₁ = expm1b_kernel_16(Val(B), r₁)
#   r₂ = fma(N_float₂, LogBo256L(Val{B}(), Float64), r₂)
#   js₁ = vpermi2pd(inds₁, TABLE_EXP_64_0, TABLE_EXP_64_1)
#   small_part₁ = vfmadd(js₁, expr₁, js₁)
#   small_part₂ = vfmadd(js₂, expm1b_kernel(Val{B}(), r₂), js₂)
#   res₁ = vscalef(small_part₁, 0.0625*N_float₁)
#   res₂ = vscalef(small_part₂, 0.00390625*N_float₂)
#   return VecUnroll((res₁, res₂))
# end

# @inline _vexp(x, ::True) = vexp2( 1.4426950408889634 * x, True() )
# @inline _vexp(x, ::True) = vexp2( mul_ieee(1.4426950408889634, x), True() )
# @inline _vexp10(x, ::True) = vexp2( 3.321928094887362 * x, True() )
# @inline _vexp(x) = _vexp(x, has_feature(Val(:x86_64_avx512f)))
# @inline _vexp10(x) = _vexp10(x, has_feature(Val(:x86_64_avx512f)))

@inline vexp(x::AbstractSIMD, ::True) = vexp_avx512(x, Val(ℯ))
@inline vexp2(x::AbstractSIMD, ::True) = vexp_avx512(x, Val(2))
@inline vexp10(x::AbstractSIMD, ::True) = vexp_avx512(x, Val(10))
# @inline vexp(x::AbstractSIMD{W,Float32}, ::True) where {W} = vexp_generic(x, Val(ℯ))
# @inline vexp2(x::AbstractSIMD{W,Float32}, ::True) where {W} = vexp_generic(x, Val(2))
# @inline vexp10(x::AbstractSIMD{W,Float32}, ::True) where {W} = vexp_generic(x, Val(10))

@inline Base.exp(v::AbstractSIMD{W}) where {W} = vexp(float(v))
@inline Base.exp2(v::AbstractSIMD{W}) where {W} = vexp2(float(v))
@inline Base.exp10(v::AbstractSIMD{W}) where {W} = vexp10(float(v))
@static if (Base.libllvm_version ≥ v"11") & ((Sys.ARCH === :x86_64) | (Sys.ARCH === :i686))
    @inline vexp(v::AbstractSIMD) = vexp(v, has_feature(Val(:x86_64_avx512f)))
    @inline vexp2(v::AbstractSIMD) = vexp2(v, has_feature(Val(:x86_64_avx512f)))
    @inline vexp10(v::AbstractSIMD) = vexp10(v, has_feature(Val(:x86_64_avx512f)))
else
    @inline vexp(v::AbstractSIMD) = vexp(v, False())
    @inline vexp2(v::AbstractSIMD) = vexp2(v, False())
    @inline vexp10(v::AbstractSIMD) = vexp10(v, False())
end
@inline vexp(v::Union{Float32,Float64}) = vexp(v, False())
@inline vexp2(v::Union{Float32,Float64}) = vexp2(v, False())
@inline vexp10(v::Union{Float32,Float64}) = vexp10(v, False())
@inline vexp(v::AbstractSIMD{2,Float32}) = vexp(v, False())
@inline vexp2(v::AbstractSIMD{2,Float32}) = vexp2(v, False())
@inline vexp10(v::AbstractSIMD{2,Float32}) = vexp10(v, False())

####################################################################################################
################################## Non-AVX512 implementation #######################################
####################################################################################################

@inline function vexp_generic_core(x::Union{Float64,AbstractSIMD{<:Any,Float64}}, ::Val{B}) where {B}
    N_float = muladd(x, LogBo256INV(Val{B}(), Float64), MAGIC_ROUND_CONST(Float64))
    N = target_trunc(reinterpret(UInt64, N_float))
    N_float = N_float - MAGIC_ROUND_CONST(Float64)
    r = fast_fma(N_float, LogBo256U(Val{B}(), Float64), x, fma_fast())
    r = fast_fma(N_float, LogBo256L(Val{B}(), Float64), r, fma_fast())
    # @show (N & 0x000000ff) % Int
    js = vload(VectorizationBase.zero_offsets(stridedpointer(J_TABLE)), (N & 0x000000ff,))
    k = N >>> 0x00000008
    small_part = reinterpret(UInt64, vfmadd(js, expm1b_kernel(Val{B}(), r), js))
    # return reinterpret(Float64, small_part), r, k, N_float, js
    twopk = (k % UInt64) << 0x0000000000000034
    res = reinterpret(Float64, twopk + small_part)
    return res
end
@inline function vexp_generic(x::Union{Float64,AbstractSIMD{<:Any,Float64}}, ::Val{B}) where {B}
    res = vexp_generic_core(x, Val{B}())
    res = ifelse(x >= MAX_EXP(Val{B}(), Float64), Inf, res)
    res = ifelse(x <= MIN_EXP(Val{B}(), Float64), 0.0, res)
    res = ifelse(isnan(x), x, res)
    return res
end

@inline function vexp_generic_core(x::Union{Float32,AbstractSIMD{<:Any,Float32}}, ::Val{B}) where {B}
    N_float = vfmadd(x, LogBINV(Val{B}(), Float32), MAGIC_ROUND_CONST(Float32))
    N = reinterpret(UInt32, N_float)
    N_float = (N_float - MAGIC_ROUND_CONST(Float32))

    r = fast_fma(N_float, LogBU(Val{B}(), Float32), x, fma_fast())
    r = fast_fma(N_float, LogBL(Val{B}(), Float32), r, fma_fast())

    small_part = reinterpret(UInt32, expb_kernel(Val{B}(), r))
    twopk = N << 0x00000017
    res = reinterpret(Float32, twopk + small_part)
    return res
end
@inline function vexp_generic(x::Union{Float32,AbstractSIMD{<:Any,Float32}}, ::Val{B}) where {B}
    res = vexp_generic_core(x, Val{B}())
    res = ifelse(x >= MAX_EXP(Val{B}(), Float32), Inf32, res)
    res = ifelse(x <= MIN_EXP(Val{B}(), Float32), 0.0f0, res)
    res = ifelse(isnan(x), x, res)
    return res
end
for (func, base) in (:vexp2=>Val(2), :vexp=>Val(ℯ), :vexp10=>Val(10))
    @eval @inline $func(x, ::False) = vexp_generic(x, $base)
end

####################################################################################################
#################################### LOG HELPERS ###################################################
####################################################################################################

# TODO: move these back to log.jl when when the log implementations there are good & well tested enough to use.

@generated function vgetexp(v::Vec{W,T}) where {W,T<:Union{Float32,Float64}}
    bits = 8W*sizeof(T)
    bits ∈ (128,256,512) || throw(ArgumentError("Vectors are $bits bits, but only 128, 256, and 512 bits are supported."))
    ltyp = LLVM_TYPES[T]
    vtyp = "<$W x $ltyp>"
    dors = T === Float64 ? 'd' : 's'
    instr = "$vtyp @llvm.x86.avx512.mask.getexp.p$(dors).$bits"
    mtyp = W == 16 ? "i16" : "i8"
    if bits == 512
        decl = "declare $instr($vtyp, $vtyp, $mtyp, i32)"
        instrs = "%res = call $instr($vtyp %0, $vtyp undef, $mtyp -1, i32 12)\nret $vtyp %res"
    else
        decl = "declare $instr($vtyp, $vtyp, $mtyp)"
        instrs = "%res = call $instr($vtyp %0, $vtyp undef, $mtyp -1)\nret $vtyp %res"
    end
    arg_syms = [:(data(v))]
    llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T}}), vtyp, [vtyp], arg_syms)
end
@generated function vgetmant(v::Vec{W,T}) where {W,T<:Union{Float32,Float64}}
    bits = 8W*sizeof(T)
    bits ∈ (128,256,512) || throw(ArgumentError("Vectors are $bits bits, but only 128, 256, and 512 bits are supported."))
    ltyp = LLVM_TYPES[T]
    vtyp = "<$W x $ltyp>"
    dors = T === Float64 ? 'd' : 's'
    instr = "$vtyp @llvm.x86.avx512.mask.getmant.p$(dors).$bits"
    mtyp = W == 16 ? "i16" : "i8"
    if bits == 512
        decl = "declare $instr($vtyp, i32, $vtyp, $mtyp, i32)"
        instrs = "%res = call $instr($vtyp %0, i32 11, $vtyp undef, $mtyp -1, i32 12)\nret $vtyp %res"
    else
        decl = "declare $instr($vtyp, i32, $vtyp, $mtyp)"
        instrs = "%res = call $instr($vtyp %0, i32 11, $vtyp undef, $mtyp -1)\nret $vtyp %res"
    end
    arg_syms = [:(data(v))]
    llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T}}), vtyp, [vtyp], arg_syms)
end
@generated function vgetmant(v::Vec{W,T}, ::Union{StaticInt{N},Val{N}}) where {W,T<:Union{Float32,Float64},N}
    bits = 8W*sizeof(T)
    bits ∈ (128,256,512) || throw(ArgumentError("Vectors are $bits bits, but only 128, 256, and 512 bits are supported."))
    ltyp = LLVM_TYPES[T]
    vtyp = "<$W x $ltyp>"
    dors = T === Float64 ? 'd' : 's'
    instr = "$vtyp @llvm.x86.avx512.mask.getmant.p$(dors).$bits"
    mtyp = W == 16 ? "i16" : "i8"
    if bits == 512
        decl = "declare $instr($vtyp, i32, $vtyp, $mtyp, i32)"
        instrs = "%res = call $instr($vtyp %0, i32 $N, $vtyp undef, $mtyp -1, i32 12)\nret $vtyp %res"
    else
        decl = "declare $instr($vtyp, i32, $vtyp, $mtyp)"
        instrs = "%res = call $instr($vtyp %0, i32 $N, $vtyp undef, $mtyp -1)\nret $vtyp %res"
    end
    arg_syms = [:(data(v))]
    llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T}}), vtyp, [vtyp], arg_syms)
end
@generated function vgetmant12(v::Vec{W,T}) where {W,T<:Union{Float32,Float64}}
    bits = 8W*sizeof(T)
    bits ∈ (128,256,512) || throw(ArgumentError("Vectors are $bits bits, but only 128, 256, and 512 bits are supported."))
    ltyp = LLVM_TYPES[T]
    vtyp = "<$W x $ltyp>"
    dors = T === Float64 ? 'd' : 's'
    instr = "$vtyp @llvm.x86.avx512.mask.getmant.p$(dors).$bits"
    mtyp = W == 16 ? "i16" : "i8"
    if bits == 512
        decl = "declare $instr($vtyp, i32, $vtyp, $mtyp, i32)"
        instrs = "%res = call $instr($vtyp %0, i32 12, $vtyp undef, $mtyp -1, i32 8)\nret $vtyp %res"
    else
        decl = "declare $instr($vtyp, i32, $vtyp, $mtyp)"
        instrs = "%res = call $instr($vtyp %0, i32 12, $vtyp undef, $mtyp -1)\nret $vtyp %res"
    end
    arg_syms = [:(data(v))]
    llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T}}), vtyp, [vtyp], arg_syms)
end

@generated function vroundscale(v::Vec{W,T}, ::Union{Val{N},StaticInt{N}}) where {W,T,N}
    bits = 8W*sizeof(T)
    bits ∈ (128,256,512) || throw(ArgumentError("Vectors are $bits bits, but only 128, 256, and 512 bits are supported."))
    ltyp = LLVM_TYPES[T]
    vtyp = "<$W x $ltyp>"
    dors = T === Float64 ? 'd' : 's'
    instr = "$vtyp @llvm.x86.avx512.mask.rndscale.p$(dors).$bits"
    mtyp = W == 16 ? "i16" : "i8"
    if bits == 512
        decl = "declare $instr($vtyp, i32, $vtyp, $mtyp, i32)"
        instrs = "%res = call $instr($vtyp %0, i32 $N, $vtyp undef, $mtyp -1, i32 4)\nret $vtyp %res"
    else
        decl = "declare $instr($vtyp, i32, $vtyp, $mtyp)"
        instrs = "%res = call $instr($vtyp %0, i32 $N, $vtyp undef, $mtyp -1)\nret $vtyp %res"
    end
    arg_syms = [:(data(v))]
    llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T}}), vtyp, [vtyp], arg_syms)
end

@inline vgetexp(v::VecUnroll) = VecUnroll(fmap(vgetexp, getfield(v, :data)))
@inline vgetmant(v::VecUnroll) = VecUnroll(fmap(vgetmant, getfield(v, :data)))
@inline vgetmant(v::VecUnroll,::Union{StaticInt{N},Val{N}}) where {N} = VecUnroll(fmap(vgetmant, getfield(v, :data), StaticInt{N}()))
@inline Base.significand(v::VecUnroll, ::True) = VecUnroll(fmap(vgetmant12, getfield(v, :data)))
@inline vroundscale(v::VecUnroll,::Union{StaticInt{N},Val{N}}) where {N} = VecUnroll(fmap(vroundscale, getfield(v, :data), StaticInt{N}()))
@inline Base.significand(v::Vec, ::True) = vgetmant12(v)
@inline Base.significand(v::AbstractSIMDVector, ::True) = vgetmant12(Vec(v))

mask_exponent(::Val{Float64}) = 0x000f_ffff_ffff_ffff
set_exponent(::Val{Float64}) = 0x3ff0_0000_0000_0000

mask_exponent(::Val{Float32}) = 0x007fffff
set_exponent(::Val{Float32}) = 0x3f800000

mask_exponent(::Val{Float16}) = 0x07ff
set_exponent(::Val{Float16}) = 0x3c00

@inline function Base.significand(v::AbstractSIMD{W,T}, ::False) where {W,T}
  reinterpret(T, (reinterpret(Base.uinttype(T), v) & mask_exponent(Val(T))) | set_exponent(Val(T)))
end
@inline Base.exponent(v::Vec, ::True)= vgetexp(v)
@inline Base.exponent(v::AbstractSIMDVector, ::True)= vgetexp(Vec(v))
@inline Base.exponent(v::VecUnroll, ::True)= VecUnroll(fmap(vgetexp, getfield(v, :data)))
@static if VERSION ≥ v"1.7.0-beta"
  using Base: inttype
else
  @inline inttype(::Type{T}) where {T} = signed(Base.uinttype(T))
end
@inline function Base.exponent(v::AbstractSIMD{W,T}, ::False) where {W,T}
  I = inttype(T)
  vshift = reinterpret(I, v) >> (Base.significand_bits(T)%I)
  e = ((vshift % Int) & Base.exponent_raw_max(T)) - I(Base.exponent_bias(T))
  convert(T, e % Int32)
end

@inline Base.significand(v::AbstractSIMD{W,T}) where {W,T<:Union{Float32,Float64}} = significand(v, has_feature(Val(:x86_64_avx512f)))
@inline Base.exponent(v::AbstractSIMD{W,T}) where {W,T<:Union{Float32,Float64}} = exponent(v, has_feature(Val(:x86_64_avx512f)))
@inline Base.significand(v::AbstractSIMD{W,Float16}) where {W} = significand(v, False())
@inline Base.exponent(v::AbstractSIMD{W,Float16}) where {W} = exponent(v, False())
@inline Base.ldexp(v::AbstractSIMD, e::AbstractSIMD) = vscalef(v, e, has_feature(Val(:x86_64_avx512f)))

# @inline function vexp2_v2(x::AbstractSIMD{8,Float64})
#     x16 = 16.0*x
#     # x8 = 8x
#     r =  vsreduce(x16, Val(0)) * 0.0625
#     @fastmath  begin
#     m = x - r
#     # m + r = x16, r ∈ (-0.5,0.5]
#     # m/16 + r/16 = x, r ∈ (-1/32, 1/32]
#     # we must now vreduce `mfrac`
#     # return m
#     end
#     # expr = expm1b_kernel_4(Val(2), r)
#     expr = expm1b_kernel_5(Val(2), r)
#     inds = convert(UInt64, vsreduce(m, Val(1)) * 16.0)
#     # inds = (reinterpret(UInt64, mfrac) >> 0x000000000000002d) & 0x000000000000000f
#     # inds = (reinterpret(UInt64, mfrac) >> 0x0000000000000035) & 0x000000000000000f
#     # @show r m mfrac reinterpret(UInt64, m) reinterpret(UInt64, mfrac)
#     js = vpermi2pd(inds, TABLE_EXP_64_0, TABLE_EXP_64_1)
#     # return r, mfrac, js
#     # @show js m 16mfrac r mfrac inds%Int ((inds>>1)%Int)
#     # js = 1.0
#     small_part = vfmadd(js, expr, js)
#     vscalef(small_part, m)#, r, mfrac, js, inds
# end

# for (func, base) in (:vexp2=>Val(2), :vexp=>Val(ℯ), :vexp10=>Val(10))
#     @eval begin
#         @inline function $func(x::AbstractSIMD{W,Float64}) where {W}
#             N_float = muladd(x, LogBo256INV($base, Float64), MAGIC_ROUND_CONST(Float64))
#             N = target_trunc(reinterpret(UInt64, N_float))
#             N_float = N_float - MAGIC_ROUND_CONST(Float64)
#             r = fast_fma(N_float, LogBo256U($base, Float64), x, fma_fast())
#             r = fast_fma(N_float, LogBo256L($base, Float64), r, fma_fast())
#             js = vload(VectorizationBase.zero_offsets(stridedpointer(J_TABLE)), (N & 0x000000ff,))
#             k = N >>> 0x00000008
#             small_part = vfmadd(js, expm1b_kernel($base, r), js)
#             twopk = (k % UInt64) << 0x0000000000000034
#             @show N_float k twopk small_part
#             return small_part
#             # res = reinterpret(Float64, twopk + small_part)
#             # return res
#             vscalef(small_part, N_float)
#         end
#     end
# end


# [
#  [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],
#  [         :                :          :        :      :          :           ],
# ]

#! format: off

parr = [
 ["state",     "ρ", 12, 12, 12, 12],
 ["state", "ρu[1]",  0,  0,  0,  0],
 ["state", "ρu[2]",  0,  0,  0,  0],
 ["state", "ρu[3]", 12,  8, 12, 12],
 ["state",    "ρθ", 12, 10, 12, 12],
]

second_order_flat = (
[
 [ "state",     "ρ",   9.95252314022507689195e-01,  9.99992856011554298590e-01,  9.98330419819817738158e-01,  1.48639562654353791886e-03 ],
 [ "state", "ρu[1]",  -6.36093867411581375298e-15,  8.28629359748938561747e-15,  5.16470747993988398929e-16,  1.17216775586980107502e-15 ],
 [ "state", "ρu[2]",  -3.70797682276813316617e-15,  9.42818503691220985643e-15,  6.07164354178963622023e-16,  1.23143801722553508699e-15 ],
 [ "state", "ρu[3]",  -1.65133743521589989450e-03,  5.29367075491514133403e-09, -8.40309050919211576736e-04,  4.66618546037470796790e-04 ],
 [ "state",    "ρθ",  -9.95249493245247940365e+00,  1.99973806108046075331e-05, -4.98740538889376860965e+00,  2.91965149708573168397e+00 ],
],
 parr,
)

second_order = (
[
 [ "state",     "ρ",   9.95252314022507689195e-01,  9.99992856011554298590e-01,  9.98330419819817738158e-01,  1.48639562654353791886e-03 ],
 [ "state", "ρu[1]",  -6.36142448250610767485e-15,  8.28656656189152269018e-15,  5.17120118507278235699e-16,  1.17302082244252352324e-15 ],
 [ "state", "ρu[2]",  -3.70772154095122499196e-15,  9.42925797553789671672e-15,  6.06575884435408245721e-16,  1.23026363135339100359e-15 ],
 [ "state", "ρu[3]",  -1.65133743521588883564e-03,  5.29367075398690075732e-09, -8.40309050919211468315e-04,  4.66618546037470417320e-04 ],
 [ "state",    "ρθ",  -9.95249493245247940365e+00,  1.99973806108054952236e-05, -4.98740538889376860965e+00,  2.91965149708573168397e+00 ],
],
 parr,
)

fourth_order_flat = (
[
 [ "state",     "ρ",   9.95377495534709000324e-01,  9.99992951378667060958e-01,  9.98321272635789513927e-01,  1.50722816639464090097e-03 ],
 [ "state", "ρu[1]",  -1.22866254290737312252e-14,  1.95946879007561421956e-14,  1.73741725325975513295e-15,  4.23083086870931785383e-15 ],
 [ "state", "ρu[2]",  -1.63058197687820133665e-14,  2.71134092821328864807e-14,  1.75302916450504673999e-15,  4.33921909208334743316e-15 ],
 [ "state", "ρu[3]",  -1.66252523985498286418e-03,  5.55021884647818839491e-08, -8.13842490777055664608e-04,  4.76414940515918383830e-04 ],
 [ "state",    "ρθ",  -9.95373884733410818626e+00, -4.05672848247591079102e-07, -4.98722877855243940104e+00,  2.97859288054288384728e+00 ],
],
 parr,
)

fourth_order = (
[
 [ "state",     "ρ",   9.95377495534709000324e-01,  9.99992951378667060958e-01,  9.98321272635789513927e-01,  1.50722816639464068413e-03 ],
 [ "state", "ρu[1]",  -1.22778750186973996879e-14,  1.95893093780676565332e-14,  1.73744809621522011733e-15,  4.23044883848200063451e-15 ],
 [ "state", "ρu[2]",  -1.62568523313266733927e-14,  2.71071963864711276060e-14,  1.75189805678757112174e-15,  4.33763908149213610479e-15 ],
 [ "state", "ρu[3]",  -1.66252523985503642377e-03,  5.55021884639224168229e-08, -8.13842490777055447768e-04,  4.76414940515918817511e-04 ],
 [ "state",    "ρθ",  -9.95373884733410818626e+00, -4.05672848142820462126e-07, -4.98722877855243940104e+00,  2.97859288054288384728e+00 ],
],
 parr,
)

#! format: on

refVals = (; second_order, fourth_order, second_order_flat, fourth_order_flat)


export GADM, NUTS, makeregions, makeregions_nuts, makeoffshoreregions, saveregions, loadregions,
        saveregions_global, subregions, all_gadm_subregions

abstract type RegionType end

struct GADM{T} <: RegionType
    parentregions::Vector{T}
    subregionnames::NTuple{N,T} where N
end
GADM(regionnames::T...) where T = GADM(T[], regionnames)
GADM(parentregions::Vector{T}, subregionnames::T...) where T = GADM(parentregions, subregionnames)

struct NUTS{T} <: RegionType
    subregionnames::NTuple{N,T} where N
end
NUTS(regionnames::T...) where T = NUTS(regionnames)

const NOREGION = typemax(Int16)

function saveregions(regionname, regiondefinitionarray; autocrop=true, bbox=[-90 -180; 90 180])
    land = JLD.load(in_datafolder("landcover.jld"), "landcover")
    if !all(bbox .== [-90 -180; 90 180])
        autocrop = false         # ignore supplied autocrop option if user changed bbox
    end
    saveregions(regionname, regiondefinitionarray, land, autocrop, bbox)
end

function saveregions(regionname, regiondefinitionarray, landcover, autocrop, bbox)
    regions, regiontype = makeregions(regiondefinitionarray; allowmixed=(regionname=="Europe_background"))
    if autocrop
        # get indexes of the bounding box containing onshore region data with 6% of padding
        lonrange, latrange = getbboxranges(regions)
        padding = round(Int, maximum(size(regions[lonrange,latrange])) * 0.06)
        lonrange, latrange = getbboxranges(regions, padding)
    else
        latrange, lonrange = bbox2ranges(roundbbox(bbox,100), 100)          # TO DO: remove hardcoded raster density
    end
    landcover = landcover[lonrange, latrange]
    regions = regions[lonrange, latrange]

    if regionname != "Global_GADM0" && regionname != "Europe_background"
        if regiontype == :NUTS
            println("\nNUTS region definitions detected (using Europe_background region file)...")
            europeregions = loadregions("Europe_background")[1][lonrange, latrange]
            regions[(regions.==0) .& (europeregions.>0)] .= NOREGION
        elseif regiontype == :GADM
            println("\nGADM region definitions detected (using Global_GADM0 region file)...")
            globalregions = loadregions("Global_GADM0")[1][lonrange, latrange]
            regions[(regions.==0) .& (globalregions.>0)] .= NOREGION
        end
    end

    # Find the closest region pixel for all non-region pixels (land and ocean)
    println("\nAllocate non-region pixels to the nearest region (for offshore wind)...")
    territory = regions[feature_transform(regions.>0)]

    # Allocate ocean and lake pixels to the region with the closest land region.
    # Even VERY far offshore pixels will be allocated to whatever region is nearest, but
    # those areas still won't be available for offshore wind power because of the
    # requirement to be close enough to the electricity grid (or rather the grid proxy).
    offshoreregions = territory .* (landcover .== 0)

    if regionname != "Global_GADM0" && regionname != "Europe_background"
        # Allocate land pixels with region==0 to the closest land region.
        # This ensures that the regions dataset is pixel-compatible with the landcover dataset.
        regions = territory .* (landcover .> 0)
    end

    println("\nSaving regions and offshoreregions...")
    regionlist = Symbol.(regiondefinitionarray[:,1])

    JLD.save(in_datafolder("regions_$regionname.jld"), "regions", regions, "offshoreregions", offshoreregions,
                "regionlist", regionlist, "lonrange", lonrange, "latrange", latrange, compress=true)
end

function saveregions_global(; args...)
    println("\nCreating a global GADM region file to identify countries and land areas later...")
    g = readdlm(in_datafolder("gadmfields.csv"), ',', skipstart=1)
    gadm0 = unique(string.(g[:,2]))
    regiondefinitionarray = [gadm0 GADM.(gadm0)]
    saveregions("Global_GADM0", regiondefinitionarray; args..., autocrop=false)
    println("Global GADM region file saved.")

    println("\nCreating a 'background' NUTS region file to identify non-European land areas later...")
    regiondefinitionarray = [NUTS_Europe; non_NUTS_Europe]
    saveregions("Europe_background", regiondefinitionarray; args..., autocrop=false)
    println("\nEurope_background region file saved.")

    println("\nCreating a region file for the 44 countries with synthetic demand training data...")
    regiondefinitionarray = [syntheticdemandregions GADM.(syntheticdemandregions)]
    saveregions("SyntheticDemandRegions", regiondefinitionarray; args...)
    println("\nSyntheticDemandRegions file saved.")
end

function loadregions(regionname)
    jldopen(in_datafolder("regions_$regionname.jld"), "r") do file
        return read(file, "regions"), read(file, "offshoreregions"), read(file, "regionlist"),
                    read(file, "lonrange"), read(file, "latrange")
    end
end

function makeregions(regiondefinitionarray; allowmixed=false)
    regionnames, nutsdef, gadmdef = splitregiondefinitions(regiondefinitionarray)
    use_nuts, use_gadm = !all(isempty.(nutsdef)), !all(isempty.(gadmdef))
    regiontype =  (use_gadm && !use_nuts) ? :GADM  :
                  (use_nuts && !use_gadm) ? :NUTS  :
                  (use_nuts && use_gadm) ? :MIXED  :  :WEIRD
    !allowmixed && regiontype==:MIXED && error("Sorry, mixed NUTS & GADM definitions are not supported yet.")
    region = zeros(Int16, (36000,18000))    # hard code size for now
    if use_nuts
        nuts, subregionnames = read_nuts()
        makeregions_nuts!(region, nuts, subregionnames, nutsdef)
    end
    if use_gadm
        gadm, subregionnames = read_gadm()
        makeregions_gadm!(region, gadm, subregionnames, gadmdef)
    end
    return region, regiontype
end

function regions2matlab(gisregion)
    regions, offshoreregions, regionlist, lonrange, latrange = loadregions(gisregion)
    matopen(in_datafolder("regions_$gisregion.mat"), "w", compress=true) do file
        write(file, "regions", regions)
        write(file, "offshoreregions", offshoreregions)
        write(file, "regionlist", string.(regionlist))
        write(file, "lonrange", collect(lonrange))
        write(file, "latrange", collect(latrange))
    end
end

function splitregiondefinitions(regiondefinitionarray)
    regionnames = regiondefinitionarray[:,1]
    regiondefinitions = [regdef isa Tuple ? regdef : (regdef,) for regdef in regiondefinitionarray[:,2]]
    nutsdef = [Tuple(rd for rd in regdef if rd isa NUTS) for regdef in regiondefinitions]
    gadmdef = [Tuple(rd for rd in regdef if rd isa GADM) for regdef in regiondefinitions]
    return regionnames, nutsdef, gadmdef
end

function makeregions_gadm!(region, gadm, subregionnames, regiondefinitions)
    println("Making region index matrix...")
    regionlookup = build_inverseregionlookup(regiondefinitions)
    rows, cols = size(region)
    updateprogress = Progress(cols, 1)
    for c in randperm(cols)
        for r = 1:rows
            gadm_uid = gadm[r,c]
            (gadm_uid == 0 || gadm_uid == 78413 || region[r,c] > 0) && continue    # ignore Caspian Sea (weirdly classified as a region in GADM)
            reg0, reg1, reg2 = subregionnames[gadm_uid,:]
            regid = lookup_regionnames(regionlookup, reg0, reg1, reg2)
            if regid > 0
                region[r,c] = regid
            end
        end
        next!(updateprogress)
    end
end

function makeregions_nuts!(region, nuts, subregionnames, regiondefinitions)
    println("Making region index matrix...")
    regionlookup = Dict(r => i for (i,tuptup) in enumerate(regiondefinitions)
                                    for ntup in tuptup for r in ntup.subregionnames)
    rows, cols = size(region)
    updateprogress = Progress(cols, 1)
    for c in randperm(cols)
        for r = 1:rows
            nuts_id = nuts[r,c]
            (nuts_id == 0 || region[r,c] > 0) && continue
            reg = subregionnames[nuts_id]
            while length(reg) >= 2
                regid = get(regionlookup, reg, 0)
                if regid > 0
                    region[r,c] = regid
                    break
                end
                reg = reg[1:end-1]
            end
        end
        next!(updateprogress)
    end
end

function lookup_regionnames(regionlookup, reg0, reg1, reg2)
    v = get(regionlookup, (reg0, "*", "*"), 0)
    v > 0 && return v
    v = get(regionlookup, (reg0, reg1, "*"), 0)
    v > 0 && return v
    return get(regionlookup, (reg0, reg1, reg2), 0)
end

function build_inverseregionlookup(regiondefinitions)
    d = Dict{Tuple{String,String,String}, Int}()
    for reg = 1:length(regiondefinitions)
        for regdef in regiondefinitions[reg]
            parentregions, subregionnames = regdef.parentregions, regdef.subregionnames
            regions = ["*", "*", "*"]
            regions[1:length(parentregions)] = parentregions
            for s in subregionnames
                regions[length(parentregions)+1] = s
                d[regions...] = reg
            end
        end
    end
    return d
end

function getsubregions(regtype::Type{GADM}, regionnames)
    gadm = readdlm(in_datafolder("gadmfields.csv"), ',', skipstart=1)[:, 2:4]
    for reg in regionnames
        gadm = gadm[gadm[:,1].==reg, 2:end]
    end
    return sort(unique(gadm[:,1]))
end

function getsubregions(regtype::Type{NUTS}, regionname)
    length(regionname) > 1 && error("Give only one string argument to match beginning of NUTS region names.")
    nuts = readdlm(in_datafolder("nutsfields.csv"), ',', skipstart=1)[:, 3]
    if isempty(regionname)
        return sort(unique(first.(nuts, 2)))
    else
        return sort(nuts[startswith.(nuts, regionname)])
    end
end

function subregions(regtype::Type{T} where T <: RegionType, regionnames::String...)
    reglist = join(regionnames, ", ")
    selected = string.(getsubregions(regtype, regionnames))
    selectedlist = join(selected, ", ")
    # isempty(regionnames) && println("Showing top level $regtype regions:")
    # println("$regtype($reglist): $selectedlist")
    return selected
end

function all_gadm_subregions(country::AbstractString, level::Int)
    level == 1 && return [country GADM(country)]
    return all_gadm_subregions(country, subregions(GADM, country), level)
end

function all_gadm_subregions(country::AbstractString, reg2::AbstractString, level::Int)
    reg2 == "" && return [country GADM(country)]
    level == 2 && return [reg2 GADM([country], reg2)]
    return vcat(all_gadm_subregions.(country, reg2, subregions(GADM, country, reg2))...)
end

function all_gadm_subregions(country::AbstractString, reg2::AbstractString, reg3::AbstractString)
    reg3 == "" && return [reg2 GADM([country], reg2)]
    return [reg3 GADM([country, reg2], reg3)]
end

all_gadm_subregions(countries::AbstractArray, level::Int) =
        vcat(all_gadm_subregions.(countries, level)...)

all_gadm_subregions(country::AbstractString, regions2::AbstractArray, level::Int) =
        vcat(all_gadm_subregions.(country, regions2, level)...)


# NOTE: This code is copied from: https://github.com/JuliaPOMDP/LocalApproximationValueIteration.jl
module LocalApproximationPolicyEvaluation
    using Random
    using Printf
    using POMDPs
    using POMDPModelTools
    using POMDPSimulators
    using GridInterpolations
    using LocalFunctionApproximation
    using Distributions
    import POMDPs: Solver, solve, Policy, action, value
    using POMDPLinter: @POMDP_require, @req, @subreq, @warn_requirements

    export LocalPolicyEvalSolver, LocalPolicyEvalPolicy, action_and_prob

    mutable struct LocalPolicyEvalSolver{I<:LocalFunctionApproximator, RNG<:AbstractRNG} <: Solver
        interp::I # Will be copied over by value to each policy
        action_probability::Function # a function with signature action_probability(mdp, s, a)
        max_iterations::Int64 # max number of iterations
        belres::Float64 # the Bellman Residual
        verbose::Bool # Whether to print while solving or not
        rng::RNG # Seed if req'd
        is_mdp_generative::Bool # Whether to treat underlying MDP model as generative
        n_generative_samples::Int64 # If underlying model generative, how many samples to use
    end

    # Default constructor
    function LocalPolicyEvalSolver(interp::I, action_probability::Function;
                                                    max_iterations::Int64=100, belres::Float64=1e-3,
                                                    verbose::Bool=false, rng::RNG=Random.GLOBAL_RNG,
                                                    is_mdp_generative::Bool=false, n_generative_samples::Int64=0) where {I<:LocalFunctionApproximator, RNG<:AbstractRNG}
        return LocalPolicyEvalSolver(interp, action_probability, max_iterations, belres, verbose, rng, is_mdp_generative, n_generative_samples)
    end

    # Unparameterized constructor just for getting requirements
    function LocalPolicyEvalSolver()
        throw(ArgumentError("LocalPolicyEvalSolver needs a LocalFunctionApproximator object for construction!"))
    end

    # NOTE : We work directly with the value function
    # And extract actions at the end by using the interpolation object
    mutable struct LocalPolicyEvalPolicy{I<:LocalFunctionApproximator, RNG<:AbstractRNG} <: Policy
        interp::I # General approximator to be used in VI
        action_probability::Function # a function with signature action_probability(mdp, s, a)
        action_map::Vector # Maps the action index to the concrete action type
        mdp::Union{MDP,POMDP} # Uses the model for indexing in the action function
        is_mdp_generative::Bool # (Copied from solver.is_mdp_generative)
        n_generative_samples::Int64 # (Copied from solver.n_generative_samples)
        rng::RNG # (Copied from solver.rng)
    end

    # The policy can be created using the MDP and solver information
    # The policy's function approximation object (interp) is obtained by deep-copying over the
    # solver's interp object. The other policy parameters are also obtained from the solver
    function LocalPolicyEvalPolicy(mdp::Union{MDP,POMDP},
                                                    solver::LocalPolicyEvalSolver)
        return LocalPolicyEvalPolicy(deepcopy(solver.interp),solver.action_probability, ordered_actions(mdp),mdp,
                                                      solver.is_mdp_generative,solver.n_generative_samples,solver.rng)
    end

    @POMDP_require solve(solver::LocalPolicyEvalSolver, mdp::Union{MDP,POMDP}) begin

        P = typeof(mdp)
        S = statetype(P)
        A = actiontype(P)
        @req discount(::P)
        @subreq ordered_actions(mdp)

        @req actionindex(::P, ::A)
        @req actions(::P, ::S)
        as = actions(mdp)
        a = first(as)

        @req convert_s(::Type{S}, ::AbstractVector{Float64}, ::P)
        @req convert_s(::Type{Vector{Float64}}, ::S, ::P)

        # Have different requirements depending on whether solver MDP is generative or explicit
        if solver.is_mdp_generative
            @req gen(::DDNOut{(:sp, :r)}, ::P, ::S, ::A, ::typeof(solver.rng))
        else
            @req transition(::P, ::S, ::A)
            pts = get_all_interpolating_points(solver.interp)
            pt = first(pts)
            ss = POMDPs.convert_s(S, pt, mdp)
            dist = transition(mdp, ss, a)
            D = typeof(dist)
            @req support(::D)
        end

    end

    function POMDPs.solve(solver::LocalPolicyEvalSolver, mdp::Union{MDP,POMDP})

        @warn_requirements solve(solver, mdp)

        # Ensure that generative model has a non-zero number of samples
        if solver.is_mdp_generative
            @assert solver.n_generative_samples > 0
        end

        # Solver parameters
        max_iterations = solver.max_iterations
        belres = solver.belres
        discount_factor = discount(mdp)

        # Initialize the policy
        policy = LocalPolicyEvalPolicy(mdp,solver)

        total_time = 0.0
        iter_time = 0.0

        # Get attributes of interpolator
        # Since the policy object is created by the solver, it directly
        # modifies the value of the interpolator of the created policy
        num_interps = n_interpolating_points(policy.interp)
        interp_points = get_all_interpolating_points(policy.interp)
        interp_values = get_all_interpolating_values(policy.interp)

        # Obtain the vector of states by converting the corresponding
        # vector of interpolation points/samples to the state type
        # using the user-provided convert_s function
        S = statetype(typeof(mdp))
        interp_states = Vector{S}(undef, num_interps)
        for (i, pt) in enumerate(interp_points)
            interp_states[i] = POMDPs.convert_s(S, pt, mdp)
        end

        # State transition dictionary
        next_state_dict = Dict()

        # Outer loop for Value Iteration
        for i = 1 : max_iterations
            residual::Float64 = 0.0
            iter_time = @elapsed begin

            for (istate,s) in enumerate(interp_states)
                sub_aspace = actions(mdp, s)

                if isterminal(mdp, s)
                    interp_values[istate] = 0.0
                else
                    old_util = interp_values[istate]
                    total_util = 0

                    for a in sub_aspace
                        u::Float64 = 0.0

                        # Do bellman backup based on generative / explicit model
                        if solver.is_mdp_generative
                            # Generative Model
                            for j in 1:solver.n_generative_samples
                                sp_point, r, isTerm = 0, 0, false
                                if solver.n_generative_samples ==1 && haskey(next_state_dict, (s,a))
                                    sp_point, r, isTerm = next_state_dict[(s,a)]
                                else
                                    sp, r = gen(mdp, s, a, solver.rng)
                                    isTerm = isterminal(mdp, sp)
                                    sp_point = POMDPs.convert_s(Vector{Float64}, sp, mdp)
                                    solver.n_generative_samples == 1  && (next_state_dict[(s,a)] = (sp_point, r, isTerm))
                                end

                                u += r

                                # Only interpolate sp if it is non-terminal
                                if !isTerm
                                    u += discount_factor*compute_value(policy.interp, sp_point)
                                end
                            end
                            u = u / solver.n_generative_samples
                        else
                            # Explicit Model
                            dist = transition(mdp, s, a)
                            for (sp, p) in weighted_iterator(dist)
                                p == 0.0 ? continue : nothing
                                r = reward(mdp, s, a, sp)
                                u += p*r

                            # Only interpolate sp if it is non-terminal
                                if !isterminal(mdp, sp)
                                    sp_point = POMDPs.convert_s(Vector{Float64}, sp, mdp)
                                    u += p * (discount_factor*LocalFunctionApproximation.compute_value(policy.interp, sp_point))
                                end
                            end # next-states
                        end
                        total_util += solver.action_probability(mdp, s, a)*u
                    end #action

                    # Update this interpolant value
                    interp_values[istate] = total_util
                    util_diff = abs(total_util - old_util)
                    util_diff > residual ? (residual = util_diff) : nothing
                end
            end #state

            end #time
            total_time += iter_time
            solver.verbose ? @printf("[Iteration %-4d] residual: %10.3G | iteration runtime: %10.3f ms, (%10.3G s total)\n", i, residual, iter_time*1000.0, total_time) : nothing
            residual < belres ? break : nothing

        end #main
        return policy
    end

    function POMDPs.value(policy::LocalPolicyEvalPolicy, s::S) where S
        # Call the conversion function on the state to get the corresponding vector
        # That represents the point at which to interpolate the function
        s_point = POMDPs.convert_s(Vector{Float64}, s, policy.mdp)
        val = LocalFunctionApproximation.compute_value(policy.interp, s_point)
        return val
    end

    # Not explicitly stored in policy - extract from value function interpolation
    POMDPs.action(policy::LocalPolicyEvalPolicy, s, rng= Random.GLOBAL_RNG) = action_and_prob(policy, s, rng)[1]

    function Distributions.logpdf(policy::LocalPolicyEvalPolicy, s, a)
        mdp = policy.mdp
        us = [value(policy, s, a)*policy.action_probability(mdp, s, a) for a in actions(policy.mdp, s)]
        if sum(us) == 0
            us = [policy.action_probability(mdp, s, a) for a in actions(policy.mdp, s)]
        end
        p = value(policy, s, a)*policy.action_probability(mdp, s, a) / sum(us)
        return (p == 0) ? log(policy.action_probability(mdp, s, a)) : log(p)
    end

    Distributions.logpdf(policy::LocalPolicyEvalPolicy, h::SimHistory) = sum([logpdf(policy, s, a) for (s,a) in eachstep(h, (:s, :a))])

    function action_and_prob(policy::LocalPolicyEvalPolicy, s, rng= Random.GLOBAL_RNG)
        mdp = policy.mdp
        us = [value(policy, s, a)*policy.action_probability(mdp, s, a) for a in actions(policy.mdp, s)]
        if sum(us) == 0
            us = [policy.action_probability(mdp, s, a) for a in actions(policy.mdp, s)]
        end
        us /= sum(us)
        ai = rand(rng, Categorical(us))
        policy.action_map[ai], us[ai]
    end

    # GlobalApproximationFailureProbe the action-value for some state-action pair
    # This is also used in the above function
    function POMDPs.value(policy::LocalPolicyEvalPolicy, s::S, a::A) where {S,A}

        mdp = policy.mdp
        discount_factor = discount(mdp)

        u::Float64 = 0.0

        # As in solve(), do different things based on whether
        # mdp is generative or explicit
        if policy.is_mdp_generative
            for j in 1:policy.n_generative_samples
                sp, r = gen(mdp, s, a, policy.rng)
                sp_point = POMDPs.convert_s(Vector{Float64}, sp, mdp)
                u += r
                if !isterminal(mdp, sp)
                    u += discount_factor*LocalFunctionApproximation.compute_value(policy.interp, sp_point)
                end
            end
            u = u / policy.n_generative_samples
        else
            dist = transition(mdp, s, a)
            for (sp, p) in weighted_iterator(dist)
                p == 0.0 ? continue : nothing
                r = reward(mdp, s, a, sp)
                u += p*r

                # Only interpolate sp if it is non-terminal
                if !isterminal(mdp, sp)
                    sp_point = POMDPs.convert_s(Vector{Float64}, sp, mdp)
                    u += p*(discount_factor*LocalFunctionApproximation.compute_value(policy.interp, sp_point))
                end
            end
        end

        return u
    end

end


using Pkg
Pkg.add("JSON3")

import JSON3
using Sockets

struct Coordinate
    x::Float64
    y::Float64
    z::Float64
end

function calc(text)
    jobj = JSON3.read(text)
    coordinates = jobj["coordinates"]
    len = length(coordinates)
    x = y = z = 0

    for coord in coordinates
        x += coord["x"]
        y += coord["y"]
        z += coord["z"]
    end

    Coordinate(x / len, y / len, z / len)
end

function notify(msg)
    try
        socket = connect("localhost", 9001)
        write(socket, msg)
        close(socket)
    catch
        # standalone usage
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    right = Coordinate(2.0, 0.5, 0.25)
    for v in [
        """{"coordinates":[{"x":2.0,"y":0.5,"z":0.25}]}""",
        """{"coordinates":[{"y":0.5,"x":2.0,"z":0.25}]}""",
    ]
        left = calc(v)
        if left != right
            println(stderr, "$(left) != $(right)")
            exit(1)
        end
    end

    text = open("/tmp/1.json") do file
        read(file, String)
    end

    notify("Julia (JSON3)\t$(getpid())")
    results = calc(text)
    notify("stop")

    println(results)
end