using Test, ASE, NeighbourLists, JuLIP
using LinearAlgebra: norm

h1("Misc ASE Tests")

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

h3("Check that the cubic unit cell for Al is reproduced correctly")
a0 = 2.025
p = [ [0.0;0.0;0.0] [0.0;a0;a0] [a0;0.0;a0] [a0;a0;0.0] ]
at = bulk("Al", cubic=true)
println(@test (positions(at) |> mat) == p)

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

h2("Check neighbourlist without periodicity")
# TODO: implement a test with periodicity as well???
at = bulk("Al", cubic=true) * 3
set_pbc!(at, (false,false,false))
rcut = 1.7 * a0
nlist = neighbourlist(at, rcut)

@info "   ... assemble neighbour list ..."
# create a neighbourlist via a naive double-loop
simple = zeros(length(at), length(at))
X = positions(at)
for n = 2:length(at), m = 1:n-1
   if norm(X[m]-X[n]) <= rcut
      simple[n,m] = simple[m,n] = 1
   end
end

@info "   ... check the bond-iterator ... "
pass_bonds_test = true
for (i,j,r,R) in pairs(nlist)
   if !( (simple[i,j] == 1) && (abs(norm(X[i]-X[j]) - r) < 1e-12) &&
         (norm(X[j]-X[i] - R) < 1e-12)  )
      global pass_bonds_test = false
      break
   end
   # switch the flag
   simple[i,j] = -1
end

h3("check that no false bonds have been found")
println(@test pass_bonds_test)
h3("check that all bonds have been found")
println(@test maximum(simple) == 0)

# revert to original
simple *= -1
h3("   ... check the site iterator ... ")
pass_site_test = true
for (i,j,r,R) in sites(nlist)
   for n = 1:length(j)
      if simple[i,j[n]] != 1
         global pass_site_test = false
         break
      end
      simple[i,j[n]] = -1
   end
   if maximum(simple[i,:]) != 0
      global pass_site_test = false
      break
   end
end
println(@test pass_site_test)
println(@test maximum(simple) == 0)

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

# h3("Checking `***_data`, `***_array`, `***_info`, `***_transient`")
#
# at = bulk("Cu") * 3
# N = length(at)
# # set an array and test that it is read back correctly
# z = rand(N)
# set_array!(at, "z", z)
# @test get_array(at, "z") == z
# # set some info and test reading
# i = "some info"
# set_info!(at, "i", i)
# @test get_info(at, "i") == i
# # ***_transient should be tested automatically via the calculators.
# # test the has_***
# @test has_array(at, "z")
# @test has_info(at, "i")
# @test !has_array(at, "i")
# @test !has_info(at, "z")

# println("Checking momenta and velocites")
# p = rand(3,N) |> vecs
# set_momenta!(at, p)
# @test isapprox(momenta(at) |> mat, p |> mat)
# v = p ./ masses(at)
# @test isapprox(velocities(at) |> mat, v |> mat)
# set_velocities!(at, v/2.0)
# @test isapprox(velocities(at) |> mat, v/2.0 |> mat)
# @test isapprox(momenta(at) |> mat, p/2.0 |> mat)
#
# println("Checking chemical_symbols")
# @test chemical_symbols(at) == ["Cu" for i in 1:length(at)]

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

# h3("testing static_neighbourlist")
# nlist = static_neighbourlist(at, rnn("Cu") * 1.1)
# rattle!(at, 0.05)
# nlist1 = static_neighbourlist(at, rnn("Cu") * 1.1)
# @test nlist === nlist1

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

h3("testing read and write")
fname = "test.xyz"
at = bulk("Cu") * 2
write_xyz(fname, at)
at2 = read_xyz(fname)
println(@test positions(at) == positions(at2))
rm(fname)

# ---------
h3("Conversion between JuLIP.Atoms and ASEAtoms")
at1 = bulk(:Si)          # JuLIP.Atoms
at2 = ASEAtoms(at1)      # ASE.ASEAtoms
@show cell(at2)
at3 = Atoms(at2)         # JuLIP.Atoms
println(@test at1 == at3)

# ---------
h3("Conversion between JuLIP.Atoms and ASEAtoms, Multi-species")
at1 = bulk(:Al) * 3         # JuLIP.Atoms
at1.Z[1:(length(at1)÷2)] .= atomic_number(:Ti)
at1.M[:] .= 1.0 .+ rand(length(at1))
at2 = ASEAtoms(at1)      # ASE.ASEAtoms
@show cell(at2)
at3 = Atoms(at2)         # JuLIP.Atoms
println(@test at1 == at3)


import Base: IteratorSize, HasLength, IsInfinite, length

struct FiniteTimeTrigger <: AbstractFiniteTrigger
    time_fire_at::Dates.Time
    n::Int
end

struct InfiniteTimeTrigger <: AbstractInfiniteTrigger
    time_fire_at::Dates.Time
end

"""
    TimeTrigger(t::Dates.Time[, n=number_of_times])

A trigger which should trigger a job daily at a given time.

# Optional parameter
- `n=1`: trigger once
- `n=-1` (default): trigger every day indefinitely
- `n=value`: trigger just a number of times
"""
function TimeTrigger(time_fire_at; n=-1)
    if n < 0
        InfiniteTimeTrigger(time_fire_at)
    else
        FiniteTimeTrigger(time_fire_at, n)
    end
end

"""
    Trigger(t::Dates.Time[, n=number_of_times])

Return an `TimeTrigger` which should trigger a job daily at a given time (once, a finite number of times or indefinitely).
"""
Trigger(t::Dates.Time; kwargs...) = TimeTrigger(t; kwargs...)

function get_next_dt_fire(trigger::Union{FiniteTimeTrigger,InfiniteTimeTrigger}, dt_previous_fire, dt_now)
    d = Date(dt_now)
    t = Dates.Time(dt_now)
    if trigger.time_fire_at > t
        d + trigger.time_fire_at
    else
        d + trigger.time_fire_at + Dates.Day(1)
    end
end


data = [10, 50, 30]
bar(data; width=1, fill="hotpink")


#=
Functions for generating a single replication of a particular experiment
Dependencies: util.jl, sgd.jl, reformulation.jl, random_forests_po.jl, validation_set.jl
=#

using Parameters, DataFrames, Statistics, Random, Gurobi

## Type for strong the results of an experiment replication
@with_kw mutable struct replication_results
    SPOplus_spoloss_test::Union{Float64, Missing} = missing # SPO loss of SPO+ SGD on test set
    LS_spoloss_test::Union{Float64, Missing} = missing # SPO loss of LS SGD on test set
    SSVM_spoloss_test::Union{Float64, Missing} = missing # SPO loss of SSVM-Hamming SGD on test set
    RF_spoloss_test::Union{Float64, Missing} = missing # SPO loss of random forests on test set
    Absolute_spoloss_test::Union{Float64, Missing} = missing # SPO loss of ell_1 on test set
    Huber_spoloss_test::Union{Float64, Missing} = missing  # SPO loss of Huber on test set
    Baseline_spoloss_test::Union{Float64, Missing} = missing # SPO loss of baseline that returns mean of training data, on test set
    SPOplus_hammingloss_test::Union{Float64, Missing} = missing # Hamming loss of SPO+ SGD on test set
    LS_hammingloss_test::Union{Float64, Missing} = missing # Hamming loss of LS SGD on test set
    SSVM_hammingloss_test::Union{Float64, Missing} = missing # Hamming loss of SSVM SGD on test set
    RF_hammingloss_test::Union{Float64, Missing} = missing # Hamming loss of RF on test set
    Absolute_hammingloss_test::Union{Float64, Missing} = missing # Hamming loss of ell_1 on test set
    Huber_hammingloss_test::Union{Float64, Missing} = missing  # Hamming loss of Huber on test set
    Baseline_hammingloss_test::Union{Float64, Missing} = missing # Hamming loss of baseline on test set
    zstar_avg_test::Union{Float64, Missing} = missing # Average of z^*(c_i) on testing set
end

"""
Make a blank data frame for storing the results of the experiment.
"""
function make_blank_complete_df()
    results_df = DataFrame(
        grid_dim = Int64[],
        n_train = Int64[],
        n_holdout = Int64[],
        n_test = Int64[],
        p_features = Int64[],
        polykernel_degree = Int64[],
        polykernel_noise_half_width = Float64[],
        SPOplus_spoloss_test = Union{Float64, Missing}[], # SPO loss of SPO+ SGD on test set
        LS_spoloss_test = Union{Float64, Missing}[],  # SPO loss of LS SGD on test set
        SSVM_spoloss_test = Union{Float64, Missing}[], # SPO loss of SSVM-Hamming SGD on test set
        RF_spoloss_test = Union{Float64, Missing}[], # SPO loss of random forests on test set
        Absolute_spoloss_test = Union{Float64, Missing}[], # SPO loss of ell_1 on test set
        Huber_spoloss_test = Union{Float64, Missing}[],  # SPO loss of Huber on test set
        Baseline_spoloss_test = Union{Float64, Missing}[], # SPO loss of baseline that returns mean of training data, on test set
        SPOplus_hammingloss_test = Union{Float64, Missing}[], # Hamming loss of SPO+ SGD on test set
        LS_hammingloss_test = Union{Float64, Missing}[], # Hamming loss of LS SGD on test set
        SSVM_hammingloss_test = Union{Float64, Missing}[], # Hamming loss of SSVM SGD on test set
        RF_hammingloss_test = Union{Float64, Missing}[], # Hamming loss of RF on test set
        Absolute_hammingloss_test = Union{Float64, Missing}[], # Hamming loss of ell_1 on test set
        Huber_hammingloss_test = Union{Float64, Missing}[],  # Hamming loss of Huber on test set
        Baseline_hammingloss_test = Union{Float64, Missing}[], # Hamming loss of baseline on test set
        zstar_avg_test = Union{Float64, Missing}[]) # Average of z^*(c_i) on testing set

    return results_df
end

