@testset "setup object" begin
    param = :complexConjug
    θ₀ = [10.0, -5.0, 5.0, 0.0, 3.0]
    P˟ = BSI.FitzhughDiffusion(param, θ₀...)
    obs = [[0.0, 0.0], [-0.3, -0.2], [0.4, 0.2]]
    tt = [0.0, 1.0, 1.5]
    P̃ = [BSI.FitzhughDiffusionAux(param, θ₀..., tt[1], obs[1], tt[2], obs[2]),
         BSI.FitzhughDiffusionAux(param, θ₀..., tt[2], obs[2], tt[3], obs[3])]
    setup = BSI.MCMCSetup(P˟, P̃, PartObs())

    @testset "initialisation" begin
        @test setup.P˟ == P˟
        @test setup.P̃ == P̃
        @test setup.blocking == BSI.NoBlocking()
        @test setup.blocking_params == ([], 0.1, BSI.NoChangePt())
        @test !any(values(setup.setup_completion))
        @test @suppress !BSI.check_if_complete(setup, [:obs])
        @test @suppress !BSI.check_if_complete(setup, [:imput])
        @test @suppress !BSI.check_if_complete(setup, [:tkern])
        @test @suppress !BSI.check_if_complete(setup, [:prior])
        @test @suppress !BSI.check_if_complete(setup, [:mcmc])
        @test @suppress !BSI.check_if_complete(setup, [:solv])
    end

    L = [1. 0.; 0. 1.]
    Σ = [0.5 0.0; 0.0 1.0]
    BSI.set_observations!(setup, [L, L], [Σ, 2*Σ], obs, tt)

    @testset "setting observations" begin
        @test setup.Ls == [L, L]
        @test setup.Σs == [Σ, 2*Σ]
        @test setup.obs == obs
        @test setup.obs_times == tt
        @test setup.fpt == [nothing, nothing]
        @test @suppress !BSI.check_if_complete(setup, [:imput])
        @test @suppress !BSI.check_if_complete(setup, [:tkern])
        @test @suppress !BSI.check_if_complete(setup, [:prior])
        @test @suppress !BSI.check_if_complete(setup, [:mcmc])
        @test @suppress !BSI.check_if_complete(setup, [:solv])
        @test @suppress BSI.check_if_complete(setup, [:obs])
    end

    dt = 0.01
    τ(t₀,T) = (x) ->  t₀ + (x-t₀) * (2-(x-t₀)/(T-t₀))
    BSI.set_imputation_grid!(setup, dt)

    @testset "setting imputation grid" begin
        @test setup.dt == dt
        @test setup.τ(tt[1], tt[2])(0.5*(tt[1]+tt[2])) == τ(tt[1], tt[2])(0.5*(tt[1]+tt[2]))
        @test @suppress !BSI.check_if_complete(setup, [:tkern])
        @test @suppress !BSI.check_if_complete(setup, [:prior])
        @test @suppress !BSI.check_if_complete(setup, [:mcmc])
        @test @suppress !BSI.check_if_complete(setup, [:solv])
        @test @suppress BSI.check_if_complete(setup, [:obs, :imput])
    end

    t_kernels = [BSI.RandomWalk([0.002, 0.1], [true, true]),
                 BSI.RandomWalk([0.2, 1.0], 2)]
    ρ = 0.5
    param_updt = true
    updt_coord_true = (Val((true,true,false,false,false)),
                       Val((false,true,true,false,false)))
    updt_coord = ((1,2),(2,3))
    updt_type=(BSI.MetropolisHastingsUpdt(),
               BSI.ConjugateUpdt())
    set_transition_kernels!(setup, t_kernels, ρ, param_updt, updt_coord,
                            updt_type)
    @testset "setting transition kernels" begin
        @test setup.t_kernel == t_kernels
        @test setup.ρ == ρ
        @test setup.param_updt == param_updt
        @test setup.updt_coord == updt_coord_true
        @test setup.updt_type == updt_type
        @test !BSI.check_if_adapt(setup.adaptive_prop)
        @test @suppress !BSI.check_if_complete(setup, [:prior])
        @test @suppress !BSI.check_if_complete(setup, [:mcmc])
        @test @suppress !BSI.check_if_complete(setup, [:solv])
        @test @suppress BSI.check_if_complete(setup, [:obs, :imput, :tkern])
    end

    @testset "reformat updt_coord object" begin
        @test BSI.reformat_updt_coord((1,3),[1,2,3]) == (Val((true,false,true)),)
        @test BSI.reformat_updt_coord(nothing,[1,2,3]) == (Val((true,)),)
        @test BSI.reformat_updt_coord([(1,2),(2,3)],[1,2,3]) == (Val{(true, true, false)}(), Val{(false, true, true)}())
        @test_throws AssertionError BSI.reformat_updt_coord([4],[1,2,3])
    end

    priors = BSI.Priors((BSI.ImproperPrior(), BSI.ImproperPrior()))
    x0_prior = BSI.KnownStartingPt(obs[1])
    set_priors!(setup, priors, x0_prior)
    @testset "setting priors" begin
        @test setup.priors == priors
        @test setup.x0_prior == x0_prior
        @test @suppress !BSI.check_if_complete(setup, [:mcmc])
        @test @suppress !BSI.check_if_complete(setup, [:solv])
        @test @suppress BSI.check_if_complete(setup, [:obs, :imput, :tkern, :prior])
    end

    num_mcmc_steps = 100
    BSI.set_mcmc_params!(setup, num_mcmc_steps)
    @testset "setting  mcmc parameters" begin
        @test setup.num_mcmc_steps == num_mcmc_steps
        @test isnan(setup.save_iter)
        @test isnan(setup.verb_iter)
        @test setup.skip_for_save == 1
        @test setup.warm_up == 0
        @test @suppress !BSI.check_if_complete(setup, [:solv])
        @test @suppress BSI.check_if_complete(setup, [:obs, :imput, :tkern, :prior,
                                                  :mcmc])
    end

    BSI.set_solver!(setup)
    @testset "setting solver" begin
        @test setup.solver == BSI.Ralston3()
        @test setup.change_pt == BSI.NoChangePt()
        @test @suppress BSI.check_if_complete(setup)
    end

    @testset "determining data type" begin
        @test BSI.determine_data_type(setup) == (SArray{Tuple{2},Float64,1,2},
                                             Float64)
    end

    BSI.prepare_containers!(setup)
    @testset "setting internal containers" begin
        @test setup.Wnr == Wiener{Float64}()
        @test eltype(setup.WW) == SamplePath{Float64}
        @test length(setup.WW) == 2
        @test eltype(setup.XX) == SamplePath{SArray{Tuple{2},Float64,1,2}}
        @test length(setup.XX) == 2
        @test setup.Ls == [SMatrix{2,2}(L), SMatrix{2,2}(L)]
        @test setup.Σs == [SMatrix{2,2}(Σ), SMatrix{2,2}(2*Σ)]
        @test setup.obs == map(x->SVector{2}(x), obs)
    end

    BSI.initialise!(Float64, setup)
    @testset "initialisation of proposal law" begin
        @test length(setup.P) == 2
        @test typeof(setup.P[1].Pt) == typeof(setup.P̃[1])
        @test setup.P[1].Target == setup.P˟
    end
end


"""
$(TYPEDEF)
"""
struct NonlinearSolution{T,N,uType,R,P,A,O,uType2} <: AbstractNonlinearSolution{T,N}
  u::uType
  resid::R
  prob::P
  alg::A
  retcode::Symbol
  original::O
  left::uType2
  right::uType2
end

const SteadyStateSolution = NonlinearSolution

function build_solution(prob::AbstractNonlinearProblem,
                        alg,u,resid;calculate_error = true,
                        retcode = :Default,
                        original = nothing,
                        left = nothing,
                        right = nothing,
                        kwargs...)

  T = eltype(eltype(u))
  N = length((size(prob.u0)...,))

  NonlinearSolution{T,N,typeof(u),typeof(resid),
                    typeof(prob),typeof(alg),typeof(original),typeof(left)}(
                    u,resid,prob,alg,retcode,original,left,right)
end

function sensitivity_solution(sol::AbstractNonlinearSolution,u)
  T = eltype(eltype(u))
  N = length((size(sol.prob.u0)...,))

  NonlinearSolution{T,N,typeof(u),typeof(sol.resid),
                    typeof(sol.prob),typeof(sol.alg),
                    typeof(sol.original),typeof(sol.left)}(
                    u,sol.resid,sol.prob,sol.alg,sol.retcode,sol.original,sol.left,sol.right)
end


module SinKing

include("./types.jl")
include("./network.jl")
include("./signals.jl")

module AgentParts
include("./agent_parts/iz_neuron.jl")
# include("./agent_parts/conductance_injection.jl")
# include("./agent_parts/markram_transmitter.jl")
include("./agent_parts/lif_neuron.jl")
include("./agent_parts/dc.jl")
end

module Agents
#include("./agents/iz_cond_mrkrm.jl")
include("./agents/lif_simple.jl")
include("./agents/dc_source.jl")
include("./agents/delayer.jl")
end

end # module


module Gabi

#Included Files
include(joinpath(Pkg.dir(),"Gabi", "src","models","svm","LibSvm.jl"))
include(joinpath(Pkg.dir(),"Gabi","src","models","svm","SvmModel.jl"))
include(joinpath(Pkg.dir(),"Gabi","src","models","knn","knn.jl"))
# Exports
export Svm
export predict
export train!
export SvmType
export SvmKernel
export Knn
export normalize
end # module


module Fields

export SpectralElementField, assemble, data

using CUDA
using Printf

using ClimateMachine.Mesh.Grids: DiscontinuousSpectralElementGrid
using ClimateMachine.MPIStateArrays: MPIStateArray
using ..Domains: AbstractDomain

#####
##### SpectralElementField
#####

struct SpectralElementField{E, D, G, A, V, R}
    elements::E
    domain::D
    grid::G
    data::A
    realdata::V
    realelems::R # Required to build realdata from data
end

"""
    SpectralElementField(domain::RectangularDomain, state::MPIStateArray, variable_index::Int)

Returns a Cartesian `view` into `state.realdata[:, variable_index, :]`,
assuming that `state.realdata` lives on `RectangularDomain`.

`SpectralElementField.elements` is a three-dimensional array of `RectangularElements`.
"""
function SpectralElementField(
    domain::AbstractDomain,
    grid::DiscontinuousSpectralElementGrid,
    state::MPIStateArray,
    variable_index::Int,
)

    data = view(state.data, :, variable_index, :)
    realdata = view(state.realdata, :, variable_index, :)
    realelems = state.realelems

    return SpectralElementField(domain, grid, data, realdata, realelems)
end

const SEF = SpectralElementField

Base.@propagate_inbounds Base.getindex(field::SEF, i, j, k) =
    field.elements[i, j, k]
Base.size(field::SEF) = size(field.elements)
Base.eltype(field::SEF) = eltype(field.realdata)

Base.maximum(f, field::SEF) = maximum([maximum(f, el) for el in field.elements])
Base.minimum(f, field::SEF) = minimum([minimum(f, el) for el in field.elements])

Base.maximum(field::SEF) = maximum([maximum(el) for el in field.elements])
Base.minimum(field::SEF) = minimum([minimum(el) for el in field.elements])

Base.show(io::IO, field::SEF{E}) where {E} =
    print(io, "SpectralElementField{$(E.name.wrapper)}")

#####
##### Domain-specific stuff
#####

include("rectangular_element.jl")
include("rectangular_spectral_element_fields.jl")

#####
##### CPU and GPU-friendly element assembly
#####

"""
    assemble(u::SpectralElementField)

Assemble `u.elements` into a single `element::eltype(u)`, averaging shared nodes.
"""
assemble(u::SpectralElementField) = assemble(u.elements)

"""
    assemble_data(u::SpectralElementField{<:CuArray})

Assemble the data in `u.elements` into a single `Array`, averaging shared nodes.
"""
function assemble(u::SpectralElementField{<:CuArray})

    domain = u.domain
    index = u.variable_index

    cpudata = Array(u.realdata)

    u_cpu = SpectralElementField(cpudata, domain)

    return assemble(u_cpu)
end

end # module


   0:2e-3:2π    .|>d->(P=
 fill(5<<11,64 ,25);z=8cis(
d)sin(.46d);P[ 64,:].=10;for
r=0:98,c=0 :5^3 x,y=@.mod(2-
$reim((.016c-r/49im-1-im)z),
 4)-2;4-x^2>√2(y+.5-√√x^2)^
  2&&(P[c÷2+1,r÷4+1]|=Int(
    ")*,h08H¨"[4&4c+1+r&
      3])-40)end;print(
       "\e[H\e[1;31m",
         join(Char.(
            P)))
             );


using Random
using Distributions
#using FunctionWrappers

const Age=UInt8
abstract type AbstractSimParams end

abstract type InfectionModulation end

abstract type AbstractOutsideCases end

include("params/age_coupling.jl")
include("params/households.jl")
include("params/friendship.jl")
include("params/modulations.jl")
include("params/progression.jl")
include("params/hospital.jl")
include("params/phonetracing.jl")
include("params/spreading.jl")
include("params/strains.jl")
include("params/outside_cases.jl")
include("params/screening.jl")

struct SimParams <: AbstractSimParams
  household_ptrs::Vector{Tuple{PersonIdx,PersonIdx}}  # (i1,i2) where i1 and i2 are the indices of first and last member of the household

  ages::Vector{Age}
  genders::BitVector

  progressions::Vector{Progression} # not sure if progressions should be there

  strain_table::StrainTable

  constant_kernel_param::Float32
  household_kernel_param::Float32
  age_coupling_params::Union{Nothing, AgeCouplingParams} # nothing if kernel not active
  hospital_kernel_params::Union{Nothing, HospitalInfectionParams}  # nothing if hospital kernel not active

  hospital_detections::Bool
  mild_detection_prob::Float64

  backward_tracing_prob::Float32
  backward_detection_delay::TimeDiff

  forward_tracing_prob::Float32
  forward_detection_delay::TimeDiff

  quarantine_length::Float32
  testing_time::TimeDiff

  phone_tracing_params::Union{Nothing, PhoneTracingParams}

  infection_modulation_function::Union{Nothing,InfectionModulation}
  screening_params::Union{Nothing, ScreeningParams}
  spreading_params::Union{Nothing, SpreadingParams}
