# This file defines some struct types used in the SMLMData package.
"""
    SMLD

 SMLD is the highest level abstract type in SMLMData.
"""
abstract type SMLD end

"""
    SMLD2D

Structure containing 2D localization data

    # Fields
    - connectID: integer ID associating localizations
    - x: x position
    - y: y position
    - σ_x: standard error uncertainty in x
    - σ_y: standard error uncertainty in y
    - photons: total photons
    - σ_photons: standard error uncertainty in σ_photons
    - bg: fit fluorescence background (photons/pixel)
    - σ_bg: standard error uncertainty in bg
    - framenum: integer framenumber within a dataset
    - datasetnum: integer dataset number
    - datasize::Vector{Int}: size of image area
    - nframes: integer frames per dataset
    - ndatasets: number of dataasets
    - datafields: bookeeping - do not modify
"""
mutable struct SMLD2D <: SMLD
    connectID::Vector{Int}
    x::Vector{Float64}
    y::Vector{Float64}
    σ_x::Vector{Float64}
    σ_y::Vector{Float64}
    photons::Vector{Float64}
    σ_photons::Vector{Float64}
    bg::Vector{Float64}
    σ_bg::Vector{Float64}
    framenum::Vector{Int}
    datasetnum::Vector{Int}
    datasize::Vector{Int}
    nframes::Int
    ndatasets::Int
    datafields::NTuple{11, Symbol}
    SMLD2D() = new()
end

"""
    SMLD2D(nlocalizations::Int)

Constructor to generate an empty `smld` with a specific size.

# Description
This is a constructor for the SMLD2D struct which allows you to populate the
structure with undefined values for a predefined number of localizations.
"""
function SMLD2D(nlocalizations::Int)
    smld = SMLD2D()
    smld.x = zeros(Float64, nlocalizations)
    smld.y = zeros(Float64, nlocalizations)
    smld.σ_x = zeros(Float64, nlocalizations)
    smld.σ_y = zeros(Float64, nlocalizations)
    smld.bg = zeros(Float64, nlocalizations)
    smld.σ_bg = zeros(Float64, nlocalizations)
    smld.photons = zeros(Float64, nlocalizations)
    smld.σ_photons = zeros(Float64, nlocalizations)
    smld.connectID = collect(1:nlocalizations)
    smld.framenum = zeros(Int, nlocalizations)
    smld.datasetnum = zeros(Int, nlocalizations)
    smld.nframes = 0
    smld.ndatasets = 0
    smld.datasize = [0; 0]
    smld.datafields = (:connectID, :x, :y, :σ_x, :σ_y,
        :photons, :σ_photons, :bg, :σ_bg, :framenum, :datasetnum)

    return smld
end

"""
    SMLD2D(data::DataFrames.DataFrame)

Constructor to generate an `smld` from a data frame.

# Description
This is a constructor for the SMLD2D struct which allows you to populate the
structure with data defined in the dataframe `data`. The intention is that a
.csv table can be organized with localizations on each row as
`[datasetnum, framenum, x, y, σ_x, σ_y]`, loaded using the CSV package, and
placed into a dataframe with the DataFrames package.
"""
function SMLD2D(data::DataFrames.DataFrame)
    smld = SMLD2D()
    smld.connectID = collect(1:size(data)[1])
    smld.datasetnum = Int.(data[:, 1])
    smld.framenum = Int.(data[:, 2])
    smld.x = Float64.(data[:, 3])
    smld.y = Float64.(data[:, 4])
    smld.σ_x = Float64.(data[:, 5])
    smld.σ_y = Float64.(data[:, 6])
    smld.photons = Float64.(data[:, 7])
    smld.σ_photons = Float64.(data[:, 8])
    smld.bg = Float64.(data[:, 9])
    smld.σ_bg = Float64.(data[:, 10])
    smld.nframes = Int(maximum(smld.framenum))
    smld.ndatasets = Int(Base.length(unique(smld.datasetnum)))
    smld.datasize = [ceil(maximum(data[:, 3]) - 0.5);
                     ceil(maximum(data[:, 4]) - 0.5)]
    smld.datafields = (:connectID, :x, :y, :σ_x, :σ_y,
        :photons, :σ_photons, :bg, :σ_bg, :framenum, :datasetnum)

    return smld
end

"""
    Base.length(smld::SMLMData.SMLD2D)

Returns the number of localizations in `smld`.
"""
function Base.length(smld::SMLMData.SMLD2D)
    # Determine the length of `smld` by counting the number of frames.
    return Base.length(smld.framenum)
end

"""
    Base.getindex(smld::SMLMData.SMLD2D, ind)

Returns the requested localization from `smld`.
"""
function Base.getindex(smld::SMLMData.SMLD2D, ind)
    return SMLMData.isolatesmld(smld, ind)
end


# This file was generated, do not modify it. # hide
RandomForestRegressor = @load RandomForestRegressor pkg=DecisionTree
rfr = RandomForestRegressor()

rfr_m = machine(rfr, X, y);


using Test
using Distributed

addprocs(2)
using IncrementalInference
@everywhere using IncrementalInference

@testset "test multiprocess solveTree!" begin

fg = generateGraph_Kaess()

solveTree!(fg)

end


# Authors: Ali Siahkoohi, alisk@gatech.edu
# Date: March 2022

using DrWatson
@quickactivate :ConditionalNFs4Imaging

using InvertibleNetworks
using HDF5
using Random
using Statistics
using ProgressMeter
using PyPlot
using Seaborn
using LinearAlgebra
using Flux
using PyCall: @py_str

font_prop, sfmt = sef_plot_configs()
args = read_config("test_amortized_imaging.json")
args = parse_input_args(args)

py"""
from scipy.signal import hilbert
import numpy as np

def normalize_std(mu, sigma):
    analytic_mu = hilbert(mu, axis=1)
    return sigma*np.abs(analytic_mu)/(np.abs(analytic_mu)**2 + 5e+2), analytic_mu
"""
normalize_std(mu, sigma) = py"normalize_std"(mu, sigma)

max_epoch = args["max_epoch"]
lr = args["lr"]
lr_step = args["lr_step"]
batchsize = args["batchsize"]
n_hidden = args["n_hidden"]
depth = args["depth"]
sim_name = args["sim_name"]
if args["epoch"] == -1
    args["epoch"] = args["max_epoch"]
end
epoch = args["epoch"]

# Define raw data directory
mkpath(datadir("training-data"))
data_path = datadir("training-data", "training-pairs.h5")

# Download the dataset into the data directory if it does not exist
if isfile(data_path) == false
    run(`wget https://www.dropbox.com/s/53u8ckb9aje8xv4/'
        'training-pairs.h5 -q -O $data_path`)
end

# Load seismic images and create training and testing data
file = h5open(data_path, "r")
X_train = file["dm"][:, :, :, :]
Y_train = file["rtm"][:, :, :, :]

nx, ny, nc, nsamples = size(X_train)

AN_x = ActNorm(nsamples)
AN_y = ActNorm(nsamples)
X_train = AN_x.forward(X_train)
Y_train = AN_y.forward(Y_train)

# Dimensions after wavelet squeeze to increase no. of channels
nx = Int(nx / 2)
ny = Int(ny / 2)
n_in = Int(nc * 4)

# Create network
CH = NetworkConditionalHINT(n_in, n_hidden, depth)

# Loading the experiment—only network weights and training loss
loaded_keys = load_experiment(args, ["Params", "fval", "fval_eval", "train_idx"])
Params = loaded_keys["Params"]
fval = loaded_keys["fval"]
fval_eval = loaded_keys["fval_eval"]
train_idx = loaded_keys["train_idx"]
put_params!(CH, convert(Array{Any,1}, Params))

# test data pairs
idx = shuffle(setdiff(1:nsamples, train_idx))[1]
X_fixed = wavelet_squeeze(X_train[:, :, :, idx:idx])
Y_fixed = wavelet_squeeze(Y_train[:, :, :, idx:idx])

# Now select single fixed sample from all Ys
Zy_fixed = CH.forward_Y(Y_fixed)

# Draw new Zx, while keeping Zy fixed
n_samples = 1000
X_post = zeros(Float32, nx, ny, n_in, n_samples)
CH = CH |> gpu

test_batchsize = 4
test_loader = Flux.DataLoader(
    (randn(Float32, nx, ny, n_in, n_samples), repeat(Zy_fixed, 1, 1, 1, n_samples)),
    batchsize = test_batchsize,
    shuffle = false,
)

p = Progress(length(test_loader))

for (itr, (X, Y)) in enumerate(test_loader)
    Base.flush(Base.stdout)
    counter = (itr - 1) * test_batchsize + 1

    X = X |> gpu
    Y = Y |> gpu

    X_post[:, :, :, counter:(counter+size(X)[4]-1)] = (CH.inverse(X, Y)[1] |> cpu)
    ProgressMeter.next!(p)
end

X_post = wavelet_unsqueeze(X_post)
X_post = AN_x.inverse(X_post)
X_post[1:10, :, :, :] .= 0.0f0

# Some stats
X_post_mean = mean(X_post; dims = 4)
X_post_std = std(X_post; dims = 4)

X_fixed = wavelet_unsqueeze(X_fixed)
Y_fixed = wavelet_unsqueeze(Y_fixed)
Zy_fixed = wavelet_unsqueeze(Zy_fixed)

X_fixed = AN_x.inverse(X_fixed)
Y_fixed = AN_y.inverse(Y_fixed)
Y_fixed[1:10, :, :, :] .= 0.0f0

save_dict = @strdict max_epoch epoch lr lr_step batchsize n_hidden depth sim_name
save_path = plotsdir(sim_name, savename(save_dict; digits = 6))

spacing = [20.0, 12.5]
extent = [0.0, size(X_fixed, 1) * spacing[1], size(X_fixed, 2) * spacing[2], 0.0] / 1e3

signal_to_noise(xhat, x) = -20.0 * log(norm(x - xhat) / norm(x)) / log(10.0)

snr_list = []
for j = 1:n_samples
    push!(snr_list, signal_to_noise(X_post[:, :, :, j], X_fixed[:, :, :, 1]))
end
X_post_mean_snr = signal_to_noise(X_post_mean[:, :, :, 1], X_fixed[:, :, :, 1])
Y_fixed_snr = signal_to_noise(Y_fixed[:, :, :, 1], X_fixed[:, :, :, 1])

# Training loss
fig = figure("training logs", figsize = (7, 2.5))
if epoch == args["max_epoch"]
    plot(
        range(0, epoch, length = length(fval_eval)),
        fval_eval,
        color = "#addbc2",
        label = "validation loss",
    )
else
    plot(
        range(0, epoch, length = length(fval_eval[1:findfirst(fval_eval .== 0.0f0)-1])),
        fval_eval[1:findfirst(fval_eval .== 0.0f0)-1],
        color = "#addbc2",
        label = "validation loss",
    )
end
ticklabel_format(axis = "y", style = "sci", useMathText = true)
title("Negative log-likelihood")
ylabel("Validation objective")
xlabel("Epochs")
wsave(joinpath(save_path, "log.png"), fig)
close(fig)

# Plot the true model
fig = figure("x", figsize = (7.68, 4.8))
imshow(
    X_fixed[:, :, 1, 1],
    vmin = -1.5e3,
    vmax = 1.5e3,
    aspect = 1,
    cmap = "Greys",
    resample = true,
    interpolation = "lanczos",
    filterrad = 1,
    extent = extent,
)
title(L"High-fidelity image, $\mathbf{x}$")
colorbar(fraction = 0.03, pad = 0.01, format = sfmt)
grid(false)
xlabel("Horizontal distance (km)")
ylabel("Depth (km)")
safesave(joinpath(save_path, string(idx), "true_model.png"), fig)
close(fig)