"""
Make a new data frame that contains one row of results for an experiment instance that was just run.
"""
function build_complete_row(grid_dim,
    n_train, n_holdout, n_test,
    p_features, polykernel_degree, polykernel_noise_half_width, results_struct)

    df = make_blank_complete_df()

    push!(df, [
        grid_dim,
        n_train,
        n_holdout,
        n_test,
        p_features,
        polykernel_degree,
        polykernel_noise_half_width,
        results_struct.SPOplus_spoloss_test,
        results_struct.LS_spoloss_test,
        results_struct.SSVM_spoloss_test,
        results_struct.RF_spoloss_test,
        results_struct.Absolute_spoloss_test,
        results_struct.Huber_spoloss_test,
        results_struct.Baseline_spoloss_test,
        results_struct.SPOplus_hammingloss_test,
        results_struct.LS_hammingloss_test,
        results_struct.SSVM_hammingloss_test,
        results_struct.RF_hammingloss_test,
        results_struct.Absolute_hammingloss_test,
        results_struct.Huber_hammingloss_test,
        results_struct.Baseline_hammingloss_test,
        results_struct.zstar_avg_test])

    return df
end

"""
Run one instance of the shortest path experiment. Uses the reformulation approach.
`num_lambda` is the number of lambdas used on the grid for each method with regularization.
Returns a replication_results struct.
"""
function shortest_path_replication(grid_dim,
    n_train, n_holdout, n_test,
    p_features, polykernel_degree, polykernel_noise_half_width;
    num_lambda = 10, lambda_max = missing, lambda_min_ratio = 0.0001, regularization = :ridge,
    gurobiEnvOracle = Gurobi.Env(), gurobiEnvReform = Gurobi.Env(), different_validation_losses = false,
    include_rf = true)

    # Get oracle
    sources, destinations = convert_grid_to_list(grid_dim, grid_dim)
    d_feasibleregion = length(sources)
    sp_oracle = sp_flow_jump_setup(sources, destinations, 1, grid_dim^2; gurobiEnv = gurobiEnvOracle)
    sp_graph = shortest_path_graph(sources = sources, destinations = destinations,
        start_node = 1, end_node = grid_dim^2, acyclic = true)

    # Generate Data
    B_true = rand(Bernoulli(0.5), d_feasibleregion, p_features)

    (X_train, c_train) = generate_poly_kernel_data_simple(B_true, n_train, polykernel_degree, polykernel_noise_half_width)
    (X_validation, c_validation) = generate_poly_kernel_data_simple(B_true, n_holdout, polykernel_degree, polykernel_noise_half_width)
    (X_test, c_test) = generate_poly_kernel_data_simple(B_true, n_test, polykernel_degree, polykernel_noise_half_width)

    # Add intercept in the first row of X
    X_train = vcat(ones(1,n_train), X_train)
    X_validation = vcat(ones(1,n_holdout), X_validation)
    X_test = vcat(ones(1,n_test), X_test)

    # Get Hamming labels
    (z_train, w_train) = oracle_dataset(c_train, sp_oracle)
    (z_validation, w_validation) = oracle_dataset(c_validation, sp_oracle)
    (z_test, w_test) = oracle_dataset(c_test, sp_oracle)

    c_ham_train = ones(d_feasibleregion, n_train) - w_train
    c_ham_validation = ones(d_feasibleregion, n_holdout) - w_validation
    c_ham_test = ones(d_feasibleregion, n_test) - w_test

    # Put train + validation together
    X_both = hcat(X_train, X_validation)
    c_both = hcat(c_train, c_validation)
    c_ham_both = hcat(c_ham_train, c_ham_validation)
    train_ind = collect(1:n_train)
    validation_ind = collect((n_train + 1):(n_train + n_holdout))

    # Set validation losses
    if different_validation_losses
        spo_plus_val_loss = :spo_loss
        ls_val_loss = :least_squares_loss
        ssvm_val_loss = :hamming_loss
        absolute_val_loss = :absolute_loss
        huber_val_loss = :huber_loss
    else
        spo_plus_val_loss = :spo_loss
        ls_val_loss = :spo_loss
        ssvm_val_loss = :spo_loss
        absolute_val_loss = :spo_loss
        huber_val_loss = :spo_loss
    end

    ### Algorithms ###

    # SPO+
    best_B_SPOplus, best_lambda_SPOplus = validation_set_alg(X_both, c_both, sp_oracle; sp_graph = sp_graph,
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :sp_spo_plus_reform, validation_loss = spo_plus_val_loss),
        path_alg_parms = reformulation_path_parms(num_lambda = num_lambda, lambda_max = lambda_max, regularization = regularization,
            gurobiEnv = gurobiEnvReform, lambda_min_ratio = lambda_min_ratio, algorithm_type = :SPO_plus))

    # Least squares
    best_B_leastSquares, best_lambda_leastSquares = validation_set_alg(X_both, c_both, sp_oracle; sp_graph = sp_graph,
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :ls_jump, validation_loss = ls_val_loss),
        path_alg_parms = reformulation_path_parms(num_lambda = num_lambda, lambda_max = lambda_max, regularization = regularization,
            gurobiEnv = gurobiEnvReform, lambda_min_ratio = lambda_min_ratio, algorithm_type = :leastSquares))

    # SSVM Hamming
    best_B_SSVM, best_lambda_SSVM = validation_set_alg(X_both, c_ham_both, sp_oracle; sp_graph = sp_graph,
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :sp_spo_plus_reform, validation_loss = ssvm_val_loss),
        path_alg_parms = reformulation_path_parms(num_lambda = num_lambda, lambda_max = lambda_max, regularization = regularization,
            gurobiEnv = gurobiEnvReform, lambda_min_ratio = lambda_min_ratio, algorithm_type = :SSVM_hamming))

    # RF
    if include_rf
        rf_mods = train_random_forests_po(X_train, c_train;
            rf_alg_parms = rf_parms())
    end

    # Absolute
    best_B_Absolute, best_lambda_Absolute = validation_set_alg(X_both, c_both, sp_oracle; sp_graph = sp_graph,
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :ls_jump, validation_loss = absolute_val_loss),
        path_alg_parms = reformulation_path_parms(num_lambda = num_lambda, lambda_max = lambda_max, regularization = regularization,
            po_loss_function = :absolute, gurobiEnv = gurobiEnvReform, lambda_min_ratio = lambda_min_ratio, algorithm_type = :Absolute))

    # Huber
    best_B_fake, best_lambda_fake = validation_set_alg(X_both, c_both, sp_oracle; sp_graph = sp_graph,
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :ls_jump, validation_loss = ls_val_loss),
        path_alg_parms = reformulation_path_parms(lambda_max = 0.0001, regularization = regularization,
            num_lambda = 2, gurobiEnv = gurobiEnvReform, algorithm_type = :Huber_LS_fake))

    fake_ls_list = abs.(vec(c_train - best_B_fake*X_train))
    delta_from_fake_ls = median(fake_ls_list)

    best_B_Huber, best_lambda_Huber = validation_set_alg(X_both, c_both, sp_oracle; sp_graph = sp_graph,
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :ls_jump, validation_loss = huber_val_loss),
        path_alg_parms = reformulation_path_parms(num_lambda = num_lambda, lambda_max = lambda_max, regularization = regularization, po_loss_function = :huber,
            huber_delta = delta_from_fake_ls, lambda_min_ratio = lambda_min_ratio, gurobiEnv = gurobiEnvReform, algorithm_type = :Huber))

    # Baseline
    c_bar_train = mean(c_train, dims=2)

    ### Populate final results ###
    final_results = replication_results()

    final_results.SPOplus_spoloss_test = spo_loss(best_B_SPOplus, X_test, c_test, sp_oracle)
    final_results.LS_spoloss_test = spo_loss(best_B_leastSquares, X_test, c_test, sp_oracle)
    final_results.SSVM_spoloss_test = spo_loss(best_B_SSVM, X_test, c_test, sp_oracle)

    if include_rf
        rf_preds_test = predict_random_forests_po(rf_mods, X_test)
        final_results.RF_spoloss_test = spo_loss(Matrix(1.0I, d_feasibleregion, d_feasibleregion), rf_preds_test, c_test, sp_oracle)
    else
        final_results.RF_spoloss_test = missing
    end

    final_results.Absolute_spoloss_test = spo_loss(best_B_Absolute, X_test, c_test, sp_oracle)
    final_results.Huber_spoloss_test = spo_loss(best_B_Huber, X_test, c_test, sp_oracle)

    c_bar_test_preds = zeros(d_feasibleregion, n_test)
    for i = 1:n_test
        c_bar_test_preds[:, i] = c_bar_train
    end
    final_results.Baseline_spoloss_test = spo_loss(Matrix(1.0I, d_feasibleregion, d_feasibleregion), c_bar_test_preds, c_test, sp_oracle)

    # Now Hamming
    final_results.SPOplus_hammingloss_test = spo_loss(best_B_SPOplus, X_test, c_ham_test, sp_oracle)
    final_results.LS_hammingloss_test = spo_loss(best_B_leastSquares, X_test, c_ham_test, sp_oracle)
    final_results.SSVM_hammingloss_test = spo_loss(best_B_SSVM, X_test, c_ham_test, sp_oracle)

    if include_rf
        final_results.RF_hammingloss_test = spo_loss(Matrix(1.0I, d_feasibleregion, d_feasibleregion), rf_preds_test, c_ham_test, sp_oracle)
    else
        final_results.RF_hammingloss_test = missing
    end

    final_results.Absolute_hammingloss_test = spo_loss(best_B_Absolute, X_test, c_ham_test, sp_oracle)
    final_results.Huber_hammingloss_test = spo_loss(best_B_Huber, X_test, c_ham_test, sp_oracle)
    final_results.Baseline_hammingloss_test = spo_loss(Matrix(1.0I, d_feasibleregion, d_feasibleregion), c_bar_test_preds, c_ham_test, sp_oracle)

    final_results.zstar_avg_test = mean(z_test)

    return final_results
end