end

numindividuals(params::SimParams) = length(params.household_ptrs)
straindata(params::SimParams, strain::StrainKind) = getdata(params.strain_table, strain)

progressionof(params::SimParams, person_id::Integer) = params.progressions[person_id]
severityof(params::SimParams, person_id::Integer) = progressionof(params, person_id).severity
householdof(params::SimParams, person_id::Integer) = UnitRange(params.household_ptrs[person_id]...)
age(params::SimParams, person_id::Integer) = params.ages[person_id]
gender(params::SimParams, person_id::Integer) = params.genders[person_id]

socialcompetence(params::SimParams, person_id::Integer) =
  nothing!=params.friendship_kernel_params && socialcompetence(params.friendship_kernel_params, person_id)

ishealthcare(params::SimParams, person_id::Integer) =
  nothing!=params.hospital_kernel_params && ishealthcare(params.hospital_kernel_params, person_id)
uses_phone_tracing(params::SimParams, person_id::Integer) =
  nothing!=params.phone_tracing_params && uses_phone_tracing(params.phone_tracing_params, person_id)

spreading(params::SimParams, person_id::Integer) = isnothing(params.spreading_params) ? 1.0 : spreading(params.spreading_params, person_id)

function load_params(rng=MersenneTwister(0);
        population::DataFrame,
        infection_modulation_name::Union{Nothing,AbstractString}=nothing,
        infection_modulation_params::NamedTuple=NamedTuple{}(),
        kwargs...
        )

  individuals_df::DataFrame = population

  num_individuals = individuals_df |> nrow

  dist_incubation_time = LogNormal(1.3669786931887833, 0.5045104580676582)
  dist_symptom_onset_time = Gamma(0.8738003969079596, 2.9148873266517685)
  dist_hospitalization_time = Gamma(1.1765988120148885, 2.6664347368236787)
  dist_mild_recovery_time = Uniform(11, 17)
  dist_severe_recovery_time = Uniform(4*7, 8*7)
  dist_death_time = LogNormal(2.610727719719777, 0.44476420066780653)

  progressions = Vector{Progression}(undef, num_individuals);
  resample!(rng, progressions, individuals_df.age,
    dist_incubation_time,
    dist_symptom_onset_time,
    dist_hospitalization_time,
    dist_mild_recovery_time,
    dist_severe_recovery_time,
    dist_death_time
  )

  infection_modulation_function = isnothing(infection_modulation_name) ? nothing : make_infection_modulation(infection_modulation_name; infection_modulation_params...)

  make_params(
    rng;
    individuals_df=individuals_df,
    progressions=progressions,
    infection_modulation_function=infection_modulation_function,
    kwargs...
  )
end

function make_params(
  rng::AbstractRNG=MersenneTwister(0);
  individuals_df::DataFrame,
  progressions::AbstractArray{Progression},

  infection_modulation_function=nothing,

  screening_params::Union{Nothing,ScreeningParams}=nothing,

  hospital_kernel_param::Float64=0.0,
  healthcare_detection_prob::Float64=0.8,
  healthcare_detection_delay::Float64=1.0,

  constant_kernel_param::Float64=1.0,
  household_kernel_param::Float64=1.0,

  hospital_detections::Bool=true,
  mild_detection_prob::Float64=0.0,

  backward_tracing_prob::Float64=0.0,
  backward_detection_delay::Float64=1.0,

  forward_tracing_prob::Float64=0.0,
  forward_detection_delay::Float64=1.0,

  quarantine_length::Float64=14.0,
  testing_time::Float64=1.0,

  phone_tracing_usage::Real=0.0,
  phone_detection_delay::Real=0.25,
  phone_tracing_usage_by_household::Bool=false,

  spreading_alpha::Union{Nothing,Real}=nothing,
  spreading_x0::Real=1,
  spreading_truncation::Real=Inf,

  british_strain_multiplier::Real=1.70,
  delta_strain_multiplier::Real=1.7*1.5,

  age_coupling_thresholds::Union{Nothing, AbstractArray{T} where T<:Real}=nothing,
  age_coupling_weights::Union{Nothing, AbstractMatrix{T} where T<:Real}=nothing,
  age_coupling_use_genders::Bool=false,
  age_coupling_param::Union{Nothing, Real}=nothing,
)
  sort!(individuals_df, :household_index)

  num_individuals = individuals_df |> nrow

  @assert num_individuals == length(progressions)

  household_ptrs = make_household_ptrs(individuals_df.household_index)

  strain_table = make_strains(british_multiplier=british_strain_multiplier, delta_multiplier=delta_strain_multiplier)

  age_coupling_kernel_params =
    if nothing === age_coupling_weights && nothing === age_coupling_thresholds && nothing === age_coupling_param; nothing
    elseif age_coupling_weights !== nothing && age_coupling_thresholds !== nothing
      @assert minimum(individuals_df.age) >= 0
      @assert maximum(individuals_df.age) * (age_coupling_use_genders+1) < typemax(GroupIdx)
      AgeCouplingParams(
        individuals_df.age,
        age_coupling_use_genders === nothing ? individuals_df.gender : nothing,
        age_coupling_thresholds, age_coupling_weights,
        age_coupling_param
        )
    else error("age couplig params not fully given")
    end

  hospital_kernel_params =
    if 0 == hospital_kernel_param; nothing
    elseif 0.0 < hospital_kernel_param;
      HospitalInfectionParams(
        individuals_df.ishealthcare,
        (UInt32(1):UInt32(num_individuals))[individuals_df.ishealthcare],
        hospital_kernel_param,
        healthcare_detection_prob,
        healthcare_detection_delay
      )
    else error("hospital_kernel_param must be postive or 0, got $hospital_kernel_param")
    end

  phone_tracing_params =
    if 0 == phone_tracing_usage; nothing
    elseif 0.0 < phone_tracing_usage <= 1.0
      phone_tracing_usage_by_household ?
        PhoneTracingParams(rng, num_individuals, phone_tracing_usage, phone_detection_delay, 1, household_ptrs) :
        PhoneTracingParams(rng, num_individuals, phone_tracing_usage, phone_detection_delay, 1)
    else error("tracing_app_usage must be nonnegative, got $phone_tracing_usage")
    end

  spreading_params =
    if nothing === spreading_alpha; nothing
    elseif 0.0 < spreading_alpha
      SpreadingParams(rng, num_individuals, alpha=spreading_alpha, x0=spreading_x0, truncation=spreading_truncation)
    else error("spreading_alpha must be larger than 0, got $spreading_alpha")
    end

  params = SimParams(
    household_ptrs,
    individuals_df.age,
    individuals_df.gender,

    progressions,

    strain_table,

    constant_kernel_param,
    household_kernel_param,
    age_coupling_kernel_params,
    hospital_kernel_params,

    hospital_detections,
    mild_detection_prob,

    backward_tracing_prob,
    backward_detection_delay,

    forward_tracing_prob,
    forward_detection_delay,

    quarantine_length, # quarantine length
    testing_time, # testing time

    phone_tracing_params,
    infection_modulation_function,
    screening_params,
    spreading_params
  )
  params
end

function saveparams(dict, p::SimParams)
  dict["num_individuals"] = numindividuals(p)
  dict["constant/kernel_param"] = p.constant_kernel_param
  dict["household/kernel_param"] = p.household_kernel_param
  dict["quarantine/duration"] = p.quarantine_length
  dict["detections/hospital"] = p.hospital_detections
  dict["detections/mild"] = p.mild_detection_prob

  dict["infection_modulation"] = p.infection_modulation_function

  dict["tracing/backward_prob"] = p.backward_tracing_prob
  dict["tracing/backward_detection_delay"] = p.backward_detection_delay
  dict["tracing/forward_tracing_prob"] = p.forward_tracing_prob
  dict["tracing/forward_detection_delay"] = p.forward_detection_delay
  dict["tracing/testing_time"] = p.testing_time

  saveparams(dict, p.progressions, "progressions/")
  nothing===p.hospital_kernel_params || saveparams(dict, p.hospital_kernel_params, "hospital/")
  nothing===p.phone_tracing_params || saveparams(dict, p.phone_tracing_params, "phone_tracing/")
  nothing===p.spreading_params || saveparams(dict, p.spreading_params, "spreading")
  nothing
end


# Script executed by Biowulf to run code

@everywhere include("/home/carsonc/StochasticGene/src/StochasticGene.jl")

include("/home/carsonc/StochasticGene/runfiles/scriptfunctions.jl")

@time fit_rna(parse(Int,ARGS[1]),ARGS[2],ARGS[3],parse(Int,ARGS[4]),parse(Float64,ARGS[5]),ARGS[6],ARGS[7],ARGS[8],ARGS[9],ARGS[10],parse(Int,ARGS[11]),parse(Bool,ARGS[12]),parse(Bool,ARGS[13]),parse(Bool,ARGS[14]))

# Arguments
# 1: nchains
# 2: gene
# 3: cell
# 4: cond
# 5: G
# 6: maxtime
# 7: infolder
# 8: resultfolder
# 9: datafolder
# 10: inlabel
# 11: label
# 12: nsets (number of rate sets)
# 13: runcycle (bool)
# 14: transient (bool)


function get_nous1()
m = Model()