# Plot the observed data
fig = figure("y", figsize = (7.68, 4.8))
imshow(
    Y_fixed[:, :, 1, 1],
    vmin = -1.5e6,
    vmax = 1.5e6,
    aspect = 1,
    cmap = "Greys",
    resample = true,
    interpolation = "lanczos",
    filterrad = 1,
    extent = extent,
)
title("RTM image")
colorbar(fraction = 0.03, pad = 0.01, format = sfmt)
grid(false)
xlabel("Horizontal distance (km)")
ylabel("Depth (km)")
safesave(joinpath(save_path, string(idx), "observed_data.png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("x_cm", figsize = (7.68, 4.8))
imshow(
    X_post_mean[:, :, 1, 1],
    vmin = -1.5e3,
    vmax = 1.5e3,
    aspect = 1,
    cmap = "Greys",
    resample = true,
    interpolation = "lanczos",
    filterrad = 1,
    extent = extent,
)
title("Conditional mean estimate")
colorbar(fraction = 0.03, pad = 0.01, format = sfmt)
grid(false)
xlabel("Horizontal distance (km)")
ylabel("Depth (km)")
safesave(joinpath(save_path, string(idx), "conditional_mean.png"), fig)
close(fig)

# Plot the pointwise standard deviation
fig = figure("x_std", figsize = (7.68, 4.8))
imshow(
    X_post_std[:, :, 1, 1],
    vmin = 40,
    vmax = 250.0,
    aspect = 1,
    cmap = "OrRd",
    resample = true,
    interpolation = "lanczos",
    filterrad = 1,
    extent = extent,
    norm = matplotlib.colors.LogNorm(),
)
title(L"Pointwise standard deviation, $\sigma$")
cp = colorbar(fraction = 0.03, pad = 0.01)
grid(false)
xlabel("Horizontal distance (km)")
ylabel("Depth (km)")
safesave(joinpath(save_path, string(idx), "pointwise_std.png"), fig)
close(fig)

for ns = 1:10
    fig = figure("x_cm", figsize = (7.68, 4.8))
    imshow(
        X_post[:, :, 1, ns],
        vmin = -1.5e3,
        vmax = 1.5e3,
        aspect = 1,
        cmap = "Greys",
        resample = true,
        interpolation = "lanczos",
        filterrad = 1,
        extent = extent,
    )
    title(L"Posterior sample, $\mathbf{x} \sim p(\mathbf{x} \mid \mathbf{y})$")
    colorbar(fraction = 0.03, pad = 0.01, format = sfmt)
    grid(false)
    xlabel("Horizontal distance (km)")
    ylabel("Depth (km)")
    safesave(joinpath(save_path, string(idx), "sample.png"), fig)
    close(fig)
end

# Plot the pointwise standard deviation
normalized_std, analytic_mu = normalize_std(X_post_mean[:, :, 1, 1], X_post_std[:, :, 1, 1])

fig = figure("x_std", figsize = (7.68, 4.8))
imshow(
    normalized_std,
    vmin = 3.0f-2,
    vmax = 1.2f0,
    aspect = 1,
    cmap = "OrRd",
    resample = true,
    interpolation = "hermite",
    filterrad = 1,
    extent = extent,
    norm = matplotlib.colors.LogNorm(),
)
title("Normalized pointwise standard deviation")
colorbar(fraction = 0.03, pad = 0.01, format = sfmt)
grid(false)
xlabel("Horizontal distance (km)")
ylabel("Depth (km)")
wsave(joinpath(save_path, string(idx), "normalized_pointwise_std.png"), fig)
close(fig)

println("SNR of conditional mean: ", X_post_mean_snr)
println("SNR of RTM image: ", Y_fixed_snr)


# This file is a part of ForwardDiffPullbacks.jl, licensed under the MIT License (MIT).

# ToDo: Use ProjectTo in rrules (requires ChainRulesCore >= v0.10.11).

struct FwdDiffPullbackThunk{F<:Base.Callable,T<:Tuple,i,U<:Any} <: ChainRulesCore.AbstractThunk
    f::F
    xs::T
    ΔΩ::U
end

function FwdDiffPullbackThunk(f::F, xs::T, ::Val{i}, ΔΩ::U) where {F<:Base.Callable,T<:Tuple,i,U<:Any}
    FwdDiffPullbackThunk{F,T,i,U}(f, xs, ΔΩ)
end

@inline function ChainRulesCore.unthunk(tnk::FwdDiffPullbackThunk{F,T,i,U}) where {F,T,i,U}
    forwarddiff_fwd_back(tnk.f, tnk.xs, Val(i), unthunk(tnk.ΔΩ))
end

# ToDo: Remove (obsolete with ChainRulesCore >= v0.10.):
(tnk::FwdDiffPullbackThunk)() = ChainRulesCore.unthunk(tnk)

Base.@generated function _forwarddiff_pullback_thunks(f::Base.Callable, xs::NTuple{N,Any}, ΔΩ::Any) where N
    Expr(:tuple, NoTangent(), (:(ForwardDiffPullbacks.FwdDiffPullbackThunk(f, xs, Val($i), ΔΩ)) for i in 1:N)...)
end

g_ΔΩ = nothing

@inline function ChainRulesCore.rrule(wrapped_f::WithForwardDiff, xs::Vararg{Any,N}) where N
    # @info "RUN ChainRulesCore.rrule(wrapped_f::WithForwardDiff, xs) with N = $N"
    f = wrapped_f.f
    y = f(xs...)
    with_fwddiff_pullback(ΔΩ) = begin
        #@info "XXXXXXXXXX" ΔΩ
        #global g_ΔΩ = ΔΩ
        _forwarddiff_pullback_thunks(f, xs, ΔΩ)
    end
    return y, with_fwddiff_pullback
end

struct FwdDiffBCPullbackThunk{F<:Base.Callable,T<:Tuple,i,U} <: ChainRulesCore.AbstractThunk
    f::F
    Xs::T
    ΔΩA::U
end

function FwdDiffBCPullbackThunk(f::F, Xs::T, ::Val{i}, ΔΩA::U) where {F<:Base.Callable,T<:Tuple,i,U}
    FwdDiffBCPullbackThunk{F,T,i,U}(f, Xs, ΔΩA)
end

@inline function ChainRulesCore.unthunk(tnk::FwdDiffBCPullbackThunk{F,T,i,U}) where {F,T,i,U}
    forwarddiff_bc_fwd_back(tnk.f, tnk.Xs, Val(i), unthunk(tnk.ΔΩA))
end

# ToDo: Remove (obsolete with ChainRulesCore >= v0.10.):
(tnk::FwdDiffBCPullbackThunk)() = ChainRulesCore.unthunk(tnk)

Base.@generated function _forwarddiff_bc_pullback_thunks(f::Base.Callable, Xs::NTuple{N,Any}, ΔΩA::Any) where N
    Expr(:tuple, NoTangent(), NoTangent(), (:(ForwardDiffPullbacks.FwdDiffBCPullbackThunk(f, Xs, Val($i), ΔΩA)) for i in 1:N)...)
end

function ChainRulesCore.rrule(::typeof(Base.broadcasted), wrapped_f::WithForwardDiff, Xs::Vararg{Any,N}) where N
    # @info "RUN ChainRulesCore.rrule(Base.broadcasted, wrapped_f::WithForwardDiff, Xs) with N = $N"
    f = wrapped_f.f
    y = broadcast(f, Xs...)
    bc_with_fwddiff_pullback(ΔΩA) = _forwarddiff_bc_pullback_thunks(f, Xs, ΔΩA)
    return y, bc_with_fwddiff_pullback
end


function utils_unit_tests()
    generate_symmetric_distance_matrix_tests()
    set_indices_tests()
    build_computed_data_tests()
    create_singleton_location_groups_tests()
    preprocess_instance_unit_tests()
    create_intersection_tws_tests()
end

function create_intersection_tws_tests()
    tws_1 = [RVRP.Range(10.0, 20.0), RVRP.Range(35.0, 40.0)]
    tws_2 = [RVRP.Range(15.0, 25.0), RVRP.Range(28.0, 36.0)]
    intersection = RVRP.ranges_intersection(tws_1, tws_2)
    @test intersection == [RVRP.Range(15.0, 20.0), RVRP.Range(35.0, 36.0)]

    tws_1 = [RVRP.Range(10.0, 100.0), RVRP.Range(110.0, 130.0)]
    tws_2 = [RVRP.Range(9.0, 25.0), RVRP.Range(28.0, 36.0), RVRP.Range(88.0, 120.0)]
    intersection = RVRP.ranges_intersection(tws_1, tws_2)
    @test intersection == [RVRP.Range(10.0, 25.0), RVRP.Range(28.0, 36.0), RVRP.Range(88.0, 100.0), RVRP.Range(110.0, 120.0)]
end

function generate_symmetric_distance_matrix_tests()
    xs = [rand(1:20) for i in 1:5]
    ys = [rand(1:20) for i in 1:5]
    matrix = RVRP.generate_symmetric_distance_matrix(xs, ys)
    @test size(matrix) == (5,5)
    for i in 1:5
        @test matrix[i,i] == 0.0
        for j in i+1:5
            @test matrix[i,j] == matrix[j,i]
        end
    end
end

function set_indices_tests()
    data = RVRP.RvrpInstance()
    data.locations = [RVRP.Location(
        id = "", index = rand(-100:-1)
    ) for i in 1:10]
    RVRP.set_indices(data)
    for loc_idx in 1:length(data.locations)
        @test data.locations[loc_idx].index == loc_idx
    end
end

function build_computed_data_tests()
    data = RVRP.RvrpInstance()
    data.locations = [RVRP.Location(
        id = string("loc_", i), index = rand(-100:-1)
    ) for i in 1:9]
    data.vehicle_categories = [RVRP.VehicleCategory(
        id = string("vc_", i)
    ) for i in 1:9]
    data.vehicle_sets = [RVRP.HomogeneousVehicleSet(
        id = string("vs_", i)
    ) for i in 1:12]
    data.product_specification_classes = [RVRP.ProductSpecificationClass(
        id = string("psc_", i)
    ) for i in 1:5]
    data.requests = [RVRP.Request(
        id = string("req_", i)
    ) for i in 10:13]
    RVRP.set_indices(data)
    data.location_groups = RVRP.create_singleton_location_groups(data.locations)
    data.vehicle_categories[1].capacities.of_vehicle = Dict{String,Float64}(
        "weird_name_1" => 10, "wolow_34" => 12, "volume" => 15
    )
    data.vehicle_categories[2].capacities.of_vehicle = Dict{String,Float64}(
        "weird_name_1" => 10, "wolow_34" => 12, "new_name_1" => 15
    )
    data.product_specification_classes[1].capacity_consumptions = Dict{String,Tuple{Float64,Float64}}(
        "weird_name_1" => (1,10), "new_name_2" => (1,12), "wolow_34" => (1,15)
    )
    data.vehicle_categories[1].properties.of_vehicle = Dict{String,Float64}(
        "prop_1" => 10, "prop_34" => 12, "prop_3" => 15
    )
    data.vehicle_categories[2].properties.of_vehicle = Dict{String,Float64}(
        "prop_1" => 10, "prop_34" => 12, "new_prop_1" => 15
    )
    data.product_specification_classes[1].property_requirements = Dict{String,Float64}(
        "prop_1" => 10, "new_prop_2" => 12, "prop_34" => 15
    )
    computed_data = RVRP.build_computed_data(data)
    for (k,v) in computed_data.location_id_2_index
        @test v >= 1
        @test v <= 9
    end
    for (k,v) in computed_data.location_group_id_2_index
        @test v >= 1
        @test v <= 9
    end
    for (k,v) in computed_data.vehicle_category_id_2_index
        @test v >= 1
        @test v <= 9
    end
    for (k,v) in computed_data.vehicle_set_id_2_index
        @test v >= 1
        @test v <= 12
    end
    for (k,v) in computed_data.product_specification_class_id_2_index
        @test v >= 1
        @test v <= 5
    end
    for (k,v) in computed_data.request_id_2_index
        @test v >= 1
        @test v <= 4
    end
    for (k,v) in computed_data.capacity_id_2_index
        @test v >= 1
        @test v <= 5
    end
    for (k,v) in computed_data.property_id_2_index
        @test v >= 1
        @test v <= 5
    end
end

function create_singleton_location_groups_tests()
    data = RVRP.RvrpInstance()
    data.locations = [RVRP.Location(
        index = rand(-100:-1)
    ) for i in 1:10]
    def_loc_groups = RVRP.create_singleton_location_groups(data.locations)
    for i in 1:length(data.locations)
        @test string(data.locations[i].id, "_loc_group") == def_loc_groups[i].id
    end
end

function preprocess_instance_unit_tests()
    data = RVRP.RvrpInstance()
    RVRP.preprocess_instance(data)
    @test length(data.locations) == 0
    @test length(data.location_groups) == 0
    @test length(data.product_compatibility_classes) == 1
    @test data.product_compatibility_classes[1].id == "default_id"
    @test length(data.product_sharing_classes) == 1
    @test data.product_sharing_classes[1].id == "default_id"
    @test length(data.product_specification_classes) == 1
    @test data.product_specification_classes[1].id == "default_id"
    @test length(data.requests) == 0
    @test length(data.vehicle_categories) == 1
    @test data.vehicle_categories[1].id == "default_id"
    @test length(data.vehicle_sets) == 0
end


using Random, Distributions, DelimitedFiles
Random.seed!()

# Here I define the global quantities used by all pivot MCMC run which are the pivot moves
const pivot_moves = Array{Int,3}(zeros(47, 2, 3)) # Tensor that defines all the pivot moves of the octahedral symmetry group
const perms = [[1 2 3]; [1 3 2]; [2 1 3]; [2 3 1]; [3 1 2]; [3 2 1]]
const tr_signs = [[1 1 1]; [1 1 -1]; [1 -1 1]; [1 -1 -1]; [-1 1 1]; [-1 1 -1]; [-1 -1 1]; [-1 -1 -1]]
p_moves_count = 1
for i in 1:6
    for j in 1:8
        if i==1 && j==1
            continue
        else
            pivot_moves[p_moves_count, 1, :] .= perms[i,:]
            pivot_moves[p_moves_count, 2, :] .= tr_signs[j,:]

            global p_moves_count += 1
        end
    end
end

#########################################################################
#########################################################################

# This structure contains all I need to compute the Hamiltonian
struct Magnetic_polymer{T1<:Int, T2<:AbstractFloat}
    n_mono::T1
    inv_temp::T2
    J::T2
    alpha::T2
    energy::T2
    spins::Array{T1}
    fields::Array{T2}
    coord::Array{T1,2}         # Coordinates of the monomers
    trial_coord::Array{T1,2}
    hash_saw::Dict{Int64, Int64}
    neighbours::Array{T1, 2}
    trial_neighbours::Array{T1, 2}
end

function Magnetic_polymer(n_mono::T1, inv_temp::T2, J::T2, alpha::T2) where {T1<:Int, T2<:AbstractFloat}
    Magnetic_polymer(n_mono, inv_temp, J, alpha, 0.0, zeros(Int, n_mono), zeros(n_mono), zeros(Int, 3,n_mono),
    zeros(Int, 3,n_mono), Dict{Int, Int}(), zeros(Int, 7,n_mono), zeros(Int, 7,n_mono))
end

function initialize_poly!(poly::Magnetic_polymer, spin_values::Array{Int})
    poly.spins .= rand(spin_values, poly.n_mono)
    #poly.fields .= (rand(poly.n_mono).*2) .-1
    for i_mono in 1:poly.n_mono
        poly.coord[1, i_mono] = i_mono-1
    end
    for i_mono in 1:poly.n_mono
        for j in 1:3
            poly.trial_coord[j, i_mono] = poly.coord[j, i_mono]
        end
    end
end

function initialize_poly!(poly::Magnetic_polymer)
    initialize_poly!(poly,[0,1])
end

#The following method of the "initialize_poly!()" function allows initialization from file
function initialize_poly!(poly::Magnetic_polymer, file_name::String)
    io = open(file_name, "r")
    data = readdlm(io,Int64)
    close(io)
    for i_mono in 1:poly.n_mono
        for j in 1:3
            poly.coord[j,i_mono] = data[i_mono,j]
        end
        poly.spins[i_mono] = data[i_mono, 4]
    end
end

function set_fields!(polymers::Array{Magnetic_polymer}, ff::Array{Float64})
    for polymer in polymers
        for i_mono in 1:polymer.n_mono
            polymer.fields[i_mono] = ff[i_mono]
        end
    end
end

# This structure stores series of observables I'm computing along the MCMC
struct MC_data{T1<:Int, T2<:AbstractFloat}
    n_steps::T1
    energies::Array{T2}
    magnetization::Array{T2}
    rg2::Array{T2}
end

function MC_data(n_steps::Int)
    MC_data(n_steps, zeros(n_steps), zeros(n_steps), zeros(n_steps))
end

#########################################################################
#########################################################################

function try_pivot!(k::Int, move::Int, coo::Array{Int,2}, t_coo::Array{Int,2}, n_mono::Int)
    for i_mono in 1:k
        for j in 1:3
            #pol.trial_coord[j, i_mono] = pol.coord[j, i_mono]
            t_coo[j, i_mono] = coo[j, i_mono]
        end
    end
    for i_mono in k+1:n_mono
        for j in 1:3
            t_coo[j,i_mono] = (pivot_moves[move,2,j]*(coo[pivot_moves[move,1,j],i_mono]
                                         - coo[pivot_moves[move,1,j],k]) + coo[j,k])
        end
    end
end

#########################################################################

function checksaw!(k::Int, pol::Magnetic_polymer, a::Int)
    L = max(k, pol.n_mono-k) + 1
    is_saw = true
    t = 0
    while(t<L && is_saw)
        if k+t <= pol.n_mono
            saw_key = pol.trial_coord[1,k+t]*a^2 + pol.trial_coord[2,k+t]*a + pol.trial_coord[3,k+t]
            if haskey(pol.hash_saw, saw_key)
                is_saw = false
            else
                pol.hash_saw[saw_key] = k+t
            end
        end
        if k-t >= 1 && t!=0
            saw_key = pol.trial_coord[1,k-t]*a^2 + pol.trial_coord[2,k-t]*a + pol.trial_coord[3,k-t]
            if haskey(pol.hash_saw, saw_key)
                is_saw = false
            else
                pol.hash_saw[saw_key] = k-t
            end
        end
        t += 1
    end
    return is_saw
end

#########################################################################
## poi quando si faranno le mosse locali non si userà questa funnzione perché
## non ci sarà bisogno di ricalcolare TUTTA la lista dei vicini
function compute_neighbours!(coo::Array{Int,2}, near::Array{Int,2}, dic::Dict{Int,Int}, a::Int, n_mono::Int)
    for i_mono in 1:n_mono
        near[1, i_mono] = 0
        for j in 2:7
            near[j, i_mono] = -1
        end
        saw_key_base = coo[1,i_mono]*a^2 + coo[2,i_mono]*a + coo[3,i_mono]
        for j in 1:3
            for k in -1:2:1
                #coo[j,i_mono] += k
                #saw_key = coo[1,i_mono]*a^2 + coo[2,i_mono]*a + coo[3,i_mono]
                saw_key = saw_key_base + k*a^(3-j)
                if haskey(dic, saw_key)
                    near[1, i_mono] += 1
                    n_neigh = near[1, i_mono]+1
                    near[n_neigh, i_mono] = dic[saw_key]
                end
                #coo[j,i_mono] -= k
            end
        end
    end
end

#########################################################################

function gyration_radius_squared(pol::Magnetic_polymer)
    rcm = zeros(3)
    rg2 = 0.0
    for i in 1:pol.n_mono
        for j in 1:3
            #rg2 += sum((view(pol.coord,:,i) .- rcm).^2)
            rg2 += pol.coord[j,i]^2
            rcm[j] += pol.coord[j,i]
        end
    end
    r = 0.0
    for j in 1:3
        r += rcm[j]^2
    end
    return rg2/pol.n_mono - r/(pol.n_mono)^2
end

#########################################################################

function compute_new_energy(pol::Magnetic_polymer, near::Array{Int,2}, ff::Array{Float64}, s::Array{Int})
    ene = 0.0
    ene_J = 0.0
    for i_mono in 1:pol.n_mono
        ene -= ff[i_mono] * s[i_mono]
        for j in 1:near[1, i_mono]
            ene_J -= s[i_mono] * s[near[j+1,i_mono]] * pol.J * 0.5
        end
    end
    return ene*(1-pol.alpha) + ene_J*pol.alpha
end

#########################################################################
#########################################################################
### In this section I there are the functions to perform local moves on the self-avoiding walk

function single_bead_flip!(t_coo::Array{Int,2}, dic::Dict{Int, Int}, n_mono::Int, a::Int)
    mono = rand(2:n_mono-1)
    dist = 0
    for j in 1:3
        dist += (t_coo[j,mono+1]-t_coo[j,mono-1])^2
    end
    if dist == 2
        saw_key = 0
        for j in 1:3
            saw_key += (t_coo[j,mono-1]+t_coo[j,mono+1]-t_coo[j,mono])*a^(3-j)
        end

        if !haskey(dic,saw_key)
            delete!(dic, t_coo[1,mono]*a^2 + t_coo[2,mono]*a + t_coo[3,mono])
            for j in 1:3
                t_coo[j,mono] = t_coo[j,mono-1]+t_coo[j,mono+1]-t_coo[j,mono]
            end
            dic[saw_key] = mono
        end
    end
end

#########################################################################

function crankshaft_180!(t_coo::Array{Int,2}, dic::Dict{Int, Int}, n_mono::Int, a::Int)
    mono = rand(1:n_mono-3)
    dist = 0
    for j in 1:3
        dist += (t_coo[j,mono+3]-t_coo[j,mono])^2
    end
    if dist==1
        saw_key1 = 0
        saw_key2 = 0
        for j in 1:3
            saw_key1 += (2*t_coo[j,mono] - t_coo[j,mono+1])*a^(3-j)
            saw_key2 += (2*t_coo[j,mono+3] - t_coo[j,mono+2])*a^(3-j)
        end
        if !haskey(dic,saw_key1) && !haskey(dic,saw_key2)
            delete!(dic, t_coo[1,mono+1]*a^2 + t_coo[2,mono+1]*a + t_coo[3,mono+1])
            delete!(dic, t_coo[1,mono+2]*a^2 + t_coo[2,mono+2]*a + t_coo[3,mono+2])
            for j in 1:3
                t_coo[j,mono+1] = 2*t_coo[j,mono] - t_coo[j,mono+1]
                t_coo[j,mono+2] = 2*t_coo[j,mono+3] - t_coo[j,mono+2]
            end
            dic[saw_key1] = mono+1
            dic[saw_key2] = mono+2
        end
    end
end

#########################################################################

function crankshaft_90_270!(t_coo::Array{Int,2}, dic::Dict{Int, Int}, n_mono::Int, a::Int, new_coord1::Array{Int}, new_coord2::Array{Int})
    mono = rand(1:n_mono-3)
    dist = 0
    axis = 0
    for j in 1:3
        dist += (t_coo[j,mono+3]-t_coo[j,mono])^2
    end
    if dist == 1
        orient = 0
        for j in 1:3
            if (t_coo[j,mono+3]-t_coo[j,mono]) != 0
                axis = j
                orient = t_coo[j,mono+3]-t_coo[j,mono]
                break
            end
        end
        p=0
        t=0
        if axis==1
            p=2
            t=rand((3,2))
        elseif axis==2
            p=6
            t=rand((2,5))
        elseif axis==3
            p=3
            t=rand((5,3))
        else
            println("Invalid Rotation Axis!")
        end

        for j in 1:3
            new_coord1[j] = t_coo[j,mono] + tr_signs[t,j]*(t_coo[perms[p,j],mono+1]-t_coo[perms[p,j],mono])*orient
            new_coord2[j] = t_coo[j,mono+3] + tr_signs[t,j]*(t_coo[perms[p,j],mono+2]-t_coo[perms[p,j],mono+3])*orient
        end

        saw_key1 = new_coord1[1]*a^2 + new_coord1[2]*a + new_coord1[3]
        saw_key2 = new_coord2[1]*a^2 + new_coord2[2]*a + new_coord2[3]
        if !haskey(dic,saw_key1) && !haskey(dic,saw_key2)
            delete!(dic, t_coo[1,mono+1]*a^2 + t_coo[2,mono+1]*a + t_coo[3,mono+1])
            delete!(dic, t_coo[1,mono+2]*a^2 + t_coo[2,mono+2]*a + t_coo[3,mono+2])
            for j in 1:3
                t_coo[j,mono+1] = new_coord1[j]
                t_coo[j,mono+2] = new_coord2[j]
            end
            dic[saw_key1] = mono+1
            dic[saw_key2] = mono+2
        end
    end
end

#########################################################################
#########################################################################
function spins_MC!(pol::Magnetic_polymer, n_flips::Int, ff::Array{Float64}, s::Array{Int}, near::Array{Int,2})
    delta_tot = 0.0
    for i_flip in 1:n_flips
        flip_candidate = rand(1:pol.n_mono)
        local_ene = 0.0
        local_ene -= ff[flip_candidate] * s[flip_candidate] * (1-pol.alpha)
        for j in 1:near[1,flip_candidate]
            local_ene -= pol.J * pol.alpha * s[flip_candidate] * s[near[j+1,flip_candidate]]
        end

        #s[flip_candidate]==1 ? trial_spin_value=-1 : trial_spin_value=1
        s[flip_candidate]==1 ? trial_spin_value=0 : trial_spin_value=1

        trial_local_ene = 0.0
        trial_local_ene -= ff[flip_candidate] * trial_spin_value * (1-pol.alpha)
        for j in 1:near[1,flip_candidate]
            trial_local_ene -= pol.J * pol.alpha * trial_spin_value * s[near[j+1,flip_candidate]]
        end

        delta_ene = trial_local_ene - local_ene
        acc = 0.0

        delta_ene<=0 ? acc=1.0 : acc=exp(-delta_ene*pol.inv_temp)

        if acc > rand()
            delta_tot += delta_ene
            s[flip_candidate] = trial_spin_value
        end
    end
    return delta_tot
end

#########################################################################
#########################################################################

#= This is the version of MC_run where I don't save the configurations of the spins (only the last one)
   This is useful if I just want to perform the simulation and I am interested only in storing the
    values of some observables along the markov chain. If instead I want to carry out some simulator
    based inference task, I will also need to save spin configurations. Expoiting julia's
    multiple dispatch feature I can always use the same "MC_run!()" function, the specific method
    is inferred by the compiler depending on the list of argument types. "MC_run_save!()" is the function
    that save the configurations and it will also be made into a MC_run!() module=#

function MC_run_base!(pol::Magnetic_polymer, traj::MC_data, start::Int, finish::Int)
    for i_mono in 1:pol.n_mono
        for j in 1:3
            pol.trial_coord[j,i_mono] = pol.coord[j,i_mono]
        end
    end
    new_coord1 = zeros(Int,3)
    new_coord2 = zeros(Int,3)
    n_acc = 0
    n_pivots = 0
    #=
      I represent the possible positions of the monomers in the cubic lattice
      as seeing the vector of coordinates as the representation in base "hash_base".
      e.g. (xi,yi,zi) --> xi*hash_base^2 + yi*hash_base + zi*1
      Since the polymer is fixed in the origin the coordinates are ≤ n_mono in modulus, so if I take hash_base
      >= 2*n_mono + 1, the mapping I think it should be 1-to-1. This speeds up the code significantly as the
      hashing of integer valued keys for the dictionary containing a certain polymer config. is much
      cheaper than hashing keys which are arrays sliced from a matrix
    =#
    hash_base = 2*pol.n_mono + 1

    checksaw!(1, pol, hash_base)
    compute_neighbours!(pol.coord, pol.neighbours, pol.hash_saw, hash_base, pol.n_mono)
    energy = compute_new_energy(pol,pol.neighbours,pol.fields, pol.spins)

    traj.energies[start] = energy
    traj.magnetization[start] = mean(pol.spins)
    traj.rg2[start] = gyration_radius_squared(pol)

    empty!(pol.hash_saw)

    for i_step in (start+1):finish
        pivot = rand(2:pol.n_mono-1)
        p_move = rand(1:47)
        if i_step%50000 == 0
            println("Pivot: ", pivot)
            println("Attmpted move: ", p_move)
            println("step number: ", i_step, "\n")
        end
        try_pivot!(pivot, p_move, pol.coord, pol.trial_coord, pol.n_mono)
        still_saw = checksaw!(pivot, pol, hash_base)

        # If the pivot move was unsuccessful try next move on the previous configuration
        if !still_saw
            empty!(pol.hash_saw)
            for i_mono in 1:pol.n_mono
                for j in 1:3
                    pol.trial_coord[j,i_mono] = pol.coord[j,i_mono]
                end
                pol.hash_saw[pol.trial_coord[1,i_mono]*hash_base^2+pol.trial_coord[2,i_mono]*hash_base+pol.trial_coord[3,i_mono]] = i_mono
            end
        else
            n_pivots += 1
        end

        for i_local in 1:pol.n_mono
            mv = rand(1:4)
            if mv==1
                single_bead_flip!(pol.trial_coord,pol.hash_saw,pol.n_mono,hash_base)
            elseif mv==2
                crankshaft_180!(pol.trial_coord,pol.hash_saw,pol.n_mono,hash_base)
            else
                crankshaft_90_270!(pol.trial_coord,pol.hash_saw,pol.n_mono,hash_base,new_coord1,new_coord2)
            end
        end

        compute_neighbours!(pol.trial_coord, pol.trial_neighbours, pol.hash_saw, hash_base, pol.n_mono)
        trial_energy = compute_new_energy(pol,pol.trial_neighbours,pol.fields, pol.spins)
        delta_energy = trial_energy - energy
        acc = 0.0
        delta_energy<=0 ? acc=1.0 : acc=exp(-delta_energy*pol.inv_temp)
        if acc > rand()
            for i_mono in 1:pol.n_mono
                for j in 1:3
                    pol.coord[j,i_mono] = pol.trial_coord[j,i_mono]
                end
            end
            for i_mono in 1:pol.n_mono
                for j in 1:7
                    pol.neighbours[j,i_mono] = pol.trial_neighbours[j,i_mono]
                end
            end
            energy = trial_energy
            n_acc += 1
        end

        energy += spins_MC!(pol, pol.n_mono,pol.fields,pol.spins,pol.neighbours) ## spins_MC! return the total energy variation of the accepted spin n_flips
        traj.rg2[i_step] = gyration_radius_squared(pol)
        traj.magnetization[i_step] = mean(pol.spins)
        traj.energies[i_step] = energy

        # This next if ... end can be removed. I kept as a sanity check in an early stage
        # when I had problems with the computation of the nearest neighbours
        #=if !isinteger(energy)
            println("Non integer energy:  ", energy)
            println(pol.coord)
            println(pol.neighbours .- pol.trial_neighbours)
            println(pol.trial_neighbours)
            break
        end=#

        empty!(pol.hash_saw)
    end
end

#########################################################################
#########################################################################
#########################################################################
#########################################################################
#########################################################################
#########################################################################

# This function is like MC_run_base! but also save configurations
function MC_run_save!(pol::Magnetic_polymer, traj::MC_data, start::Int, finish::Int, spins_configs::Matrix{Int}, sample_lag::Int, n_samples::Int)
    for i_mono in 1:pol.n_mono
        for j in 1:3
            pol.trial_coord[j,i_mono] = pol.coord[j,i_mono]
        end
    end
    new_coord1 = zeros(Int,3)
    new_coord2 = zeros(Int,3)
    n_acc = 0
    n_pivots = 0

    hash_base = 2*pol.n_mono + 1

    checksaw!(1, pol, hash_base)
    compute_neighbours!(pol.coord, pol.neighbours, pol.hash_saw, hash_base, pol.n_mono)
    energy = compute_new_energy(pol,pol.neighbours,pol.fields, pol.spins)

    traj.energies[start] = energy
    traj.magnetization[start] = mean(pol.spins)
    traj.rg2[start] = gyration_radius_squared(pol)

    empty!(pol.hash_saw)

    if start%sample_lag == 0
        i_sample = div(start,sample_lag)
        for i_mono in 1:pol.n_mono
            spins_configs[i_mono, i_sample]  = pol.spins[i_mono]
        end
    end

    for i_step in (start+1):min(finish,sample_lag*n_samples)
        pivot = rand(2:pol.n_mono-1)
        p_move = rand(1:47)
        if i_step%50000 == 0
            println("Pivot: ", pivot)
            println("Attmpted move: ", p_move)
            println("step number: ", i_step, "\n")
        end
        try_pivot!(pivot, p_move, pol.coord, pol.trial_coord, pol.n_mono)
        still_saw = checksaw!(pivot, pol, hash_base)

        # If the pivot move was unsuccessful try next move on the previous configuration
        if !still_saw
            empty!(pol.hash_saw)
            for i_mono in 1:pol.n_mono
                for j in 1:3
                    pol.trial_coord[j,i_mono] = pol.coord[j,i_mono]
                end
                pol.hash_saw[pol.trial_coord[1,i_mono]*hash_base^2+pol.trial_coord[2,i_mono]*hash_base+pol.trial_coord[3,i_mono]] = i_mono
            end
        else
            n_pivots += 1
        end

        for i_local in 1:pol.n_mono
            mv = rand(1:4)
            if mv==1
                single_bead_flip!(pol.trial_coord,pol.hash_saw,pol.n_mono,hash_base)
            elseif mv==2
                crankshaft_180!(pol.trial_coord,pol.hash_saw,pol.n_mono,hash_base)
            else
                crankshaft_90_270!(pol.trial_coord,pol.hash_saw,pol.n_mono,hash_base,new_coord1,new_coord2)
            end
        end

        compute_neighbours!(pol.trial_coord, pol.trial_neighbours, pol.hash_saw, hash_base, pol.n_mono)
        trial_energy = compute_new_energy(pol,pol.trial_neighbours,pol.fields, pol.spins)
        delta_energy = trial_energy - energy
        acc = 0.0
        delta_energy<=0 ? acc=1.0 : acc=exp(-delta_energy*pol.inv_temp)
        if acc > rand()
            for i_mono in 1:pol.n_mono
                for j in 1:3
                    pol.coord[j,i_mono] = pol.trial_coord[j,i_mono]
                end
            end
            for i_mono in 1:pol.n_mono
                for j in 1:7
                    pol.neighbours[j,i_mono] = pol.trial_neighbours[j,i_mono]
                end
            end
            energy = trial_energy
            n_acc += 1
        end

        energy += spins_MC!(pol, pol.n_mono,pol.fields,pol.spins,pol.neighbours) ## spins_MC! return the total energy variation of the accepted spin n_flips
        traj.rg2[i_step] = gyration_radius_squared(pol)
        traj.magnetization[i_step] = mean(pol.spins)
        traj.energies[i_step] = energy

        #=if !isinteger(energy)
            println("Non integer energy:  ", energy)
            println(pol.coord)
            println(pol.neighbours .- pol.trial_neighbours)
            println(pol.trial_neighbours)
            break
        end=#
        empty!(pol.hash_saw)
        if i_step%sample_lag == 0
            i_sample = div(i_step,sample_lag)
            for i_mono in 1:pol.n_mono
                spins_configs[i_mono, i_sample]  = pol.spins[i_mono]
            end
        end
    end
end

#########################################################################
#########################################################################
#########################################################################
# Multiple Markov Chains

function MMC_run_base!(polymers::Array{Magnetic_polymer}, trajs::Array{MC_data},
                    n_strides::Int, stride::Int, inv_temps::Array{Float64})
    n_temps = length(inv_temps)
    accepted_swaps = 0
    temp_coord = zeros(Int,3)

    for i_temp in 1:n_temps
        MC_run!(polymers[i_temp], trajs[i_temp],1,stride)
    end

    for i_strides in 2:n_strides
        swap = rand(1:n_temps-1)
        delta_ene = (inv_temps[swap] - inv_temps[swap+1]) * (trajs[swap+1].energies[(i_strides-1)*stride] - trajs[swap].energies[(i_strides-1)*stride])
        alpha_swap = 0.0
        delta_ene<=0 ? alpha_swap=1.0 : alpha_swap=exp(-delta_ene)
        if alpha_swap >= rand()
            for i_mono in 1:polymers[1].n_mono
                for j in 1:3
                   temp_coord[j] = polymers[swap].coord[j,i_mono]
                   polymers[swap].coord[j,i_mono] = polymers[swap+1].coord[j,i_mono]
                   polymers[swap+1].coord[j,i_mono] =  temp_coord[j]
                end
                temp_spin = polymers[swap].spins[i_mono]
                polymers[swap].spins[i_mono] = polymers[swap+1].spins[i_mono]
                polymers[swap+1].spins[i_mono] = temp_spin
            end
            accepted_swaps += 1
        end

        for i_temp in 1:n_temps
            MC_run!(polymers[i_temp], trajs[i_temp],(i_strides-1)*stride+1,i_strides*stride)
        end
    end
    #println("Accepted_swaps: ", accepted_swaps)
end

function MMC_run_save!(polymers::Array{Magnetic_polymer}, trajs::Array{MC_data},
                    n_strides::Int, stride::Int, inv_temps::Array{Float64}, spins_configs::Matrix{Int}, sample_lag::Int, n_samples::Int)
    n_temps = length(inv_temps)
    accepted_swaps = 0
    temp_coord = zeros(Int,3)

    MC_run!(polymers[1], trajs[1],1,stride, spins_configs,sample_lag, n_samples)
    # Only the lowest temperature is the system we're using for our likelihood approx
    # the chains at higher temps are only used to "fluidify" the chain of interest
    for i_temp in 2:n_temps
        MC_run!(polymers[i_temp], trajs[i_temp],1,stride, sample_lag, n_samples)
    end

    for i_strides in 2:n_strides
        swap = rand(1:n_temps-1)
        delta_ene = (inv_temps[swap] - inv_temps[swap+1]) * (trajs[swap+1].energies[i_strides*stride] - trajs[swap].energies[i_strides*stride])
        alpha_swap = 0.0
        delta_ene<=0 ? alpha_swap=1.0 : alpha_swap=exp(-delta_ene)
        if alpha_swap >= rand()
            for i_mono in 1:polymers[1].n_mono
                for j in 1:3
                   temp_coord[j] = polymers[swap].coord[j,i_mono]
                   polymers[swap].coord[j,i_mono] = polymers[swap+1].coord[j,i_mono]
                   polymers[swap+1].coord[j,i_mono] =  temp_coord[j]
                end
                temp_spin = polymers[swap].spins[i_mono]
                polymers[swap].spins[i_mono] = polymers[swap+1].spins[i_mono]
                polymers[swap+1].spins[i_mono] = temp_spin
            end
            accepted_swaps += 1
        end
        MC_run!(polymers[1], trajs[1],(i_strides-1)*stride+1,i_strides*stride, spins_configs,sample_lag, n_samples)
        # Only the lowest temperature is the system we're using for our likelihood approx
        # the chains at higher temps are only used to "fluidify" the chain of interest
        for i_temp in 2:n_temps
            MC_run!(polymers[i_temp], trajs[i_temp],(i_strides-1)*stride+1,i_strides*stride, sample_lag, n_samples)
        end
    end
    #println("Accepted_swaps: ", accepted_swaps)
end


function ⪇(x::F, y::F) where {F <: AbstractFloat}
    return x < y && eps(F) < abs(x - y)
end

function ⪉(x::F, y::F) where {F <: AbstractFloat}
    return x < y && x ≉  y
end


abstract type AbstractVLAE end

function elbo(m::AbstractVLAE, x::AbstractArray{T,4}) where T
    # encoder pass - KL divergence
    μzs_σzs = _encoded_mu_vars(m, x)
    zs = map(y->rptrick(y...), μzs_σzs)
    kldl = sum(map(y->Flux.mean(kld(y...)), μzs_σzs))

    # decoder pass - logpdf
    lpdf = Flux.mean(logpdf(m, x, zs...))

    -kldl + lpdf
end

function _encoded_mu_vars(m::AbstractVLAE, x)
    nl = length(m.e)

    h = x
    mu_vars = map(1:nl) do i
        h = m.e[i](h)
        h, μz, σz = _extract_mu_var(m.f, h, i)
        (μz, σz)
    end
    mu_vars
end

function _extract_mu_var(f, h::AbstractArray{T,4}, i) where T
    nch = floor(Int,size(h,3)/2)
    μz, σz = mu_var(f[i](h[:,:,(nch+1):end,:]))
    h = h[:,:,1:nch,:]
    h, μz, σz
end

function _extract_mu_var(f, h::AbstractArray{T,2}, i) where T
    nch = floor(Int,size(h,1)/2)
    μz, σz = mu_var(f[i](h[(nch+1):end,:]))
    h = h[1:nch,:]
    h, μz, σz
end

function _decoder_out(m::AbstractVLAE, zs...)
    nl = length(m.d)
    @assert length(zs) == nl

    h = m.g[1](zs[end])
    # now, propagate through the decoder
    for i in 1:nl-1
        h1 = m.d[i](h)
        h2 = m.g[i+1](zs[end-i])
        h = _cat_arrays(h1, h2)
    end
    h = m.d[end](h)
end
function _decoded_mu_var(m::AbstractVLAE, zs...)
    h = _decoder_out(m, zs...)
    μx, σx = mu_var(h)
end
function _decoded_mu_var1(m::AbstractVLAE, zs...)
    h = _decoder_out(m, zs...)
    μx, σx = mu_var1(h)
end
_cat_arrays(h1::AbstractArray{T,4}, h2::AbstractArray{T,4}) where T = cat(h1, h2, dims=3)
_cat_arrays(h1::AbstractArray{T,2}, h2::AbstractArray{T,2}) where T = cat(h1, h2, dims=1)

# encoding functions
function _batched_encs(f, m, x, batchsize, args...; kwargs...)
    local encs
    @suppress begin
        encs = map(y->f(gpu(m), gpu(y), args...; kwargs...),
            Flux.Data.DataLoader(x, batchsize=batchsize))
    end
    [cat([y[i] for y in encs]..., dims=2) for i in 1:length(encs[1])]
end
function _batched_encs(f, m, x::AbstractArray{T,2}, batchsize, args...; kwargs...) where T
    local encs
    @suppress begin
        encs = map(y->cpu(f(m, y, args...; kwargs...)),
            Flux.Data.DataLoader(x, batchsize=batchsize))
    end
    [cat([y[i] for y in encs]..., dims=2) for i in 1:length(encs[1])]
end
function _encode_ith(m::AbstractVLAE, x, i::Int; mean=false)
    μzs_σzs = _encoded_mu_vars(m, x)
    if mean
        return [μzs_σzs[i][1]]
    else
        return [rptrick(μzs_σzs[i]...)]
    end
end
_encode_ith_batched(m::AbstractVLAE, x, i::Int; mean=false, batchsize=128) =
    _batched_encs(_encode_ith, m, x, batchsize, i; mean=mean)[1]
function _encode_all(m::AbstractVLAE, x; mean=false)
    if mean
        return Tuple(map(y->y[1], _encoded_mu_vars(m, x)))
    else
        return Tuple(map(y->rptrick(y...), _encoded_mu_vars(m, x)))
    end
end

"""
    encode(m::AbstractVLAE, x, [i::Int]; mean=false, batchsize=128)

Encoding in the i-th latent layer. If i is not given, all encodings are returned instead.
"""
encode(m::AbstractVLAE, x, i::Int; batchsize::Int=128, mean=false) =
    _encode_ith_batched(m , x, i; batchsize=batchsize, mean=mean)
encode(m::AbstractVLAE, x; batchsize::Int=128, mean=false) =
    _batched_encs(_encode_all, m, x, batchsize; mean=mean)
encode_mean(args...; kwargs...) = encode(args...; mean=true, kwargs...)
encode_all(m::AbstractVLAE, x; kwargs...) =
     encode(m, x; kwargs...)

# decoding functions
function decode(m::AbstractVLAE, zs...)
    # this is for 2D data
    if length(m.xdim) == 1
        μx, σx = _decoded_mu_var(m, zs...)
        return rptrick(μx, σx)
    end
    # this is for images
    if m.xdist == Val(:gaussian)
        μx, σx = _decoded_mu_var1(m, zs...)
        return devectorize(rptrick(μx, σx), m.xdim...)
    else # bernoulli
        return _decoder_out(m, zs...)
    end
end

reconstruct(m::AbstractVLAE, x) = decode(m, encode_all(m, x)...)

function reconstruction_probability(m::AbstractVLAE, x::AbstractArray{T,2}) where T
    zs = map(y->rptrick(y...), _encoded_mu_vars(m, x))
    return -logpdf(m, x, zs...)
end
function reconstruction_probability(m::AbstractVLAE, x::AbstractArray{T,4}) where T
    gx = gpu(x)
    zs = map(y->rptrick(y...), _encoded_mu_vars(m, gx))
    return -logpdf(m, gx, zs...)
end
reconstruction_probability(m::AbstractVLAE, x, L::Int) = mean([reconstruction_probability(m,x) for _ in 1:L])
function reconstruction_probability(m::AbstractVLAE, x, L::Int, batchsize::Int)
    if size(x, ndims(x)) == 0
        return Float32[]
    else
        return vcat(map(b->cpu(reconstruction_probability(m, b, L)), Flux.Data.DataLoader(x, batchsize=batchsize))...)
    end
end

function generate(m::AbstractVLAE, n::Int)
    nl = length(m.e)
    zs = [randn(Float32, m.zdim, n) for _ in 1:nl]

    if is_gpu(m)
        zs = gpu(zs)
    end

    cpu(decode(m, zs...))
end

is_gpu(m::AbstractVLAE) = typeof(m.f[1][2].W) <: Flux.CUDA.CuArray


function aggregate(obj::AggregatableStatistic, arg0::Iterable)
    return jcall(obj, "aggregate", void, (Iterable,), arg0)
end

function aggregate(obj::AggregatableStatistic, arg0::Object)
    return jcall(obj, "aggregate", void, (Object,), arg0)
end

function aggregate(obj::AggregatableStatistic, arg0::Vector{Object})
    return jcall(obj, "aggregate", void, (Vector{Object},), arg0)
end


function set_passive_or_contact_points(point_types::Array{PointType, 3}, potential::Array{T, 3},
                        grid::CylindricalGrid{T}, obj, pot::T, use_nthreads::Int = 1) where {T}
    if !isnan(pot)
        @onthreads 1:use_nthreads for iz in workpart(axes(potential, 3), 1:use_nthreads, Base.Threads.threadid())
            @inbounds for iφ in axes(potential, 2)
                for ir in axes(potential, 1)
                    pt::CylindricalPoint{T} = CylindricalPoint{T}( grid.axes[1].ticks[ir], grid.axes[2].ticks[iφ], grid.axes[3].ticks[iz] )
                    if pt in obj
                        potential[ ir, iφ, iz ] = pot
                        point_types[ ir, iφ, iz ] = zero(PointType)
                    end
                end
            end
        end
    end
    nothing
end

function set_point_types_and_fixed_potentials!(point_types::Array{PointType, 3}, potential::Array{T, 3},
        grid::CylindricalGrid{T}, det::SolidStateDetector{T};
        weighting_potential_contact_id::Union{Missing, Int} = missing,
        use_nthreads::Int = Base.Threads.nthreads(),
        not_only_paint_contacts::Val{NotOnlyPaintContacts} = Val{true}(),
        paint_contacts::Val{PaintContacts} = Val{true}())::Nothing where {T <: SSDFloat, NotOnlyPaintContacts, PaintContacts}

    @onthreads 1:use_nthreads for iz in workpart(axes(potential, 3), 1:use_nthreads, Base.Threads.threadid())
        @inbounds for iφ in axes(potential, 2)
            for ir in axes(potential, 1)
                pt::CylindricalPoint{T} = CylindricalPoint{T}( grid.axes[1].ticks[ir], grid.axes[2].ticks[iφ], grid.axes[3].ticks[iz] )
                if in(pt, det.semiconductor)
                    point_types[ ir, iφ, iz ] += pn_junction_bit
                end
            end
        end
    end
    isEP = ismissing(weighting_potential_contact_id)
    if !ismissing(det.passives)
        for passive in det.passives
            pot::T = isEP ? passive.potential : (isnan(passive.potential) ? passive.potential : zero(T))
            set_passive_or_contact_points(point_types, potential, grid, passive.geometry, pot, use_nthreads)
        end
    end
    if NotOnlyPaintContacts
        for contact in det.contacts
            pot::T = isEP ? contact.potential : contact.id == weighting_potential_contact_id
            set_passive_or_contact_points(point_types, potential, grid, contact.geometry, pot, use_nthreads)
        end
    end
    if PaintContacts
        for contact in det.contacts
            pot::T = isEP ? contact.potential : contact.id == weighting_potential_contact_id
            fs = ConstructiveSolidGeometry.surfaces(contact.geometry)
            for face in fs
                paint!(point_types, potential, face, contact.geometry, pot, grid)
            end
        end
    end
    nothing
end

function fill_ρimp_ϵ_ρfix(ρ_eff_imp_tmp::Array{T}, ϵ::Array{T}, ρ_eff_fix_tmp::Array{T},
    ::Type{Cylindrical}, mpz::Vector{T}, mpφ::Vector{T}, mpr::Vector{T}, axr::Vector{T}, use_nthreads::Int, obj) where {T}
    @inbounds begin
        @onthreads 1:use_nthreads for iz in workpart(axes(ϵ, 3), 1:use_nthreads, Base.Threads.threadid())
            pos_z::T = mpz[iz]
            for iφ in axes(ϵ, 2)
                pos_φ::T = mpφ[iφ]
                for ir in axes(ϵ, 1)
                    pos_r::T = mpr[ir]
                    if (ir == 1 && axr[1] == 0) pos_r = axr[2] * 0.5 end
                    pt::CylindricalPoint{T} = CylindricalPoint{T}(pos_r, pos_φ, pos_z)
                    if pt in obj
                        ρ_eff_imp_tmp[ir, iφ, iz]::T, ϵ[ir, iφ, iz]::T, ρ_eff_fix_tmp[ir, iφ, iz]::T = get_ρimp_ϵ_ρfix(pt, obj)
                    end
                end
            end
        end
    end
    nothing
end

function PotentialCalculationSetup(det::SolidStateDetector{T}, grid::CylindricalGrid{T},
                medium::NamedTuple = material_properties[materials["vacuum"]],
                potential_array::Union{Missing, Array{T, 3}} = missing,
                imp_scale::Union{Missing, Array{T, 3}} = missing;
                weighting_potential_contact_id::Union{Missing, Int} = missing,
                point_types = missing,
                use_nthreads::Int = Base.Threads.nthreads(),
                sor_consts = (1.0, 1.0),
                not_only_paint_contacts::Bool = true, paint_contacts::Bool = true)::PotentialCalculationSetup{T} where {T}
    r0_handling::Bool = typeof(grid.axes[1]).parameters[2] == :r0
    only_2d::Bool = length(grid.axes[2]) == 1 ? true : false
    @assert grid.axes[1][1] == 0 "Something is wrong. R-axis has `:r0`-boundary handling but first tick is $(axr[1]) and not 0."

    is_weighting_potential::Bool = !ismissing(weighting_potential_contact_id)
    depletion_handling::Bool = is_weighting_potential && !ismissing(point_types)

    @inbounds begin
        begin # Geometrical weights of the Axes
            nr = size(grid)[1]

            # R-axis
            axr::Vector{T} = collect(grid.axes[1]) # real grid points/ticks -> the potential at these ticks are going to be calculated
            r_inv::Vector{T} = inv.(axr)
            if r0_handling r_inv[1] = inv(axr[2] * 0.5) end
            r_ext::Vector{T} = get_extended_ticks(grid.axes[1])
            Δr_ext::Vector{T} = diff(r_ext)
            Δr_ext_inv::Vector{T} = inv.(Δr_ext)
            mpr::Vector{T} = midpoints(r_ext)
            Δmpr::Vector{T} = diff(mpr)
            Δmpr_inv::Vector{T} = inv.(Δmpr)
            Δmpr_squared::Vector{T} = T(0.5) .* ((mpr[2:end].^2) .- (mpr[1:end-1].^2))
            if r0_handling
                Δmpr_squared[1] = T(0.5) * (mpr[2]^2)
                Δr_ext_inv[1] = 0 # -> left weight @r=0 becomes 0 through this
            end
            Δhmprr::Vector{T} = mpr[2:end] - axr # distances between midpoints and real grid points (half distances -> h), needed for weights wrr, wrl, ... & dV (volume element)
            Δhmprl::Vector{T} = axr - mpr[1:end - 1]
            wrr::Vector{T} = Δmpr_inv .* Δhmprr # weights for epislon_r adding
            wrl::Vector{T} = Δmpr_inv .* Δhmprl
            if r0_handling wrl[1]::T = 1 - wrr[1] end
            wr::Array{T, 2} = zeros(T, 6, length(wrr))
            wr[1, :] = wrr
            wr[2, :] = wrl
            wr[3, :] = r_inv .* Δmpr
            wr[4, :] = Δr_ext_inv[2:end] .* mpr[2:end]
            wr[5, :] = Δr_ext_inv[1:length(wrr)] .* mpr[1:length(wrr)]
            wr[6, :] = Δmpr_squared
            gw_r::GeometricalRadialAxisWeights{T} = GeometricalRadialAxisWeights{T}( wr ) # Weights needed for Field Simulation loop
            r_ext_mid::T = (r_ext[end - 1] - r_ext[2]) / 2
            grid_boundary_factor_r_left::T = abs((r_ext[2] - r_ext_mid) / (r_ext[1] - r_ext_mid))
            grid_boundary_factor_r_right::T = r0_handling ? abs(r_ext[end - 1] / r_ext[end]) : abs((r_ext[end - 1] - r_ext_mid) / (r_ext[end] - r_ext_mid))

            # φ-axis
            axφ::Vector{T} = collect(grid.axes[2]) # real grid points/ticks -> the potential at these ticks are going to be calculated
            φ_ext::Vector{T} = get_extended_ticks(grid.axes[2])

            Δφ_ext::Vector{T} = diff(φ_ext)
            Δφ_ext_inv::Vector{T} = inv.(Δφ_ext)
            mpφ::Vector{T} = midpoints(φ_ext)
            Δmpφ::Vector{T} = diff(mpφ)
            Δmpφ_inv::Vector{T} = inv.(Δmpφ)
            Δhmpφr::Vector{T} = mpφ[2:end] - axφ # distances between midpoints and real grid points (half distances -> h), needed for weights wrr, wrl, ... & dV (volume element)
            Δhmpφl::Vector{T} = axφ - mpφ[1:end - 1]
            wφr::Vector{T} = Δmpφ_inv .* Δhmpφr # weights for epislon_r adding
            wφl::Vector{T} = Δmpφ_inv .* Δhmpφl
            wφ::Array{T, 2} = zeros(T, 4, length(wφr) + 1)
            wφ[1, 1:length(wφr)] = wφr
            wφ[2, 1:length(wφr)] = wφl
            wφ[3, 1:length(wφr)] = Δmpφ
            wφ[4, :] = Δφ_ext_inv
            gw_φ::GeometricalAzimutalAxisWeights{T} = GeometricalAzimutalAxisWeights{T}( wφ ) # Weights needed for Field Simulation loop
            φ_ext_mid::T = (φ_ext[end - 1] - φ_ext[2]) / 2
            grid_boundary_factor_φ_right::T = abs((φ_ext[end - 1] - φ_ext_mid) / (φ_ext[end] - φ_ext_mid))
            grid_boundary_factor_φ_left::T = abs((φ_ext[2] - φ_ext_mid) / (φ_ext[1] - φ_ext_mid))

            # Z-axis
            axz::Vector{T} = collect(grid.axes[3]) # real grid points/ticks -> the potential at these ticks are going to be calculated
            z_ext::Vector{T} = get_extended_ticks(grid.axes[3])
            Δz_ext::Vector{T} = diff(z_ext)
            Δz_ext_inv::Vector{T} = inv.(Δz_ext)
            mpz::Vector{T} = midpoints(z_ext)
            Δmpz::Vector{T} = diff(mpz)
            Δmpz_inv::Vector{T} = inv.(Δmpz)
            Δhmpzr::Vector{T} = mpz[2:end] - axz # distances between midpoints and real grid points (half distances -> h), needed for weights wrr, wrl, ... & dV (volume element)
            Δhmpzl::Vector{T} = axz - mpz[1:end - 1]
            wzr::Vector{T} = Δmpz_inv .* Δhmpzr # weights for epislon_r adding
            wzl::Vector{T} = Δmpz_inv .* Δhmpzl
            wz::Array{T, 2} = zeros(T, 4, length(wzr) + 1)
            wz[1, 1:length(wzr)] = wzr
            wz[2, 1:length(wzr)] = wzl
            wz[3, 1:length(wzr)] = Δmpz
            wz[4, :] = Δz_ext_inv
            gw_z::GeometricalCartesianAxisWeights{T} = GeometricalCartesianAxisWeights{T}( wz ) # Weights needed for Field Simulation loop
            z_ext_mid::T = (z_ext[end - 1] - z_ext[2]) / 2
            grid_boundary_factor_z_right::T = abs((z_ext[end - 1] - z_ext_mid) / (z_ext[end] - z_ext_mid))
            grid_boundary_factor_z_left::T = abs((z_ext[2] - z_ext_mid) / (z_ext[1] - z_ext_mid))

            geom_weights::NTuple{3, AbstractGeometricalAxisWeights{T}} = (gw_r, gw_φ, gw_z) # Weights needed for Field Simulation loop
            grid_boundary_factors::NTuple{3, NTuple{2, T}} = ((grid_boundary_factor_r_left, grid_boundary_factor_r_right), (grid_boundary_factor_φ_left, grid_boundary_factor_φ_right), (grid_boundary_factor_z_left, grid_boundary_factor_z_right))
        end

        bias_voltages::Vector{T} = if length(det.contacts) > 0
            [i.potential for i in det.contacts]
        else
            T[0]
        end
        minimum_applied_potential::T = minimum(bias_voltages)
        maximum_applied_potential::T = maximum(bias_voltages)
        bias_voltage::T = maximum_applied_potential - minimum_applied_potential
        sor_slope = (sor_consts[2] .- sor_consts[1]) / (nr - 1 )
        sor_const::Vector{T} = T[ sor_consts[1] + (i - 1) * sor_slope for i in 1:nr]

        medium_ϵ_r::T = medium.ϵ_r
        ϵ = fill(medium_ϵ_r, length(mpr), length(mpφ), length(mpz))
        ρ_eff_imp_tmp = zeros(T, length(mpr), length(mpφ), length(mpz))
        ρ_eff_fix_tmp = zeros(T, length(mpr), length(mpφ), length(mpz))
        fill_ρimp_ϵ_ρfix(ρ_eff_imp_tmp, ϵ, ρ_eff_fix_tmp, Cylindrical, mpz, mpφ, mpr, axr, use_nthreads, det.semiconductor)
        if !ismissing(det.passives)
            for passive in det.passives
                fill_ρimp_ϵ_ρfix(ρ_eff_imp_tmp, ϵ, ρ_eff_fix_tmp, Cylindrical, mpz, mpφ, mpr, axr, use_nthreads, passive)
            end
        end
        if depletion_handling
            for iz in axes(ϵ, 3)
                pos_z = mpz[iz]
                for iφ in axes(ϵ, 2)
                    pos_φ = mpφ[iφ]
                    for ir in axes(ϵ, 1)
                        pos_r = mpr[ir]
                        if (ir == 1 && axr[1] == 0) pos_r = axr[2] * 0.5 end
                        pt::CylindricalPoint{T} = CylindricalPoint{T}(pos_r, pos_φ, pos_z)
                        ig = find_closest_gridpoint(pt, point_types.grid)
                        if is_undepleted_point_type(point_types.data[ig...])
                            ϵ[ir, iφ, iz] *= scaling_factor_for_permittivity_in_undepleted_region(det.semiconductor) * (1 - imp_scale[ig...])
                        elseif is_fixed_point_type(point_types.data[ig...])
                            ϵ[ir, iφ, iz] *= scaling_factor_for_permittivity_in_undepleted_region(det.semiconductor)
                        end
                    end
                end
            end
        end

        ϵ0_inv::T = inv(ϵ0)
        ρ_eff_imp_tmp *= ϵ0_inv
        ρ_eff_fix_tmp *= ϵ0_inv

        volume_weights::Array{T, 4} = RBExtBy2Array(T, grid)
        q_eff_imp::Array{T, 4} = RBExtBy2Array(T, grid)
        q_eff_fix::Array{T, 4} = RBExtBy2Array(T, grid)
        for iz in range(2, stop = length(z_ext) - 1)
            inz::Int = iz - 1
            irbz::Int = rbidx(inz)
            for iφ in range(2, stop = length(φ_ext) - 1)
                inφ::Int = iφ - 1
                for ir in range(2, stop = length(r_ext) - 1)
                    inr::Int = ir - 1;

                    rbi::Int = iseven(inr + inφ + inz) ? rb_even::Int : rb_odd::Int
                    # rbinds = irbz, iφ, ir, rbi

                    ρ_imp_cell::T = 0
                    ρ_eff_fix_cell::T = 0
                    if !is_weighting_potential
                        if inr > 1
                            ρ_imp_cell += ρ_eff_imp_tmp[ ir,  iφ,  iz] * wzr[inz] * wrr[inr] * wφr[inφ]
                            ρ_imp_cell += ρ_eff_imp_tmp[ ir,  iφ, inz] * wzl[inz] * wrr[inr] * wφr[inφ]
                            ρ_imp_cell += ρ_eff_imp_tmp[ ir, inφ,  iz] * wzr[inz] * wrr[inr] * wφl[inφ]
                            ρ_imp_cell += ρ_eff_imp_tmp[ ir, inφ, inz] * wzl[inz] * wrr[inr] * wφl[inφ]

                            ρ_imp_cell += ρ_eff_imp_tmp[inr,  iφ,  iz] * wzr[inz] * wrl[inr] * wφr[inφ]
                            ρ_imp_cell += ρ_eff_imp_tmp[inr,  iφ, inz] * wzl[inz] * wrl[inr] * wφr[inφ]
                            ρ_imp_cell += ρ_eff_imp_tmp[inr, inφ,  iz] * wzr[inz] * wrl[inr] * wφl[inφ]
                            ρ_imp_cell += ρ_eff_imp_tmp[inr, inφ, inz] * wzl[inz] * wrl[inr] * wφl[inφ]

                            ρ_eff_fix_cell += ρ_eff_fix_tmp[ ir,  iφ,  iz] * wzr[inz] * wrr[inr] * wφr[inφ]
                            ρ_eff_fix_cell += ρ_eff_fix_tmp[ ir,  iφ, inz] * wzl[inz] * wrr[inr] * wφr[inφ]
                            ρ_eff_fix_cell += ρ_eff_fix_tmp[ ir, inφ,  iz] * wzr[inz] * wrr[inr] * wφl[inφ]
                            ρ_eff_fix_cell += ρ_eff_fix_tmp[ ir, inφ, inz] * wzl[inz] * wrr[inr] * wφl[inφ]

                            ρ_eff_fix_cell += ρ_eff_fix_tmp[inr,  iφ,  iz] * wzr[inz] * wrl[inr] * wφr[inφ]
                            ρ_eff_fix_cell += ρ_eff_fix_tmp[inr,  iφ, inz] * wzl[inz] * wrl[inr] * wφr[inφ]
                            ρ_eff_fix_cell += ρ_eff_fix_tmp[inr, inφ,  iz] * wzr[inz] * wrl[inr] * wφl[inφ]
                            ρ_eff_fix_cell += ρ_eff_fix_tmp[inr, inφ, inz] * wzl[inz] * wrl[inr] * wφl[inφ]
                        else
                            ρ_imp_cell += ρ_eff_imp_tmp[ ir,  iφ,  iz] * wzr[inz] * 0.5 #wφr[inφ]
                            ρ_imp_cell += ρ_eff_imp_tmp[ ir,  iφ, inz] * wzl[inz] * 0.5 #wφr[inφ]
                            ρ_imp_cell += ρ_eff_imp_tmp[ ir, inφ,  iz] * wzr[inz] * 0.5 #wφl[inφ]
                            ρ_imp_cell += ρ_eff_imp_tmp[ ir, inφ, inz] * wzl[inz] * 0.5 #wφl[inφ]

                            ρ_eff_fix_cell += ρ_eff_fix_tmp[ ir,  iφ,  iz] * wzr[inz] * 0.5 #wφr[inφ]
                            ρ_eff_fix_cell += ρ_eff_fix_tmp[ ir,  iφ, inz] * wzl[inz] * 0.5 #wφr[inφ]
                            ρ_eff_fix_cell += ρ_eff_fix_tmp[ ir, inφ,  iz] * wzr[inz] * 0.5 #wφl[inφ]
                            ρ_eff_fix_cell += ρ_eff_fix_tmp[ ir, inφ, inz] * wzl[inz] * 0.5 #wφl[inφ]
                        end
                    end
                    if inr > 1
                        wrr_eps::T = ϵ[  ir,  iφ, inz + 1] * wφr[inφ] * wzr[inz]
                        wrr_eps   += ϵ[  ir, inφ, inz + 1] * wφl[inφ] * wzr[inz]
                        wrr_eps   += ϵ[  ir,  iφ, inz ]    * wφr[inφ] * wzl[inz]
                        wrr_eps   += ϵ[  ir, inφ, inz ]    * wφl[inφ] * wzl[inz]
                        # # left weight in r: wrr
                        wrl_eps::T = ϵ[ inr,  iφ, inz + 1] * wφr[inφ] * wzr[inz]
                        wrl_eps   += ϵ[ inr, inφ, inz + 1] * wφl[inφ] * wzr[inz]
                        wrl_eps   += ϵ[ inr,  iφ, inz ]    * wφr[inφ] * wzl[inz]
                        wrl_eps   += ϵ[ inr, inφ, inz ]    * wφl[inφ] * wzl[inz]
                        # right weight in φ: wφr
                        wφr_eps::T = ϵ[ inr,  iφ, inz + 1] * wrl[inr] * wzr[inz]
                        wφr_eps   += ϵ[  ir,  iφ, inz + 1] * wrr[inr] * wzr[inz]
                        wφr_eps   += ϵ[ inr,  iφ, inz ]    * wrl[inr] * wzl[inz]
                        wφr_eps   += ϵ[  ir,  iφ, inz ]    * wrr[inr] * wzl[inz]
                        # left weight in φ: wφl
                        wφl_eps::T = ϵ[ inr, inφ, inz + 1] * wrl[inr] * wzr[inz]
                        wφl_eps   += ϵ[  ir, inφ, inz + 1] * wrr[inr] * wzr[inz]
                        wφl_eps   += ϵ[ inr, inφ, inz ]    * wrl[inr] * wzl[inz]
                        wφl_eps   += ϵ[  ir, inφ, inz ]    * wrr[inr] * wzl[inz]
                        # right weight in z: wzr
                        wzr_eps::T = ϵ[  ir,  iφ, inz + 1] * wrr[inr] * wφr[inφ]
                        wzr_eps   += ϵ[  ir, inφ, inz + 1] * wrr[inr] * wφl[inφ]
                        wzr_eps   += ϵ[ inr,  iφ, inz + 1] * wrl[inr] * wφr[inφ]
                        wzr_eps   += ϵ[ inr, inφ, inz + 1] * wrl[inr] * wφl[inφ]
                        # left weight in z: wzr
                        wzl_eps::T = ϵ[  ir,  iφ,    inz ] * wrr[inr] * wφr[inφ]
                        wzl_eps   += ϵ[  ir, inφ,    inz ] * wrr[inr] * wφl[inφ]
                        wzl_eps   += ϵ[ inr,  iφ,    inz ] * wrl[inr] * wφr[inφ]
                        wzl_eps   += ϵ[ inr, inφ,    inz ] * wrl[inr] * wφl[inφ]

                        volume_weight::T = wrr_eps * Δr_ext_inv[ ir] * mpr[ ir] * Δmpφ[inφ] * Δmpz[inz]
                        volume_weight += wrl_eps * Δr_ext_inv[inr] * mpr[inr] * Δmpφ[inφ] * Δmpz[inz]
                        volume_weight += wφr_eps * r_inv[inr] * Δφ_ext_inv[ iφ] * Δmpr[inr] * Δmpz[inz]
                        volume_weight += wφl_eps * r_inv[inr] * Δφ_ext_inv[inφ] * Δmpr[inr] * Δmpz[inz]
                        volume_weight += wzr_eps * Δz_ext_inv[ iz] * Δmpφ[inφ] * Δmpr_squared[inr]
                        volume_weight += wzl_eps * Δz_ext_inv[inz] * Δmpφ[inφ] * Δmpr_squared[inr]

                        volume_weights[ irbz, iφ, ir, rbi ] = inv(volume_weight)

                        dV::T = Δmpz[inz] * Δmpφ[inφ] * Δmpr_squared[inr]
                        q_eff_imp[ irbz, iφ, ir, rbi ] = dV * ρ_imp_cell
                        q_eff_fix[ irbz, iφ, ir, rbi ] = dV * ρ_eff_fix_cell
                    else
                        wrr_eps = ϵ[  ir,  iφ, inz + 1] * 0.5 * wzr[inz]
                        wrr_eps   += ϵ[  ir, inφ, inz + 1] * 0.5 * wzr[inz]
                        wrr_eps   += ϵ[  ir,  iφ, inz ]    * 0.5 * wzl[inz]
                        wrr_eps   += ϵ[  ir, inφ, inz ]    * 0.5 * wzl[inz]
                        # # left weight in r: wrr
                        wrl_eps = ϵ[ inr,  iφ, inz + 1] * 0.5 * wzr[inz]
                        wrl_eps   += ϵ[ inr, inφ, inz + 1] * 0.5 * wzr[inz]
                        wrl_eps   += ϵ[ inr,  iφ, inz ]    * 0.5 * wzl[inz]
                        wrl_eps   += ϵ[ inr, inφ, inz ]    * 0.5 * wzl[inz]
                        # right weight in φ: wφr
                        wφr_eps = ϵ[ inr,  iφ, inz + 1] * wrl[inr] * wzr[inz]
                        wφr_eps   += ϵ[  ir,  iφ, inz + 1] * wrr[inr] * wzr[inz]
                        wφr_eps   += ϵ[ inr,  iφ, inz ]    * wrl[inr] * wzl[inz]
                        wφr_eps   += ϵ[  ir,  iφ, inz ]    * wrr[inr] * wzl[inz]
                        # left weight in φ: wφl
                        wφl_eps = ϵ[ inr, inφ, inz + 1] * wrl[inr] * wzr[inz]
                        wφl_eps   += ϵ[  ir, inφ, inz + 1] * wrr[inr] * wzr[inz]
                        wφl_eps   += ϵ[ inr, inφ, inz ]    * wrl[inr] * wzl[inz]
                        wφl_eps   += ϵ[  ir, inφ, inz ]    * wrr[inr] * wzl[inz]
                        # right weight in z: wzr
                        wzr_eps = ϵ[  ir,  iφ, inz + 1] * wrr[inr] * 0.5
                        wzr_eps   += ϵ[  ir, inφ, inz + 1] * wrr[inr] * 0.5
                        wzr_eps   += ϵ[ inr,  iφ, inz + 1] * wrl[inr] * 0.5
                        wzr_eps   += ϵ[ inr, inφ, inz + 1] * wrl[inr] * 0.5
                        # left weight in z: wzr
                        wzl_eps = ϵ[  ir,  iφ,    inz ] * wrr[inr] * 0.5
                        wzl_eps   += ϵ[  ir, inφ,    inz ] * wrr[inr] * 0.5
                        wzl_eps   += ϵ[ inr,  iφ,    inz ] * wrl[inr] * 0.5
                        wzl_eps   += ϵ[ inr, inφ,    inz ] * wrl[inr] * 0.5

                        volume_weight = wrr_eps * Δr_ext_inv[ ir] * mpr[ ir] * 2π * Δmpz[inz]
                        volume_weight += wrl_eps * Δr_ext_inv[inr] * mpr[inr] * Δmpφ[inφ] * Δmpz[inz]
                        volume_weight += wφr_eps * r_inv[inr] * 0.15915494f0 * Δmpr[inr] * Δmpz[inz]
                        volume_weight += wφl_eps * r_inv[inr] * 0.15915494f0 * Δmpr[inr] * Δmpz[inz]
                        volume_weight += wzr_eps * Δz_ext_inv[ iz] * 2π * Δmpr_squared[inr]
                        volume_weight += wzl_eps * Δz_ext_inv[inz] * 2π * Δmpr_squared[inr]

                        volume_weights[ irbz, iφ, ir, rbi ] = inv(volume_weight)

                        dV = Δmpz[inz] * 2π * Δmpr_squared[inr]
                        q_eff_imp[ irbz, iφ, ir, rbi ] = dV * ρ_imp_cell
                        q_eff_fix[ irbz, iφ, ir, rbi ] = dV * ρ_eff_fix_cell
                    end
                end
            end
        end
        potential = ismissing(potential_array) ? zeros(T, size(grid)...) : potential_array
        point_types = ones(PointType, size(grid)...)
        set_point_types_and_fixed_potentials!( point_types, potential, grid, det,
                weighting_potential_contact_id = weighting_potential_contact_id,
                use_nthreads = use_nthreads,
                not_only_paint_contacts = Val(not_only_paint_contacts),
                paint_contacts = Val(paint_contacts)  )
        rbpotential  = RBExtBy2Array( potential, grid )
        rbpoint_types = RBExtBy2Array( point_types, grid )
    end # @inbounds

    pcs = PotentialCalculationSetup(
        grid,
        rbpotential,
        rbpoint_types,
        volume_weights,
        q_eff_imp,
        ones(T, size(rbpoint_types)),
        q_eff_fix,
        ϵ,
        broadcast(gw -> gw.weights, geom_weights),
        sor_const,
        bias_voltage,
        maximum_applied_potential,
        minimum_applied_potential,
        grid_boundary_factors
     )

    apply_boundary_conditions!(pcs, Val{ true}(), only_2d ? Val{true}() : Val{false}()) # even points
    apply_boundary_conditions!(pcs, Val{false}(), only_2d ? Val{true}() : Val{false}()) # odd
    return pcs
end

Grid(pcs::PotentialCalculationSetup) = pcs.grid

function ElectricPotentialArray(pcs::PotentialCalculationSetup{T, Cylindrical, 3, Array{T, 3}})::Array{T, 3} where {T}
    pot::Array{T, 3} = Array{T, 3}(undef, size(pcs.grid))
    for iz in axes(pot, 3)
        irbz::Int = rbidx(iz)
        for iφ in axes(pot, 2)
            irbφ::Int = iφ + 1
            idxsum::Int = iz + iφ
            for ir in axes(pot, 1)
                irbr::Int = ir + 1
                rbi::Int = iseven(idxsum + ir) ? rb_even::Int : rb_odd::Int
                pot[ir, iφ, iz] = pcs.potential[ irbz, irbφ, irbr, rbi ]
            end
        end
    end
    for iz in axes(pot, 3)
        p_r0 = mean(pot[1,:,iz])
        pot[1,:,iz] .= p_r0
    end
    return pot
end

function ImpurityScaleArray(pcs::PotentialCalculationSetup{T, Cylindrical, 3, Array{T, 3}})::Array{T, 3} where {T}
    s::Array{T, 3} = Array{T, 3}(undef, size(pcs.grid))
    for iz in axes(s, 3)
        irbz::Int = rbidx(iz)
        for iφ in axes(s, 2)
            irbφ::Int = iφ + 1
            idxsum::Int = iz + iφ
            for ir in axes(s, 1)
                irbr::Int = ir + 1
                rbi::Int = iseven(idxsum + ir) ? rb_even::Int : rb_odd::Int
                s[ir, iφ, iz] = pcs.imp_scale[ irbz, irbφ, irbr, rbi ]
            end
        end
    end
    return s
end

function PointTypeArray(pcs::PotentialCalculationSetup{T, Cylindrical, 3, Array{T, 3}})::Array{PointType, 3} where {T}
    point_types::Array{PointType, 3} = zeros(PointType, size(pcs.grid))
    for iz in axes(point_types, 3)
        irbz::Int = rbidx(iz)
        for iφ in axes(point_types, 2)
            irbφ::Int = iφ + 1
            idxsum::Int = iz + iφ
            for ir in axes(point_types, 1)
                irbr::Int = ir + 1
                rbi::Int = iseven(idxsum + ir) ? rb_even::Int : rb_odd::Int
                point_types[ir, iφ, iz] = pcs.point_types[irbz, irbφ, irbr, rbi ]
            end
        end
    end
    return point_types
end

function EffectiveChargeDensityArray(pcs::PotentialCalculationSetup{T, Cylindrical, 3, Array{T, 3}})::Array{T} where {T}
    ρ::Array{T, 3} = zeros(T, size(pcs.grid))
    for iz in axes(ρ, 3)
        irbz::Int = rbidx(iz)
        Δmpz::T = pcs.geom_weights[3][3, iz]
        for iφ in axes(ρ, 2)
            irbφ::Int = iφ + 1
            idxsum::Int = iz + iφ
            Δmpφ::T = pcs.geom_weights[2][3, iφ]
            Δmpzφ::T = Δmpz * Δmpφ
            for ir in axes(ρ, 1)
                irbr::Int = ir + 1
                rbi::Int = iseven(idxsum + ir) ? rb_even::Int : rb_odd::Int
                dV::T = pcs.geom_weights[1][6, ir] * Δmpzφ  #Δmpz[inz] * Δmpφ[inφ] * Δmpr_squared[inr]
                if ir == 1
                    dV = dV * 2π / Δmpφ
                end
                ρ[ir, iφ, iz] = pcs.q_eff_imp[irbz, irbφ, irbr, rbi ] / dV
            end
        end
    end
    return ρ
end

function FixedEffectiveChargeDensityArray(pcs::PotentialCalculationSetup{T, Cylindrical, 3, Array{T, 3}})::Array{T} where {T}
    ρ::Array{T, 3} = zeros(T, size(pcs.grid))
    for iz in axes(ρ, 3)
        irbz::Int = rbidx(iz)
        Δmpz::T = pcs.geom_weights[3][3, iz]
        for iφ in axes(ρ, 2)
            irbφ::Int = iφ + 1
            idxsum::Int = iz + iφ
            Δmpφ::T = pcs.geom_weights[2][3, iφ]
            Δmpzφ::T = Δmpz * Δmpφ
            for ir in axes(ρ, 1)
                irbr::Int = ir + 1
                rbi::Int = iseven(idxsum + ir) ? rb_even::Int : rb_odd::Int
                dV::T = pcs.geom_weights[1][6, ir] * Δmpzφ  #Δmpz[inz] * Δmpφ[inφ] * Δmpr_squared[inr]
                if ir == 1
                    dV = dV * 2π / Δmpφ
                end
                ρ[ir, iφ, iz] = pcs.q_eff_fix[irbz, irbφ, irbr, rbi ] / dV
            end
        end
    end
    return ρ
end

function DielectricDistributionArray(pcs::PotentialCalculationSetup{T, S, 3, Array{T, 3}})::Array{T, 3} where {T, S}
    return pcs.ϵ_r
end


#=
  Original Author: Anastasia Yendiki

  Copyright © 2022 The General Hospital Corporation (Boston, MA) "MGH"

  Terms and conditions for use, reproduction, distribution and contribution
  are found in the 'FreeSurfer Software License Agreement' contained
  in the file 'LICENSE' found in the FreeSurfer distribution, and here:

  https://surfer.nmr.mgh.harvard.edu/fswiki/FreeSurferSoftwareLicense

  Reporting: freesurfer@nmr.mgh.harvard.edu
=#

using ColorTypes, ImageView

export view

"""
    view(mri::MRI, plane::Char='a')

View an `MRI` structure in a slice viewer, along the specified `plane` ('a':
axial; 's': sagittal; 'c': coronal).
"""
function view(mri::MRI, plane::Char='a')

  # Find which axes of the volume correspond to the specified viewing plane
  ax = view_axes(mri.vox2ras, plane)

  (ax1, ax2) = abs.(ax)
  (flip1, flip2) = (ax .< 0)

  if mri.ispermuted
    (ax1 == 1) && (ax1 = 2)
    (ax1 == 2) && (ax1 = 1)
    (ax2 == 1) && (ax2 = 2)
    (ax2 == 2) && (ax2 = 1)
  end

  # Find the through-plane axis of the volume
  ax3 = setdiff(1:3, [ax1, ax2])[1]

  # Set names and colors of axis labels
  if plane == 'a'			# Axial plane: A->P, R->L
    label1 = ["A", "P"]
    color1 = julia_green

    label2 = ["R", "L"]
    color2 = julia_red
  elseif plane == 's'			# Sagittal plane: S->I, P->A
    label1 = ["S", "I"]
    color1 = julia_blue

    label2 = ["P", "A"]
    color2 = julia_green
  else					# Coronal plane: S->I, R->L
    label1 = ["S", "I"]
    color1 = julia_blue

    label2 = ["R", "L"]
    color2 = julia_red
  end

  # Convert volume to RGB/Gray array
  rgb = vol_to_rgb(mri.vol)

  # Set the size of the canvas that the volume will be displayed in
  csize = size(rgb)[[ax1, ax2]]
  csize = (300, Int.(round.(300 * csize[2] / csize[1])))

  # Display volume
  gui = ImageView.imshow(rgb, axes=(ax1,ax2), flipy=flip1, flipx=flip2,
                              canvassize=csize, name=mri.fspec)

  # Move display to middle slice
  push!(gui["roi"]["slicedata"].signals[1], Int(size(mri.vol, ax3)/2))

  # Add annotation for image axes
  fsize = round(5 * maximum(mri.volsize) / 256)

  ImageView.annotate!(gui,
            AnnotationText(size(mri.vol, ax2)*.5, size(mri.vol, ax1)*.05,
            label1[1], color=color1, fontsize=fsize))
  ImageView.annotate!(gui,
            AnnotationText(size(mri.vol, ax2)*.5, size(mri.vol, ax1)*.95,
            label1[2], color=color1, fontsize=fsize))
  ImageView.annotate!(gui,
            AnnotationText(size(mri.vol, ax2)*.05, size(mri.vol, ax1)*.5,
            label2[1], color=color2, fontsize=fsize))
  ImageView.annotate!(gui,
            AnnotationText(size(mri.vol, ax2)*.95, size(mri.vol, ax1)*.5,
            label2[2], color=color2, fontsize=fsize))

  # Add annotation for b-value and gradient vector
  blabel = Vector{String}(undef, 0)

  if !isempty(mri.bval)
    blabel = "b=" .* string.(Int.(round.(mri.bval)))
  end

  if !isempty(mri.bvec)
    blabel = blabel .* "\ng=[" .*
                       string.(round.(mri.bvec[:,1] * 100)/100) .* "," .*
                       string.(round.(mri.bvec[:,2] * 100)/100) .* "," .*
                       string.(round.(mri.bvec[:,3] * 100)/100) .* "]"
  end

  for ib = 1:length(blabel)
    ImageView.annotate!(gui,
              AnnotationText(size(mri.vol, ax2)*.1, size(mri.vol, ax1)*.1,
              blabel[ib], t = ib,
              color=RGB(1,1,1), fontsize=fsize, halign="left"))
  end

  return gui
end


# © 2016 Massachusetts Institute of Technology.  See LICENSE file for details.
elementwise(::ANode) = false

for i in [:.+, :.-, :.*, :./, :.^]
  @eval elementwise(::ANode{$(QuoteNode(i))}) = true
end

for i in [:+, :*]
  @eval elementwise(::ANode{$(QuoteNode(Symbol("c", i)))}) = true
end

for i in [:/, :^]
  @eval elementwise(::ANode{$(QuoteNode(Symbol("c", i)))}) = true
  @eval elementwise(::ANode{$(QuoteNode(Symbol(i, "c")))}) = true
end

# Create a devectorized kernel that evaluates node n
# taking the nodes [inp] as input.  Returns a Symbol.
function devec_kernel(n::ANode, inp::Vector{ANode})
  b = :(for i in 1:length(out) end)

  symtab = Dict{ANode,Any}()
  symtab[n] = :(out[i])
  for i in 1:length(inp)
    symtab[inp[i]] = :(input[$i][i])
  end

  ready = ANode[n]
  done = Set{ANode}([ready;  inp])
  while !isempty(ready)
    i = pop!(ready)
    out = get!(gensym, symtab, i)
    arg = map(x->get!(gensym, symtab, x), i.input)

    if sym(i) == Symbol("c*")
      unshift!(b.args[2].args, :($out = *($(i.arg), $(arg...))))
    elseif sym(i) == Symbol("c+")
      unshift!(b.args[2].args, :($out = +($(i.arg), $(arg...))))
    elseif sym(i) == Symbol("c^")
      unshift!(b.args[2].args, :($out = ^($(i.arg), $(arg...))))
    elseif sym(i) == Symbol("^c")
      unshift!(b.args[2].args, :($out = ^($(arg...), $(i.arg))))
    else
      unshift!(b.args[2].args, :($out = $(sym(i))($(arg...))))
    end

    for j in i.input
      if !(j in done)
        push!(done, j)
        push!(ready, j)
      end
    end
  end

  #@show b

  s = gensym("devec")
  @eval @inline do_forward!(n::ANode{$(QuoteNode(s))}, out, input...) = @fastmath @inbounds $b
  return s
end

function add_color(colors::Dict{ANode,Set{ANode}}, node::ANode, c::ANode)
  set = get!(Set{ANode}, colors, node)
  if c in set return end
  push!(set, c)
  for i in node.input add_color(colors, i, c) end
  for i in node.control add_color(colors, i, c) end
  return
end

function intern_values{K,V}(d::Dict{K,V})
  cache = Dict{V,V}()
  for (k,v) in d
    if haskey(cache, v)
      d[k] = cache[v]
    else
      cache[v] = v
    end
  end
  return d
end

devectorize(n::ANode...) = devectorize(collect(ANode, n))
function devectorize(n::Vector{ANode})
  nodes = connected_nodes(n)
  outputs = Dict{ANode,Vector{ANode}}()
  for i in nodes
    get!(Vector{ANode}, outputs, i)
    for j in i.input push!(get!(Vector{ANode}, outputs, j), i) end
    for j in i.control push!(get!(Vector{ANode}, outputs, j), i) end
  end

  colors = Dict{ANode,Set{ANode}}()
  for i in n
    add_color(colors, i, i)
  end
  intern_values(colors)

  inelligible = Set{ANode}(n)
  for (i,output) in outputs
    c = colors[i]
    if !elementwise(i)
      #println("!elementwise $i")
      push!(inelligible, i)
    elseif !all(elementwise, output)
      #println("!elementwise output $i")
      push!(inelligible, i)
    elseif any(j->colors[j]!==c, output)
      #println("multicolor output $i")
      push!(inelligible, i)
    end
  end

  replace = Dict{ANode,ANode}()

  for i in collect(inelligible)
    s = copy(inelligible)
    delete!(s, i)
    a = setdiff(connected_nodes(i, s), inelligible)
    inp = Set{ANode}(i.input)
    for j in a
      for k in j.input
        push!(inp, k)
      end
    end
    setdiff!(inp, a)

    if !isempty(a)
      #@show i
      #@show a
      #@show inp

      d = ANode(devec_kernel(i, collect(inp)), (inp...), size(i))
      for j in outputs[i]
        for m in 1:length(j.input)
          if j.input[m] === i
            j.input[m] = d
          end
        end

        for m in 1:length(j.control)
          if j.control[m] === i
            j.control[m] = d
          end
        end
      end

      delete!(inelligible, i)
      push!(inelligible, d)
      replace[i] = d
      #println()
    end
  end

  return map(x->get(replace, x, x), n)
end


"""
# PolyFields.jl
Welcome to PolyFields.jl!
"""
module PolyFields

#==============================================================================#
# Dependencies
#==============================================================================#

import Statistics: mean, std
import LinearAlgebra: norm, dot, det, mul!, inv!, lu!
import Dates

using Printf
using JSON
using Parameters
using RecipesBase

using Random
using StaticArrays
using Interpolations
using FFTW

#==============================================================================#
# Exports
#==============================================================================#

# Types
export AbstractSystem, AbstractSpecies, AbstractConstraint, AbstractPropagator, AbstractInteraction, AbstractFieldUpdater
export FieldGrid, PseudoSpectralSolver, FFTHolder, UnitCell
export Monomer, Point, Homopolymer, Diblock, Multiblock
export FloryInteraction, EdwardsInteraction
export Ensemble, Canonical, Grand, Compressibility
export FieldSystem, SCFTOptions, SCFTResults
export DomainUpdater, EulerUpdater, AndersonUpdater, MomentumUpdater

# Methods
export make_basis, basis_matrix, basis_dimensions
export add_monomer!, add_species!, add_cell!, add_interaction!, set_interaction!, validate
export density!, muphi!, potentials!, potential!, uniform_fields!, residuals!
export free_energy, free_energy_bulk, scfstress, scfstress!, field_error, stress_error
export step!, scft!, fieldinit!, savefields, loadfields, scft!
export energy, energy_bulk, monomer_fraction, chain_grid
export updater_with_system

#==============================================================================#
# Load files
#==============================================================================#

# General
include("types.jl")
include("utils.jl")

# Cell and system
include("system/fft.jl")
include("system/mde.jl")
include("system/monomer.jl")
include("system/basis.jl")
include("system/cell.jl")
include("system/compress.jl")
include("system/system.jl")

# Species architectures
include("species/base.jl")
include("species/point.jl")
include("species/multiblock.jl")

# # Interaction classes
include("interactions/base.jl")
include("interactions/flory.jl")
include("interactions/edwards.jl")

# # Field updaters
include("updaters/base.jl")
include("updaters/euler.jl")
include("updaters/anderson.jl")
include("updaters/momentum.jl")
include("updaters/domain.jl")

include("fileio.jl")
include("simulate.jl")
# # include("viz.jl")

end


# import Pkg
# Pkg.add("YAML")
import YAML

const CONVENTIONAL_WEIGHT = Dict{Int,Float64}(
    1 => 1.008,
    3 => 6.94,
    5 => 10.81,
    6 => 12.011,
    7 => 14.007,
    8 => 15.999,
    12 => 24.305,
    14 => 28.085,
    16 => 32.06,
    17 => 35.45,
    35 => 79.904,
    81 => 204.38
)

const path = joinpath(
    dirname(@__FILE__), "../assets/raw/nist_atomic_weights20200410.txt")
const dest = joinpath(
    dirname(@__FILE__), "../assets/const/atomicweights.yaml")
const revdest = joinpath(
    dirname(@__FILE__), "../assets/const/symboltonumber.yaml")

function run()
    data = []
    revdict = Dict{String,Int}()
    lines = eachline(path)
    state = nothing
    next = iterate(lines, state)
    while true
        rcd = Dict()
        iso = Dict()
        while true
            (line, state) = next
            line == "" && break
            k, v = split(line, " = ")
            if k == "Atomic Number"
                rcd["Number"] = parse(Int, v)
            elseif k == "Atomic Symbol"
                rcd["Symbol"] = v
            elseif k == "Mass Number"
                iso["Number"] = parse(Int, v)
            elseif k == "Relative Atomic Mass"
                m = match(r"(.*?)\((.*?)#?\)", v)
                value = parse(Float64, m[1])
                unc = parse(Float64, string(
                    replace(m[1], r"[0-9]" => "0")[1:end-length(m[2])], m[2]))
                iso["Mass"] = value
                iso["MassUncertainty"] = unc
            elseif k == "Isotopic Composition"
                m = match(r"(.*?)\((.*?)#?\)", v)
                if m !== nothing
                    value = parse(Float64, m[1])
                    unc = parse(Float64, string(
                        replace(m[1], r"[0-9]" => "0")[1:end-length(m[2])], m[2]))
                elseif tryparse(Float64, v) !== nothing
                    value = parse(Float64, v)
                    unc = 0.0
                else
                    value = NaN
                    unc = NaN
                end
                iso["Composition"] = value
                iso["CompositionUncertainty"] = unc
            elseif k == "Standard Atomic Weight"
                m1 = match(r"\[(.*?),(.*?)\]", v)
                m2 = match(r"(.*?)\((.*?)\)", v)
                m3 = match(r"\[(.*?)\]", v)
                if m1 !== nothing
                    rcd["Weight"] = CONVENTIONAL_WEIGHT[rcd["Number"]]
                    rcd["WeightType"] = "Interval"
                    rcd["WeightLower"] = parse(Float64, m1[1])
                    rcd["WeightHigher"] = parse(Float64, m1[2])
                elseif m2 !== nothing
                    rcd["Weight"] = parse(Float64, m2[1])
                    unc = parse(Float64, string(
                        replace(m2[1], r"[0-9]" => "0")[1:end-length(m2[2])], m2[2]))
                    rcd["WeightType"] = "Uncertainty"
                    rcd["WeightUncertainty"] = unc
                elseif m3 !== nothing
                    rcd["Weight"] = parse(Int, m3[1])
                    rcd["WeightType"] = "MostStable"
                else
                    rcd["Weight"] = NaN
                    rcd["WeightType"] = "NotAvailable"
                end
            elseif k == "Notes"
                rcd["Notes"] = []
                v != " " && append!(rcd["Notes"], split(v, ","))
            end
            next = iterate(lines, state)
            next === nothing && break
        end
        revdict[rcd["Symbol"]] = rcd["Number"]
        if length(data) + 1 == rcd["Number"]
            if !isempty(data)
                e = data[end]
                # Monoisotopic mass
                comp = [
                    r["Composition"] for r in e["Isotopes"]
                    if r["Composition"] !== NaN]
                if !isempty(comp)
                    val, i = findmax(comp)
                    e["Monoisotopic"] = e["Isotopes"][i]["Mass"]
                    e["MonoisotopicUncertainty"
                        ] = e["Isotopes"][i]["MassUncertainty"]
                else
                    e["Monoisotopic"] = NaN
                    e["MonoisotopicUncertainty"] = NaN
                end
            end
            rcd["Isotopes"] = [iso]
            push!(data, rcd)
            # New entry
            if next === nothing
                e = data[end]
                # Monoisotopic mass
                comp = [
                    r["Composition"] for r in e["Isotopes"]
                    if r["Composition"] !== NaN]
                if !isempty(comp)
                    val, i = findmax(comp)
                    e["Monoisotopic"] = e["Isotopes"][i]["Mass"]
                    e["MonoisotopicUncertainty"
                        ] = e["Isotopes"][i]["MassUncertainty"]
                else
                    e["Monoisotopic"] = NaN
                    e["MonoisotopicUncertainty"] = NaN
                end
                break
            end
        else
            if rcd["Symbol"] != data[end]["Symbol"]
                iso["Symbol"] = rcd["Symbol"]
            end
            push!(data[end]["Isotopes"], iso)
        end
        next = iterate(lines, state)
    end
    println(length(data))
    YAML.write_file(dest, data)
    YAML.write_file(revdest, revdict)
end

run()


module TestScikitLearnWrapper

include(joinpath("..", "fixture_learners.jl"))
using .FixtureLearners
nfcp = NumericFeatureClassification()

using FactCheck

importall Orchestra.Types
importall Orchestra.Transformers.ScikitLearnWrapper
using PyCall
@pyimport sklearn.neighbors as NN
@pyimport random as RAN

function skl_fit_and_transform!(learner::Learner, problem::MLProblem, seed=1)
  RAN.seed(seed)
  srand(seed)
  return fit_and_transform!(learner, problem, seed)
end

function backend_fit_and_transform!(sk_learner, seed=1)
  RAN.seed(seed)
  srand(seed)
  sk_learner[:fit](nfcp.train_instances, nfcp.train_labels)
  return collect(sk_learner[:predict](nfcp.test_instances))
end

function behavior_check(learner::Learner, sk_learner)
  # Predict with Orchestra learner
  orchestra_predictions = skl_fit_and_transform!(learner, nfcp)

  # Predict with original backend learner
  original_predictions = backend_fit_and_transform!(sk_learner)

  # Verify same predictions
  @fact orchestra_predictions => original_predictions
end

facts("scikit-learn learners") do
  context("SKLLearner gives same results as its backend") do
    learner_names = collect(keys(ScikitLearnWrapper.learner_dict))
    for learner_name in learner_names
      sk_learner = ScikitLearnWrapper.learner_dict[learner_name]()
      impl_options = Dict()

      fragile_learners = [
        "AdaBoostClassifier",
        "BaggingClassifier",
        "ExtraTreesClassifier",
        "GradientBoostingClassifier",
        "RandomForestClassifier",
        "SGDClassifier",
        "RadiusNeighborsClassifier",
        "RidgeClassifierCV"
      ]

      learner = SKLLearner({
        :learner => learner_name,
        :impl_options => impl_options
      })

      if !in(learner_name, fragile_learners)
        behavior_check(learner, sk_learner)
      else
        skl_fit_and_transform!(learner, nfcp)
        @fact 1 => 1
      end
    end
  end
end

end # module