"""
Run multiple replications of the shortest path experiment and return a data frame.
`rng_seed` and `num_trials` are integers
`n_train_vec`, `polykernel_degree_vec`, and `polykernel_noise_half_width_vec` are
vectors so that we try all combinations of these values.
"""
function shortest_path_multiple_replications(rng_seed, num_trials, grid_dim,
    n_train_vec, n_test,
    p_features, polykernel_degree_vec, polykernel_noise_half_width_vec;
    num_lambda = 10, lambda_max = missing, lambda_min_ratio = 0.0001, holdout_percent = 0.25, regularization = :ridge,
    different_validation_losses = false,
    include_rf = true)

    Random.seed!(rng_seed)

    big_results_df = make_blank_complete_df()

    envOracle = setup_gurobi_env(method_type = :default, use_time_limit = false)
    envReform = setup_gurobi_env(method_type = :method3, use_time_limit = false)

    for n_train in n_train_vec
        for polykernel_degree in polykernel_degree_vec
            for polykernel_noise_half_width in polykernel_noise_half_width_vec
                println("Moving on to n_train = $n_train, polykernel_degree = $polykernel_degree, polykernel_noise_half_width = $polykernel_noise_half_width")
                for trial = 1:num_trials
                    println("Current trial is $trial")
                    n_holdout = round(Int, holdout_percent*n_train)

                    current_results = shortest_path_replication(grid_dim,
                        n_train, n_holdout, n_test,
                        p_features, polykernel_degree, polykernel_noise_half_width;
                        num_lambda = num_lambda, lambda_max = lambda_max, lambda_min_ratio = lambda_min_ratio, regularization = regularization,
                        gurobiEnvOracle = envOracle, gurobiEnvReform = envReform,
                        different_validation_losses = different_validation_losses,
                        include_rf = include_rf)

                    current_df_row = build_complete_row(grid_dim,
                        n_train, n_holdout, n_test,
                        p_features, polykernel_degree, polykernel_noise_half_width, current_results)

                    big_results_df = vcat(big_results_df, current_df_row)
                end
            end
        end
    end

    return big_results_df
end

### Portfolio stuff ###

"""
Run one instance of the portfolio experiment. Uses the SGD approach.
`num_lambda` is the number of lambdas used on the grid for each method with regularization.
Returns a replication_results struct.
"""
function portfolio_replication(num_assets, num_factors, n_train, n_holdout, n_test, p_features,
    n_sigmoid_polydegree, noise_multiplier_tau;
    num_lambda = 10, lambda_max = :practical, lambda_min_ratio = 0.0001,
    gurobiEnvOracle = Gurobi.Env(), gurobiEnvReform = Gurobi.Env(), different_validation_losses = false,
    data_type = :sigmoid)

    # Get oracle
    factor_size = 0.0025*noise_multiplier_tau
    sigma_noise = 0.01*noise_multiplier_tau

    L_mat = 2*factor_size*rand(num_assets, num_factors) - factor_size*ones(num_assets, num_factors)
    Sigma = L_mat*L_mat' + (sigma_noise^2)*Matrix(1.0I, num_assets, num_assets)

    avg_vec = ones(num_assets)/num_assets
    gamma = avg_vec'*Sigma*avg_vec
    gamma = 2.25*gamma

    port_oracle = portfolio_simplex_jump_setup(Sigma, gamma; gurobiEnv = gurobiEnvOracle)

    # Generate Data
    return_intercept = 0
    if data_type == :rbf_kernel
        sigma_kernel = sqrt(p_features)
        # Rademacher instead of Bernoulli
        B_true = 2*rand(Bernoulli(0.5), num_assets, p_features) - ones(num_assets, p_features)

        (X_train, c_train) = generate_rbf_kernel_returns_data(B_true, n_train, sigma_kernel, L_mat, sigma_noise, return_intercept)
        (X_validation, c_validation) = generate_rbf_kernel_returns_data(B_true, n_holdout, sigma_kernel, L_mat, sigma_noise, return_intercept)
        (X_test, c_test) = generate_rbf_kernel_returns_data(B_true, n_test, sigma_kernel, L_mat, sigma_noise, return_intercept)
    elseif data_type == :sigmoid
        bernoulli_prob = 0.5

        (X_all, c_all) = generate_sigmoid_returns_data(p_features, num_assets, n_sigmoid_polydegree, n_train + n_holdout + n_test, bernoulli_prob,
        	L_mat, sigma_noise, return_intercept)
        c_all = c_all ./ n_sigmoid_polydegree

        X_train = X_all[:, 1:n_train]
        c_train = c_all[:, 1:n_train]
        X_validation = X_all[:, (n_train + 1):(n_train + n_holdout)]
        c_validation = c_all[:, (n_train + 1):(n_train + n_holdout)]
        X_test = X_all[:, (n_train + n_holdout + 1):(n_train + n_holdout + n_test)]
        c_test = c_all[:, (n_train + n_holdout + 1):(n_train + n_holdout + n_test)]
    elseif data_type == :poly_kernel
        #### zzz need to do properly tomorrow
        B_true = rand(Bernoulli(0.5), num_assets, p_features)
        #B_true = rand(num_assets, p_features)
        polykernel_degree = n_sigmoid_polydegree

        (X_train, c_train) = generate_poly_kernel_returns_data(B_true, n_train, polykernel_degree,
        	L_mat, sigma_noise, return_intercept)
        (X_validation, c_validation) = generate_poly_kernel_returns_data(B_true, n_holdout, polykernel_degree,
        	L_mat, sigma_noise, return_intercept)
        (X_test, c_test) = generate_poly_kernel_returns_data(B_true, n_test, polykernel_degree,
        	L_mat, sigma_noise, return_intercept)
    else
        error("Enter a valid data generation type")
    end

    #println("Train returns are:")
    #println(-c_train[:,1])

    # Add intercept in the first row of X
    X_train = vcat(ones(1,n_train), X_train)
    X_validation = vcat(ones(1,n_holdout), X_validation)
    X_test = vcat(ones(1,n_test), X_test)

    # Put train + validation together
    X_both = hcat(X_train, X_validation)
    c_both = hcat(c_train, c_validation)
    train_ind = collect(1:n_train)
    validation_ind = collect((n_train + 1):(n_train + n_holdout))

    # Set validation losses
    if different_validation_losses
        spo_plus_val_loss = :spo_loss
        ls_val_loss = :least_squares_loss
        ssvm_val_loss = :hamming_loss
        absolute_val_loss = :absolute_loss
        huber_val_loss = :huber_loss
    else
        spo_plus_val_loss = :spo_loss
        ls_val_loss = :spo_loss
        ssvm_val_loss = :spo_loss
        absolute_val_loss = :spo_loss
        huber_val_loss = :spo_loss
    end

    ### Algorithms ###

    # SPO+
    best_B_SPOplus, best_lambda_SPOplus = validation_set_alg(X_both, c_both, port_oracle;
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :spo_plus_sgd, resolve_sgd = false, validation_loss = spo_plus_val_loss),
        path_alg_parms = sgd_path_parms(num_lambda = num_lambda, lambda_max = lambda_max, lambda_min_ratio = lambda_min_ratio,
            obj_accuracy = 0.001, step_type = :long_dynamic, verbose = true))

    # Least squares
    best_B_leastSquares, best_lambda_leastSquares = validation_set_alg(X_both, c_both, port_oracle;
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :ls_jump, validation_loss = ls_val_loss),
        path_alg_parms = reformulation_path_parms(num_lambda = num_lambda, lambda_max = lambda_max, regularization = :ridge,
            gurobiEnv = gurobiEnvReform, lambda_min_ratio = lambda_min_ratio, algorithm_type = :leastSquares))

    # RF
    rf_mods = train_random_forests_po(X_train, c_train;
        rf_alg_parms = rf_parms())

    # Absolute
    best_B_Absolute, best_lambda_Absolute = validation_set_alg(X_both, c_both, port_oracle;
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :ls_jump, validation_loss = absolute_val_loss),
        path_alg_parms = reformulation_path_parms(num_lambda = num_lambda, lambda_max = lambda_max, regularization = :ridge,
            po_loss_function = :absolute, gurobiEnv = gurobiEnvReform, lambda_min_ratio = lambda_min_ratio, algorithm_type = :Absolute))

    # Huber
    best_B_fake, best_lambda_fake = validation_set_alg(X_both, c_both, port_oracle;
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :ls_jump, validation_loss = ls_val_loss),
        path_alg_parms = reformulation_path_parms(lambda_max = 0.0001, regularization = :ridge,
            num_lambda = 2, gurobiEnv = gurobiEnvReform, algorithm_type = :Huber_LS_fake))

    fake_ls_list = abs.(vec(c_train - best_B_fake*X_train))
    delta_from_fake_ls = median(fake_ls_list)

    best_B_Huber, best_lambda_Huber = validation_set_alg(X_both, c_both, port_oracle;
        train_ind = train_ind, validation_ind = validation_ind,
        val_alg_parms = val_parms(algorithm_type = :ls_jump, validation_loss = huber_val_loss),
        path_alg_parms = reformulation_path_parms(num_lambda = num_lambda, lambda_max = lambda_max, regularization = :ridge, po_loss_function = :huber,
            huber_delta = delta_from_fake_ls, lambda_min_ratio = lambda_min_ratio, gurobiEnv = gurobiEnvReform, algorithm_type = :Huber))

    # Baseline
    c_bar_train = mean(c_train, dims=2)

    ### Populate final results ###
    final_results = replication_results()

    final_results.SPOplus_spoloss_test = spo_loss(best_B_SPOplus, X_test, c_test, port_oracle)
    final_results.LS_spoloss_test = spo_loss(best_B_leastSquares, X_test, c_test, port_oracle)

    rf_preds_test = predict_random_forests_po(rf_mods, X_test)
    final_results.RF_spoloss_test = spo_loss(Matrix(1.0I, num_assets, num_assets), rf_preds_test, c_test, port_oracle)

    final_results.Absolute_spoloss_test = spo_loss(best_B_Absolute, X_test, c_test, port_oracle)
    final_results.Huber_spoloss_test = spo_loss(best_B_Huber, X_test, c_test, port_oracle)

    c_bar_test_preds = zeros(num_assets, n_test)
    for i = 1:n_test
        c_bar_test_preds[:, i] = c_bar_train
    end
    final_results.Baseline_spoloss_test = spo_loss(Matrix(1.0I, num_assets, num_assets), c_bar_test_preds, c_test, port_oracle)

    (z_test, w_test) = oracle_dataset(c_test, port_oracle)
    final_results.zstar_avg_test = mean(z_test)

    return final_results
end