# ----- Variables ----- #
@variable(m, objvar)
x_Idx = Any[2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
@variable(m, x[x_Idx])
b_Idx = Any[50, 51]
@variable(m, b[b_Idx], Bin)
JuMP.set_lower_bound(x[16], 0.0)
JuMP.set_lower_bound(x[14], 0.0)
JuMP.set_lower_bound(x[38], 0.0)
JuMP.set_lower_bound(x[42], 0.0)
JuMP.set_lower_bound(x[22], 0.0)
JuMP.set_lower_bound(x[2], 0.0)
JuMP.set_lower_bound(x[9], 0.0)
JuMP.set_lower_bound(x[8], 0.0)
JuMP.set_lower_bound(x[43], 0.0)
JuMP.set_lower_bound(x[36], 0.0)
JuMP.set_lower_bound(x[4], 0.0)
JuMP.set_lower_bound(x[32], 0.0)
JuMP.set_lower_bound(x[27], 0.0)
JuMP.set_lower_bound(x[3], 0.0)
JuMP.set_lower_bound(x[25], 0.0)
JuMP.set_lower_bound(x[30], 0.0)
JuMP.set_lower_bound(x[11], 0.0)
JuMP.set_lower_bound(x[29], 0.0)
JuMP.set_lower_bound(x[5], 0.0)
JuMP.set_lower_bound(x[37], 0.0)
JuMP.set_lower_bound(x[24], 0.0)
JuMP.set_lower_bound(x[41], 0.0)
JuMP.set_lower_bound(x[18], 0.0)
JuMP.set_lower_bound(x[7], 0.0)
JuMP.set_lower_bound(x[13], 0.0)
JuMP.set_lower_bound(x[21], 0.0)
JuMP.set_lower_bound(x[10], 0.0)
JuMP.set_lower_bound(x[26], 0.0)
JuMP.set_lower_bound(x[45], 0.0)
JuMP.set_lower_bound(x[12], 0.0)
JuMP.set_lower_bound(x[40], 0.0)
JuMP.set_lower_bound(x[44], 0.0)
JuMP.set_lower_bound(x[31], 0.0)
JuMP.set_lower_bound(x[33], 0.0)
JuMP.set_lower_bound(x[28], 0.0)
JuMP.set_lower_bound(x[35], 0.0)
JuMP.set_lower_bound(x[6], 0.0)
JuMP.set_lower_bound(x[17], 0.0)
JuMP.set_lower_bound(x[23], 0.0)
JuMP.set_lower_bound(x[34], 0.0)
JuMP.set_lower_bound(x[19], 0.0)
JuMP.set_lower_bound(x[20], 0.0)
JuMP.set_lower_bound(x[39], 0.0)
JuMP.set_lower_bound(x[15], 0.0)
JuMP.set_upper_bound(x[2], 300.0)
JuMP.set_upper_bound(x[3], 300.0)
JuMP.set_upper_bound(x[4], 300.0)
JuMP.set_upper_bound(x[5], 300.0)
JuMP.set_upper_bound(x[6], 300.0)
JuMP.set_upper_bound(x[7], 300.0)
JuMP.set_upper_bound(x[8], 300.0)
JuMP.set_upper_bound(x[9], 300.0)
JuMP.set_upper_bound(x[10], 300.0)
JuMP.set_upper_bound(x[11], 300.0)
JuMP.set_upper_bound(x[12], 300.0)
JuMP.set_upper_bound(x[13], 300.0)
JuMP.set_upper_bound(x[14], 300.0)
JuMP.set_upper_bound(x[15], 300.0)
JuMP.set_upper_bound(x[16], 300.0)
JuMP.set_upper_bound(x[17], 300.0)
JuMP.set_upper_bound(x[18], 300.0)
JuMP.set_upper_bound(x[19], 300.0)
JuMP.set_upper_bound(x[22], 100.0)
JuMP.set_upper_bound(x[23], 100.0)
JuMP.set_upper_bound(x[24], 100.0)
JuMP.set_upper_bound(x[25], 100.0)
JuMP.set_upper_bound(x[26], 100.0)
JuMP.set_upper_bound(x[27], 100.0)
JuMP.set_upper_bound(x[28], 1.0)
JuMP.set_upper_bound(x[29], 1.0)
JuMP.set_upper_bound(x[30], 1.0)
JuMP.set_upper_bound(x[31], 1.0)
JuMP.set_upper_bound(x[32], 1.0)
JuMP.set_upper_bound(x[33], 1.0)
JuMP.set_upper_bound(x[34], 1.0)
JuMP.set_upper_bound(x[35], 1.0)
JuMP.set_upper_bound(x[36], 1.0)
JuMP.set_upper_bound(x[37], 1.0)
JuMP.set_upper_bound(x[38], 1.0)
JuMP.set_upper_bound(x[39], 1.0)
JuMP.set_upper_bound(x[40], 1.0)
JuMP.set_upper_bound(x[41], 1.0)
JuMP.set_upper_bound(x[42], 1.0)
JuMP.set_upper_bound(x[43], 1.0)
JuMP.set_upper_bound(x[44], 1.0)
JuMP.set_upper_bound(x[45], 1.0)
JuMP.set_lower_bound(x[46], 0.85)
JuMP.set_upper_bound(x[46], 1.0)
JuMP.set_lower_bound(x[47], 0.85)
JuMP.set_upper_bound(x[47], 1.0)
JuMP.set_lower_bound(x[48], 0.85)
JuMP.set_upper_bound(x[48], 1.0)
JuMP.set_lower_bound(x[49], 0.85)
JuMP.set_upper_bound(x[49], 1.0)

# ----- Constraints ----- #
@NLconstraint(m, e1, -((0.0042656*x[29]-0.0005719*x[28]+0.0093514*x[46]+0.0077308*x[47]-0.0139904)*x[4]+(0.0016371*x[31]+0.0288996*x[32]+0.0338147*x[48]+0.0373349*x[49]-0.0661588)*x[5])+objvar-0.23947*b[50]-0.75835*b[51] == 0.0)
@constraint(m, e2, x[2]+x[3]+x[20]+x[21] == 300.0)
@constraint(m, e3, x[6]-x[12]-x[13] == 0.0)
@constraint(m, e4, x[7]-x[11]-x[14]-x[15] == 0.0)
@constraint(m, e5, x[8]-x[10]-x[16]-x[17] == 0.0)
@constraint(m, e6, x[9]-x[18]-x[19] == 0.0)
@NLconstraint(m, e7, -x[10]*x[40]-0.333333333333333*x[2]+x[22] == 0.0)
@NLconstraint(m, e8, -x[10]*x[41]-0.333333333333333*x[2]+x[23] == 0.0)
@NLconstraint(m, e9, -x[10]*x[42]-0.333333333333333*x[2]+x[24] == 0.0)
@NLconstraint(m, e10, -x[11]*x[37]-0.333333333333333*x[3]+x[25] == 0.0)
@NLconstraint(m, e11, -x[11]*x[38]-0.333333333333333*x[3]+x[26] == 0.0)
@NLconstraint(m, e12, -x[11]*x[39]-0.333333333333333*x[3]+x[27] == 0.0)
@NLconstraint(m, e13, -(x[6]*x[34]+x[7]*x[37])+x[22] == 0.0)
@NLconstraint(m, e14, -(x[6]*x[35]+x[7]*x[38])+x[23] == 0.0)
@NLconstraint(m, e15, -(x[6]*x[36]+x[7]*x[39])+x[24] == 0.0)
@NLconstraint(m, e16, -(x[8]*x[40]+x[9]*x[43])+x[25] == 0.0)
@NLconstraint(m, e17, -(x[8]*x[41]+x[9]*x[44])+x[26] == 0.0)
@NLconstraint(m, e18, -(x[8]*x[42]+x[9]*x[45])+x[27] == 0.0)
@NLconstraint(m, e19, x[22]*x[46]-x[6]*x[34] == 0.0)
@NLconstraint(m, e20, x[23]*x[47]-x[7]*x[38] == 0.0)
@NLconstraint(m, e21, x[26]*x[48]-x[8]*x[41] == 0.0)
@NLconstraint(m, e22, x[27]*x[49]-x[9]*x[45] == 0.0)
@NLconstraint(m, e23, x[12]*x[34]+x[14]*x[37]+x[16]*x[40]+x[18]*x[43]+0.333333333333333*x[20] == 30.0)
@NLconstraint(m, e24, x[12]*x[35]+x[14]*x[38]+x[16]*x[41]+x[18]*x[44]+0.333333333333333*x[20] == 50.0)
@NLconstraint(m, e25, x[12]*x[36]+x[14]*x[39]+x[16]*x[42]+x[18]*x[45]+0.333333333333333*x[20] == 30.0)
@NLconstraint(m, e26, x[13]*x[34]+x[15]*x[37]+x[17]*x[40]+x[19]*x[43]+0.333333333333333*x[21] == 70.0)
@NLconstraint(m, e27, x[13]*x[35]+x[15]*x[38]+x[17]*x[41]+x[19]*x[44]+0.333333333333333*x[21] == 50.0)
@NLconstraint(m, e28, x[13]*x[36]+x[15]*x[39]+x[17]*x[42]+x[19]*x[45]+0.333333333333333*x[21] == 70.0)
@NLconstraint(m, e29, x[4]*x[28]-x[22] == 0.0)
@NLconstraint(m, e30, x[4]*x[29]-x[23] == 0.0)
@NLconstraint(m, e31, x[4]*x[30]-x[24] == 0.0)
@NLconstraint(m, e32, x[5]*x[31]-x[25] == 0.0)
@NLconstraint(m, e33, x[5]*x[32]-x[26] == 0.0)
@NLconstraint(m, e34, x[5]*x[33]-x[27] == 0.0)
@constraint(m, e35, x[34]+x[35]+x[36] == 1.0)
@constraint(m, e36, x[37]+x[38]+x[39] == 1.0)
@constraint(m, e37, x[40]+x[41]+x[42] == 1.0)
@constraint(m, e38, x[43]+x[44]+x[45] == 1.0)
@constraint(m, e39, x[28]+x[29]+x[30] == 1.0)
@constraint(m, e40, x[31]+x[32]+x[33] == 1.0)
@constraint(m, e41, x[36] == 0.0)
@constraint(m, e42, x[43] == 0.0)
@constraint(m, e43, x[4]-300*b[50] <= 0.0)
@constraint(m, e44, x[5]-300*b[51] <= 0.0)

# ----- Objective ----- #
@objective(m, Min, objvar)
return m
end


function _update_till_convergence!( pssrb::PotentialSimulationSetupRB{T, S, 3},
                                    convergence_limit::T,
                                    device_array_type::Type{DAT};
                                    n_iterations_between_checks = 500,
                                    depletion_handling::Val{depletion_handling_enabled} = Val{false}(),
                                    only2d::Val{only_2d} = Val{false}(),
                                    is_weighting_potential::Val{_is_weighting_potential} = Val{false}(),
                                    use_nthreads::Int = Base.Threads.nthreads(),
                                    max_n_iterations::Int = 10_000, # -1
                                    verbose::Bool = true
                                )::T where {T, S, depletion_handling_enabled, only_2d, _is_weighting_potential, DAT <: GPUArrays.AbstractGPUArray}
    device = get_device(DAT)
    N_grid_points = prod(size(pssrb.potential)[1:3] .- 2)
    kernel = get_sor_kernel(S, device)
    @showprogress for i in 1:max_n_iterations
        update_even_points = true
        wait(kernel(
            pssrb.potential, pssrb.point_types, pssrb.volume_weights, pssrb.q_eff_imp, pssrb.q_eff_fix, pssrb.ϵ_r,
            pssrb.geom_weights, pssrb.sor_const, update_even_points, depletion_handling_enabled, _is_weighting_potential, only_2d,
            ndrange = N_grid_points
        ))
        apply_boundary_conditions!(pssrb, Val(update_even_points), only2d)
        update_even_points = false
        wait(kernel(
            pssrb.potential, pssrb.point_types, pssrb.volume_weights, pssrb.q_eff_imp, pssrb.q_eff_fix, pssrb.ϵ_r,
            pssrb.geom_weights, pssrb.sor_const, update_even_points, depletion_handling_enabled, _is_weighting_potential, only_2d,
            ndrange = N_grid_points
        ))
        apply_boundary_conditions!(pssrb, Val(update_even_points), only2d)
    end
    return 0
end

get_sor_kernel(::Type{Cylindrical}, args...) = sor_cyl_gpu!(args...)
get_sor_kernel(::Type{Cartesian},   args...) = sor_car_gpu!(args...)

function get_device end


"""
SARSA-λ
-------

See learner/sarsaLlearner.jl for the core implementation.

get_all_puzzles(), get_all_episodes_in_puzzle(), get_data_by_puzzle_episode()
are in rl_data_toolsforalgos.jl
"""
function getLL_SarsaLambda(data::Array{Int,2}, alpha::Float64, gamma::Float64,
                            lambda::Float64, temp::Float64;
                            reward_magnitude=1., action_cost = 0., init_Q=0.,
                            initQSA=nothing, use_replacing_trace=true)
    data_stats, nr_states, nr_actions, nr_samples = get_data_stats_and_dimensions(data)
    if data_stats.nr_subjects != 1
        error("Invalid argument: getLL_SarsaLambda supports only one subject")
    end

    LL_unbiased = 0.

    action_selection_probs = zeros(Float64, nr_samples)

    # signals:
    RPE_record = zeros(Float64, nr_samples)
    Qsa_record = zeros(Float64, nr_samples) # trial-by-trial action value
    Vs_record = zeros(Float64, nr_samples)  # trial-by-trial max(Q(s,:))
    finalQSA_record = Dict{Int,Any}() # end of Episode, Q(sa) estimate
    RPE_QV_record = zeros(Float64, nr_samples) # the RPE defined as  Q(s_t, a_t)  -  gamma*V(s_t+1), where V(s_t+1) = max(Q(s_t+1, :))

    # Record various values (not used eventually in paper)
    state_on_counter_V1 = 1; state_on_counter_V2 = 1; state_on_counter_V3 = 1
    total_nr_of_state_on = nr_samples + data_stats.nr_episodes  # actions + nr_goal states (no action at goal)
    StateValue_V1_record = zeros(Float64, total_nr_of_state_on)  # value = max(Q(s,:)), 0 at goal state
    StateValue_V2_record = zeros(Float64, total_nr_of_state_on)  # r + gamma*V_next
    StateValue_V3_record = zeros(Float64, total_nr_of_state_on)  # V AFTER update. Start state: V before update

    PolicyParametersImbalance_record = zeros(Float64, nr_samples)
    PolicyImbalance_record = zeros(Float64, nr_samples)
    PolicyParametersImbalanceAllstates_record = Array{Array{Float64,1}, 1}(undef, nr_samples)
    [PolicyParametersImbalanceAllstates_record[i] = zeros(Float64, nr_states) for i in 1:nr_samples]
    PolicyParametersDiffChosenVsUnchosen_record = zeros(Float64, nr_samples)
    PolicyDiffChosenVsUnchosen_record = zeros(Float64, nr_samples)

    current_puzzle_id = -1
    # Puzzle = independent block of the experiment.
    # The code was made for general purposes, in this experiment we
    # have only ONE puzzle.
    is_start_state = true
    isendofexperiment = false
    modulatingfactor = 0.
    learner = []
    for i = 1:nr_samples
        if current_puzzle_id != data[i,Int(c_puzzle)]
            current_puzzle_id = data[i,Int(c_puzzle)]
            # if new puzzle -> reset learning
            if initQSA == nothing
                learner = SarsaL(nr_states = nr_states,
                                    nr_actions = nr_actions,
                                    α = alpha, γ = gamma,
                                    λ = lambda, initvalue = init_Q,
                                    use_replacing_trace = use_replacing_trace
                                    )
            else
                learner = SarsaL(nr_states = nr_states,
                                    nr_actions = nr_actions,
                                    α = alpha,
                                    γ = gamma,
                                    λ = lambda,
                                    QSA = copy(initQSA[current_puzzle_id]),
                                    use_replacing_trace = use_replacing_trace
                                    )
            end
        end
        state = data[i, Int(c_state)]
        action = data[i, Int(c_action)]
        reward::Float64= data[i, Int(c_reward)]
        reward *= reward_magnitude
        next_state = data[i, Int(c_nextState)]

        Qsa_record[i] = copy(learner.QSA[state, action])
        Vs_record[i] = copy(maximum(learner.QSA[state, :]))

        StateValue_V1_record[state_on_counter_V1] = maximum(learner.QSA[state, :])
        state_on_counter_V1 += 1
        StateValue_V2_record[state_on_counter_V2] = learner.γ*maximum(learner.QSA[state, :])
        state_on_counter_V2 += 1

        # and if goal state:
        if reward>0
            StateValue_V1_record[state_on_counter_V1]  = 0
            state_on_counter_V1 += 1
            StateValue_V2_record[state_on_counter_V2] = reward #+ action_cost
            state_on_counter_V2 += 1
        end

        if is_start_state
            StateValue_V3_record[state_on_counter_V3] = maximum(learner.QSA[state, :])
            state_on_counter_V3 += 1
        end

        p_actions = getSoftmaxProbs(learner.QSA[state,:], temp, [0., 0.])
        action_selection_probs[i] = copy(p_actions[action])
        LL_unbiased += log(p_actions[action])

        PolicyParametersImbalance_record[i] = abs(learner.QSA[state, 1] - learner.QSA[state, 2])
        PolicyImbalance_record[i] = abs(p_actions[1] - p_actions[2])
        PolicyParametersImbalanceAllstates_record[i] = abs.(learner.QSA[:, 1] - learner.QSA[:, 2])
        PolicyParametersDiffChosenVsUnchosen_record[i] = learner.QSA[state, action] - learner.QSA[state, mod(action,2)+1]
        PolicyDiffChosenVsUnchosen_record[i] = p_actions[action] - p_actions[mod(action,2)+1]

        # Sort out if end of episode or end of experiment
        next_action = 0
        if reward==0 #non goal state
            rpe_qv = learner.γ * maximum(learner.QSA[next_state, :]) - learner.QSA[state, action]
            if nr_samples > i
                next_action = data[i+1,Int(c_action)]
            else
                # println("i:$i \nS:$state | a:$action | R:$reward | S':$next_state")
                isendofexperiment = true
            end
        end

        update!(learner, state, action, next_state, next_action, reward,
                isendofexperiment, modulatingfactor)

        if reward>0; rpe_qv = copy(learner.rpe[1]); end

        RPE_record[i] = copy(learner.rpe[1])
        RPE_QV_record[i] = copy(rpe_qv)

        finalQSA_record[current_puzzle_id] = copy(learner.QSA)

        StateValue_V3_record[state_on_counter_V3] = maximum(learner.QSA[state, :])
        state_on_counter_V3 += 1
        is_start_state = reward>0 # next state will be a start state
    end

    maxLL = LL_unbiased
    trial_by_trial_p_sa = (action_selection_probs)

    trace_type = use_replacing_trace ? "ReplTrc" : "EligTrc"
    algo = "Sarsa-Lambda ($trace_type)"

    param_v = [alpha, gamma, lambda, temp,reward_magnitude, action_cost, init_Q]
    param_str = "alpha=$(round(alpha,digits=4)), gamma=$(round(gamma,digits=4)),
                lambda=$(round(lambda,digits=4)), temp=$(round(temp,digits=4)),
                reward_magnitude=$(round(reward_magnitude,digits=4)),
                action_cost=$(round(action_cost,digits=4)),
                initQSA=$(repr(initQSA)), use_replacing_trace=$use_replacing_trace"

    results_dict = Dict{String,Any}("trialByTrial_Qsa" => Qsa_record,
                                    "final_QSA" => finalQSA_record,
                                    "trialByTrial_Vs" => Vs_record,
                                    "RPE_QV" => RPE_QV_record,
                                    "StateValue_V1_record" => StateValue_V1_record,
                                    "StateValue_V2_record" => StateValue_V2_record,
                                    "StateValue_V3_record" => StateValue_V3_record,
                                    "PolicyParametersImbalance" => PolicyParametersImbalance_record,
                                    "PolicyImbalance" => PolicyImbalance_record,
                                    "PolicyParametersImbalanceAllstates" => PolicyParametersImbalanceAllstates_record,
                                    "PolicyParametersDiffChosenVsUnchosen" => PolicyParametersDiffChosenVsUnchosen_record,
                                    "PolicyDiffChosenVsUnchosen" => PolicyDiffChosenVsUnchosen_record)

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


# Main archetype function
function buildarchetypes(project)
    stockdata = project.dat["master"]

    # Identify different plan shapes, based on vertex count
    L = length.(stockdata.dims)
    vertexcount(L,i) = findall(L.==i)
    ishapes = Dict()
    push!(ishapes,"4vert"=>vertexcount(L,4))
    push!(ishapes,"6vert"=>vertexcount(L,6))
    push!(ishapes,"8vert"=>vertexcount(L,8))

    # Prepare new columns for archetype fields
    L = nrow(stockdata)
    stockdata[!,:width] = fill!(Array{Union{Missing,Float64}}(undef,L),missing)
    stockdata[!,:depth] = fill!(Array{Union{Missing,Float64}}(undef,L),missing)
    stockdata[!,:tol] = fill!(Array{Union{Missing,Float64}}(undef,L),missing)
    stockdata[!,:tor] = fill!(Array{Union{Missing,Float64}}(undef,L),missing)
    stockdata[!,:tp] = fill!(Array{Union{Missing,Float64}}(undef,L),missing)
    stockdata[!,:bc1] = fill!(Array{Union{Missing,String}}(undef,L),missing)

    stockdata = assess4vertplans(stockdata,ishapes["4vert"])
    return project
end

# 4 vertex dwellings
function assess4vertplans(stockdata,ishape)
    elevpair(verts) = [round(mean(verts[[1,3]]),digits=3),round(mean(verts[[2,4]]),digits=3)]
    plans_dims = elevpair.(stockdata.dims[ishape])

    plans = DataFrame(
        :width => maximum.(plans_dims),
        :depth => minimum.(plans_dims)
    )

    adiab_odd(verts) = sum(iseven.(findall(verts.==true)))<sum(isodd.(findall(verts.==true)))
    adiab_even(verts) = sum(iseven.(findall(verts.==true)))>sum(isodd.(findall(verts.==true)))

    i_even = findall(adiab_even.(stockdata.adiabBool[ishape]))
    i_odd = findall(adiab_odd.(stockdata.adiabBool[ishape]))

    dim_i(dims,i) = dims[i]
    plans.width[i_even] .= dim_i.(plans_dims[i_even],1)
    plans.width[i_odd] .= dim_i.(plans_dims[i_odd],2)
    plans.depth[i_even] .= dim_i.(plans_dims[i_even],2)
    plans.depth[i_odd] .= dim_i.(plans_dims[i_odd],1)

    stockdata.width[ishape] .= plans.width
    stockdata.depth[ishape] .= plans.depth
    stockdata.tol[ishape] .= 0.
    stockdata.tor[ishape] .= 0.
    stockdata.tp[ishape] .= 0.

    function findadjacencies(adiab_array)
        i_all = findall(adiab_array)
        adjacencies = missing
        if isempty(i_all)
            adjacencies = "EEEE"
        elseif length(i_all) == 1
            adjacencies = "EEEA"
        elseif length(i_all) == 2
            if diff(i_all) == [2]
                adjacencies = "EEAA"
            else
                adjacencies = "EAEA"
            end
        elseif length(i_all) == 3
            adjacencies = "EAAA"
        end
        return adjacencies
    end
    stockdata.bc1[ishape] = findadjacencies.(stockdata.adiabBool[ishape])

    return stockdata
end

function abovebelowadj(stockdata)
    # Prep array
    stockdata[!,:bc2] = fill!(Array{Union{Missing,String}}(undef,nrow(stockdata)),missing)
    # Apply External (E) roof and Ground (G) adjacencies to all single-UPRN footprints
    # ... THIS OVERLOOKS FLATS ABOVE NON-RES,BUT THESE ARE PICKED UP BELOW
    stockdata.bc2[stockdata.nUPRNs_undersameroof.==1] .= "EG"

    findall_arr(val,arr) = findall(arr.==val)
    i_1to4UPRN = vcat(findall_arr.((stockdata.nUPRNs_undersameroof,),1:4)...)
    i_nonmissing = findall(ismissing.(stockdata.EPC_Dwelling_Type).==false)
    findoccursin_arr(val,arr) = i_nonmissing[findall(occursin.(val,arr))]
    j = findoccursin_arr.(("ground-floor ","mid-floor ","top-floor ","basement "),(lowercase.(stockdata.EPC_Dwelling_Type[i_nonmissing]),))
    stockdata.bc2[intersect(i_1to4UPRN,j[1])] .= "AG"
    stockdata.bc2[intersect(i_1to4UPRN,j[2])] .= "AA"
    stockdata.bc2[intersect(i_1to4UPRN,j[3])] .= "EA"
    stockdata.bc2[intersect(i_1to4UPRN,j[4])] .= "AG"
    return stockdata
end

function gen_archcode(stockdata,tol_geom)
    # Find intersecting non-missing rows (all data present in geom cols)
    ismissing_col(data,col) = findall(ismissing.(data[:,col]).==false)
    i_nonmisisngs = intersect(vcat(
        intersect(ismissing_col.((stockdata,),[:width,:depth,:tol,:tor,:tp]))...),
        intersect(vcat(ismissing_col.((stockdata,),[:RelH2,:RelHMax]))
    ...))
    stockdata_exrp = deepcopy(stockdata[i_nonmisisngs,:])

    # Apply archetyping geometry tolerence (arg, in metres)
    apply_geomtol(arr,tol,factor) = lpad(string(Int(factor*round(arr./tol).*tol)),3,"0")
    w = apply_geomtol.(stockdata_exrp.width,tol_geom,10)
    d = apply_geomtol.(stockdata_exrp.depth,tol_geom,10)
    tol = apply_geomtol.(stockdata_exrp.tol,tol_geom,10)
    tor = apply_geomtol.(stockdata_exrp.tor,tol_geom,10)
    tp = apply_geomtol.(stockdata_exrp.tp,tol_geom,10)
    he = apply_geomtol.(stockdata_exrp.RelH2,tol_geom/2,100)
    ha = apply_geomtol.(stockdata_exrp.RelHMax,tol_geom/2,100)

    # Compile archetypes
    compilearchetype(w,d,tol,tor,tp,he,ha,bc1,bc2) =
        "w$w-d$d-tol$tol-tor$tor-tp$tp-he$he-ha$ha-adj$bc1$bc2"
    archcode = compilearchetype.(w,d,tol,tor,tp,he,ha,
        stockdata_exrp.bc1,stockdata_exrp.bc2)

    return archcode
end


module TuringGLM

import Base: ==, length
using LinearAlgebra: I
using Statistics: mean, median, std
using StatsBase: mad, coefnames
using StatsModels: ConstantTerm, FormulaTerm, FunctionTerm, Term
using StatsModels: hasintercept, response, term
using MixedModels: MixedModel, RandomEffectsTerm

# Different modelmatrix
using MixedModels: MixedModels
using StatsModels: StatsModels

# Tables API stuff
using Tables: Tables
using TableOperations: TableOperations

# LazyArrays
using LazyArrays: @~, LazyArray

# Distributions
using Distributions: UnivariateDistribution

using Reexport: @reexport

@reexport begin
    using Turing
    using StatsModels: @formula
    using Distributions:
        Bernoulli,
        Binomial,
        Gamma,
        InverseGaussian,
        NegativeBinomial,
        Normal,
        Poisson,
        TDist
end

include("utils.jl")
include("data_constructors.jl")
include("priors.jl")
include("turing_model.jl")

export turing_model
export CustomPrior, DefaultPrior

end # module


function abs(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "abs", RealFieldElement, (RealFieldElement,), arg0)
end

function abs(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "abs", jdouble, (jdouble,), arg0)
end

function abs(::Type{FastMath}, arg0::jfloat)
    return jcall(FastMath, "abs", jfloat, (jfloat,), arg0)
end

function abs(::Type{FastMath}, arg0::jint)
    return jcall(FastMath, "abs", jint, (jint,), arg0)
end

function abs(::Type{FastMath}, arg0::jlong)
    return jcall(FastMath, "abs", jlong, (jlong,), arg0)
end

function acos(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "acos", RealFieldElement, (RealFieldElement,), arg0)
end

function acos(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "acos", jdouble, (jdouble,), arg0)
end

function acosh(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "acosh", RealFieldElement, (RealFieldElement,), arg0)
end

function acosh(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "acosh", jdouble, (jdouble,), arg0)
end

function add_exact(::Type{FastMath}, arg0::jint, arg1::jint)
    return jcall(FastMath, "addExact", jint, (jint, jint), arg0, arg1)
end

function add_exact(::Type{FastMath}, arg0::jlong, arg1::jlong)
    return jcall(FastMath, "addExact", jlong, (jlong, jlong), arg0, arg1)
end

function asin(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "asin", RealFieldElement, (RealFieldElement,), arg0)
end

function asin(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "asin", jdouble, (jdouble,), arg0)
end

function asinh(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "asinh", RealFieldElement, (RealFieldElement,), arg0)
end

function asinh(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "asinh", jdouble, (jdouble,), arg0)
end

function atan(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "atan", RealFieldElement, (RealFieldElement,), arg0)
end

function atan(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "atan", jdouble, (jdouble,), arg0)
end

function atan2(::Type{FastMath}, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(FastMath, "atan2", RealFieldElement, (RealFieldElement, RealFieldElement), arg0, arg1)
end

function atan2(::Type{FastMath}, arg0::jdouble, arg1::jdouble)
    return jcall(FastMath, "atan2", jdouble, (jdouble, jdouble), arg0, arg1)
end

function atanh(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "atanh", RealFieldElement, (RealFieldElement,), arg0)
end

function atanh(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "atanh", jdouble, (jdouble,), arg0)
end

function cbrt(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "cbrt", RealFieldElement, (RealFieldElement,), arg0)
end

function cbrt(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "cbrt", jdouble, (jdouble,), arg0)
end

function ceil(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "ceil", RealFieldElement, (RealFieldElement,), arg0)
end

function ceil(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "ceil", jdouble, (jdouble,), arg0)
end

function copy_sign(::Type{FastMath}, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(FastMath, "copySign", RealFieldElement, (RealFieldElement, RealFieldElement), arg0, arg1)
end

function copy_sign(::Type{FastMath}, arg0::RealFieldElement, arg1::jdouble)
    return jcall(FastMath, "copySign", RealFieldElement, (RealFieldElement, jdouble), arg0, arg1)
end

function copy_sign(::Type{FastMath}, arg0::jdouble, arg1::jdouble)
    return jcall(FastMath, "copySign", jdouble, (jdouble, jdouble), arg0, arg1)
end

function copy_sign(::Type{FastMath}, arg0::jfloat, arg1::jfloat)
    return jcall(FastMath, "copySign", jfloat, (jfloat, jfloat), arg0, arg1)
end

function cos(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "cos", RealFieldElement, (RealFieldElement,), arg0)
end

function cos(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "cos", jdouble, (jdouble,), arg0)
end

function cosh(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "cosh", RealFieldElement, (RealFieldElement,), arg0)
end

function cosh(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "cosh", jdouble, (jdouble,), arg0)
end

function decrement_exact(::Type{FastMath}, arg0::jint)
    return jcall(FastMath, "decrementExact", jint, (jint,), arg0)
end

function decrement_exact(::Type{FastMath}, arg0::jlong)
    return jcall(FastMath, "decrementExact", jlong, (jlong,), arg0)
end

function exp(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "exp", RealFieldElement, (RealFieldElement,), arg0)
end

function exp(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "exp", jdouble, (jdouble,), arg0)
end

function expm1(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "expm1", RealFieldElement, (RealFieldElement,), arg0)
end

function expm1(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "expm1", jdouble, (jdouble,), arg0)
end

function floor(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "floor", RealFieldElement, (RealFieldElement,), arg0)
end

function floor(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "floor", jdouble, (jdouble,), arg0)
end

function floor_div(::Type{FastMath}, arg0::jint, arg1::jint)
    return jcall(FastMath, "floorDiv", jint, (jint, jint), arg0, arg1)
end

function floor_div(::Type{FastMath}, arg0::jlong, arg1::jint)
    return jcall(FastMath, "floorDiv", jlong, (jlong, jint), arg0, arg1)
end

function floor_div(::Type{FastMath}, arg0::jlong, arg1::jlong)
    return jcall(FastMath, "floorDiv", jlong, (jlong, jlong), arg0, arg1)
end

function floor_mod(::Type{FastMath}, arg0::jint, arg1::jint)
    return jcall(FastMath, "floorMod", jint, (jint, jint), arg0, arg1)
end

function floor_mod(::Type{FastMath}, arg0::jlong, arg1::jint)
    return jcall(FastMath, "floorMod", jint, (jlong, jint), arg0, arg1)
end

function floor_mod(::Type{FastMath}, arg0::jlong, arg1::jlong)
    return jcall(FastMath, "floorMod", jlong, (jlong, jlong), arg0, arg1)
end

function fma(::Type{FastMath}, arg0::jdouble, arg1::jdouble, arg2::jdouble)
    return jcall(FastMath, "fma", jdouble, (jdouble, jdouble, jdouble), arg0, arg1, arg2)
end

function fma(::Type{FastMath}, arg0::jfloat, arg1::jfloat, arg2::jfloat)
    return jcall(FastMath, "fma", jfloat, (jfloat, jfloat, jfloat), arg0, arg1, arg2)
end

function get_exponent(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "getExponent", jint, (jdouble,), arg0)
end

function get_exponent(::Type{FastMath}, arg0::jfloat)
    return jcall(FastMath, "getExponent", jint, (jfloat,), arg0)
end

function hypot(::Type{FastMath}, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(FastMath, "hypot", RealFieldElement, (RealFieldElement, RealFieldElement), arg0, arg1)
end

function hypot(::Type{FastMath}, arg0::jdouble, arg1::jdouble)
    return jcall(FastMath, "hypot", jdouble, (jdouble, jdouble), arg0, arg1)
end

function i_e_e_eremainder(::Type{FastMath}, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(FastMath, "IEEEremainder", RealFieldElement, (RealFieldElement, RealFieldElement), arg0, arg1)
end

function i_e_e_eremainder(::Type{FastMath}, arg0::RealFieldElement, arg1::jdouble)
    return jcall(FastMath, "IEEEremainder", RealFieldElement, (RealFieldElement, jdouble), arg0, arg1)
end

function i_e_e_eremainder(::Type{FastMath}, arg0::jdouble, arg1::jdouble)
    return jcall(FastMath, "IEEEremainder", jdouble, (jdouble, jdouble), arg0, arg1)
end

function increment_exact(::Type{FastMath}, arg0::jint)
    return jcall(FastMath, "incrementExact", jint, (jint,), arg0)
end

function increment_exact(::Type{FastMath}, arg0::jlong)
    return jcall(FastMath, "incrementExact", jlong, (jlong,), arg0)
end

function log(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "log", RealFieldElement, (RealFieldElement,), arg0)
end

function log(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "log", jdouble, (jdouble,), arg0)
end

function log(::Type{FastMath}, arg0::jdouble, arg1::jdouble)
    return jcall(FastMath, "log", jdouble, (jdouble, jdouble), arg0, arg1)
end

function log10(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "log10", RealFieldElement, (RealFieldElement,), arg0)
end

function log10(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "log10", jdouble, (jdouble,), arg0)
end

function log1p(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "log1p", RealFieldElement, (RealFieldElement,), arg0)
end

function log1p(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "log1p", jdouble, (jdouble,), arg0)
end

function main(::Type{FastMath}, arg0::Vector{JString})
    return jcall(FastMath, "main", void, (Vector{JString},), arg0)
end

function max(::Type{FastMath}, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(FastMath, "max", RealFieldElement, (RealFieldElement, RealFieldElement), arg0, arg1)
end

function max(::Type{FastMath}, arg0::jdouble, arg1::jdouble)
    return jcall(FastMath, "max", jdouble, (jdouble, jdouble), arg0, arg1)
end

function max(::Type{FastMath}, arg0::jfloat, arg1::jfloat)
    return jcall(FastMath, "max", jfloat, (jfloat, jfloat), arg0, arg1)
end

function max(::Type{FastMath}, arg0::jint, arg1::jint)
    return jcall(FastMath, "max", jint, (jint, jint), arg0, arg1)
end

function max(::Type{FastMath}, arg0::jlong, arg1::jlong)
    return jcall(FastMath, "max", jlong, (jlong, jlong), arg0, arg1)
end

function min(::Type{FastMath}, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(FastMath, "min", RealFieldElement, (RealFieldElement, RealFieldElement), arg0, arg1)
end

function min(::Type{FastMath}, arg0::jdouble, arg1::jdouble)
    return jcall(FastMath, "min", jdouble, (jdouble, jdouble), arg0, arg1)
end

function min(::Type{FastMath}, arg0::jfloat, arg1::jfloat)
    return jcall(FastMath, "min", jfloat, (jfloat, jfloat), arg0, arg1)
end

function min(::Type{FastMath}, arg0::jint, arg1::jint)
    return jcall(FastMath, "min", jint, (jint, jint), arg0, arg1)
end

function min(::Type{FastMath}, arg0::jlong, arg1::jlong)
    return jcall(FastMath, "min", jlong, (jlong, jlong), arg0, arg1)
end

function multiply_exact(::Type{FastMath}, arg0::jint, arg1::jint)
    return jcall(FastMath, "multiplyExact", jint, (jint, jint), arg0, arg1)
end

function multiply_exact(::Type{FastMath}, arg0::jlong, arg1::jint)
    return jcall(FastMath, "multiplyExact", jlong, (jlong, jint), arg0, arg1)
end

function multiply_exact(::Type{FastMath}, arg0::jlong, arg1::jlong)
    return jcall(FastMath, "multiplyExact", jlong, (jlong, jlong), arg0, arg1)
end

function multiply_full(::Type{FastMath}, arg0::jint, arg1::jint)
    return jcall(FastMath, "multiplyFull", jlong, (jint, jint), arg0, arg1)
end

function multiply_high(::Type{FastMath}, arg0::jlong, arg1::jlong)
    return jcall(FastMath, "multiplyHigh", jlong, (jlong, jlong), arg0, arg1)
end

function next_after(::Type{FastMath}, arg0::jdouble, arg1::jdouble)
    return jcall(FastMath, "nextAfter", jdouble, (jdouble, jdouble), arg0, arg1)
end

function next_after(::Type{FastMath}, arg0::jfloat, arg1::jdouble)
    return jcall(FastMath, "nextAfter", jfloat, (jfloat, jdouble), arg0, arg1)
end

function next_down(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "nextDown", jdouble, (jdouble,), arg0)
end

function next_down(::Type{FastMath}, arg0::jfloat)
    return jcall(FastMath, "nextDown", jfloat, (jfloat,), arg0)
end

function next_up(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "nextUp", jdouble, (jdouble,), arg0)
end

function next_up(::Type{FastMath}, arg0::jfloat)
    return jcall(FastMath, "nextUp", jfloat, (jfloat,), arg0)
end

function pow(::Type{FastMath}, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(FastMath, "pow", RealFieldElement, (RealFieldElement, RealFieldElement), arg0, arg1)
end

function pow(::Type{FastMath}, arg0::RealFieldElement, arg1::jdouble)
    return jcall(FastMath, "pow", RealFieldElement, (RealFieldElement, jdouble), arg0, arg1)
end

function pow(::Type{FastMath}, arg0::RealFieldElement, arg1::jint)
    return jcall(FastMath, "pow", RealFieldElement, (RealFieldElement, jint), arg0, arg1)
end

function pow(::Type{FastMath}, arg0::jdouble, arg1::jdouble)
    return jcall(FastMath, "pow", jdouble, (jdouble, jdouble), arg0, arg1)
end

function pow(::Type{FastMath}, arg0::jdouble, arg1::jint)
    return jcall(FastMath, "pow", jdouble, (jdouble, jint), arg0, arg1)
end

function pow(::Type{FastMath}, arg0::jdouble, arg1::jlong)
    return jcall(FastMath, "pow", jdouble, (jdouble, jlong), arg0, arg1)
end

function random(::Type{FastMath})
    return jcall(FastMath, "random", jdouble, ())
end

function rint(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "rint", RealFieldElement, (RealFieldElement,), arg0)
end

function rint(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "rint", jdouble, (jdouble,), arg0)
end

function round(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "round", jlong, (RealFieldElement,), arg0)
end

function round(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "round", jlong, (jdouble,), arg0)
end

function round(::Type{FastMath}, arg0::jfloat)
    return jcall(FastMath, "round", jint, (jfloat,), arg0)
end

function scalb(::Type{FastMath}, arg0::RealFieldElement, arg1::jint)
    return jcall(FastMath, "scalb", RealFieldElement, (RealFieldElement, jint), arg0, arg1)
end

function scalb(::Type{FastMath}, arg0::jdouble, arg1::jint)
    return jcall(FastMath, "scalb", jdouble, (jdouble, jint), arg0, arg1)
end

function scalb(::Type{FastMath}, arg0::jfloat, arg1::jint)
    return jcall(FastMath, "scalb", jfloat, (jfloat, jint), arg0, arg1)
end

function signum(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "signum", RealFieldElement, (RealFieldElement,), arg0)
end

function signum(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "signum", jdouble, (jdouble,), arg0)
end

function signum(::Type{FastMath}, arg0::jfloat)
    return jcall(FastMath, "signum", jfloat, (jfloat,), arg0)
end

function sin(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "sin", RealFieldElement, (RealFieldElement,), arg0)
end

function sin(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "sin", jdouble, (jdouble,), arg0)
end

function sin_cos(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "sinCos", FieldSinCos, (RealFieldElement,), arg0)
end

function sin_cos(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "sinCos", SinCos, (jdouble,), arg0)
end

function sinh(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "sinh", RealFieldElement, (RealFieldElement,), arg0)
end

function sinh(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "sinh", jdouble, (jdouble,), arg0)
end

function sqrt(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "sqrt", RealFieldElement, (RealFieldElement,), arg0)
end

function sqrt(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "sqrt", jdouble, (jdouble,), arg0)
end

function subtract_exact(::Type{FastMath}, arg0::jint, arg1::jint)
    return jcall(FastMath, "subtractExact", jint, (jint, jint), arg0, arg1)
end

function subtract_exact(::Type{FastMath}, arg0::jlong, arg1::jlong)
    return jcall(FastMath, "subtractExact", jlong, (jlong, jlong), arg0, arg1)
end

function tan(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "tan", RealFieldElement, (RealFieldElement,), arg0)
end

function tan(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "tan", jdouble, (jdouble,), arg0)
end

function tanh(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "tanh", RealFieldElement, (RealFieldElement,), arg0)
end

function tanh(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "tanh", jdouble, (jdouble,), arg0)
end

function to_degrees(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "toDegrees", RealFieldElement, (RealFieldElement,), arg0)
end

function to_degrees(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "toDegrees", jdouble, (jdouble,), arg0)
end

function to_int_exact(::Type{FastMath}, arg0::jlong)
    return jcall(FastMath, "toIntExact", jint, (jlong,), arg0)
end

function to_radians(::Type{FastMath}, arg0::RealFieldElement)
    return jcall(FastMath, "toRadians", RealFieldElement, (RealFieldElement,), arg0)
end

function to_radians(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "toRadians", jdouble, (jdouble,), arg0)
end

function ulp(::Type{FastMath}, arg0::jdouble)
    return jcall(FastMath, "ulp", jdouble, (jdouble,), arg0)
end

function ulp(::Type{FastMath}, arg0::jfloat)
    return jcall(FastMath, "ulp", jfloat, (jfloat,), arg0)
end


## Contributed by @doublestrong

## Need 6 npz file and timing

using Caesar
# using RoME, Distributions
# using GraphPlot
using RoMEPlotting
Gadfly.set_default_plot_size(35cm,20cm)

# using Rebugger
# using Gadfly
using NPZ
using MAT

# # add multiproc support
using Distributed
addprocs(3)
using Caesar
@everywhere using Caesar, RoME

# # 90% precompile
for i in 1:5
  warmUpSolverJIT()
end

##

# # [OPTIONAL] add more julia processes to speed up inference
# using Distributed
# nprocs() < 3 ? addprocs(4-nprocs()) : nothing

# @everywhere using RoME
N = 1000
# start with an empty factor graph object
fg = initfg()

# fieldnames(typeof(getSolverParams(fg)))
# getSolverParams(fg).graphinit = true  # init on graph first then solve on tree (default)
# getSolverParams(fg).graphinit = false # init and solve on tree
getSolverParams(fg).useMsgLikelihoods = true
getSolverParams(fg).N = 150

# getSolverParams(fg).N = N
# # Add the first pose :x0
# addVariable!(fg, :x0, Pose2)
# # Add a few more poses
# for i in 1:5
#   addVariable!(fg, Symbol("x$(i)"), Pose2)
# end
# # Add landmarks
# for i in 1:2
#   addVariable!(fg, Symbol("l$(i)"), Point2)
# end

timing = zeros(0)

## Step 0

# Add at a fixed location Prior to pin :x0 to a starting location (0,0,0)
s_time = time_ns()
addVariable!(fg, :x0, Pose2)
addVariable!(fg, :l1, Point2)
addFactor!(fg, [:x0], PriorPose2( MvNormal([0; 0; 0], diagm([0.3;0.3;0.3].^2)) ))
ppr = Pose2Point2Range(MvNormal([7.323699045815659], diagm([0.3].^2)))
addFactor!(fg, [:x0; :l1], ppr )
# Gadfly.draw(PDF("fg.pdf", 10cm, 10cm),pl)  # or PNG(...)
# ensureAllInitialized!(fg)
# dfgplot(fg)

## Perform inference
tree, smt, hist = solveTree!(fg)
e_time = time_ns()
elapsed = (e_time - s_time)/10^9
append!(timing, elapsed)
println("elapsed time: ", elapsed)

# pl = plotKDE(fg, [:x0,:l1],dims=[1;2], levels=4)
pl = plotKDE(fg, [:x0,:l1],dims=[1;2], levels=4)

##

Gadfly.draw(PDF("batch_1.pdf", 20cm, 10cm),pl)
X0 = rand(getBelief(fg, :x0) ,N)
L1 = rand(getBelief(fg, :l1) ,N)
batch1 = vcat(X0,L1)
npzwrite("batch1.npz",batch1)

# plX1 |> PNG("/tmp/testX1.png")
# X2 = rand(getBelief(fg, :x2) ,N)
# X3 = rand(getBelief(fg, :x3) ,N)
# X4 = rand(getBelief(fg, :x4) ,N)
# X5 = rand(getBelief(fg, :x5) ,N)
# res = vcat(X0,X1,L1)
# using NPZ
# npzwrite("res.npz",res)
# Gadfly.draw(PDF("batch.pdf", 10cm, 10cm),pl)  # or PNG(...)
# error("here's a breakpoint")

## Step 1
i = 1
s_time = time_ns()
addVariable!(fg, :x1, Pose2)
pp = Pose2Pose2(MvNormal([9.82039106655709;-0.040893643507463134;0.7851856362659602], diagm([0.3;0.3;0.05].^2)))
ppr = Pose2Point2Range(MvNormal([6.783355788017825], diagm([0.3].^2)))
addFactor!(fg, [Symbol("x$(i-1)"); Symbol("x$(i)")], pp )
addFactor!(fg, [Symbol("x$(i)"); :l1], ppr )

##

tree, smt, hist = solveTree!(fg, tree)

e_time = time_ns()
elapsed = (e_time - s_time)/10^9
append!(timing, elapsed)
println("elapsed time: ", elapsed)

pl = plotKDE(fg, [:x0,:x1,:l1],dims=[1;2], levels=4)

##

Gadfly.draw(PDF("batch_2.pdf", 20cm, 10cm),pl)
X0 = rand(getBelief(fg, :x0) ,N)
L1 = rand(getBelief(fg, :l1) ,N)
X1 = rand(getBelief(fg, :x1) ,N)
batch2 = vcat(X0,X1,L1)
npzwrite("batch2.npz",batch2)

# error("here's a breakpoint")
## Step 2

i = i + 1
s_time = time_ns()
addVariable!(fg, :x2, Pose2)
addVariable!(fg, :l2, Point2)
pp = Pose2Pose2(MvNormal([9.824301116341609;-0.1827443713759503;0.7586281983933238], diagm([0.3;0.3;0.05].^2)))
ppr = Pose2Point2Range(MvNormal([6.768864478545091], diagm([0.3].^2)))
addFactor!(fg, [Symbol("x$(i-1)"); Symbol("x$(i)")], pp )
addFactor!(fg, [Symbol("x$(i)"); :l2], ppr )

##

tree, smt, hist = solveTree!(fg, tree)
# tree, smt, hist = solveTree!(fg; smtasks, recordcliqs=ls(fg))

e_time = time_ns()
elapsed = (e_time - s_time)/10^9
append!(timing, elapsed)
println("elapsed time: ", elapsed)

pl = plotKDE(fg, [:x0,:x1,:x2,:l1,:l2],dims=[1;2], levels=4)

##
Gadfly.draw(PDF("batch_3.pdf", 20cm, 10cm),pl)
X0 = rand(getBelief(fg, :x0) ,N)
L1 = rand(getBelief(fg, :l1) ,N)
L2 = rand(getBelief(fg, :l2) ,N)
X1 = rand(getBelief(fg, :x1) ,N)
X2 = rand(getBelief(fg, :x2) ,N)
batch3 = vcat(X0,X1,X2,L1,L2)
npzwrite("batch3.npz",batch3)

#
## Step 3

i = i + 1
s_time = time_ns()
addVariable!(fg, :x3, Pose2)
pp = Pose2Pose2(MvNormal([9.776502885351334; -0.010587078502017132; 1.5591793408311467], diagm([0.3;0.3;0.05].^2)))
ppr = Pose2Point2Range(MvNormal([7.401417053438512], diagm([0.3].^2)))
addFactor!(fg, [Symbol("x$(i-1)"); Symbol("x$(i)")], pp )
addFactor!(fg, [Symbol("x$(i)"); :l2], ppr )

##
tree, smt, hist = solveTree!(fg, tree)
e_time = time_ns()
elapsed = (e_time - s_time)/10^9
append!(timing, elapsed)
println("elapsed time: ", elapsed)

pl = plotKDE(fg, [:x0,:x1,:x2,:x3,:l1,:l2],dims=[1;2], levels=4)

##
Gadfly.draw(PDF("batch_4.pdf", 20cm, 10cm),pl)
X0 = rand(getBelief(fg, :x0) ,N)
L1 = rand(getBelief(fg, :l1) ,N)
L2 = rand(getBelief(fg, :l2) ,N)
X1 = rand(getBelief(fg, :x1) ,N)
X2 = rand(getBelief(fg, :x2) ,N)
X3 = rand(getBelief(fg, :x3) ,N)
batch4 = vcat(X0,X1,X2,X3,L1,L2)
npzwrite("batch4.npz",batch4)

## Step 4 (:l2 has triple range measurements)
i = i + 1
s_time = time_ns()
addVariable!(fg, Symbol("x4"), Pose2)
pp = Pose2Pose2(MvNormal([9.644657137571507; 0.5847494762836476; 0.7422440549101994], Matrix(Diagonal([0.3;0.3;0.05].^2))))
ppr = Pose2Point2Range(MvNormal([7.331883435735532], Matrix(Diagonal([0.3].^2))))
addFactor!(fg, [Symbol("x$(i-1)"); Symbol("x$(i)")], pp )
addFactor!(fg, [Symbol("x$(i)"); Symbol("l2")], ppr )

##
tree, smt, hist = solveTree!(fg, tree)
e_time = time_ns()
elapsed = (e_time - s_time)/10^9
append!(timing, elapsed)
println("elapsed time: ", elapsed)

pl = plotKDE(fg, [:x0,:x1,:x2,:x3,:x4,:l1,:l2],dims=[1;2], levels=4)

##
Gadfly.draw(PDF("batch_5.pdf", 20cm, 10cm),pl)
X0 = rand(getBelief(fg, :x0) ,N)
L1 = rand(getBelief(fg, :l1) ,N)
L2 = rand(getBelief(fg, :l2) ,N)
X1 = rand(getBelief(fg, :x1),N)
X2 = rand(getBelief(fg, :x2) ,N)
X3 = rand(getBelief(fg, :x3) ,N)
X4 = rand(getBelief(fg, :x4),N)
batch5 = vcat(X0,X1,X2,X3,X4,L1,L2)
npzwrite("batch5.npz",batch5)

## Step 5 (triple range to :l1)
i = i + 1
s_time = time_ns()
addVariable!(fg, Symbol("x5"), Pose2)
pp = Pose2Pose2(MvNormal([9.725096752125593; 0.6094800276622434; 0.7750400402527422], Matrix(Diagonal([0.3;0.3;0.05].^2))))
ppr = Pose2Point2Range(MvNormal([6.476782344345081], Matrix(Diagonal([0.3].^2))))
addFactor!(fg, [Symbol("x$(i-1)"); Symbol("x$(i)")], pp )
addFactor!(fg, [Symbol("x$(i)"); Symbol("l1")], ppr )

##
tree, smt, hist = solveTree!(fg, tree)
e_time = time_ns()
elapsed = (e_time - s_time)/10^9
append!(timing, elapsed)
println("elapsed time: ", elapsed)

pl = plotKDE(fg, [:x0,:x1,:x2,:x3,:x4,:x5,:l1,:l2],dims=[1;2], levels=4)

##
Gadfly.draw(PDF("batch_6.pdf", 20cm, 10cm),pl)
X0 = rand(getBelief(fg, :x0),N)
L1 = rand(getBelief(fg, :l1),N)
L2 = rand(getBelief(fg, :l2),N)
X1 = rand(getBelief(fg, :x1),N)
X2 = rand(getBelief(fg, :x2),N)
X3 = rand(getBelief(fg, :x3),N)
X4 = rand(getBelief(fg, :x4),N)
X5 = rand(getBelief(fg, :x5),N)
batch6 = vcat(X0,X1,X2,X3,X4,X5,L1,L2)
npzwrite("batch6.npz",batch6)
npzwrite("timing.npz",timing)

#
#
# pl = plotSLAM2D(fg)
#
# # For scripting use-cases you can export the image
# Gadfly.draw(PDF("/tmp/test.pdf", 20cm, 10cm),pl)  # or PNG(...)
#
# L1 = rand(getBelief(fg, :l1) ,N)
# # L2 = rand(getBelief(fg, :l2) ,N)
# X0 = rand(getBelief(fg, :x0) ,N)
# X1 = rand(getBelief(fg, :x1) ,N)
# # X2 = rand(getBelief(fg, :x2) ,N)
# # X3 = rand(getBelief(fg, :x3) ,N)
# # X4 = rand(getBelief(fg, :x4) ,N)
# # X5 = rand(getBelief(fg, :x5) ,N)
# res = vcat(X0,X1,L1)
# using NPZ
# npzwrite("res.npz",res)
# # pl2 = Gadfly.plot(x=X0[1,:],y=X0[2,:], Geom.hexbin);
# # p1H = hstack(pl2)
# #
# # # convert to file
# # p1H |> PNG("/home/chad/Pictures/julia_x0_step0.png")
#
# # using RoMEPlotting
#
# #drawPoses(fg)
# # If you have landmarks, you can instead call
# # drawPosesLandms(fg)
# #
# # # # Draw the KDE for x0
# # # plotKDE(fg, :x0)
# # # # Draw the KDE's for x0 and x1
# # plotKDE(fg, [:x1])
# # plotKDE(fg, [:x2])
# # plotKDE(fg, [:x3])
# # plotKDE(fg, [:x4])
# # plotKDE(fg, [:x5])
# # plotKDE(fg, [:l1])
# # plotKDE(fg, [:l2])


using Test, Algorithms, Documenter
DocMeta.setdocmeta!(Algorithms, :DocTestSetup, :(using Algorithms); recursive=true)
doctest(Algorithms)
#@test Maths.abs_value(5)==10


# ==================================================================================================== #
#                                                                                #
#                                          Helper Functions                                            #
#                                                                                #
# ==================================================================================================== #
function readConfig(file)
    @info "Reading Inputs..."

    tml = TOML.parsefile(file)
    sim = tml["Simulation"]
    χfill = nothing
    rr = r"^fixed:(?P<start>\N+):(?P<stop>\N+)"

    if tml["Simulation"]["chi_unusable_fill_value"] == "0"
        χfill = zero_χ_fill
    elseif tml["Simulation"]["chi_unusable_fill_value"] == "chi_lambda"
        χfill = lambda_χ_fill
    elseif tml["Simulation"]["chi_unusable_fill_value"] == "chi"
        χfill = χ_fill
    else
        error("could not parse chi fill value")
    end
    tc_type_f = Symbol(lowercase(tml["Simulation"]["fermionic_tail_correction"]))
    if !(tc_type_f in [:nothing, :richardson, :shanks])
        error("Unrecognized tail correction type \"$(tc_type_f)\"")
    end
    tc_type_b = Symbol(lowercase(tml["Simulation"]["bosonic_tail_correction"]))
    if !(tc_type_b in [:nothing, :richardson, :shanks, :coeffs])
        error("Unrecognized tail correction type \"$(tc_type_b)\"")
    end
    λc_type = Symbol(lowercase(tml["Simulation"]["lambda_correction"]))
    if !(λc_type in [:nothing, :sp, :sp_ch])
        error("Unrecognized tail correction type \"$(λc_type)\"")
    end
    ωsum_inp = lowercase(tml["Simulation"]["bosonic_sum_range"])
    m = match(rr, ωsum_inp)
    ωsum_type = if m !== nothing
        tuple(parse.(Int, m.captures)...)
    elseif ωsum_inp in ["common", "individual", "full"]
        Symbol(ωsum_inp)
    else
        error("Could not parse bosonic_sum. Allowed input is: common, individual, full, fixed:N:M")
    end
    smoothing = Symbol(lowercase(tml["Simulation"]["omega_smoothing"]))
    if !(smoothing in [:nothing, :range, :full])
        error("Unrecognized smoothing type \"$(smoothing)\"")
    end

    λrhs_type = Symbol(lowercase(tml["Simulation"]["rhs"]))
    !(λrhs_type in [:native, :fixed, :error_comp]) && error("Could not parse rhs type for lambda correction. Options are native, fixed, error_comp.")

    dbg_full_eom_omega = (haskey(tml["Debug"], "full_EoM_omega") && tml["Debug"]["full_EoM_omega"]) ? true : false

    env = EnvironmentVars(   tml["Environment"]["inputDataType"],
                             tml["Environment"]["writeFortran"],
                             tml["Environment"]["loadAsymptotics"],
                             tml["Environment"]["inputDir"],
                             tml["Environment"]["freqFile"],
                             tml["Environment"]["inputVars"],
                             tml["Environment"]["asymptVars"],
                             tml["Environment"]["cast_to_real"],
                             String([(i == 1) ? uppercase(c) : lowercase(c)
                                     for (i, c) in enumerate(tml["Environment"]["loglevel"])]),
                             lowercase(tml["Environment"]["logfile"]),
                             tml["Environment"]["progressbar"]
                            )
    mP = ModelParameters(tml["Model"]["U"],
                            tml["Model"]["mu"],
                            tml["Model"]["beta"],
                            tml["Model"]["nden"])

    freqFilePath = ""
    if !isfile(env.freqFile)
        m = match(r"b(?<bf>\d+)f(?<ff>\d+)",env.freqFile)
        nBose = parse(Int, m[:bf])
        nFermi = parse(Int, m[:ff])
        freqFilePath = "$(@__DIR__)/../scripts/freqList.jld2"
        gen_mesh(nBose, nFermi, 0, freqFilePath)
    else
        freqFilePath = env.freqFile
    end
    JLD2.@load freqFilePath freqRed_map freqList freqList_min parents ops nFermi nBose shift base offset

    sP = SimulationParameters(nBose,nFermi,shift,
                               tc_type_f,
                               tc_type_b,
                               λc_type,
                               ωsum_type,
                               λrhs_type,
                               tml["Simulation"]["force_full_bosonic_chi"],
                               χfill,
                               tml["Simulation"]["bosonic_tail_coeffs"],
                               tml["Simulation"]["fermionic_tail_coeffs"],
                               tml["Simulation"]["usable_prct_reduction"],
                               smoothing,
                               dbg_full_eom_omega
    )
    kGrids = []
    for Nk in tml["Simulation"]["Nk"]
        kGrid = gen_kGrid(tml["Model"]["kGrid"], Nk)
        push!(kGrids, kGrid)
    end
    kGridLoc = gen_kGrid(tml["Model"]["kGrid"], 1)
    return mP, sP, env, kGrids, kGridLoc, freqRed_map, freqList, freqList_min, parents, ops, nFermi, nBose, shift, base, offset
end

function convertGF!(GF, storedInverse, storeFull)
    if storedInverse
        GF = 1 ./ GF
    end
    if storeFull
        NH = size(GF, 1)
        tmp = copy(GF)
        GF = Array{eltype(GF)}(undef, 2*NH - 1)
        GF[1:(NH-1)] = reverse(conj.(tmp[2:NH]))
        GF[NH:end] = tmp[1:NH]
    end
end

function print_chi_bubble(qList, res, simParams)
    for j in 1:size(res,1)
        print(" ========== ω = $(j-(simParams.n_iω + 1)) =============== \n")
        for k in 1:size(res,2)
            print(" ---------- ν = $(k-1) -------------- \n")
            for (qi,q) in enumerate(qList)
                @printf("   q = (%.2f,%.2f): %.2f + %.2fi\n", q[1],q[2], real(res[j, k, qi]), imag(res[j, k, qi]))
            end
        end
    end
end

function reduce_3Freq(inp, freqBox, simParams)
    n_iω = simParams.n_iω
    n_iν = simParams.n_iν
    start_ωₙ =  -(freqBox[1,1] + n_iω) + 1
    start_νₙ =  -(freqBox[2,1] + n_iν) + 2

    if (length(start_ωₙ:(start_ωₙ+2*n_iω)) != size(inp,1)) ||
       (length(start_νₙ:(start_νₙ + 2*n_iν - 1)) != size(inp,2))
        @error "Input frequency range does not match input data. This can lead to wrong results! Input size: ", size(inp),
        " got ", length(start_ωₙ:(start_ωₙ+2*n_iω)), ",", length(start_νₙ:(start_νₙ + 2*n_iν - 1))
    end
    # This selects the range of [-ωₙ, ωₙ] and [-νₙ, νₙ-1] from gamma
    return inp[start_ωₙ:(start_ωₙ+2*n_iω), start_νₙ:(start_νₙ + 2*n_iν - 1), start_νₙ:(start_νₙ +2*n_iν - 1)]
end

reduce_3Freq!(inp, freqBox, simParams) = inp = reduce_3Freq(inp, freqBox, simParams)

function convert_from_fortran(simParams, env, loadFromBak=false)
    @info "Reading Fortran Input, this can take several minutes."
        g0 = readFortranSymmGF(simParams.n_iν+2*simParams.n_iω, env.inputDir*"/g0mand", storedInverse=true)
        gImp = readFortranSymmGF(simParams.n_iν+2*simParams.n_iω, env.inputDir*"/gm_wim", storedInverse=false)
        freqBox, Γch, Γsp = readFortranΓ(env.inputDir*"/gamma_dir")
        @info "Done Reading Gamma"
        χDMFTch, χDMFTsp = readFortranχDMFT(env.inputDir*"/chi_dir")
        @info "Done Reading chi"
    Γch = -1.0 .* reduce_3Freq(Γch, freqBox, simParams)
    Γsp = -1.0 .* reduce_3Freq(Γsp, freqBox, simParams)
    χDMFTch = reduce_3Freq(χDMFTch, freqBox, simParams)
    χDMFTsp = reduce_3Freq(χDMFTsp, freqBox, simParams)
    if env.writeFortran
        @info "Writing HDF5 (vars.jdl) and Fortran (fortran_out/) output."
        writeFortranΓ("fortran_out", "gamma", simParams, Γch, Γsp)
        writeFortranΓ("fortran_out", "chi", simParams, 0.5 .* (χDMFTch .+ χDMFTsp), 0.5 .* (χDMFTch .- χDMFTsp))
    end
    if length(env.inputVars) > 0
        JLD2.@save env.inputVars Γch Γsp χDMFTch χDMFTsp gImp g0
    end
end

# ==================================================================================================== #
#                                                                                #
#                                             JLD2 Input                                               #
#                                                                                #
# ==================================================================================================== #
function readEDAsymptotics_julia(env)
    @error("Asymptotics from jld2 not implemented yet")
end

# ==================================================================================================== #
#                                                                                #
#                                            Parquet Input                                             #
#                                                                                #
# ==================================================================================================== #
function readFortranSymmGF_pq(nFreq, filename; storedInverse, storeFull=false)
    GF = redirect_stdout(open("/dev/null","w")) do
        load(filename) |> @orderby(_.__index_level_0__) |> @take(nFreq)|> @map((_.Re) + (_.Im)*1im) |> collect
    end
    convertGF!(GF, storedInverse, storeFull)
    println(GF)
    return GF
end

function readFortranΓ_pq(fileName::String)
    #in = redirect_stdout(open("/dev/null","w")) do
        #load(filename) |> @orderby(_.__index_level_0__, _.__index_level_1__)
    #end
    #ωₙ, freqBox, Γcharge0, Γspin0
    Γcharge = Array{Complex{Float64}}(undef, length(files), size(Γspin0,1), size(Γspin0,2))
    Γspin   = Array{Complex{Float64}}(undef, length(files), size(Γspin0,1), size(Γspin0,2))
    Γcharge[1,:,:] = Γcharge0
    Γspin[1,:,:]   = Γspin0
    ω_min = ωₙ
    ω_max = ωₙ

    for (i,file) in enumerate(files[2:end])
        ωₙ, _, Γcharge_new, Γspin_new = readFortran3FreqFile(dirName * "/" * file, sign=-1.0)
        ω_min = if ωₙ < ω_min ωₙ else ω_min end
        ω_max = if ωₙ > ω_max ωₙ else ω_max end
        Γcharge[i,:,:] = Γcharge_new
        Γspin[i,:,:] = Γspin_new
    end
    freqBox = [ω_min ω_max; freqBox[1,1] freqBox[1,2]; freqBox[2,1] freqBox[2,2]]
    #Γcharge = permutedims(Γcharge, [3,1,2])
    #Γspin   = permutedims(Γspin, [3,1,2])
    return freqBox, Γcharge, Γspin
end

function readFortranχDMFT_pq(filename::String)
    in = redirect_stdout(open("/dev/null","w")) do
        load(filename) |> collect
    end
    #_, _, χup, χdown = readFortran3FreqFile(dirName * "/" * files[1], sign=1.0, freqInteger=false)

    for file in files[2:end]
        _, _, χup_new, χdown_new = readFortran3FreqFile(dirName * "/" * file, sign=1.0, freqInteger=false)
        χup = cat(χup, χup_new, dims=3)
        χdown = cat(χdown, χdown_new, dims=3)
    end
    χup = permutedims(χup, [3,1,2])
    χdown   = permutedims(χdown, [3,1,2])
    χCharge = χup .+ χdown
    χSpin   = χup .- χdown
    return χCharge, χSpin
end

function convert_from_fortran_pq(simParams, env)
    println("Reading Fortran Parquet-File Input, this can take several minutes.")
    g0 = readFortranSymmGF_pq(simParams.n_iν+simParams.n_iω, env.inputDir*"/g0mand.parquet", storedInverse=true)
    gImp = readFortranSymmGF_pq(simParams.n_iν+simParams.n_iω, env.inputDir*"/gm_wim.parquet", storedInverse=false)
    χDMFTCharge, χDMFTSpin = readFortranχDMFT_pq(env.inputDir*"/vert_chi.parquet")
    freqBox, Γcharge, Γspin = readFortranΓ_pq(env.inputDir*"/GAMMA_DM_FULLRANGE.parquet")
    Γcharge = -1.0 .* reduce_3Freq(Γcharge, freqBox, simParams)
    Γspin = -1.0 .* reduce_3Freq(Γspin, freqBox, simParams)
    χDMFTCharge = reduce_3Freq(χDMFTCharge, freqBox, simParams)
    χDMFTSpin = reduce_3Freq(χDMFTSpin, freqBox, simParams)
    println("Writing HDF5 (vars.jdl) and Fortran (fortran_out/) output.")
    writeFortranΓ("fortran_out", "gamma", simParams, Γcharge, Γspin)
    writeFortranΓ("fortran_out", "chi", simParams, 0.5 .* (χDMFTCharge .+ χDMFTSpin), 0.5 .* (χDMFTCharge .- χDMFTSpin))
    JLD2.@save env.inputVars Γch Γsp χDMFTch χDMFTsp gImp g0
end

# ==================================================================================================== #
#                                                                                #
#                                            Text Input                                                #
#                                                                                #
# ==================================================================================================== #
function readFortranqωFile(filename, nDims; readq = false, data_cols = 2)
    InString = open(filename, "r") do f
        readlines(f)
    end
    start = if readq 1 else nDims+1  end
    InArr = parse.(Float64,hcat(split.(InString)...)[start:end,:])
    if readq
        qVecArr = InArr[1:3,:]
    else
        qVecArr = []
        nDims = 0
    end
    data_λ = InArr[nDims+1,:] + InArr[nDims+2,:] .* 1im
    data = InArr[nDims+3,:] + InArr[nDims+4,:] .* 1im
    return qVecArr, data_λ, data
end

function readFortranΓ(dirName::String)
    files = readdir(dirName)
    ωₙ, freqBox, Γcharge0, Γspin0 = readFortran3FreqFile(dirName * "/" * files[1], sign=-1.0)
    Γcharge = Array{Complex{Float64}}(undef, length(files), size(Γspin0,1), size(Γspin0,2))
    Γspin   = Array{Complex{Float64}}(undef, length(files), size(Γspin0,1), size(Γspin0,2))
    Γcharge[1,:,:] = Γcharge0
    Γspin[1,:,:]   = Γspin0
    ω_min = ωₙ
    ω_max = ωₙ

    for (i,file) in enumerate(files[2:end])
        ωₙ, _, Γcharge_new, Γspin_new = readFortran3FreqFile(dirName * "/" * file, sign=-1.0)
        ω_min = if ωₙ < ω_min ωₙ else ω_min end
        ω_max = if ωₙ > ω_max ωₙ else ω_max end
        Γcharge[i+1,:,:] = Γcharge_new
        Γspin[i+1,:,:] = Γspin_new
    end
    freqBox = [ω_min ω_max; freqBox[1,1] freqBox[1,2]; freqBox[2,1] freqBox[2,2]]
    Γcharge = permutedims(Γcharge, [1,3,2])
    Γspin   = permutedims(Γspin, [1,3,2])
    return freqBox, Γcharge, Γspin
end

"""
    Returns χ_DMFT[ω, ν, ν']
"""
function readFortranχDMFT(dirName::String)
    files = readdir(dirName)
    _, _, χupup, χupdo = readFortran3FreqFile(dirName * "/" * files[1], sign=1.0, freqInteger=false)
    for file in files[2:end]
        _, _, χup_new, χupdo_new = readFortran3FreqFile(dirName * "/" * file, sign=1.0, freqInteger=false)
        χupup = cat(χupup, χup_new, dims=3)
        χupdo = cat(χupdo, χupdo_new, dims=3)
    end
    χupup = permutedims(χupup, [3,2,1])
    χupdo   = permutedims(χupdo, [3,2,1])
    χCharge = χupup .+ χupdo
    χSpin   = χupup .- χupdo
    return χCharge, χSpin
end

function readFortranχlDGA(dirName::String, nDims)
    files = readdir(dirName)
    qVecs, data_λ_i, data_i = readFortranqωFile(dirName * "/" * files[1], nDims, readq = true);
    data = Array{Complex{Float64}}(undef, length(files), 2, length(data_i))
    data[1,1,:] = data_λ_i
    data[1,2,:] = data_i
    for (i,file) in enumerate(files[2:end])
        _, data_λ_i, data_i = readFortranqωFile(dirName * "/" * file, nDims, readq = false);
        data[i+1,1,:] = data_λ_i
        data[i+1,2,:] = data_i
    end
    qVecs, data
end

function readFortranBubble(dirName::String, nBose, nFermi, nQ)
    bubble = Array{Complex{Float64}}(undef, nBose, nQ, nFermi)
    files_i = readdir(dirName)
    files = [dirName * "/" * f for f  in files_i]
    iBose = 1
    for file in files[1:end]
        InString = open(file, "r") do f
            readlines(f)
        end
        InArr = parse.(Float64,hcat(split.(InString)...)[4:end,:])
        bubble[iBose, :, :] = reshape(InArr[1,:] .+ InArr[2,:].*1im, (nQ, nFermi))
        iBose += 1
    end

    bubble = permutedims(bubble, [1,3,2])
    return bubble
end

function read_anderson_parameters(file)
    content = open(file) do f
        readlines(f)
    end

    in_epsk = false
    in_tpar = false
    ϵₖ = []
    Vₖ = []
    μ = 0
    for line in content
        if "Eps(k)" == strip(line)
            in_epsk = true
            continue
        elseif "tpar(k)" == strip(line)
            in_epsk = false
            in_tpar = true
            continue
        end

        if in_epsk
            push!(ϵₖ, parse(Float64, line))
        elseif in_tpar
            # skip last line, which is mu
            if length(Vₖ) < length(ϵₖ)
                push!(Vₖ, parse(Float64, line))
            else
                if occursin("#", line)
                    μ = parse(Float64, line[1:(findfirst("#", line))[1] - 1])
                else
                    μ = parse(Float64, line)
                end
            end
        end
    end
    return convert(Array{Float64,1}, ϵₖ), convert(Array{Float64,1}, Vₖ), μ
end

function readGImp(filename; only_positive=false)
    GFString = open(filename, "r") do f
        readlines(f)
    end

    tmp = parse.(Float64,hcat(split.(GFString)...)) # Construct a 2xN array of floats (re,im as 1st index)
    tmpG = tmp[2,:] .+ tmp[3,:].*1im
    tmpiνₙ = tmp[1,:] .* 1im
    if only_positive
        GImp = tmpG
        iνₙ  = tmpiνₙ
    else
        N = 2*size(tmpG,1)
        NH = size(tmpG,1)
        GImp = zeros(Complex{Float64}, N)
        iνₙ  = zeros(Complex{Float64}, N)
        GImp[1:(NH)] = reverse(conj.(tmpG[1:NH]))
        GImp[(NH+1):N] = tmpG[1:NH]
        iνₙ[1:(NH)] = conj.(reverse(tmpiνₙ[1:(NH)]))
        iνₙ[(NH+1):N] = tmpiνₙ[1:NH]
    end
    return iνₙ, GImp
end

function readEDAsymptotics(env, mP)
    χ_asympt = readdlm(env.inputDir * "/chi_asympt")
    χchAsympt = (χ_asympt[:,2] + χ_asympt[:,4]) / (2*mP.β*mP.β);
    χspAsympt = (χ_asympt[:,2] - χ_asympt[:,4]) / (2*mP.β*mP.β);
    #= _, χup, χdo = readFortranEDχ(env.inputDir * "/chi_dir", freqInteger = false) =#
    #= χchED = χup .+ χdo =#
    #= χspED = χup .- χdo =#
    save(env.asymptVars, "chi_ch_asympt", χchAsympt, "chi_sp_asympt", χspAsympt,
         compress=true, compatible=true)
end

function readFortranSymmGF(nFreq, filename; storedInverse, storeFull=false)
    GFString = open(filename, "r") do f
        readlines(f)
    end

    if size(GFString, 1)*(1 + 1*storeFull) < nFreq
        throw(BoundsError("nFermFreq in simulation parameters too large!"))
    end

    tmp = parse.(Float64,hcat(split.(GFString)...)[2:end,:]) # Construct a 2xN array of floats (re,im as 1st index)
    tmp = tmp[1,:] .+ tmp[2,:].*1im

    if storedInverse
        tmp = 1 ./ tmp
    end

    GF = Array{Complex{Float64}}(undef, nFreq)
    if storeFull
        NH = Int(nFreq/2)
        GF[1:(NH-1)] = reverse(conj.(tmp[2:NH]))
        GF[NH:nFreq] = tmp[1:NH]
    else
        GF = tmp[1:nFreq]
    end
    return GF
end

function readFortran3FreqFile(filename; sign = 1.0, freqInteger = true)
    InString = open(filename, "r") do f
        readlines(f)
    end
    NCols = 7
    NRows = length(InString)
    lineLen = length(InString[1])
    splitLength = floor(Int64,lineLen/NCols)
    if lineLen%NCols != 0
        @warn "   ---> Warning!! Could not find fixed column width!"
    end
    #InArr = sign .* parse.(Float64,hcat(split.(InString)...)[2:end,:])
    InArr = zeros(Float64, NRows, NCols-1)
    N = Int(sqrt(NRows))
    ErrorLine = []
    for i in 1:NRows
        row =  split_n(InString[i], splitLength, lineLen)[2:end]
        for j in 1:(NCols-1)
            InArr[i,j] = try
                sign * parse(Float64,row[j])
            catch e
                NaN
            end
        end
    end
    if freqInteger
        ωₙ = Int(parse(Float64, split(InString[1])[1]))
        tmpArr1 = filter(!isnan,InArr[:,1])
        tmpArr2 = filter(!isnan,InArr[:,2])
        middleFreqBox = Int64.([minimum(tmpArr1),maximum(tmpArr1)])
        innerFreqBox  = Int64.([minimum(tmpArr2),maximum(tmpArr2)])
        freqBox = permutedims([middleFreqBox innerFreqBox], [2,1])
    else
        #print("\rWarning, non integer frequencies in "*filename*" ignored!")
        ωₙ = 0
        freqBox = []
    end

    InCol1  = reshape(InArr[:,3] .+ InArr[:,4].*1im, (N, N))
    InCol2  = reshape(InArr[:,5] .+ InArr[:,6].*1im, (N, N))
    return ωₙ, freqBox, InCol1, InCol2
end

# ==================================================================================================== #
#                                                                                #
#                                            Text Output                                                #
#                                                                                #
# ==================================================================================================== #
function writeFortranΓ(dirName::String, fileName::String, simParams, inCol1, inCol2)
    simParams.n_iν+simParams.n_iω
    if !isdir(dirName)
        mkdir(dirName)
    end
    for (ωi,ωₙ) in enumerate(-simParams.n_iω:simParams.n_iω)
        filename = dirName * "/" * fileName * lpad(ωi-1,3,"0")
        open(filename, "w") do f
            for (νi,νₙ) in enumerate(-simParams.n_iν:(simParams.n_iν-1))
                for (ν2i,ν2ₙ) in enumerate(-simParams.n_iν:(simParams.n_iν-1))
                    @printf(f, "  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f\n", ωₙ, νₙ, ν2ₙ,
                            real(inCol1[ωi, νi, ν2i]), imag(inCol1[ωi, νi, ν2i]),
                            real(inCol2[ωi, νi, ν2i]), imag(inCol2[ωi, νi, ν2i]))
                end
            end
        end
    end
end

function writeFortranΣ(dirName::String, Σ_ladder)
    res = zeros(size(Σ_ladder,1), 3)
    if !isdir(dirName)
        mkdir(dirName)
    end

    for ki in 1:size(Σ_ladder,2)
        fn = dirName * "/SELF_Q_" * lpad(ki,6,"0") * ".dat"
        open(fn, write=true) do f
            write(f, "header...\n")
            res[:,1] = (2 .*(0:size(Σ_ladder,1)-1) .+ 1) .* π ./ mP.β
            res[:,2] = real.(Σ_ladder[:,ki])
            res[:,3] = imag.(Σ_ladder[:,ki])
            writedlm(f,  rpad.(round.(res; digits=14), 22, " "), "\t")
        end
    end
end

function writeFortranχ(dirName::String, χ, χ_λ, qGrid, usable_ω)
    if !isdir(dirName)
        mkdir(dirName)
    end
    for ωi in 1:size(χ, 1)
        fn = dirName * "/chi" * lpad(ωi-1,3,"0")
        open(fn, write=true) do f
            for qi in 1:size(χ, 2)
                res = ωi in usable_ω ? χ_λ[ωi - first(usable_ω) + 1,1] : 0.0
                if length(qGrid[1]) == 3
                    @printf(f, "  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f\n",
                        qGrid[qi][1], qGrid[qi][2], qGrid[qi][3],
                        real(χ[ωi, qi]), imag(χ[ωi, qi]),
                        real(res), imag(res))
                else
                    @printf(f, "  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f\n",
                        qGrid[qi][1], qGrid[qi][2],
                        real(χ[ωi, qi]), imag(χ[ωi, qi]),
                        real(res), imag(res))
                end
            end
        end
    end
end

function writeFortranEnergies(E_Kin, E_Pot, β, dirName::String)
    if !isdir(dirName)
        mkdir(dirName)
    end
    @assert size(E_Kin) == size(E_Pot)
    νGrid = 0:(length(E_Kin)-1)
    iν_n = iν_array(β, νGrid)

    fn_DGA = dirName * "/energiesDGA.dat"
    open(fn_DGA, "w") do f
        for i in 1:length(E_Kin)
            @printf(f, "  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f\n",
                    imag(iν_n[i]), 0.0, 0.0, real(E_Kin[i]), 0.0, 0.0, 0.0, real(E_Pot[i]), 0.0)
        end
    end
end

function writeFortranEnergies(E_Kin, E_Pot, E_Kin_ED, E_Pot_ED, β, dirName::String)
    if !isdir(dirName)
        mkdir(dirName)
    end
    @assert size(E_Kin) == size(E_Pot) == size(E_Kin_ED) == size(E_Pot_ED)
    νGrid = 0:(length(E_Kin)-1)
    iν_n = iν_array(β, νGrid)

    writeFortranEnergies(E_Kin, E_Pot, β, dirName)

    fn_DMFT = dirName * "/energiesDMFT.dat"
    open(fn_DMFT, "w") do f
        for i in 1:length(E_Kin)
            @printf(f, "  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f  %18.10f\n",
                    imag(iν_n[i]), 0.0, 0.0, real(E_Kin_ED[i]), 0.0, 0.0, 0.0, real(E_Pot_ED[i]), 0.0)
        end
    end
end


"""
    conceptnet2wv(cptnet, language)

Converts a `ConceptNet` object, `cptnet` to a `WordVectors` object.
The `language` of the word embeddings has to be specified explicitly
as a `Symbol` or `Languages.Language` (Conceptnet embeddings can be
multilingual).
"""
function conceptnet2wv(cptnet::ConceptNet{L,K,E},
                       language::Languages.Language;
                      ) where {L<:Language, K<:AbstractString, E<:AbstractFloat}
    embeddings = cptnet.embeddings[language]
    vector_size = cptnet.width
    vocab_size = length(embeddings)
    vocab = Vector{String}(undef, vocab_size)
    vectors = zeros(E, vector_size, vocab_size)
    for (i, (word, wordvector)) in enumerate(embeddings)
        vocab[i] = word
        vectors[:,i] = wordvector
    end
    wv = WordVectors(vocab, vectors)
end

function conceptnet2wv(cptnet::ConceptNet{L,K,E},
                       language::Symbol;
                      ) where {L<:Language, K<:AbstractString, E<:AbstractFloat}
    conceptnet2wv(cptnet, ConceptnetNumberbatch.LANGUAGES[language])
end


# # Interfaces

# This file defines abstract super types and includes further method
# definitions for these types so we can write our algorithm without
# the actual implementations being available.
# This also avoids some import order complications.

# ## Abstract Super Types

# ### Surrogate Models
# Each surrogate model type is configured via a data structure
# that implements `SurrogateConfig`.
# Such a configuration struct should be exported by the main module.
"Abstract super type for a configuration type defining some surrogate model."
abstract type SurrogateConfig end

"Abstract super type for meta data that is used to build a model."
abstract type SurrogateMeta end

"Abstract super type for the actual surrogate models."
abstract type SurrogateModel end

"Abstract super type wrapping around an objective, its model and the surrogate meta."
abstract type AbstractSurrogateWrapper end

"Wrapper around a list of `AbstractSurrogateWrapper`s."
abstract type AbstractSurrogateContainer end

# Define a shorthand for the type of what we save in some `AbstractIterSaveable`.
const ContainerSaveable = Union{Nothing,Tuple{Vararg{<:Union{Nothing,SurrogateMeta}}}}

# ### MOPs
# Internally we use implementations of the `AbstractObjective` for
# managing the MOP objectives and evaluation.
"Abstract super type for any kind of (vector) objective."
abstract type AbstractObjective <: MOI.AbstractVectorFunction end

# Our actual problem is repersented by an `AbstractMOP`.
"""
    AbstractMOP{T}

Abstract super type for multi-objective optimization problems.
`T` is `true` if the problem is modifyable and `false elsewise.

The user should define a `MixedMOP<:AbstractMOP{true}`, see [`MixedMOP`](@ref).
"""
abstract type AbstractMOP{T} <: MOI.ModelLike end

# ### Internal Data Managment

# For internal data management we make some effort to keep everything
# structured.
# To save the intermediate results:
"Abstract super type for stuff stored in the database."
abstract type AbstractResult{XT <: VecF, YT <: VecF} end

# `AbstractIterData` is passed to functions and provides access to
# the current site and value vectors and the trust region radius.
"Abstract super type for iteration data. Implemented by `IterData`.
Has type parameters for the vector type of itaration site, value vector and trust region radius."
abstract type AbstractIterData{XT <: VecF, YT <: VecF, DT <: NumOrVecF} end
"Abstract super type for some saveable representation of `AbstractIterData`."
abstract type AbstractIterSaveable{XT <: VecF, YT <: VecF, DT <: NumOrVecF, C <: ContainerSaveable } end

# A shorthand for everything that is either nothing or an `AbstractIterSaveable`:
const NothingOrSaveable = Union{Nothing, AbstractIterSaveable}

# Everything is kept in a database:
"Abstract database super type. Implemented by `ArrayDB` and `MockDB`."
abstract type AbstractDB{R<:AbstractResult, S<:NothingOrSaveable} end

# ### Algorithm configuration.

# We have our own type to define the solution of the trust region sub-problems:
"Abstract super type for descent step configuration."
abstract type AbstractDescentConfig end

# This might be returned by the general algorithm configuration implementing `AbstractConfig`:
"Abstract super type for user configurable algorithm configuration."
abstract type AbstractConfig end

# ### Differentiation
"Super type for making (automatic) differentiation comfortable."
abstract type DiffFn end

# Method definitions for `DiffFn`s:
get_gradient( :: DiffFn, :: Vec, :: Int ) :: Vec = nothing
get_jacobian( :: DiffFn, :: Vec ) :: Mat = nothing
get_hessian( :: DiffFn, :: Vec, :: Int ) :: Mat = nothing

# ### Enums

# These codes should be availabe everywhere and that is why we
# define them here:

@enum ITER_TYPE begin
    ACCEPTABLE = 1;     # accept trial point, shrink radius
    SUCCESSFULL = 2;    # accept trial point, grow radius
    MODELIMPROVING = 3; # reject trial point, keep radius
    INACCEPTABLE = 4;   # reject trial point, shrink radius (much)
end

@enum STOP_CODE begin
    CONTINUE = 1
    MAX_ITER = 2
    BUDGET_EXHAUSTED = 3
    CRITICAL = 4
    TOLERANCE = 5
end

# ### Interface Definitions
# Most of the interfaces are defined in subfiles and we include
# them here:
include("SurrogateModelInterface.jl");
include("AbstractSurrogateContainerInterface.jl")
include("AbstractObjectiveInterface.jl");
include("AbstractMOPInterface.jl");
include("AbstractIterDataInterface.jl")
include("AbstractResultInterface.jl")
include("AbstractDBInterface.jl")
include("AbstractConfigInterface.jl")