function portfolio_multiple_replications(rng_seed, num_trials,
    num_assets, num_factors, n_train_vec, n_test, p_features,
        n_sigmoid_polydegree_vec, noise_multiplier_tau_vec;
        num_lambda = 10, lambda_max = :practical, lambda_min_ratio = 0.0001, holdout_percent = 0.25,
        different_validation_losses = false, data_type = :poly_kernel)

    Random.seed!(rng_seed)

    big_results_df = make_blank_complete_df()

    envOracle = setup_gurobi_env(method_type = :default, use_time_limit = false)
    envReform = setup_gurobi_env(method_type = :method3, use_time_limit = false)

    for n_train in n_train_vec
        for n_sigmoid_polydegree in n_sigmoid_polydegree_vec
            for noise_multiplier_tau in noise_multiplier_tau_vec
                println("Moving on to n_train = $n_train, n_sigmoid_polydegree = $n_sigmoid_polydegree, noise_multiplier_tau = $noise_multiplier_tau")
                for trial = 1:num_trials
                    println("Current trial is $trial")
                    n_holdout = round(Int, holdout_percent*n_train)

                    current_results = portfolio_replication(num_assets, num_factors, n_train, n_holdout, n_test, p_features,
                        n_sigmoid_polydegree, noise_multiplier_tau;
                        num_lambda = num_lambda, lambda_max = lambda_max, lambda_min_ratio = lambda_min_ratio,
                        gurobiEnvOracle = envOracle, gurobiEnvReform = envReform, different_validation_losses = different_validation_losses,
                        data_type = data_type)

                    current_df_row = build_complete_row(num_assets,
                        n_train, n_holdout, n_test,
                        p_features, n_sigmoid_polydegree, noise_multiplier_tau, current_results)

                    big_results_df = vcat(big_results_df, current_df_row)
                end
            end
        end
    end

    return big_results_df
end


## -------------------------  H shocks

export HcShockSwitches, HcShockSet, HcShock
export HcShockSwitchesNone, HcShockSetNone, HcShockNone
export HcShockSwitchesDown, HcShockSetDown, HcShockDown, make_hshock_switches_down, max_steps
export HcShockSwitchesLoseLearn, HcShockSetLoseLearn, HcShockLoseLearn, make_hshock_switches_lose_learn
export has_h_shocks, value_before_h_shock, shock_probs, shock_prob, shock_probs_one_college, validate_hshocks
export make_h_shock_set, make_h_shock, sim_h_shocks, pr_hprime
export make_test_hshock_switches, make_test_hshock_set, make_test_hshock

"""
	$(SIGNATURES)

Governs dh shocks for all colleges.
Each type of `HcShock` has 3 objects:
- `HcShockSwitches`: contains switches used to construct the `HcShockSet`.
- `HcShockSet`: model object that describes shocks for all colleges.
- `HcShock`: for one college.
"""

abstract type HcShockSet <: ModelObject end
abstract type HcShockSwitches <: ModelSwitches end
abstract type HcShock end

Base.show(io :: IO,  dh :: HcShockSwitches) = print(io, typeof(dh));
Base.show(io :: IO,  dh :: HcShockSet) = print(io, typeof(dh));
Base.show(io :: IO,  dh :: HcShock) = print(io, typeof(dh));

validate_hshocks(switches :: HcShockSwitches) = true;

has_h_shocks(dh :: HcShockSet) = true;
has_h_shocks(dh :: HcShock) = true;

n_colleges(switches :: HcShockSwitches) = switches.nColleges;
n_colleges(dh :: HcShockSet) = n_colleges(dh.switches);

StructLH.describe(switches :: HcShockSwitches) = settings_table(switches);
StructLH.describe(dh :: HcShockSet) = settings_table(dh);
settings_table(dh :: HcShockSet) = settings_table(dh.switches);

"""
	$(SIGNATURES)

Given value `vNextV` at start of `t+1` by `h[t+1]`, find value at end of `t` before shocks are realized. Returns scalar.
"""
function value_before_h_shock(dh :: HcShock, hStart :: F1, hEndIdx :: Integer,
    hGridNextV :: AbstractVector{F1},  vNextV :: AbstractVector{F1}) where F1

    if has_h_shocks(dh)
        # nh = length(vNextV);
        hEnd = hGridNextV[hEndIdx];
        # vOutV = similar(vNextV);
        # Loop over h grid point at END of today
        # for ih = 1 : nh
            idxV, prV = pr_hprime(dh, hEndIdx, hStart, hEnd, hGridNextV);

            vOut = 0.0;
            for (j, idx) in enumerate(idxV)
                vOut += prV[j] * vNextV[idx];
            end
            # vOutV[ih] = vOut;
        # end
    else
        # vOutV = copy(vNextV);
        vOut = vNextV[hEndIdx];
    end
    @assert all_greater(vNextV, -1e7)  "Low vNext"
    # @assert all_greater(vOutV, -1e7)  "Low values"
    return Double(vOut)
end

include("h_shocks_none.jl");
include("h_shocks_down.jl");
include("h_shocks_lose_learning.jl");

## --------  Testing

make_test_hshock_switches(; nColleges = 4, maxSteps = 2) =
    make_hshock_switches_down(nColleges, maxSteps);

make_test_hshock_set(; nColleges = 4, maxSteps = 2) =
    make_h_shock_set(ObjectId(:test),
        make_test_hshock_switches(; nColleges = nColleges, maxSteps = maxSteps));

make_test_hshock(; nColleges = 4, maxSteps = 2) =
    make_h_shock(
        make_test_hshock_set(; nColleges = nColleges, maxSteps = maxSteps),
        nColleges - 1);

# ------------


using Diversity.API

import Base.IteratorSize, Base.length, Base.IteratorEltype, Base.eltype

abstract type AbstractIterator{M <: AbstractMetacommunity} end

struct TypeIterator{M} <: AbstractIterator{M}
    metacommunity::M
    viewfn::Function
end

function TypeIterator(fn::Function, meta::M) where M <: AbstractMetacommunity
    n = ndims(fn(meta))
    if n == 2
        viewfn = (meta, i) -> view(fn(meta), i, :)
    elseif n == 1
        viewfn = (meta, i) -> view(fn(meta), i)
    else
        error("Can't iterate over types for the function of the metacommunity - $fn")
    end
    return TypeIterator(meta, viewfn)
end

function TypeIterator(meta::M) where M <: AbstractMetacommunity
    return TypeIterator(getabundance, meta)
end

import Base.iterate
function iterate(ti::TypeIterator, state = 1)
    if state > counttypes(ti.metacommunity)
        return nothing
    end
    return ti.viewfn(ti.metacommunity, state), state + 1
end

function IteratorSize(ti::Type{TypeIterator})
    return HasLength()
end

function length(ti::TypeIterator)
    return counttypes(ti.metacommunity)
end

function IteratorEltype(ti::Type{TypeIterator})
    return HasEltype()
end

function eltype(ti::TypeIterator)
    return eltype(ti.viewfn(ti.metacommunity, 1))
end

struct SubcommunityIterator{M} <: AbstractIterator{M}
    metacommunity::M
    viewfn::Function
end

function SubcommunityIterator(fn::Function,
                              meta::M) where M <: AbstractMetacommunity
    n = ndims(fn(meta))
    if n == 2
        viewfn = (meta, i) -> view(fn(meta), :, i)
    elseif n == 1
        viewfn = (meta, i) -> view(fn(meta), i)
    else
        error("Can't iterate over subcommunities for the " *
              "function of the metacommunity - $fn")
    end
    return SubcommunityIterator(meta, viewfn)
end

function SubcommunityIterator(meta::M) where M <: AbstractMetacommunity
    return SubcommunityIterator(getabundance, meta)
end

function iterate(si::SubcommunityIterator, state = 1)
    if state > countsubcommunities(si.metacommunity)
        return nothing
    end
    return si.viewfn(si.metacommunity, state), state + 1
end

function IteratorSize(si::Type{SubcommunityIterator})
    return HasLength()
end

function length(si::SubcommunityIterator)
    return countsubcommunities(si.metacommunity)
end

function IteratorEltype(si::Type{SubcommunityIterator})
    return HasEltype()
end

function eltype(si::SubcommunityIterator)
    return eltype(si.viewfn(si.metacommunity, 1))
end


#> This code implements the MLJ model interface for models in the
#> DecisionTree.jl package. It is annotated so that it may serve as a
#> template for other regressors of Deterministic type and classifiers
#> of Probabilistic type. The annotations, which begin with "#>",
#> should be removed (but copy this file first!). See also the model
#> interface specification at "doc/adding_new_models.md".

#> Note that all models need to "register" their location by setting
#> `load_path(<:ModelType)` appropriately.

module DecisionTree_

#> export the new models you're going to define (and nothing else):
export DecisionTreeClassifier, DecisionTreeRegressor

import MLJBase

#> needed for classifiers:
using CategoricalArrays

#> import package:
import ..DecisionTree # strange syntax b/s we are lazy-loading

## CLASSIFIER

"""
    DecisionTreeClassifer(; kwargs...)

CART decision tree classifier from
[https://github.com/bensadeghi/DecisionTree.jl/blob/master/README.md](https://github.com/bensadeghi/DecisionTree.jl/blob/master/README.md). Predictions
are Probabilistic.

For post-fit pruning, set `post-prune=true` and set
`min_purity_threshold` appropriately. Other hyperparameters as per
package documentation cited above.

"""
mutable struct DecisionTreeClassifier <: MLJBase.Probabilistic
    pruning_purity::Float64
    max_depth::Int
    min_samples_leaf::Int
    min_samples_split::Int
    min_purity_increase::Float64
    n_subfeatures::Float64
    display_depth::Int
    post_prune::Bool
    merge_purity_threshold::Float64
end

# keywork constructor:
#> all arguments are kwargs with a default value
function DecisionTreeClassifier(
    ; pruning_purity=1.0
    , max_depth=-1
    , min_samples_leaf=1
    , min_samples_split=2
    , min_purity_increase=0.0
    , n_subfeatures=0
    , display_depth=5
    , post_prune=false
    , merge_purity_threshold=0.9)

    model = DecisionTreeClassifier(
        pruning_purity
        , max_depth
        , min_samples_leaf
        , min_samples_split
        , min_purity_increase
        , n_subfeatures
        , display_depth
        , post_prune
        , merge_purity_threshold)

    message = MLJBase.clean!(model)       #> future proof by including these
    isempty(message) || @warn message #> two lines even if no clean! defined below

    return model
end

#> The following optional method (the fallback does nothing, returns
#> empty warning) is called by the constructor above but also by the
#> fit methods below:
function MLJBase.clean!(model::DecisionTreeClassifier)
    warning = ""
    if  model.pruning_purity > 1
        warning *= "Need pruning_purity < 1. Resetting pruning_purity=1.0.\n"
        model.pruning_purity = 1.0
    end
    if model.min_samples_split < 2
        warning *= "Need min_samples_split >= 2. Resetting min_samples_slit=2.\n"
        model.min_samples_split = 2
    end
    return warning
end

#> A required `fit` method returns `fitresult, cache, report`. (Return
#> `cache=nothinobserg` unless you are overloading `update`)
function MLJBase.fit(model::DecisionTreeClassifier
             , verbosity::Int   #> must be here (and typed!!) even if not used (as here)
             , X
             , y)

    Xmatrix = MLJBase.matrix(X)

    yplain = identity.(y) # y as plain not abstact vector
    classes_seen = unique(yplain)

    tree = DecisionTree.build_tree(yplain,
                                   Xmatrix,
                                   model.n_subfeatures,
                                   model.max_depth,
                                   model.min_samples_leaf,
                                   model.min_samples_split,
                                   model.min_purity_increase)
    if model.post_prune
        tree = DecisionTree.prune_tree(tree, model.merge_purity_threshold)
    end

    verbosity < 3 || DecisionTree.print_tree(tree, model.display_depth)

    fitresult = (tree, classes_seen)

    #> return package-specific statistics (eg, feature rankings,
    #> internal estimates of generalization error) in `report`, which
    #> should be a named tuple with the same type every call (can have
    #> empty values):

    cache = nothing
    report = NamedTuple{}()

    return fitresult, cache, report

end

MLJBase.fitted_params(::DecisionTreeClassifier, fitresult) = fitresult[1]

function MLJBase.predict(::DecisionTreeClassifier
                     , fitresult
                     , Xnew)
    Xmatrix = MLJBase.matrix(Xnew)

    tree, classes_seen = fitresult

    y_probabilities = DecisionTree.apply_tree_proba(tree, Xmatrix, classes_seen)
    return [MLJBase.UnivariateFinite(classes_seen, y_probabilities[i,:])
            for i in 1:size(y_probabilities, 1)]
end

## REGRESSOR

"""
    DecisionTreeRegressor(; kwargs...)

CART decision tree classifier from
[https://github.com/bensadeghi/DecisionTree.jl/blob/master/README.md](https://github.com/bensadeghi/DecisionTree.jl/blob/master/README.md). Predictions
are Deterministic.

For post-fit pruning, set `post-prune=true` and set
`pruning_purity_threshold` appropriately. Other hyperparameters as per
package documentation cited above.

"""

mutable struct DecisionTreeRegressor <: MLJBase.Deterministic
    pruning_purity_threshold::Float64
    max_depth::Int
    min_samples_leaf::Int
    min_samples_split::Int
    min_purity_increase::Float64
    n_subfeatures::Int
    post_prune::Bool
end

# constructor:
#> all arguments are kwargs with a default value
function DecisionTreeRegressor(;
    pruning_purity_threshold=0.0
    , max_depth=-1
    , min_samples_leaf=5
    , min_samples_split=2
    , min_purity_increase=0.0
    , n_subfeatures=0
    , post_prune=false)

    model = DecisionTreeRegressor(
       pruning_purity_threshold
       , max_depth
       , min_samples_leaf
       , min_samples_split
       , min_purity_increase
       , n_subfeatures
       , post_prune)

    message = MLJBase.clean!(model)       #> future proof by including these
    isempty(message) || @warn message #> two lines even if no clean! defined below

    return model
end

function MLJBase.clean!(model::DecisionTreeRegressor)
    warning = ""
    if model.min_samples_split < 2
        warning *= "Need min_samples_split >= 2. Resetting min_samples_slit=2.\n"
        model.min_samples_split = 2
    end
    return warning
end

function MLJBase.fit(model::DecisionTreeRegressor
             , verbosity::Int   #> must be here (and typed) even if not used (as here)
             , X
             , y)

    Xmatrix = MLJBase.matrix(X)

    fitresult = DecisionTree.build_tree(float.(y)
				   , Xmatrix
				   , model.n_subfeatures
				   , model.max_depth
				   , model.min_samples_leaf
				   , model.min_samples_split
				   , model.min_purity_increase)

    if model.post_prune
        fitresult = DecisionTree.prune_tree(fitresult,
                                            model.pruning_purity_threshold)
    end
    cache = nothing
    report = nothing

    return fitresult, cache, report
end

MLJBase.fitted_params(::DecisionTreeRegressor, fitresult) = fitresult

function MLJBase.predict(model::DecisionTreeRegressor
                     , fitresult
                     , Xnew)
    Xmatrix = MLJBase.matrix(Xnew)
    return DecisionTree.apply_tree(fitresult,Xmatrix)
end

## METADATA

DTTypes=Union{DecisionTreeClassifier,DecisionTreeRegressor}

MLJBase.package_name(::Type{<:DTTypes}) = "DecisionTree"
MLJBase.package_uuid(::Type{<:DTTypes}) = "7806a523-6efd-50cb-b5f6-3fa6f1930dbb"
MLJBase.package_url(::Type{<:DTTypes}) = "https://github.com/bensadeghi/DecisionTree.jl"
MLJBase.is_pure_julia(::Type{<:DTTypes}) = true

MLJBase.load_path(::Type{<:DecisionTreeClassifier}) = "MLJModels.DecisionTree_.DecisionTreeClassifier"
MLJBase.load_path(::Type{<:DecisionTreeRegressor}) = "MLJModels.DecisionTree_.DecisionTreeRegressor"

MLJBase.input_scitype_union(::Type{<:DecisionTreeClassifier}) = MLJBase.Continuous
MLJBase.input_scitype_union(::Type{<:DecisionTreeRegressor}) = MLJBase.Continuous

MLJBase.target_scitype_union(::Type{<:DecisionTreeClassifier}) = MLJBase.Finite
MLJBase.target_scitype_union(::Type{<:DecisionTreeRegressor}) = MLJBase.Continuous

MLJBase.input_is_multivariate(::Type{<:DecisionTreeClassifier}) = true
MLJBase.input_is_multivariate(::Type{<:DecisionTreeRegressor}) = true

end # module


using Test
using Libxc

# Wrap in an outer testset to get a full report if one test fails
@testset "Libxc.jl" begin

@testset "Metadata" begin
    @test Libxc.libxc_version ≥ v"5.1.0"
    @test Libxc.libxc_doi isa AbstractString
end

@testset "Functional list" begin
    available = available_functionals()
    @test :lda_x     in available
    @test :lda_c_vwn in available
    @test :lda_xc_teter93 in available
    @test :gga_x_pbe in available
    @test :gga_c_pbe in available
end

@testset "Functional construction" begin
    lda = Functional(:lda_x, n_spin=2)
    @test is_lda(lda)
    @test !is_gga(lda)
    @test !is_mgga(lda)
    @test lda.family == :lda
    @test sort(lda.flags) == sort([:vxc, :dim3, :fxc, :exc])
    @test lda.kind == :exchange
    @test lda.name == "Slater exchange"
    @test lda.n_spin == 2
    @test all(i in supported_derivatives(lda) for i in 0:2)
end

@testset "LDA and GGA evaluate" begin
    rho = [0.1, 0.2, 0.3, 0.4, 0.5]
    sigma = [0.2, 0.3, 0.4, 0.5, 0.6]
    result = zeros(Float64, 5)

    # LDA
    lda_x = Functional(:lda_x)
    result = evaluate(lda_x, rho=rho, derivatives=0).zk
    @test result ≈ [-0.342809, -0.431912, -0.494416, -0.544175, -0.586194] atol=1e-5

    # GGA
    gga_x = Functional(:gga_x_pbe)
    result = evaluate(gga_x, rho=rho, sigma=sigma, derivatives=0).zk
    @test result ≈ [-0.452598, -0.478878, -0.520674, -0.561428, -0.598661] atol=1e-5
end

@testset "LDA and GGA evaluate!" begin
    rho = [0.1, 0.2, 0.3, 0.4, 0.5]
    sigma = [0.2, 0.3, 0.4, 0.5, 0.6]
    result = zeros(Float64, 5)

    # LDA
    func = Functional(:lda_x)
    evaluate!(func, rho=rho, zk=result)
    @test result ≈ [-0.342809, -0.431912, -0.494416, -0.544175, -0.586194] atol=1e-5

    # GGA
    func = Functional(:gga_x_pbe)
    evaluate!(func, rho=rho, sigma=sigma, zk=result)
    @test result ≈ [-0.452598, -0.478878, -0.520674, -0.561428, -0.598661] atol=1e-5
end

@testset "LDA / GGA evaluate! with spin" begin
    rho = abs.(randn(5))
    sigma = abs.(randn(5))

    # Duplicate rho and sigma for spin = 2 tests
    rho2 = zeros(Float64, 2length(rho))
    rho2[1:2:end] = rho2[2:2:end] = rho ./ 2

    sigma2 = zeros(Float64, 3length(rho))
    sigma2[1:3:end] = sigma2[2:3:end] = sigma2[3:3:end] = sigma ./ 4

    # LSDA
    for sym in (:lda_x, :lda_c_vwn)
        E = zeros(Float64, length(rho))
        func = Functional(sym, n_spin=1)
        evaluate!(func, rho=rho, zk=E)

        E2 = zeros(Float64, length(rho))
        func = Functional(sym, n_spin=2)
        evaluate!(func, rho=rho2, zk=E2)

        @test E ≈ E2
    end

    # GGA
    for sym in (:gga_x_pbe, :gga_c_pbe)
        E = zeros(Float64, length(rho))
        func = Functional(sym, n_spin=1)
        evaluate!(func, rho=rho, sigma=sigma, zk=E)

        E2 = zeros(Float64, length(rho))
        func = Functional(sym, n_spin=2)
        evaluate!(func, rho=rho2, sigma=sigma2, zk=E2)

        @test E ≈ E2
    end
end

@testset "LDA / GGA evaluate with spin" begin
    shape = (2, 3, 4)
    rho   = reshape(abs.(randn(shape)), 1, shape...)
    sigma = reshape(abs.(randn(shape)), 1, shape...)

    # Duplicate rho and sigma for spin = 2 tests
    rho2 = 0.5vcat(rho, rho)
    sigma2 = 0.25vcat(sigma, sigma, sigma)

    # LSDA
    for sym in (:lda_x, :lda_c_vwn)
        res = evaluate(Functional(sym, n_spin=1), rho=rho, zk=zeros(shape))
        @test size(res.zk) == shape
        @test size(res.vrho) == (1, shape...)

        res2 = evaluate(Functional(sym, n_spin=2), rho=rho2)
        @test size(res2.zk) == shape
        @test size(res2.vrho) == (2, shape...)

        @test res.zk ≈ res2.zk
        @test res.vrho[1, :, :, :] ≈ res2.vrho[1, :, :, :]
        @test res.vrho[1, :, :, :] ≈ res2.vrho[2, :, :, :]
    end

    # GGA
    for sym in (:gga_x_pbe, :gga_c_pbe)
        res = evaluate(Functional(sym, n_spin=1), rho=rho, sigma=sigma)
        @test size(res.zk) == shape
        @test size(res.vrho) == (1, shape...)
        @test size(res.vsigma) == (1, shape...)

        res2 = evaluate(Functional(sym, n_spin=2), rho=rho2, sigma=sigma2)
        @test size(res2.zk) == shape
        @test size(res2.vrho) == (2, shape...)
        @test size(res2.vsigma) == (3, shape...)

        @test res.zk ≈ res2.zk
        @test res.vrho[1, :, :, :] ≈ res2.vrho[1, :, :, :]
        @test res.vrho[1, :, :, :] ≈ res2.vrho[2, :, :, :]
        @test 4res.vsigma ≈ sum(res2.vsigma, dims=1)
    end
end

@testset "mGGA evaluate! without Laplacian" begin
    rho    = [0.1, 0.2, 0.3, 0.4, 0.5]
    sigma  = [0.2, 0.3, 0.4, 0.5, 0.6]
    tau    = [0.5, 0.4, 0.3, 0.2, 0.1]
    result = zeros(Float64, 5)

    tpss_x = Functional(:mgga_x_tpss, n_spin=1)
    evaluate!(tpss_x, rho=rho, sigma=sigma, tau=tau, zk=result)
    @test result ≈ [-0.4254894, -0.4685985, -0.5341998, -0.6017734, -0.6646443] atol=1e-6
end

@testset "mGGA evaluate! with Laplacian" begin
    rho    = [0.1, 0.2, 0.3, 0.4, 0.5]
    sigma  = [0.2, 0.3, 0.4, 0.5, 0.6]
    lapl   = [-0.2, -0.3, -0.4, -0.5, -0.6]
    tau    = [0.5, 0.4, 0.3, 0.2, 0.1]
    vrho   = randn(Float64, 5)
    vsigma = randn(Float64, 5)
    vtau   = randn(Float64, 5)
    vlapl  = randn(Float64, 5)

    # TB09 is a potential functional, construct and check this
    tb09_x = Functional(:mgga_x_tb09, n_spin=1)
    @test all(i in supported_derivatives(tb09_x) for i in 1:2)
    @test !(0 in supported_derivatives(tb09_x))

    evaluate!(tb09_x; rho, sigma, lapl, tau, vrho, vsigma, vtau, vlapl)
    @test vrho ≈ [ 0.07015148, -0.4085103, -0.73564854, -0.97043302, -1.15268102] atol=1e-6
    @test iszero(vsigma)
    @test iszero(vtau)
    @test iszero(vlapl)
end

@testset "evaluate! with invalid arguments" begin
    rho = [0.1, 0.2, 0.3, 0.4, 0.5]
    func = Functional(:lda_x)

    # Derivatives ≥ 3 not compiled into libxc at the moment
    @test_throws ArgumentError evaluate(func, rho=rho, derivatives=0:3)
    @test_throws ArgumentError evaluate!(func, rho=rho, v3rho3=randn(5))

    @test_throws DimensionMismatch evaluate!(func, rho=rho, zk=randn(2))
end

@testset "Custom evaluate! dispatch" begin
    # A dummy version of Slater exchange for Float32
    function Libxc.evaluate!(func::Functional, ::Val{:lda}, rho::Array{Float32};
                             zk::Array{Float32})
        @assert func.identifier == :lda_x
        @. zk = -Float32(3/4) * cbrt(Float32(3/π) * rho)
    end
    rho = Float32[0.1, 0.2, 0.3, 0.4, 0.5]
    result = zeros(Float32, 5)
    func = Functional(:lda_x)

    res = evaluate(func, rho=rho, derivatives=0)
    @test eltype(res.zk) == Float32
    @test eltype(res.zk) == Float32
    @test res.zk ≈ [-0.342809, -0.431912, -0.494416, -0.544175, -0.586194] atol=1e-5
end

@testset "Reference extraction" begin
    lda = Functional(:lda_x)
    @test length(lda.references) == 2
    @test lda.references[1].doi  == "10.1017/S0305004100016108"
    @test lda.references[2].doi  == "10.1007/BF01340281"
end

@testset "Setting thresholds" begin
    lda = Functional(:lda_x)

    lda.density_threshold = 1e-4
    @test lda.density_threshold == 1e-4
    lda.density_threshold = 1e-6
    @test lda.density_threshold == 1e-6

    for field in (:zeta_threshold, :sigma_threshold, :tau_threshold)
        setproperty!(lda, field, 1e-4)
        @test getproperty(lda, field) == 1e-4
        setproperty!(lda, field, 1e-6)
        @test getproperty(lda, field) == 1e-6
    end
end

@testset "Read-only properties" begin
    lda = Functional(:lda_x)

    # LDA has none of these things
    allprops = [:exx_coefficient, :cam_alpha, :cam_beta, :cam_omega, :nlc_b, :nlc_C]
    for s in allprops
        @test isnothing(getproperty(lda, s))
    end

    # Global hybrid functional
    b3lyp = Functional(:hyb_gga_xc_b3lyp)
    @test b3lyp.exx_coefficient == 0.2
    @test isnothing(b3lyp.cam_alpha)
    @test isnothing(b3lyp.cam_beta)
    @test isnothing(b3lyp.nlc_b)

    # SCAN + VV10 correlation functional
    scan_vv10 = Functional(:xc_mgga_c_scan_vv10)
    @test isnothing(scan_vv10.exx_coefficient)
    @test isnothing(scan_vv10.cam_alpha)
    @test scan_vv10.nlc_b == 14.0
    @test scan_vv10.nlc_C == 0.0093
    @test !needs_laplacian(scan_vv10)

    # ωB97 functional
    ωB97 = Functional(:xc_hyb_gga_xc_wb97)
    @test isnothing(ωB97.exx_coefficient)
    @test ωB97.cam_alpha == 1.0
    @test ωB97.cam_beta  == -1.0
    @test ωB97.cam_omega == 0.4
    @test isnothing(ωB97.nlc_b)

    # ωB97x_v functional
    ωB97x_v = Functional(:xc_hyb_gga_xc_wb97x_v)
    @test isnothing(ωB97.exx_coefficient)
    @test ωB97x_v.cam_alpha == 1.0
    @test ωB97x_v.cam_beta  == -(1.0 - 0.167)
    @test ωB97x_v.cam_omega == 0.3
    @test ωB97x_v.nlc_b     == 6.0
    @test ωB97x_v.nlc_C     == 0.01
end
end  # outer wrapper


# A BonesaTuner has state for optimizing a certain algorithm on a set of
# parameters described by a ParameterSet. The parameter vectors that are tuned
# only take values in [0.0, 1.0] and it is up to the ParameterSet to translate
# that to meaningful ranges.
type BonesaTuner
  numvectors::Int64                     # Current num of param vectors in archive
  parameter_vectors::Array{Float64, 2}  # Archive of param vectors, each column is one vector
  utilities::Array{Float64, 2}          # Each column has Nu * Np utility values for one and the same parameter vector
  num_utility_values
end

# The predicted utility of a parameter vector x on problem is a weigthed, kernel
# smoothed average of the utilities of the vectors in the archive.
function predicted_utility(ba::BonesaArchive, x, problem)
end

function time_is_up(bt::BonesaTuner)
  time() - bt.start_time >= bt.max_time_seconds
end

# Start tuning parameters given a maximum time budget in hours.
function tune_parameters(bt::BonesaTuner, max_hours = 1.0)
  # Set up for measuring time left.
  bt.start_time = time()
  bt.max_time_seconds = max_hours * 60.0 * 60.0

  # Create a random set of vectors and then evaluate them on the problem set.
  create_random_parameter_vectors(bt, bt.M)
  bt.utilities = zeros(num_utilities(bt) * num_problems(bt.problem_set), bt.M)
  for(i in 1:num_param_vectors(bt))
    if time_is_up(bt)
      # Cut the size of the archive since we did not have time to eval them all... :(
      bt.parameter_vectors = bt.parameter_vectors[:,1:(i-1)]
      bt.utilities = bt.utilities[:,1:(i-1)]
      break
    end
    bt.utilities[:, i] = evaluate_param_vector(bt, bt.parameter_vectors[:,i])
  end

  # Iteratively add to archive as we learn more about which param vectors are good.
  while( !time_is_up(bt) )

    # Update constants and temp calcs needed to calc pareto strenght below

  end

  # Select the terminal set (pareto front) of the param vectors.

end

# Calc predicted utilities, variance and support as well as pareto dominance
# and strengths for all param vectors in the archive.
function calc_temp_values_for_archive(bt::BonesaTuner)
  # Update magic constant c based on size of archive
  bt.c = calc_approximate_magic_constant_c(size_of_archive(bt), numdims(bt))

  # Calc the good vectors (saved as (boolean) indicator values)
  mean_actual_utilities = mean(bt.utilities, 2) # Mean per row => mean for each utility value

end

function support_of_vector(bt::BonesaTuner, pvector)

end

# For now we use random sampling but use lhs here instead, long-term.
function create_random_parameter_vectors(bt::BonesaTuner, num_vectors)
  bt.parameter_vectors = rand(num_params(bt), num_vectors)
end

function num_utility_values(bt::BonesaTuner)
  if bt.num_utility_values == false
    # Evaluate a param vector once to know how many utility values are returned.
    res = evaluate_param_vector_on_problem(bt, bt.problem_set[1], rand(num_params(bt)))
    bt.num_utility_values = length(res) + 1 # Add one for time
  else
    bt.num_utility_values
  end
end

function evaluate_param_vector(bt, pvector)
  parameters = map_to_parameter_values(bt.parameter_set, pvector)
  nu = num_utility_values(bt)
  np = num_problems(bt.problem_set)
  utilities = zeros(np*nu)

  for(i in 1:np)
    p = bt.problem_set[i]
    start = 1 + (i-1) * nu
    tic()
    utilities[ustart:(ustart+nu-2)] = evaluate_param_vector_on_problem(bt, p, parameters)
    utilities[ustart+nu-1] = toq()
  end

  utilities
end

# Find an approximation to the BONESA magic constant c. It is the value
# that makes sum(w) == 50, where w is exp(c * normalized_distance(x, y)) for
# m uniformly, randomly sampled on the l-dimensional hypercube.
function calc_approximate_magic_constant_c(m, l, num_repeats = 5, tolerance = 1e-3)
  # Do a number of random approximations and then average
  mean([find_approximate_c_given_random_sample_of_points(m, l, tolerance) for i in 1:num_repeats])
end

function find_approximate_c_given_random_sample_of_points(m, l, tolerance = 1e-3)
  # Randomly sample m l-dimensional points
  points = rand(l, m)

  # To calc the distances we need the std dev per dimension
  std_devs = std(points, 2)

  # Now calc the distances between points. We only sample them if m is large.
  distances = zeros(Float64, m, m)
  for(i in 1:m)
    x = points[:,i]
    for(j in i:m)
      y = points[:,j]
      distances[i, j] = distances[i, j] = std_distance(x, y, std_devs)
    end
  end

  # Find a value of c by using bisection so that we are close to 50.0
  f = (c) -> average_sum_of_distances_for_c(c, distances) - 50.0
  bisection(f, -100.0, 100.0, tolerance)
end

# Find a value v that makes abs(f(v)) <= tolerance where a and b are guesses
# for v where b > a. This implementation assumes that we can find a negative
# (or positive) value of the function by stepping down (or up) => will not
# work for any function but throws an ArgumentError after many iterations of trying.
function bisection(f, a, b = 2*a, tolerance = 1e-5, max_iterations = 1e4)
  if f(a) > 0 && f(b) > 0
    a = step_value_until_changes_sign(f, a, -1.0, max_iterations, max_iterations)
  elseif f(a) < 0 && f(b) < 0
    b = step_value_until_changes_sign(f, b, 1.0, max_iterations, max_iterations)
  end

  p = (a + b)/2
  error = abs(f(p))
  iterations = 0
  while iterations < max_iterations && error > tolerance
    iterations += 1
    (f(a)*f(p) < 0) ? (b = p) : (a = p)
    p = (a + b)/2
    error = abs(f(p))
  end
  if iterations >= max_iterations
    throw(ArgumentError("Could not find a good approximation after $(iterations) iterations."))
  else
    return p
  end
end

function step_value_until_changes_sign(func, value::Float64, step_size::Float64 = 1.0, max_iterations = 1e3)
  wanted_sign = -1 * sign(func(value))
  iterations = 0
  while iterations < max_iterations && sign(func(value)) != wanted_sign
    iterations += 1
    value = value + step_size
    step_size *= 2
  end
  if iterations >= max_iterations
    throw(ArgumentError("Could not invert the sign after $(iterations) iterations."))
  else
    return value
  end
end

function average_sum_of_distances_for_c(c, distances)
  mean( exp(c * collect(triu(distances))) )
end

function std_distance(x, y, std_devs)
  sqrt(1/l * sum(((x .- y) ./ std_devs).^2))
end


module Tools

export alias, aliasRegion!, setupIndex!, getMarkup, mergeMcosts
export checkConvg, fastFill!, whsaleprice, collectΠ!
using TypeDef
using NLsolve, LoopVectorization

function fastFill!(
  A    ::Matrix{Float64},
  nRow ::Int64,
  nCol ::Int64)

  @avx for j=1:nCol
    for i=1:nRow
      A[i,j] = 0
    end
  end
end

function alias(
  dim::dimType,
  sim::simType,
  optD2::Bool=false)

  s      = sim.s
  num    = sim.num
  D      = sim.D
  ms     = sim.ms
  ∂s∂p   = sim.∂s∂p
  ∂S∂p   = sim.∂S∂p
  ∂D∂p   = sim.∂D∂p

  fill!(ms,0)

  if optD2==true
    ∂²S∂p² = sim.∂²S∂p²
    ∂²D∂p² = sim.∂²D∂p²
    @inbounds for j=1:dim.product
      D[j] = 0
      for i=1:dim.product
        ∂D∂p[i,j] = 0
        for k=1:dim.product
          ∂²D∂p²[k,i,j] = 0
        end
      end
    end
    return s,num,D,ms,∂s∂p,∂S∂p,∂²S∂p²,∂D∂p,∂²D∂p²
  else
    @inbounds for j=1:dim.product
      D[j] = 0
      for i=1:dim.product
        ∂D∂p[i,j] = 0
      end
    end
    return s,num,D,ms,∂s∂p,∂S∂p,∂D∂p
  end
end

function aliasRegion!(
  num   ::Vector{Float64},
  indx  ::indxType,
  r     ::Int64,
  dim   ::dimType,
  mat   ::matType,
  par   ::parType,
  sim   ::simType)

  fill!(num,0)
  # r   = g - minimum(indx.market) + 1 # region index
  g   = indx.market[r]
  αᵢ  = view(par.αᵢ,:,r)
  μ̃   = view(sim.μ̃,:,:,r)
  pop = mat.pop[r]
  # active = view(indx.active,:,r)
  active = findall(view(indx.active,:,r))

  return αᵢ,μ̃,pop,active
end

function setupIndex!(
  sim ::simType,
  indx::indxType,
  pdline::BitArray{1},
  dim ::dimType,
  mat ::matType,
  mask::maskType,
  t   ::Int64)

  indx.market .= dim.region*(t-1) .+ collect(1:dim.region)
  sim.μ̃       .= view(mat.μ̃,:,:,indx.market)
  sim.p       .= view(mat.p,:,t)

  anyActive = findmax(view(mat.active,:,indx.market), dims=2)[1]
  for j=1:dim.product
    indx.anyActive[j] = anyActive[j] & pdline[j]
    indx.mno[j]       = mask.mno[j]  & indx.anyActive[j]
    indx.mvno[j]      = mask.mvno[j] & indx.anyActive[j]
    for r=1:dim.region
      g = indx.market[r]
      indx.active[j,r] = mat.active[j,g] & pdline[j]
    end
  end
  for j=1:3
    for i=1:dim.mvno
      indx.mvnp[i,j] = mat.mvnp[i,j] & indx.anyActive[dim.mno+i]
    end
  end

end

function getMarkup(
  p     ::Vector{Float64},
  wprice::Vector{Float64},
  mcost ::Vector{Vector{Float64}})

  markupR = p - mcost[2]
  markupW = wprice - mcost[1]
  markupV = p - wprice - mcost[2]

  return markupR, markupW, markupV
end

function mergeMcosts(
  mcostinc,
  mcostfree,
  dim::dimType,
  indx::indxType,
  mask::maskType)

  # mcost = zeros(dim.product)
  # mcost[indx.anyActive .& .!mask.free] = mcostinc
  # mcost[indx.anyActive .& mask.free] .= mcostfree

  mcost = [zeros(dim.product) for t=1:2]
  idx = findall(indx.anyActive .& .!mask.free)
  for j in eachindex(idx)
    i = idx[j]
    if indx.mno[i]
      mcost[2][i] = mcostinc[j]
    elseif indx.mvno[i]
      mcost[2][i] = 0
      mcost[1][i] = mcostinc[j]
    end
  end
  mcost[2][mask.free] = mcostfree
  return mcost
end

function checkConvg(
  iter::Int64,
  Op  ::NLsolve.SolverResults,
  fail::Bool=false)

  exitflag = (converged(Op)==true) & (maximum(Op.zero)<100) & (fail==false)

  return exitflag
end

function whsaleprice(
  p,
  dim::dimType,
  mat::matType,
  sim::simType,
  indx::indxType)

  any    = indx.anyActive
  D1own  = sim.∂D∂p' .* mat.own
  markup = -D1own[any,any] \ sim.D[any]
  wprice = zeros(dim.product)
  wprice[any] = p[any] - markup

  return wprice
end

function collectΠ!( # Collect profits both retail & wholesale
  Π,
  eq  ::equilType,
  vInt::Bool,
  t   ::Int64,
  dim ::dimType,
  mat ::matType,
  indx::indxType,
  mask::maskType)

  # Incumbents' profit
  prodIdx = findall(indx.anyActive)
  for j in prodIdx
    if vInt==false
      i_retail = mat.firm[j,2]
      Π[i_retail] += eq.profit[j,t,2]
      if mask.mvno[j]
        i_whsale = mat.firm[j,3]
        Π[i_whsale] += eq.profit[j,t,1]
      end
    else
      i_whsale = mat.firm[j,3]
      Π[i_whsale] += eq.profit[j,t,2]
    end
  end # end of j loop

end

end


const FlatField{P,T,M} = Union{FlatS0{P,T,M},FlatS2{P,T,M},FlatS02{P,T,M}}

const FlatFieldMap{P,T,M} = Union{FlatMap{P,T,M},FlatS2Map{P,T,M},FlatS02Map{P,T,M}}
const FlatFieldFourier{P,T,M} = Union{FlatFourier{P,T,M},FlatS2{P,T,M},FlatS02Fourier{P,T,M}}

### pretty printing
@show_datatype show_datatype(io::IO, t::Type{F}) where {N,θ,∂mode,D,T,M,F<:FlatField{Flat{N,θ,∂mode,D},T,M}} =
print(io, "$(pretty_type_name(F)){$(join(N .* (1,1), "×"))$(D==1 ? "" : "(×$D)") map, $(θ)′ pixels, $(∂mode.name.name), $M}")
for F in (:FlatMap, :FlatFourier,
          :FlatQUMap, :FlatQUFourier, :FlatEBMap, :FlatEBFourier,
          :FlatIQUMap, :FlatIQUFourier, :FlatIEBMap, :FlatIEBFourier)
    @eval pretty_type_name(::Type{<:$F}) = $(string(F))
end
function Base.summary(io::IO, f::FlatField{<:Flat{N,<:Any,<:Any,D}}) where {N,D}
    print(io, "$(length(f))", (D==1 ? "" : "(×$D)"), "-element ")
    showarg(io, f, true)
end

### field info
@memoize fieldinfo(::Type{P},::Type{T}=Float32,::Type{M}=Matrix) where {Nside,θpix,∂mode,D,P<:Flat{Nside,θpix,∂mode,D},T,M} =
    (;FlatInfo(T,basetype(M),θpix,Nside,D)..., ∂mode=∂mode)
@memoize fieldinfo(::Type{F}) where {P<:Flat,T,M,F<:FlatField{P,T,M}} =
    (;fieldinfo(P,T,M)..., @namedtuple(P,M,B=basis(F),S=spin(F))...)
fieldinfo(::F) where {F<:FlatField} = fieldinfo(F)

### promotion & conversion
# note: we don't need to promote the eltype T here since that will be
# automatically handled in broadcasting
function promote(f1::F1, f2::F2) where {T1,θ1,N1,∂mode1,F1<:FlatField{<:Flat{N1,θ1,∂mode1},T1},T2,θ2,N2,∂mode2,F2<:FlatField{<:Flat{θ2,N2,∂mode2},T2}}
    B     = promote_type(basis(F1),basis(F2))
    ∂mode = promote_type(∂mode1,∂mode2)
    B(∂mode(f1)), B(∂mode(f2))
end
(::Type{∂mode})(f::F) where {∂mode<:∂modes,N,θ,D,F<:FlatS0{<:Flat{N,θ,<:Any,D}}} = basetype(F){Flat{N,θ,∂mode,D}}(fieldvalues(f)...)
(::Type{∂mode})(f::FieldTuple{B}) where {∂mode<:∂modes,B} = FieldTuple{B}(map(∂mode,f.fs))

### basis-like definitions
LenseBasis(::Type{<:FlatS0})  =    Map
LenseBasis(::Type{<:FlatS2})  =  QUMap
LenseBasis(::Type{<:FlatS02}) = IQUMap
DerivBasis(::Type{<:FlatS0{<:Flat{<:Any,<:Any,fourier∂}}})  =    Fourier
DerivBasis(::Type{<:FlatS2{<:Flat{<:Any,<:Any,fourier∂}}})  =  QUFourier
DerivBasis(::Type{<:FlatS02{<:Flat{<:Any,<:Any,fourier∂}}}) = IQUFourier
DerivBasis(::Type{<:FlatS0{<:Flat{<:Any,<:Any,map∂}}})      =    Map
DerivBasis(::Type{<:FlatS2{<:Flat{<:Any,<:Any,map∂}}})      =  QUMap
DerivBasis(::Type{<:FlatS02{<:Flat{<:Any,<:Any,map∂}}})     = IQUMap

### derivatives

## Fourier-space
# the use of @generated here is for memoization so that we don't have do the
# `prefactor * im * ...` each time. actually, since for the first two cases
# these are only 1-d arrays, that's a completely unnecessary optimization, but
# without the @generated it triggers an order-dependenant compilation bug which
# e.g. slows down LenseFlow by a factor of ~4 so we gotta keep it for now ¯\_(ツ)_/¯
@generated function broadcastable(::Type{F}, ::∇diag{1,<:Any,prefactor}) where {P,T,M,prefactor,F<:FlatFourier{P,T,M}}
    @unpack kx = fieldinfo(F)
    basetype(M)(@. prefactor * im * kx')
end
@generated function broadcastable(::Type{F}, ::∇diag{2,<:Any,prefactor}) where {P,T,M,prefactor,F<:FlatFourier{P,T,M}}
    @unpack ky, Ny = fieldinfo(F)
    basetype(M)(@. prefactor * im * ky[1:Ny÷2+1])
end
@generated function broadcastable(::Type{F}, ::∇²diag) where {P,T,M,F<:FlatFourier{P,T,M}}
    @unpack kx, ky, Ny = fieldinfo(F)
    basetype(M)(@. kx'^2 + ky[1:Ny÷2+1]^2)
end

## Map-space
function copyto!(f′::F, bc::Broadcasted{<:Any,<:Any,typeof(*),Tuple{∇diag{coord,covariant,prefactor},F}}) where {coord,covariant,prefactor,T,θ,N,D,F<:FlatMap{Flat{N,θ,map∂,D},T}}
    D!=1 && error("Gradients of batched map∂ flat maps not implemented yet.")
    f = bc.args[2]
    n,m = size(f.Ix)
    α = 2 * prefactor * fieldinfo(f).Δx
    if coord==1
        @inbounds for j=2:m-1
            @simd for i=1:n
                f′[i,j] = (f[i,j+1] - f[i,j-1])/α
            end
        end
        @inbounds for i=1:n
            f′[i,1] = (f[i,2]-f[i,end])/α
            f′[i,end] = (f[i,1]-f[i,end-1])/α
        end
    elseif coord==2
        @inbounds for j=1:n
            @simd for i=2:m-1
                f′[i,j] = (f[i+1,j] - f[i-1,j])/α
            end
            f′[1,j] = (f[2,j]-f[end,j])/α
            f′[end,j] = (f[1,j]-f[end-1,j])/α
        end
    end
    f′
end

### bandpass
HarmonicBasis(::Type{<:FlatS0}) = Fourier
HarmonicBasis(::Type{<:FlatQU}) = QUFourier
HarmonicBasis(::Type{<:FlatEB}) = EBFourier
broadcastable(::Type{F}, bp::BandPass) where {P,T,F<:FlatFourier{P,T}} = Cℓ_to_2D(P,T,bp.Wℓ)

### logdets
logdet(L::Diagonal{<:Complex,<:FlatFourier}) = batch(real(sum_kbn(nan2zero.(log.(L.diag[:Il,full_plane=true])),dims=(1,2))))
logdet(L::Diagonal{<:Real,   <:FlatMap})     = batch(real(sum_kbn(nan2zero.(log.(complex.(L.diag.Ix))),dims=(1,2))))
### traces
tr(L::Diagonal{<:Complex,<:FlatFourier}) = batch(real(sum_kbn(L.diag[:Il,full_plane=true],dims=(1,2))))
tr(L::Diagonal{<:Real,   <:FlatMap})     = batch(real(sum_kbn(complex.(L.diag.Ix),dims=(1,2))))

### misc
Cℓ_to_Cov(f::FlatField{P,T,M}, args...; kwargs...) where {P,T,M} = adapt(M, Cℓ_to_Cov(P,T,spin(f),args...; kwargs...))

function pixwin(f::FlatField)
    @unpack θpix,P,T,k = fieldinfo(f)
    Diagonal(FlatFourier{P,T}((pixwin.(θpix,k) .* pixwin.(θpix,k'))[1:end÷2+1,:]))
end

global_rng_for(::Type{<:FlatField{<:Any,<:Any,M}}) where {M} = global_rng_for(M)

"""
    fixed_white_noise(rng, F)

Like white noise but the amplitudes are fixed to unity, only the phases are
random. Currently only implemented when F is a Fourier basis. Note that unlike
[`white_noise`](@ref), fixed white-noise generated in EB and QU Fourier bases
are not statistically the same.
"""
fixed_white_noise(rng, F::Type{<:FlatFieldFourier}) =
     exp.(im .* angle.(basis(F)(white_noise(rng,F)))) .* fieldinfo(F).Nside


if get(ARGS, 1, "Array") == "CuArray"
    using CUDA
    using CUDAKernels
    ArrayType = CUDA.CuArray
    CUDA.allowscalar(false)
    get_device() = CUDADevice()
else
    ArrayType = Array
    get_device() = CPU()
end

@show ArrayType

# Re-include test_helper so we can run
# this file on its own.
include("test_helper.jl")

using KernelAbstractions

# only grab array problems
filter!(x->x.x_init isa AbstractArray, problem_list)

@kernel function solve_kernel!(
    f,
    f′,
    MethodType,
    x_init, x_lower, x_upper,
    tol,
    dst::AbstractArray{FT, N},
) where {FT, N}
    i = @index(Group, Linear)
    @inbounds begin
        # Construct the method and solve the system.
        # Store the solution in `dst`
        method = get_method(MethodType, x_init[i], x_lower[i], x_upper[i], f′)
        sol = find_zero(f, method, CompactSolution(), tol)
        dst[i] = sol.root
    end
end

@testset "CPU/GPU kernel test" begin
    n_elem = problem_size()
    device = get_device()
    work_groups = (1,)
    ndrange = (n_elem,)

    for prob in problem_list
        FT = typeof(prob.x̃)
        x_init = prob.x_init
        x_lower = prob.x_lower
        x_upper = prob.x_upper
        for MethodType in (
                    SecantMethodType(),
                    RegulaFalsiMethodType(),
                    NewtonsMethodADType(),
                    NewtonsMethodType(),
                )
            for tol in get_tolerances(FT)
                a_dst = Array{FT}(undef, n_elem)
                d_dst = ArrayType(a_dst)
                kernel! = solve_kernel!(device, work_groups)
                event = kernel!(
                    prob.f,
                    prob.f′,
                    MethodType,
                    x_init,
                    x_lower,
                    x_upper,
                    tol,
                    d_dst;
                    ndrange = ndrange
                )
                wait(device, event)

                @test all(d_dst .≈ prob.x̃)
            end
        end
    end
end