abstract type FermiAction{Dim,Dirac,fermion,gauge} end

include("./StaggeredFermiAction.jl")
include("./WilsonFermiAction.jl")
include("./DomainwallFermiAction.jl")
include("./MobiusDomainwallFermiAction.jl")
include("./WilsontypeFermiAction.jl")

function FermiAction(D::Dirac_operator{Dim},parameters_action;covneuralnet = nothing) where {NC,Dim}
    diractype = typeof(D)
    if covneuralnet ==  nothing
        hascovnet = false
    else
        hascovnet = true
    end

    if diractype <: Staggered_Dirac_operator
        return StaggeredFermiAction(D,hascovnet,covneuralnet,parameters_action)
    elseif diractype <: Wilson_Dirac_operators
        return WilsonFermiAction(D,hascovnet,covneuralnet,parameters_action)
    elseif diractype <: Domainwall_Dirac_operator
        return DomainwallFermiAction(D,hascovnet,covneuralnet)
    elseif diractype <: MobiusDomainwall_Dirac_operator
        return MobiusDomainwallFermiAction(D,hascovnet,covneuralnet)
    elseif diractype <: Wilson_GeneralDirac_operator
        return Wilson_GeneralDirac_FermiAction(D,hascovnet,covneuralnet,parameters_action)
    else
        error("Action type $diractype is not supported")
    end

end

function evaluate_FermiAction(fermi_action::FermiAction,U,ϕ::AbstractFermionfields)
    error("evaluate_FermiAction(fermi_action,U,ϕ) is not implemented in type fermi_action:$(typeof(fermi_action)), U:$(typeof(U)), and ϕ:$(typeof(ϕ)),  ")
end

function gauss_sampling_in_action!(η::AbstractFermionfields,U,fermi_action::FermiAction)
    error("gauss_sampling_in_action!(η,fermi_action) is not implemented in type η:$(typeof(η)), fermi_action:$(typeof(fermi_action))")
end

#=
 Conventional case:
det(D)^Nf = det(D^+ D)^{Nf/2}
 = int dphi dphi^* exp[- phi^* (D^+ D)^{-1} phi]
 = int dphi dphi^* exp[- phi^* D^{-1} (D^+)^{-1} phi]

 RHMC case:
det(D)^Nf =
 = int dphi dphi^* exp[- phi^* D^{-Nf} phi]
 = int dphi dphi^* exp[- phi^* D^{-Nf/2} D^{-Nf/2} phi]
=#

function sample_pseudofermions!(ϕ::AbstractFermionfields,U,fermi_action::FermiAction,ξ)
    error("sample_pseudofermions!(ϕ,fermi_action,ξ) is not implemented in type ϕ:$(typeof(ϕ)), fermi_action:$(typeof(fermi_action)), ξ:$(typeof(ξ))")
end

function calc_UdSfdU(fermi_action::FermiAction{Dim,Dirac,fermion,gauge},U::Vector{<: AbstractGaugefields},ϕ::AbstractFermionfields) where {Dim,Dirac,fermion,gauge}
    x = U[1]
    UdSfdU = Array{typeof(x),1}(undef,Dim)
    for μ=1:Dim
        UdSfdU[μ] = similar(x)
    end
    calc_UdSfdU!(UdSfdU,fermi_action,U,ϕ)
end

function calc_UdSfdU!(UdSfdU::Vector{<: AbstractGaugefields},fermi_action::FermiAction,U::Vector{<: AbstractGaugefields},ϕ::AbstractFermionfields)
    error("cald_UdSfdU!(UdSfdU,fermi_action,U) is not implemented in type UdSfdU:$(typeof(UdSfdU)), fermi_action:$(typeof(fermi_action)), U:$(typeof(U)), ϕ:$(typeof(ϕ))")
end


module Dispersal
# Use the README as the module docs
@doc read(joinpath(dirname(@__DIR__), "README.md"), String) Dispersal

using ConstructionBase,
      Colors,
      Dates,
      DimensionalData,
      Distributed,
      Distributions,
      DocStringExtensions,
      FieldDefaults,
      FieldDocTables,
      FieldMetadata,
      Flatten,
      LinearAlgebra,
      LossFunctions,
      Mixers,
      PoissonRandom,
      Reexport,
      Setfield,
      Statistics

@reexport using DynamicGrids

using LossFunctions: ZeroOneLoss
using DimensionalData: Time

import Base: getindex, setindex!, lastindex, size, length, push!

import DynamicGrids: applyrule, applyrule!, applyrule, applyrule!,
       neighbors, sumneighbors, neighborhood, setneighbor!, mapsetneighbor!,
       radius, gridsize, mask, overflow, cellsize, ruleset, inbounds, extent,
       currenttime, currenttimestep, timestep, tspan,
       buffer, WritableGridData, aux, unwrap

import ConstructionBase: constructorof

import FieldMetadata: @default, @description, @bounds, @flattenable,
                      default, description, bounds, flattenable

export AbstractDispersalKernel, DispersalKernel

export KernelFormulation, ExponentialKernel, GeometricKernel, GaussianKernel

export DistanceMethod, CentroidToCentroid, CentroidToArea, AreaToArea, AreaToCentroid

export AbstractInwardsDispersal, InwardsBinaryDispersal, InwardsPopulationDispersal,
       SwitchedInwardsPopulationDispersal

export AbstractOutwardsDispersal, OutwardsBinaryDispersal, OutwardsPopulationDispersal

export AlleeExtinction, JumpDispersal, HumanDispersal

export BatchGroups, HeirarchicalGroups

export ExactExponentialGrowth, ExactLogisticGrowth

export DiscreteGrowth, DiscreteGrowthMap, DiscreteGrowth2, DiscreteGrowth3

export MaskGrowthMap, ExactExponentialGrowthMap, ExactLogisticGrowthMap, ExactLogisticGrowthMap2, ExactLogisticGrowthMap3

export SurvLogLogistic, SurvLogLogisticMap, SurvLogLogisticMap2

export GrowthSurvLogLogisticMap3

export MaskSurvMap

export SelectionGradientSurv, SelectionGradientSurvMap

export SelectionGradient1locusSurv, SelectionGradient1locusSurvMap, LandeVariable

export DeltaAlleleFrequencySurv, DeltaAlleleFrequencySurvMap, DeltaAlleleFrequencySurvMap_noFastmath

export MatingPopulation, MatingDispersal

export LayerCopy

export Parametriser, AbstractObjective, SimpleObjective, RegionObjective, RegionOutput,
       ColorRegionFit, Accuracy

export ThreadedReplicates, DistributedReplicates, SingleCoreReplicates

export targets, predictions

const FIELDDOCTABLE = FieldDocTable((:Description, :Default, :Bounds),
                                    (description, default, bounds);
                                    truncation=(100,40,100))

# Documentation templates
@template TYPES =
    """
    $(TYPEDEF)
    $(DOCSTRING)
    """

include("rules/mixins.jl")
include("downsampling.jl")
include("layers.jl")
include("rules/kernel/common.jl")
include("rules/kernel/inwards.jl")
include("rules/kernel/outwards.jl")
include("rules/growth.jl")
include("rules/human.jl")
include("rules/jump.jl")
include("rules/allee.jl")
include("rules/survival.jl")
include("rules/alleleFrequency.jl")
include("rules/selectionGradient.jl")
include("rules/matingPopulation.jl")
include("optimisation/optimisation.jl")
include("optimisation/objectives.jl")
include("optimisation/output.jl")

end # module


@kwargs function select!(
    op::SelectUnion,
    C::GBVecOrMat,
    A::GBVecMatOrTrans;
    thunk::Union{GBScalar, Nothing} = nothing
)
    thunk === nothing && (thunk = C_NULL)
    A, desc, _ = _handletranspose(A, desc)
    if A isa GBVector && C isa GBVector
        libgb.GxB_Vector_select(C, mask, accum, op, A, thunk, desc)
    elseif A isa GBMatrix && C isa GBMatrix
        libgb.GxB_Matrix_select(C, mask, accum, op, A, thunk, desc)
    end
end

@kwargs function select(
    op::SelectUnion,
    A::GBVecMatOrTrans;
    thunk::Union{GBScalar, Nothing} = nothing
)
    C = similar(A)
    select!(op, C, A; thunk, accum, mask, desc)
    return C
end


addprocs(8)

include("../../daggen.jl")
using DagScheduler
using BenchmarkTools

isdir(".mempool") && rm(".mempool"; recursive=true)
@everywhere begin
    DagScheduler.META_IMPL[:node] = ENV["NODE_META_IMPL"]
    DagScheduler.META_IMPL[:cluster] = ENV["CLUSTER_META_IMPL"]
    DagScheduler.META_IMPL[:map_num_entries] = 1024*5
    DagScheduler.META_IMPL[:map_entry_sz] = 256
end

node1 = NodeEnv(2, getipaddr(), [3,4,5])
node2 = NodeEnv(6, getipaddr(), [7,8,9])
runenv = DagScheduler.Plugin.setrunenv(RunEnv(; nodes=[node1,node2]))

const L = 6^4
const dag2 = gen_straight_dag(ones(Int, L));

result = collect(rundag(runenv, dag2));
#@time result = rundag(runenv, dag2);
@btime collect(rundag(runenv, dag2))

DagScheduler.cleanup(runenv)


abstract type AbstractSampler
end

struct GibbsSampler <: AbstractSampler
   nsteps::Int
end

struct TemperatureDrivenSampler <: AbstractSampler
   nsteps::Int
   autocorcoeff::Float64
   betasd::Float64
end

function TemperatureDrivenSampler(;
      nsteps::Int = 5, autocorcoeff::Float64 = 0.9,
      betasd::Float64 = 0.9)

   TemperatureDrivenSampler(nsteps, autocorcoeff, betasd)
end

struct NoSampler <: AbstractSampler
end

function sample!(particles::Particles, sampler::NoSampler, rbm::AbstractRBM,
      upfactor::Float64 = 1.0, downfactor::Float64 = 1.0)
   # do nothing
end

function sample!(vmodel::M, hmodel::M, sampler::NoSampler, rbm::AbstractRBM,
      upfactor::Float64 = 1.0, downfactor::Float64 = 1.0
      ) where {M <:AbstractArray{Float64}}
   # do nothing
end

function sample!(vmodel::M, hmodel::M, sampler::AbstractSampler, rbm::AbstractRBM,
      upfactor::Float64 = 1.0, downfactor::Float64 = 1.0
      ) where {M <:AbstractArray{Float64}}

   particles = Particles(2)
   particles[1] = vmodel
   particles[2] = hmodel
   sample!(particles, sampler, rbm, upfactor, downfactor)
   nothing
end

function sample!(vmodel::M, hmodel::M, sampler::GibbsSampler, rbm::AbstractRBM,
      upfactor::Float64 = 1.0, downfactor::Float64 = 1.0
      ) where {M <:AbstractArray{Float64}}

   for i = 1:sampler.nsteps
      samplevisible!(vmodel, rbm, hmodel, downfactor)
      samplehidden!(hmodel, rbm, vmodel, upfactor)
   end
   nothing
end

function sample!(particles::Particles, sampler::GibbsSampler, rbm::AbstractRBM,
      upfactor::Float64 = 1.0, downfactor::Float64 = 1.0)

   gibbssample!(particles, rbm, sampler.nsteps)
end

function sample!(particles::Particles, sampler::TemperatureDrivenSampler, rbm::AbstractRBM,
      upfactor::Float64 = 1.0, downfactor::Float64 = 1.0)
   # TODO respect upfactor and downfactor
   gibbssample_gamma!(particles, rbm, sampler.nsteps;
         autocorcoeff = sampler.autocorcoeff, betasd = sampler.betasd)
end


using Random
using LinearAlgebra
# const ContactImplicitMPC = Main

# include(joinpath(module_dir(), "src", "solver", "interior_point.jl"))
# include(joinpath(module_dir(), "src", "solver", "mehrotra.jl"))

################################################################################
# Test Utils
################################################################################

function get_initialization(ref_traj, t)
	Random.seed!(10)
	z = deepcopy(ref_traj.z[t])
	z += rand(length(z))
	θ = deepcopy(ref_traj.θ[t])
	return z, θ
end

function interior_point_timing(ref_traj, t, ip0)
	e1 = @belapsed begin
		z, θ = get_initialization(ref_traj, t)
		ip = deepcopy(ip0)
	end
	e2 = @belapsed begin
		z, θ = get_initialization(ref_traj, t)
		ip = deepcopy(ip0)
		interior_point_solve!(ip, z, θ)
	end
	@show e1
	@show e2
	return e2 - e1
end

# function interior_point_profiling(ref_traj, t, ip0, N::Int = 100)
# 	@profiler begin
# 		for i = 1:N
# 			z, θ = get_initialization(ref_traj, t)
# 			ip = deepcopy(ip0)
# 			interior_point_solve!(ip, z, θ)
# 		end
# 	end
# 	return nothing
# end

################################################################################
# Test data
################################################################################

s = ContactImplicitMPC.get_simulation("quadruped", "flat_2D_lc", "flat")
ref_traj = deepcopy(ContactImplicitMPC.get_trajectory(s.model, s.env,
    joinpath(ContactImplicitMPC.module_dir(), "src/dynamics/quadruped/gaits/gait2.jld2"),
    load_type = :split_traj_alt))

model = s.model
env = s.env
t = 10
nz = num_var(model, env)
nθ = num_data(model)

################################################################################
# Test Interior Point on the full non linear problem
################################################################################

z1, θ1 = get_initialization(ref_traj, t)
ip1 = interior_point(z1, θ1,
	oss = OptimizationSpace13(model, env),
	r! = s.res.r!,
    rz! = s.res.rz!,
    rθ! = s.res.rθ!,
    rz = s.rz,
    rθ = s.rθ,
    opts = InteriorPointOptions(
        max_iter=30,
        r_tol=1e-8,
        κ_tol=1e-8,
		# verbose=true
		))

interior_point_solve!(ip1)
r1 = zeros(nz)
s.res.r!(r1, ip1.z, ip1.θ, 0.0)
@testset "Interior Point Non Linear" begin
	@test norm(r1, Inf) < 1e-8
	@test abs(norm(r1, Inf) - 3.66e-9) < 1e-11
	@test ip1.iterations == 10
end
# e_ip = 1e6 * interior_point_timing(ref_traj, t, ip1)

################################################################################
# Test Interior Point on the linearized problem
################################################################################

im_traj1 = ImplicitTraj(ref_traj, s;
	κ = 1e-8,
	ip_type = :interior_point,
	opts = InteriorPointOptions(
			κ_tol = 1.0e-8,
			r_tol = 1.0e-8,
			diff_sol = true,
			max_iter=100,
			ϵ_min=0.05,
			solver=:empty_solver,
			# verbose=true
			))

z1, θ1 = get_initialization(ref_traj, t)
ip1 = deepcopy(im_traj1.ip[10])
interior_point_solve!(ip1, z1, θ1)
r1 = zeros(nz)
s.res.r!(r1, ip1.z, ip1.θ, 0.0)
@testset "Interior Point Linear" begin
	@test norm(r1, Inf) < 1e-8
end

# e_ip = 1e6 * interior_point_timing(ref_traj, t, im_traj1.ip[t])

# @profiler for k = 1:15000
# 	z1, θ1 = get_initialization(ref_traj, t)
# 	ip1 = deepcopy(im_traj1.ip[t])
# 	interior_point_solve!(ip1, z1, θ1)
# end

# z1, θ1 = get_initialization(ref_traj, t)
# ip1 = deepcopy(im_traj1.ip[t])
# @code_warntype interior_point_solve!(ip1)


#=
Configuration Load and saving
=#
function load_hourglass_to_nn(nn,config_path)
    nn.weight_path = config_path
    hg=HG2(64,1,4); #Dummy Numbers
    load_hourglass(nn.weight_path,hg)
    nn.hg = hg
    nn.use_existing_weights = true
    nn.features=StackedHourglass.features(hg)

    nothing
end

#=
Save and load training file
=#
function save_training(han,mypath=string(han.paths.backup,"/labels.jld"))
    save_training(mypath,han.frame_list,han.woi,han.nn)
end

function save_training(mypath,frame_list,woi,nn)

    file = jldopen(mypath, "w")
    write(file, "frame_list",frame_list)
    write(file, "woi",woi)
    write(file, "mean_img",nn.norm.mean_img)
    write(file, "std_img",nn.norm.std_img)

    close(file)

    nothing
end

function load_training(han::Tracker_Handles,path::String)

    file = jldopen(path, "r")
    frame_list = read(file, "frame_list")
    woi = read(file, "woi")
    if typeof(woi) <: Array
        for i=1:length(frame_list)
            han.woi[frame_list[i]]=woi[i]
        end
    else # New Version
        han.woi = woi
    end

    han.frame_list = frame_list
    set_gtk_property!(han.b["labeled_frame_adj"],:upper,length(han.frame_list))

    tracked = [true for i=1:length(han.frame_list)]
    han.tracked = Dict{Int64,Bool}(zip(frame_list,tracked))

    pole_present = [false for i=1:length(han.frame_list)]
    han.pole_present=Dict{Int64,Bool}(zip(frame_list,pole_present))

    pole_loc = [zeros(Float32,2) for i=1:length(han.frame_list)]
    han.pole_loc=Dict{Int64,Array{Float32,1}}(zip(frame_list,pole_loc))

    han.nn.norm.mean_img = read(file, "mean_img")
    han.nn.norm.std_img = read(file, "std_img")

    close(file)

end


include("../imports.jl")

@testset ExtendedTestSet "`offsetcropindices`" begin

    @test offsetcropindices((50, 50), (1:100, 1:100), (0., 0.)) == (1:50, 1:50)
    @test offsetcropindices((50, 50), (1:100, 1:100), (0.5, 0.5)) == (26:75, 26:75)
    @test offsetcropindices((50, 50, 50), (1:100, 1:100, 1:100), (1., 1., 1.)) == (51:100, 51:100, 51:100)

end

@testset ExtendedTestSet "CroppedProjectiveTransform" begin

    image = Image(rand(100, 100))
    tfm = Project(Translation(20, 20))
    crop = CenterCrop((50, 50))
    @test tfm |> crop isa CroppedProjectiveTransform

    cropped = tfm |> crop
    @test_nowarn apply(cropped, image)
end


# %% imports

using PyCall
gym = pyimport("gym")

include("../../Modules/Gridcells/StateSpikeGaussGrid.jl")
include("../../Modules/Networks/net_arch.jl")
include("simviz.jl")

# %% Build the grids

include("../../Modules/Gridcells/GridGenerator.jl")

grid_cells = grid_gen()

# %% import the trained cartpole

torch = pyimport("torch")
nn    = pyimport("torch.nn")
f     = pyimport("torch.nn.functional")

@pydef mutable struct DQLnet <: nn.Module
    function __init__(self)
        pybuiltin(:super)(DQLnet, self).__init__()
        self.fc1 = nn.Linear(4,16)
        self.fc2 = nn.Linear(16,16)
        self.fc3 = nn.Linear(16,2)
    end

    function forward(self, s)
        s = f.relu(self.fc1(s))
        s = f.relu(self.fc2(s))
        s = self.fc3(s)
        return s
    end
end

Qnet = DQLnet()
Qnet.load_state_dict(torch.load("/Users/titou/Documents/PFE/RL-STDP/Julia/CartPole/saved_cartpole.net"))

# %% Function

function count_spikes(arch::Array{Int64}, spiked::Array{Int64}, layer::Int64)
    count_layer = zeros(arch[layer])
    interval = [sum(arch[1:layer-1])+1, sum(arch[1:layer])]
    spikes_in = filter(x-> interval[1] <= x <= interval[2], spiked)
    spikes_in .-= sum(arch[1:layer-1])
    count_layer[spikes_in] .+= 1
    return count_layer
end

function left_right(count_spikes::Array{Float64}, n_left::Int64, n_right::Int64)
    sl = sum(count_spikes[1:n_left])
    sr = sum(count_spikes[n_left+1:end])
    (sl > sr) ? action = 0 : action = 1
    if sl == sr
        action = rand(0:1)
    end
    return action, sl, sr
end

# %% Simulation

global s_hist = [[],[]]

function run_sim(net::Network, n_steps::Int64, render::Bool=false;
                 method::String = "snn")
    env = gym.make("CartPole-v1")
    env.seed(0)
    obs = env.reset()
    spike_history = Array{Int64}[]
    @inbounds for j in 1:n_steps
        count_tot = zeros(net.arch[end])
        spikes = input_spikes(obs, grid_cells)
        influence = 2
        @inbounds for msec in 1:10
            (msec==10) ? train=true : train=false
            spiked = step!(net, spikes, train)
            if msec >= length(net.arch)-1
                count_tot .+= influence .* count_spikes(net.arch, spiked, length(net.arch))
                influence *= 0.95
                push!(spike_history,spiked)
            end
        end
        if method == "snn"
            n_left = div(net.arch[end],2)
            n_right = net.arch[end] - n_left
            action, sl, sr = left_right(count_tot, n_left, n_right)
            push!(s_hist[1], sl)
            push!(s_hist[2], sr)
        elseif method == "trained"
            n_left = div(net.arch[end],2)
            n_right = net.arch[end] - n_left
            action = torch.argmax(Qnet(torch.tensor(obs).float())).item()
            _, sl, sr = left_right(count_tot, n_left, n_right)
            push!(s_hist[1], sl)
            push!(s_hist[2], sr)
        elseif method == "random"
            n_left = div(net.arch[end],2)
            n_right = net.arch[end] - n_left
            action = rand(0:1)
            _, sl, sr = left_right(count_tot, n_left, n_right)
            push!(s_hist[1], sl)
            push!(s_hist[2], sr)
        end
        obs_new, reward, done, _ = env.step(action)
        if render==true
            env.render()
        end
        obs = obs_new
    end
    return spike_history
end

parameters = YAML.load(open("/Users/titou/Documents/PFE/RL-STDP/Julia/Modules/cfg.yml"))

n = Network([120,60,80], parameters)
#randomize!(n)

# %% run the simulation

@time spike_history = run_sim(n, 200, method = "snn")

for _ in 1:9
    @time spike_history .= run_sim(n, 200, method = "snn")
end

# %% Plots difference
x = collect(1:length(s_hist[1]))
y = s_hist[2] .- s_hist[1]
plot(x,y,
     xlabel = "Time step cartpole",
     ylabel = "Sum spikes left minus right",
     title = "Difference between left or right",
     legend = false)

# %% Plot s_hist
x = collect(1:length(s_hist[1]))
plot(x,s_hist,
     xlabel = "Time step cartpole",
     ylabel = "Sum spikes left and right",
     lab = ["left" "right"],
     title = "Left or Right plot")

# %% spike history plot per layer
layer_spike_hist(n.arch, spike_history, 1)

# %% spike distribution for all layers
spike_histogram(n.arch, spike_history)

# %% weights visualization heatmap
weight_viz(n, "heatmap")

# %% weights visualization histogram
weight_viz(n, "histogram")

# %% visualize network connection
heatmap(n.connection)


import Dagger: @par, @spawn, spawn
import Dagger: Chunk

@everywhere begin
    checkwid() = myid()==1

    function dynamic_fib(n)
        n <= 1 && return n
        t = Dagger.spawn(dynamic_fib, n-1)
        y = dynamic_fib(n-2)
        return (fetch(t)::Int) + y
    end

    struct ProcessLockedStruct
        x::Ptr{Int} # Zero'd during serialization
    end
    (pls::ProcessLockedStruct)(x) = x+UInt(pls.x)

    struct MulProc <: Dagger.Processor
        owner::Int
    end
    MulProc() = MulProc(myid())
    Dagger.get_parent(mp::MulProc) = OSProc(mp.owner)
    Dagger.move(src::MulProc, dest::Dagger.ThreadProc, x::Function) = Base.:*
end

@testset "@par" begin
    @testset "per-call" begin
        x = 2
        a = @par x + x
        @test a isa Dagger.Thunk
        b = @par sum([x,1,2])
        c = @par a * b
        @test collect(c) == 20
    end
    @testset "block" begin
        c = @par begin
            x = 2
            a = x + x
            b = sum([x,1,2])
            c = a * b
        end
        @test x isa Int
        @test a isa Dagger.Thunk
        @test c isa Dagger.Thunk
        @test collect(c) == 20
    end
end

@testset "@spawn" begin
    @test !isassigned(Dagger.Sch.EAGER_CONTEXT)
    @testset "per-call" begin
        x = 2
        a = @spawn x + x
        @test a isa Dagger.EagerThunk
        b = @spawn sum([x,1,2])
        c = spawn(*, a, b)
        @test c isa Dagger.EagerThunk
        @test fetch(a) == 4
        @test fetch(b) == 5
        @test fetch(c) == 20
    end
    @test Dagger.Sch.EAGER_CONTEXT[] isa Context
    @testset "waiting" begin
        a = @spawn sleep(1)
        @test !isready(a)
        wait(a)
        @test isready(a)
    end
    @testset "options" begin
        s = 1
        m = true
        a = @spawn single=s checkwid()
        b = @spawn meta=m ((_a)->_a isa Dagger.Chunk)(a)
        @test fetch(a)
        @test fetch(b)
    end
    @testset "errors" begin
        @testset "independent" begin
            a = @spawn error("Test")
            wait(a)
            @test isready(a)
            @test_throws_unwrap Dagger.ThunkFailedException fetch(a)
            b = @spawn 1+2
            @test fetch(b) == 3
        end
        @testset "direct vs indirect" begin
            a = @spawn error("Test")
            b = @spawn a+1

            ex = try
                fetch(a)
            catch err
                err
            end
            ex = Dagger.Sch.unwrap_nested_exception(ex)
            ex_str = sprint(io->Base.showerror(io,ex))
            @test occursin(r"^ThunkFailedException \(Thunk.*failure\):", ex_str)
            @test occursin("Test", ex_str)
            @test !occursin("due to a failure in", ex_str)

            ex = try
                fetch(b)
            catch err
                err
            end
            ex = Dagger.Sch.unwrap_nested_exception(ex)
            ex_str = sprint(io->Base.showerror(io,ex))
            @test occursin(r"Thunk.*failure due to a failure in", ex_str)
            @test occursin("Test", ex_str)
        end
        @testset "single dependent" begin
            a = @spawn error("Test")
            b = @spawn a+2
            @test_throws_unwrap Dagger.ThunkFailedException fetch(a)
        end
        @testset "multi dependent" begin
            a = @spawn error("Test")
            b = @spawn a+2
            c = @spawn a*2
            @test_throws_unwrap Dagger.ThunkFailedException fetch(b)
            @test_throws_unwrap Dagger.ThunkFailedException fetch(c)
        end
        @testset "dependent chain" begin
            a = @spawn error("Test")
            @test_throws_unwrap Dagger.ThunkFailedException fetch(a)
            b = @spawn a+1
            @test_throws_unwrap Dagger.ThunkFailedException fetch(b)
            c = @spawn b+2
            @test_throws_unwrap Dagger.ThunkFailedException fetch(c)
        end
        @testset "single input" begin
            a = @spawn 1+1
            b = @spawn (a->error("Test"))(a)
            @test fetch(a) == 2
            @test_throws_unwrap Dagger.ThunkFailedException fetch(b)
        end
        @testset "multi input" begin
            a = @spawn 1+1
            b = @spawn 2*2
            c = @spawn ((a,b)->error("Test"))(a,b)
            @test fetch(a) == 2
            @test fetch(b) == 4
            @test_throws_unwrap Dagger.ThunkFailedException fetch(c)
        end
        @testset "diamond" begin
            a = @spawn 1+1
            b = @spawn a+1
            c = @spawn a*2
            d = @spawn ((b,c)->error("Test"))(b,c)
            @test fetch(a) == 2
            @test fetch(b) == 3
            @test fetch(c) == 4
            @test_throws_unwrap Dagger.ThunkFailedException fetch(d)
        end
    end
    @testset "remote spawn" begin
        a = fetch(Distributed.@spawnat 2 Dagger.spawn(+, 1, 2))
        @test Dagger.Sch.EAGER_INIT[]
        @test fetch(Distributed.@spawnat 2 !(Dagger.Sch.EAGER_INIT[]))
        @test a isa Dagger.EagerThunk
        @test fetch(a) == 3

        # Mild stress-test
        @test dynamic_fib(10) == 55
    end
    @testset "undefined function" begin
        # Issues #254, #255

        # only defined on head node
        @eval evil_f(x) = x

        eager_thunks = map(1:10) do i
            single = isodd(i) ? 1 : first(workers())
            Dagger.@spawn single=single evil_f(i)
        end

        errored(t) = try
            fetch(t)
            false
        catch
            true
        end
        @test any(t->errored(t), eager_thunks)
        @test any(t->!errored(t), eager_thunks)
    end
    @testset "function chunks" begin
        @testset "lazy API" begin
            a = delayed(+)(1,2)
            @test !(a.f isa Chunk)

            a = delayed(+; scope=NodeScope())(1,2)
            @test a.f isa Chunk
            @test a.f.processor isa OSProc
            @test a.f.scope isa NodeScope

            a = delayed(+; processor=Dagger.ThreadProc(1,1))(1,2)
            @test a.f isa Chunk
            @test a.f.processor isa Dagger.ThreadProc
            @test a.f.scope isa AnyScope

            a = delayed(+; processor=Dagger.ThreadProc(1,1), scope=NodeScope())(1,2)
            @test a.f isa Chunk
            @test a.f.processor isa Dagger.ThreadProc
            @test a.f.scope isa NodeScope

            @testset "Scope Restrictions" begin
                pls = ProcessLockedStruct(Ptr{Int}(42))
                ctx = Context([1, workers()...])

                # Negative test
                @test_skip !all(x->x==43, collect(ctx, delayed(vcat)([delayed(pls)(1) for i in 1:10]...)))
                # Positive tests (no serialization)
                @test all(x->x==43, collect(ctx, delayed(vcat)([delayed(pls; scope=ProcessScope())(1) for i in 1:10]...)))
                @test all(x->x==1, collect(ctx, delayed(vcat)([delayed(pls; scope=ProcessScope(first(workers())))(1) for i in 1:10]...)))
            end
            @testset "Processor Data Movement" begin
                @everywhere Dagger.add_processor_callback!(()->MulProc(), :mulproc)
                @test collect(delayed(+; processor=MulProc())(3,4)) == 12
                @everywhere Dagger.delete_processor_callback!(:mulproc)
            end
        end
        @testset "eager API" begin
            _a = Dagger.spawn(+, 1, 2; scope=NodeScope())
            a = Dagger.Sch._find_thunk(_a)
            @test a.f isa Chunk
            @test a.f.processor isa OSProc
            @test a.f.scope isa NodeScope

            _a = Dagger.spawn(+, 1, 2; processor=Dagger.ThreadProc(1,1))
            a = Dagger.Sch._find_thunk(_a)
            @test a.f isa Chunk
            @test a.f.processor isa Dagger.ThreadProc
            @test a.f.scope isa AnyScope

            _a = Dagger.spawn(+, 1, 2; processor=Dagger.ThreadProc(1,1), scope=NodeScope())
            a = Dagger.Sch._find_thunk(_a)
            @test a.f isa Chunk
            @test a.f.processor isa Dagger.ThreadProc
            @test a.f.scope isa NodeScope
        end
    end
    @testset "parent fetch child, one thread" begin
        # Issue #282

        s = p -> p == Dagger.ThreadProc(1, 1)
        f = (x) -> 10 + x
        g = (x) -> fetch(Dagger.spawn(f, x; proclist=s))
        fetch(Dagger.spawn(g, 10; proclist=s))
    end
end


abstract type AbstractLProblem{Approx} end

struct LProblem{Approx, SCB, PP} <: AbstractLProblem{Approx}
    A::Approx
    cb::SCB
    prob::PP
    rules::ConvRules
    cf_vals::CFVals
    Tmax::Float64
    tstep::Float64
    delimiter::Int64
end

struct LProblem0{Approx, SCB, PP} <: AbstractLProblem{Approx}
    A::Approx
    cb::SCB
    prob::PP
    rules::ConvRules
    cf_vals::CFVals0
    Tmax::Float64
    tstep::Float64
    delimiter::Int64
end

function build_problem(L::Lattice, M::AbstractModel, args...; kwargs...)
    A = build_Approximation(M, L, args...)
    return build_problem(A; kwargs...)
end

function build_problem(A::abstractApproximation; Tmax::Float64 = 10.0, tstep::Float64 = 2.0^-7, delimiter::Int64 = 10, links = [(:all,:all)], axes = [:x, :y, :z])
    ARHS = build_RHS_function(A)
    rules, OSET = set_CorrFuncs(A, links, axes)
    l = length(OSET.Observables)
    tspan = (0.0,Tmax*delimiter)
    t_means = collect(0.0:tstep:(Tmax*delimiter))
    t_cors = collect(0.0:tstep:Tmax)
    saved_values = SavedValues(t_means, [zeros(l) for i in eachindex(t_means)])
    u0 = randomState(A, OSET)
    prob = ODEProblem(ARHS, u0, tspan)
    cb = SavingCallback(OSET, saved_values, u0; saveat = t_means)
    #cb = SavingCallback(OSET, saved_values; save_everystep=true)
    cf_vals = CFVals(rules, t_cors)
    return LProblem(A, cb, prob, rules, cf_vals, Tmax, tstep, delimiter)
end

function build_problem0(L::Lattice, M::AbstractModel, args...; kwargs...)
    A = build_Approximation(M, L, args...)
    return build_problem0(A; kwargs...)
end

function build_problem0(A::abstractApproximation; Tmax::Float64 = 10.0, tstep::Float64 = 2.0^-7, delimiter::Int64 = 10, links = [(:all,:all)], axes = [:x, :y, :z])
    ARHS = build_RHS_function(A)
    rules, OSET = set_CorrFuncs(A, links, axes)
    l = length(OSET.Observables)
    tspan = (0.0,Tmax*delimiter)
    t_means = collect(0.0:tstep:(Tmax*delimiter))
    t_cors = collect(0.0:tstep:Tmax)
    saved_values = SavedValues(t_means, [zeros(l) for i in eachindex(t_means)])
    u0 = randomState(A, OSET)
    prob = ODEProblem(ARHS, u0, tspan)
    cb = SavingCallback(OSET, saved_values, u0; saveat = t_means)
    #cb = SavingCallback(OSET, saved_values; save_everystep=true)
    cf_vals = CFVals0(rules, t_cors)
    return LProblem0(A, cb, prob, rules, cf_vals, Tmax, tstep, delimiter)
end

struct Logger{T1,T2} end
struct Saver{T}
    filename::String
end
Saver() = Saver{:no}("")
Saver(filename::String) = Saver{:file}(filename)
function Saver(filename::String, nTrials::Int64, cf_vals0::CFVals0)
    if !isfile(filename)
        init_file(filename, nTrials, cf_vals0)
    end
    return Saver{:fullfile}(filename)
end

function init_file(filename::AbstractString, nTrials::Int64, cf_vals0::CFVals0)
    jldopen(filename, "w") do file
        gdata = g_create(file, "data")
        file["ts"] = cf_vals0.ts
        file["str_vec"] = cf_vals0.str_vec
        file["iter"] = [0]
        for i in Base.OneTo(length(cf_vals0.str_vec))
            gdata[cf_vals0.str_vec[i]] = zeros(nTrials, size(cf_vals0.data, 1))
        end
    end
end

@inline save_data(s::Saver{:no}, cf_vals::CFVals, iter = 0) = nothing
@inline save_data(s::Saver{:file}, cf_vals::CFVals, iter) = save_data(s.filename, cf_vals, iter)
@inline save_data(s::Saver{:no}, cf_vals::CFVals0, iter) = nothing
@inline save_data(s::Saver{:fullfile}, cf_vals0::CFVals0, iter) = save_data!(s.filename, cf_vals0, iter)
function save_data(filename::AbstractString, cf_vals::CFVals, iter)
    jldopen(filename, "w") do file
        gdata = g_create(file, "data")
        gerrors = g_create(file, "errors")
        for i in Base.OneTo(length(cf_vals.str_vec))
            write(gdata, cf_vals.str_vec[i], cf_vals.data[i])
            write(gerrors, cf_vals.str_vec[i], cf_vals.errors[i])
        end
        write(file, "ts", cf_vals.ts)
        write(file, "str_vec", cf_vals.str_vec)
        write(file, "nTrials", iter)
    end
end
function save_data!(filename::AbstractString, cf_vals0::CFVals0, iter)
    jldopen(filename, "r+") do file
        gdata = file["data"]
        for i in Base.OneTo(length(cf_vals0.str_vec))
            gdata[cf_vals0.str_vec[i]][iter,:] = cf_vals0.data[i]
        end
        file["iter"][1] = iter
    end
end

function load_data(filename::AbstractString)
    jdata = load(filename)
    data = VectorOfArray([jdata["data"][key] for key in jdata["str_vec"]])
    errors = VectorOfArray([jdata["errors"][key] for key in jdata["str_vec"]])
    return CFVals(jdata["str_vec"], jdata["ts"], data, errors), jdata["nTrials"]
end
function load_data!(cf_vals::CFVals, filename::AbstractString)
    jdata = load(filename)
    for i in eachindex(jdata["str_vec"])
        key = jdata["str_vec"][i]
        cf_vals.data[:,i] = jdata["data"][key]
        cf_vals.errors[:,i] = jdata["errors"][key]
    end
    return jdata["nTrials"]
end
@inline load_data!(cf_vals::CFVals, saver::Saver{:file}) = load_data!(cf_vals, saver.filename)
@inline load_data!(cf_vals::CFVals0, saver::Saver{:fullfile}) = nothing

function load_data0(filename::AbstractString)
    file = jldopen(filename, "r")
    ts = read(file, "ts")
    str_vec = read(file, "str_vec")
    iter = read(file, "iter")[1]
    gdata = file["data"]
    datas = Dict{String, Matrix{Float64}}()
    for str in str_vec
        datas[str] = gdata[str][1:iter,:]
    end
    close(file)
    return CFArray(str_vec, ts, datas)
end

abstract type AbstractLog end
struct CMDLog <: AbstractLog end

@inline output(log_o::CMDLog, args...) = (print(args...);print("\n"))
@inline Base.close(log_o::CMDLog) = nothing

struct FileLog
    status::Int64
    offset::Float64
    filename::String
end

function FileLog(filename::AbstractString)
    if isfile(filename)
        io = open(filename, "r+")
        lc = countlines(io)
        if lc>0
            skip(io, -64)
            line = split(readline(io))
            status = parse(Int64, line[1])
            offset = status*parse(Float64, line[3])
            seekend(io)
        else
            status = 0
            offset = 0.0
        end
    else
        io = open(filename, "w")
        status = 0
        offset = 0.0
    end
    close(io)
    FileLog(status, offset, filename)
end

@inline output(log_o::FileLog, args...) = (io = open(log_o.filename, "a"); print(io, args...);print(io, "\n"); close(io))

@inline get_status(log_o::CMDLog) = 0
@inline get_status(log_o::FileLog) = log_o.status
@inline get_offset(log_o::CMDLog) = 0.0
@inline get_offset(log_o::FileLog) = log_o.offset

@inline get_string(nTrials::Int64, thread::Int64, alg::OrdinaryDiffEqAlgorithm) = @sprintf("nTr%d Alg(%s) thr%d", nTrials, alg, thread)

function get_path_string(A::abstractApproximation)
    return get_string(A.L) * "/" * get_string(A) * "/"
end
function get_path_string(A::CompositeApproximation)
    s1 = get_string(A.A1.L) * @sprintf(" + UnSp:%s", A.A2.L.cell_name)
    return s1 * "/" * get_string(A.A1) * "/"
end

function set_Logging(logger::Logger, A::abstractApproximation, path_part::AbstractString, filename, args...)
    path = set_path(logger, A, path_part)
    save_o = set_Saver(path, filename, args...)
    log_o = set_Logger(logger, path, filename)
    return log_o, save_o
end

@inline set_path(logger::Logger{:no, :cmd}, args...) = nothing
function set_path(logger::Logger, A::abstractApproximation, path_part::AbstractString)
    path = get_prefix(logger) * get_path_string(A) * path_part
    mkpath(path)
    return path
end

get_prefix(logger::Logger{:local}) = "./Data/"
get_prefix(logger::Logger{:cluster}) = ENV["WORK_DIR"] * "/Data/"

@inline set_Saver(path::Void, args...) = Saver()
set_Saver(path::AbstractString, filename::AbstractString, args...) = Saver(path * filename * ".jld", args...)
set_Saver(path::AbstractString, filename::Vector{T}, args...) where {T<:AbstractString} = [Saver(path * fname * ".jld", args...) for fname in filename]

@inline set_Logger(logger::Logger{T, :cmd}, args...) where {T} = CMDLog()
@inline set_Logger(logger::Logger{T, :cmd}, path::AbstractString, filename::AbstractString) where {T} = CMDLog()
@inline set_Logger(logger::Logger{T, :cmd}, path::AbstractString, filename::Vector{T1}) where {T, T1<:AbstractString} = CMDLog()
set_Logger(logger::Logger, path::AbstractString, filename::AbstractString) = FileLog(path * filename * ".log")
set_Logger(logger::Logger, path::AbstractString, filename::Vector{T}) where {T<:AbstractString} = FileLog(path * filename[1] * ".log")

set_Logging(LP::LProblem, nTrials::Int64, args...) = set_Logging(args...)
set_Logging(LP::LProblem0, nTrials::Int64, args...) = set_Logging(args..., nTrials, LP.cf_vals)

function simulate(LP::AbstractLProblem, nTrials::Int64, alg::OrdinaryDiffEqAlgorithm, logger::Logger{T1,T2} = Logger{:no, :cmd}(), thread::Int64 = 1, interval::Int64 = 1, RNG::AbstractRNG=Base.GLOBAL_RNG; abstol_i = 1e-14, reltol_i = 1e-6, kwargs...) where {T1,T2}
    assert(nTrials > 0)
    saved_values = LP.cb.affect!.saved_values
    rules = LP.rules
    cf_vals = LP.cf_vals
    A = LP.A
    OSET = LP.cb.affect!.save_func
    filename = (@sprintf("Tm%.2f dt%.3f del%d Obs%s/", LP.Tmax, LP.tstep, LP.delimiter, get_string(LP.rules.str_vec)), get_string(nTrials, thread, alg))
    log_o, save_o = set_Logging(LP, nTrials, logger, A, filename...)
    iter = get_status(log_o)
    offset = get_offset(log_o)
    if iter > 0
        load_data!(cf_vals, save_o)
    end
    if iter == nTrials && (T1 == :local || T1 == :cluster)
        nTrials = 2*nTrials
        #close(log_o)
        filename = (@sprintf("Tm%.2f dt%.3f del%d Obs%s/", LP.Tmax, LP.tstep, LP.delimiter, get_string(LP.rules.str_vec)), get_string(nTrials, thread, alg))
        log_o, save_o = set_Logging(LP, nTrials, logger, A, filename...)
    end
    iter += 1
    t1 = time()
    t0 = t1-offset
    integrator = init(LP.prob, alg; save_everystep = false, callback = LP.cb, abstol = abstol_i, reltol = reltol_i, kwargs...)
    reinit!(integrator, randomState(A,OSET,RNG))
    solve!(integrator)
    calculateCorrFunc!(A, saved_values, rules, cf_vals, iter)
    save_data(save_o, cf_vals, iter)
    t2 = time()
    output(log_o, @sprintf("%10d %20.5f %20.5f %10d", iter, t2-t1, (t2-t0)/iter, nTrials - iter))
    iter += 1
    while iter <= nTrials
        t1 = time()
        reinit!(integrator, randomState(A,OSET,RNG))
        solve!(integrator)
        calculateCorrFunc!(A, saved_values, rules, cf_vals, iter)
        t2 = time()
        if iter%interval == 0
            save_data(save_o, cf_vals, iter)
            output(log_o, @sprintf("%10d %20.5f %20.5f %10d", iter, t2-t1, (t2-t0)/iter, nTrials - iter))
        end
        iter += 1
    end
    save_data(save_o, cf_vals, iter)
    output(log_o, @sprintf("%10d %20.5f %20.5f %10d", nTrials, t2-t0, (t2-t0)/(iter-1), 0))
    #close(log_o)
    return cf_vals
end

export Logger

function parallel_simulate(LP::AbstractLProblem, nTrials::Int64, alg::OrdinaryDiffEqAlgorithm, logger::Logger{T1,T2} = Logger{:no, :cmd}(), interval::Int64 = 1, offset::Int64 = 0; kwargs...) where {T1,T2}
    RNGarray = initRNG(nprocs())
    @sync begin
        for i in workers()
            @spawnat i simulate(LP,nTrials, alg, logger, i+offset, interval, RNGarray[i]; kwargs...)
        end
        @spawnat 1 simulate(LP,nTrials, alg, logger, 1+offset, interval, RNGarray[1]; kwargs...)
    end
end

export parallel_simulate

function parse_dir(dirname::AbstractString = pwd())
    blacklist = ["aggregate.jld", "aggregate1.jld", "aggregate2.jld"]
    corfs = Vector{CFVals}()
    Ns = Vector{Int64}()
    for filename in readdir(dirname)
        file = joinpath(dirname, filename)
        if !(filename in blacklist) && isfile(file) && file[end-3:end] == ".jld"
            corf, Ntr = load_data(file)
            push!(corfs, corf)
            push!(Ns, Ntr)
        end
    end
    return corfs, Ns
end

function parse_dir(dirlist::Tuple{AbstractString,Vararg{AbstractString}})
    blacklist = ["aggregate.jld", "aggregate1.jld", "aggregate2.jld"]
    corfs = Vector{CFVals}()
    Ns = Vector{Int64}()
    for dirname in dirlist
        for filename in readdir(dirname)
            file = joinpath(dirname, filename)
            if !(filename in blacklist) && isfile(file) && file[end-3:end] == ".jld"
                corf, Ntr = load_data(file)
                push!(corfs, corf)
                push!(Ns, Ntr)
            end
        end
    end
    return corfs, Ns
end

function aggregate(dirname::AbstractString = pwd(), normalize::Val{T} = Val{true}()) where {T}
    corfs, Ns = parse_dir(dirname)
    return aggregate(corfs, Ns, dirname, "aggregate.jld", normalize)
end

function aggregate(dirlist::Tuple{AbstractString,Vararg{AbstractString}}, destdir::AbstractString = pwd(), normalize::Val{T} = Val{true}()) where {T}
    corfs, Ns = parse_dir(dirlist)
    mkpath(destdir)
    return aggregate(corfs, Ns, destdir, "aggregate.jld", normalize)
end

get_norm_coef(coef::Float64, ::Val{true}) = coef
get_norm_coef(coef::Float64, ::Val{false}) = 1.0

function aggregate(corfs::Vector{CFVals}, Ns::Vector{Int64}, dirname::AbstractString = pwd(), agrname = "aggregate.jld", normalize::Val{T} = Val{true}()) where {T}
    aggregate = zeros(first(corfs))
    N_total = sum(Ns)
    for i in eachindex(corfs)
        aggregate.data   .= aggregate.data  .+ corfs[i].data .* Ns[i]
        aggregate.errors .= aggregate.errors .+ corfs[i].errors
    end
    aggregate.data .= aggregate.data .* (1.0/N_total)
    for i = 1:length(corfs)
        aggregate.errors .= aggregate.errors .+ (corfs[i].data - aggregate.data).^2 .* Ns[i]
    end
    divisor = 1.0/(N_total * (N_total - 1))
    for i in 1:length(aggregate.data)
        coef = 1.0/aggregate.data[1,i]
        aggregate.data[:,i] .= aggregate.data[:,i] .* get_norm_coef(coef, normalize)
        aggregate.errors[:, i] .= sqrt.(aggregate.errors[:,i] .* divisor) .* get_norm_coef(coef, normalize)
    end
    save_data(joinpath(dirname, agrname), aggregate, N_total)
    return aggregate, N_total
end

function separate(Ns::Vector{Int64})
    i = length(Ns)
    lm = div(i,2)
    l1 = Vector{Int64}()
    l2 = Vector{Int64}()
    sum1 = 0
    sum2 = 0
    while (length(l1)<lm) && (length(l2)<lm)
        el = Ns[i]
        if sum1>sum2
            sum2 += el
            push!(l2, i)
        else
            sum1 += el
            push!(l1, i)
        end
        i -=1
    end
    if sum1 > sum2
        append!(l2, 1:i)
    else
        append!(l1, 1:i)
    end
    return l1, l2
end

function aggregate2(dirname::AbstractString = pwd(), normalize::Val{T} = Val{true}()) where {T}
    corfs, Ns = parse_dir(dirname)
    l1, l2 = separate(Ns)
    agr1, Ntot1 = aggregate(corfs[l1], Ns[l1], dirname, "aggregate1.jld", normalize)
    agr2, Ntot2 = aggregate(corfs[l2], Ns[l2], dirname, "aggregate2.jld", normalize)
    return agr1, Ntot1, agr2, Ntot2
end

function aggregate2(dirlist::Tuple{AbstractString,Vararg{AbstractString}}, destdir::AbstractString = pwd(), normalize::Val{T} = Val{true}()) where {T}
    corfs, Ns = parse_dir(dirlist)
    l1, l2 = separate(Ns)
    mkpath(destdir)
    agr1, Ntot1 = aggregate(corfs[l1], Ns[l1], destdir, "aggregate1.jld", normalize)
    agr2, Ntot2 = aggregate(corfs[l2], Ns[l2], destdir, "aggregate2.jld", normalize)
    return agr1, Ntot1, agr2, Ntot2
end


module ChemostatPlots

import Chemostat
const Ch = Chemostat
const ChU = Chemostat.Utils
const ChLP = Chemostat.LP
import LinearAlgebra: normalize
using Plots
using Distributions

include("plot_marginal.jl")
include("plot_projection2D.jl")

end


# Scalar moment file currently does not exist. Oksana needs to make it.

## Read scalar moments from a single dedicated file
# I currently created this as a dummy file
function read_scalar_moments(ds :: DataSettings)
    fPath = joinpath(data_dir(ds), "scalar_moments.csv");
    csvFile = CSV.File(fPath, header = [:name, :value],
        delim = ',', comment = commentStr);
    return csvFile |> DataFrame
end

"""
	$(SIGNATURES)

Read one scalar moments from the scalar moments file.
"""
function read_scalar_moment(ds :: DataSettings, name :: String)
    df = read_scalar_moments(ds);
    outV = df[df.name .== name, :value];
    @assert length(outV) == 1
    outVal = outV[1];
    return outVal :: Float64
end

## Fraction entering college
function frac_enter(ds :: DataSettings)
    fracEnter = read_all_from_delim_file(raw_entry_gpa_parental(ds));
    cnt = read_all_from_delim_file(raw_entry_gpa_parental(ds; momentType = :count));
    @assert check_float(fracEnter, lb = 0.45, ub = 0.65);
    ses = (fracEnter * (1.0 - fracEnter) / cnt) ^ 0.5;
    return fracEnter, ses, cnt
end

## Graduation rate (conditional on entry)
function grad_rate(ds :: DataSettings)
    gradRate = read_all_from_delim_file(raw_grad_rate_qual_gpa(ds));
    cnt = read_all_from_delim_file(raw_grad_rate_qual_gpa(ds; momentType = :count));
    @assert check_float(gradRate, lb = 0.3, ub = 0.7);
    ses = (gradRate * (1.0 - gradRate) / cnt) ^ 0.5;
    return gradRate, ses, cnt
end

## Correlation HS GPA / parental income
# in which units? +++
function corr_gpa_yp(ds :: DataSettings)
    corr = read_scalar_moment(ds, "corrGpaYp");
    @assert corr > 0.3  &&  corr < 0.8
    return corr, 0.0, 1000
    # return ScalarDeviation{Double}(name = :corrGpaYp, dataV = corr, modelV = 0.0,
    #     wtV = 1.0 ./ corr,
    #     shortStr = "corrGpaYp", longStr = "Correlation HSgpa/parental background")
end

# Number of college entrants in the sample
# Useful for computing std errors
function n_entrants(ds :: DataSettings)
    cntM = read_matrix_by_xy(count_file(ds, :fracEnter_gpM));
    cnt = round(Int, sum(cntM));
    @assert 4000 < cnt < 7_000
    return cnt
end

## Average study time
#=
We have NLSY79 moments in xls, but not 97 moments.
This is based on Babcock/Marks 2011, table 2, HERI and NSSE
For full time 4 year students. So likely a bit high

Setting `wtV = 1/dataV` makes this a percentage deviation.
=#
# function dev_mean_study_time(ds :: DataSettings)
#     dataV = mean_study_time(ds; modelUnits = true);
#     return ScalarDeviation{Double}(name = :meanStudyTime, dataV = dataV, modelV = dataV,
#         wtV = 1 ./ dataV,
#         shortStr = "meanStudyTime", longStr = "Mean study time, hours per week")
# end

# Mean study time (hours per week)
# function mean_study_time(; modelUnits :: Bool = true)
#     mst = 25.0;
#     if modelUnits
#         mst = hours_per_week_to_mtu(mst);
#     end
#     return mst
# end

## Average work hours in college
#=
No good basis right now +++
=#
# function mean_work_time(; modelUnits :: Bool = true)
#     mwt = 15.0;
#     if modelUnits
#         mwt = hours_per_week_to_mtu(mwt);
#     end
#     return mwt
# end

## Wage earned in college (per hour worked)
#=
Based on xls table "Average hourly wage" which does not seem to vary by quality
=#
# function college_wage(ds :: DataSettings; modelUnits :: Bool = true)
#     wage = 7.0;
#     if modelUnits
#         wage = dollars_data_to_model(wage, :perHour);
#     end
#     return wage
# end

# function mean_course_load(; modelUnits :: Bool = true)
#     mcl = 10;
#     if !modelUnits
#         mcl = mcl * dataCoursesPerCourse;
#     end
#     return mcl
# end

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


module GeoData

# Use the README as the module docs
@doc let
    path = joinpath(dirname(@__DIR__), "README.md")
    include_dependency(path)
    read(path, String)
end GeoData

using Adapt,
      Dates,
      Missings,
      Mmap,
      ProgressMeter,
      RecipesBase,
      Reexport,
      Requires

@reexport using DimensionalData, GeoFormatTypes

const DD = DimensionalData

using Base: tail, @propagate_inbounds

using DimensionalData: StandardIndices

export AbstractGeoArray, MemGeoArray, DiskGeoArray, GeoArray

export AbstractGeoStack, MemGeoStack, DiskGeoStack, DiskStack, GeoStack

export AbstractGeoSeries, GeoSeries

export Projected, Mapped

export Band, Lat, Lon, Vert, GeoXDim, GeoYDim, GeoZDim

export missingval, boolmask, missingmask, replace_missing,
       aggregate, aggregate!, disaggregate, disaggregate!

export crs, mappedcrs, mappedindex, mappedbounds, projectedindex, projectedbounds

export geoarray, stack, series

const Lon = X
const Lat = Y
const Vert = Z
const GeoXDim = XDim
const GeoYDim = YDim
const GeoZDim = ZDim

# DimensionalData documentation urls
const DDdocs = "https://rafaqz.github.io/DimensionalData.jl/stable/api"
const DDdimdocs = joinpath(DDdocs, "#DimensionalData.Dimension")
const DDarraydocs = joinpath(DDdocs, "#DimensionalData.AbstractDimensionalArray")
const DDabssampleddocs = joinpath(DDdocs, "#DimensionalData.AbstractSampled")
const DDsampleddocs = joinpath(DDdocs, "#DimensionalData.Sampled")
const DDlocusdocs = joinpath(DDdocs, "#DimensionalData.Locus")
const DDselectordocs = joinpath(DDdocs, "#DimensionalData.Selector")
const DDtidocs = joinpath(DDdocs, "#DimensionalData.Ti")

include("mode.jl")
include("dimensions.jl")
include("array.jl")
include("stack.jl")
include("series.jl")
include("utils.jl")
include("aggregate.jl")
include("methods.jl")
include("open.jl")
include("read.jl")
include("sources/grd.jl")
include("show.jl")
include("plotrecipes.jl")
include("convenience.jl")

function __init__()
    @require HDF5="f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" begin
        # This section is for sources that rely on HDF5, not simply any HDF5.
        include("sources/smap.jl")
    end
    @require NCDatasets="85f8d34a-cbdd-5861-8df4-14fed0d494ab" begin
        include("sources/ncdatasets.jl")
    end
    @require ArchGDAL="c9ce4bd3-c3d5-55b8-8973-c0e20141b8c3" begin
        include("resample.jl")
        include("reproject.jl")
        include("sources/gdal.jl")
    end
    @require RasterDataSources="3cb90ccd-e1b6-4867-9617-4276c8b2ca36" begin
        include("sources/rasterdatasources.jl")
    end
end

end


# Script for processing BATSRUS explicit shocktube test weak scaling runs.
#
# The runlogs are stored in Google Drive, ScalingTests.
#
# Hongyang Zhou, hyzhou@umich.edu 05/09/2019, last modified 09/03/2020
include("Utility.jl")

using PyPlot
#using Plots
#pyplot()

Machine = "Frontera"

if Machine == "Bluewaters"

   # New results after modifying the message passing
   res = postprocess("results/gnu_BATL_update")
   #res = postprocess("results/gnu")

   result_0 = sort( res[res.Thread.==0,:],(:Proc))
   result_1 = sort( res[res.Thread.==1,:],(:Proc))
   result_2 = sort( res[res.Thread.==2,:],(:Proc))
   result_4 = sort( res[res.Thread.==4,:],(:Proc))
   result_8 = sort( res[res.Thread.==8,:],(:Proc))
   result_16 = sort(res[res.Thread.==16,:],(:Proc))
   result_32 = sort(res[res.Thread.==32,:],(:Proc))

   time_0 = copy(result_0.Time)
   time_1 = copy(result_1.Time)
   time_2 = copy(result_2.Time)
   time_4 = copy(result_4.Time)
   time_8 = copy(result_8.Time)
   time_16 = copy(result_16.Time)
   time_32 = copy(result_32.Time)

   append!(time_0,fill(NaN,size(time_32)[1]-size(time_0)[1]))
   append!(time_1,fill(NaN,size(time_32)[1]-size(time_1)[1]))
   append!(time_2,fill(NaN,size(time_32)[1]-size(time_2)[1]))
   append!(time_4,fill(NaN,size(time_32)[1]-size(time_4)[1]))
   append!(time_8,fill(NaN,size(time_32)[1]-size(time_8)[1]))
   append!(time_16,fill(NaN,size(time_32)[1]-size(time_16)[1]))

   # Calculate the number of cell updates per second for 2nd order scheme
   cell_size = sort(8^3 .* 256 .* result_32.Thread .* result_32.Proc .* 2)

   speed_0 = cell_size ./ time_0
   speed_1 = cell_size ./ time_1
   speed_2 = cell_size ./ time_2
   speed_4 = cell_size ./ time_4
   speed_8 = cell_size ./ time_8
   speed_16 = cell_size ./ time_16
   speed_32 = cell_size ./ time_32

   proc = result_32.Proc .* result_32.Thread

   speed = hcat(speed_0, speed_1, speed_2, speed_4, speed_8, speed_16, speed_32)
   speed_ref = speed_0[1] ./ proc[1] .* proc

   fig = plt.figure(figsize=(18,5))
   # For 1,2,4 threads
   ax1 = plt.subplot(121)

   plt.plot(proc[1:7], speed[1:7,:]/1e9, "-o", lw=1, markersize=3, alpha=0.8)
   plt.plot(proc[1:7], speed_ref[1:7]/1e9, ":")
   plt.legend(loc="upper left",
         ("pure MPI", "1 thread", "2 threads", "4 threads", "8 threads",
         "16 threads", "32 threads", "linear ref."))
   plt.grid(true)
   ax1.xaxis.set_major_locator(plt.MaxNLocator(5))
   plt.title("BATSRUS Weak Scaling, Explicit Scheme", y=1.03, fontsize=16)
   plt.xlabel("number of cores", fontsize=16)
   plt.ylabel(L"giga cell ($10^9$) updates per second", fontsize=16)

   # For 1,2,4,8,16,32 threads
   ax2 = plt.subplot(122)
   plt.plot(proc, speed/1e9, "-o", markersize=3, alpha=0.8)
   plt.plot(proc, speed_ref/1e9, ":")
   plt.legend(loc="upper left",
         ("pure MPI", "1 thread", "2 threads", "4 threads","8 threads",
         "16 threads", "32 threads", "linear ref."))
   plt.grid(true)
   ax2.xaxis.set_major_locator(plt.MaxNLocator(5))
   plt.title("BATSRUS Weak Scaling, Explicit Scheme", y=1.03, fontsize=16)
   plt.xlabel("number of cores", fontsize=16)
   plt.ylabel(L"giga cell ($10^9$) updates per second", fontsize=16)
   #plt.text(0.9,-0.08, "number of cores", transform=ax2.transAxes)

   #=
   mylabel = ["1 thread" "16 threads" "32 threads" "" "" ""]

   plot(proc[1:end],speed[1:end,:], seriestype=[:line, :scatter],
      markersize=4, lw=2, markercolor=[:red :green :purple],
      markerstrokewidth = 0.1,
      linecolor=[:red :green :purple],
      xscale=:identity, yscale=:identity, legend=:best,
      xlim=[10^2, 1.4e5], ylim=[10^6,1.8e9],
      label=mylabel,
      xlabel="number of cores",
      ylabel="cell updates per second",
      title="BATSRUS Hybrid Weak Scaling",
      xtickfontsize=16, ytickfontsize=16,
      legendfontsize=16, axis=:equal,
      size=(900,600))

   plot!(proc[1:end],speed_ref[1:end], linestyle=:dashdot, linecolor=:black,
         label="linear ref.")
   =#

   #=
   speed = hcat(speed_1,speed_2,speed_4,speed_8)
   mylabel = ["1 thread" "2 threads" "4 threads" "8 threads" "" "" "" ""]

   plot(proc[1:6],speed[1:6,:], seriestype=[:line, :scatter],
      markersize=4, lw=2, markercolor=[1 2 3 4],
      markerstrokewidth = 0.1,
      linecolor=[1 2 3 4],
      xscale=:identity, yscale=:identity, legend=:best,
      label=mylabel,
      xlabel="number of cores",
      ylabel="cell updates per second",
      xtickfontsize=16, ytickfontsize=16,
      legendfontsize=16, axis=:equal,
      reuse=false)
   =#

   #using StatsPlots

   #=
   @df results plot(:Thread, :Time, xticks=[1,2,4,8,16,32],
      seriestype=:scatter,
      title="Timing")

   @df results plot!(:Thread, :Time, xticks=[1,2,4,8,16,32],
      seriestype=:line,
      title="Timing")
   =#

   #plotly()

   #=
   @df results plot(:Thread, :Time, xticks=[1,2,4,8,16,32], group=(:Compiler),
      seriestype=[:line,:scatter],
      markersize=0.5 .* log2.(:Proc),
      xlabel="thread",
      ylabel="execution time [s]",
      title="Timing")
   =#

elseif Machine == "Stampede"

   # New results after modifying the message passing
   res = postprocess("results/intel_stampede")

   result_1 = sort( res[res[:Thread].==1,:],(:Proc))
   result_2 = sort( res[res[:Thread].==2,:],(:Proc))
   result_4 = sort( res[res[:Thread].==4,:],(:Proc))
   result_8 = sort( res[res[:Thread].==8,:],(:Proc))
   result_12 = sort( res[res[:Thread].==12,:],(:Proc))
   result_16 = sort(res[res[:Thread].==16,:],(:Proc))

   time_1 = result_1[:Time]
   time_2 = result_2[:Time]
   time_4 = result_4[:Time]
   time_8 = result_8[:Time]
   time_12 = result_12[:Time]
   time_16 = result_16[:Time]

   append!(time_1,fill(NaN,size(time_16)[1]-size(time_1)[1]))
   append!(time_2,fill(NaN,size(time_16)[1]-size(time_2)[1]))
   append!(time_4,fill(NaN,size(time_16)[1]-size(time_4)[1]))
   append!(time_8,fill(NaN,size(time_16)[1]-size(time_8)[1]))
   append!(time_12,fill(NaN,size(time_16)[1]-size(time_12)[1]))
   append!(time_16,fill(NaN,size(time_16)[1]-size(time_16)[1]))

   # Calculate the number of cell updates per second for 2nd order scheme
   cell_size = sort(8^3 .* 256 .* result_16[:Thread] .* result_16[:Proc] .* 2)

   speed_1 = cell_size ./ time_1
   speed_2 = cell_size ./ time_2
   speed_4 = cell_size ./ time_4
   speed_8 = cell_size ./ time_8
   speed_12 = cell_size ./ time_12
   speed_16 = cell_size ./ time_16

   proc = result_16[:Proc] .* result_16[:Thread]
   #=
   speed = hcat(speed_1,speed_16,speed_32)

   mylabel = ["1 thread" "16 threads" "32 threads" "" "" ""]

   plot(proc,speed, seriestype=[:line, :scatter],
      markersize=4, lw=2, markercolor=[:red :green :purple],
      markerstrokewidth = 0.1,
      linecolor=[:red :green :purple],
      xscale=:identity, yscale=:identity, legend=:best,
      xlim=[10^2, 5.5e5], ylim=[10^6,4.5e9],
      label=mylabel,
      xlabel="number of cores",
      ylabel="cell updates per second",
      title="BATSRUS Hybrid Weak Scaling",
      xtickfontsize=16, ytickfontsize=16,
      legendfontsize=16, axis=:equal)

   speed_ref = speed_1[1] ./ proc[1] .* proc
   plot!(proc,speed_ref, linestyle=:dashdot, linecolor=:black,
         label="linear ref.")
   =#
   speed = hcat(speed_1,speed_2,speed_4,speed_8,speed_12,speed_16)
   speed_ref = speed_1[1] ./ proc[1] .* proc

   plt.figure(figsize=(5,5))
   plt.plot(proc,speed/1e9,"-o",markersize=3)
   plt.plot(proc,speed_ref/1e9,"--")
   plt.legend(loc="upper left",
         ("1 thread", "2 threads", "4 threads", "8 threads", "12 threads",
         "16 threads","linear ref."))
   plt.grid(true)
   plt.title("BATSRUS Weak Scaling, Explicit Scheme",y=1.03)
   plt.xlabel("number of cores")
   plt.ylabel(L"giga cell ($10^9$) updates per second")

   #=
   mylabel = ["1 thread", "2 threads", "4 threads", "8 threads", "16 threads",
         "", "", "", "", ""]
   mylabel = reshape(mylabel,1,:)

   plot(proc,speed[:,:], seriestype=[:line, :scatter],
      markersize=4, lw=2, markercolor=[1 2 3 4 5],
      markerstrokewidth = 0.1,
      linecolor=[1 2 3 4 5],
      xscale=:identity, yscale=:identity, legend=:best,
      label=mylabel,
      xlabel="number of cores",
      ylabel="cell updates per second",
      xtickfontsize=16, ytickfontsize=16,
      legendfontsize=16, axis=:equal,
      reuse=false)

   plot!(proc,speed_ref, linestyle=:dashdot, linecolor=:black,
         label="linear ref.")
   =#
elseif Machine == "Frontera"

   res = postprocess("RESULT_gnu_Expl")

   result_1 = sort( res[res.Thread.==1,:],(:Proc))
   result_2 = sort( res[res.Thread.==2,:],(:Proc))
   result_4 = sort( res[res.Thread.==4,:],(:Proc))
   result_7 = sort( res[res.Thread.==7,:],(:Proc))
   result_14 = sort( res[res.Thread.==14,:],(:Proc))
   result_28 = sort( res[res.Thread.==28,:],(:Proc))
   result_56 = sort(res[res.Thread.==56,:],(:Proc))

   time_1 = copy(result_1.Time)
   time_2 = copy(result_2.Time)
   time_4 = copy(result_4.Time)
   time_7 = copy(result_7.Time)
   time_14 = copy(result_14.Time)
   time_28 = copy(result_28.Time)
   time_56 = copy(result_56.Time)

   append!(time_1,fill(NaN,size(time_56)[1]-size(time_1)[1]))
   append!(time_2,fill(NaN,size(time_56)[1]-size(time_2)[1]))
   append!(time_4,fill(NaN,size(time_56)[1]-size(time_4)[1]))
   append!(time_7,fill(NaN,size(time_56)[1]-size(time_7)[1]))
   append!(time_14,fill(NaN,size(time_56)[1]-size(time_14)[1]))
   append!(time_28,fill(NaN,size(time_56)[1]-size(time_28)[1]))
   append!(time_56,fill(NaN,size(time_56)[1]-size(time_56)[1]))

   # Calculate the number of cell updates per second for 2nd order scheme
   cell_size = sort(8^3 .* 224 .* result_56.Thread .* result_56.Proc .* 2)

   speed_1 = cell_size ./ time_1
   speed_2 = cell_size ./ time_2
   speed_4 = cell_size ./ time_4
   speed_7 = cell_size ./ time_7
   speed_14 = cell_size ./ time_14
   speed_28 = cell_size ./ time_28
   speed_56 = cell_size ./ time_56

   proc = result_56.Proc .* result_56.Thread

   speed = hcat(speed_1, speed_2, speed_4, speed_7, speed_14, speed_28, speed_56)
   speed_ref = speed_1[1] ./ proc[1] .* proc

   fig = plt.figure(figsize=(6,5))

   ax1 = plt.subplot(111)

   plt.plot(proc, speed/1e9, "-o", lw=1, markersize=3, alpha=0.8)
   plt.plot(proc, speed_ref/1e9, ":")
   plt.legend(loc="upper left",
         ("1 thread", "2 thread", "4 thread", "7 threads", "14 threads",
         "28 threads", "56 threads", "linear ref."))
   plt.grid(true)
   ax1.xaxis.set_major_locator(plt.MaxNLocator(5))
   plt.title("BATSRUS Weak Scaling, Explicit Scheme", y=1.03, fontsize=16)
   plt.xlabel("number of cores", fontsize=16)
   plt.ylabel(L"giga cell ($10^9$) updates per second", fontsize=16)

end


# This file is a part of Julia. License is MIT: http://julialang.org/license

# Benchmark implementing the board logic for the game of go and
# exercising it by playing random games. Derived from
# http://www.lysator.liu.se/~gunnar/gtp/brown-1.0.tar.gz
import Base.getindex

const EMPTY = 0
const WHITE = 1
const BLACK = 2

# Function to find the opposite color.
other_color(color::Int) = WHITE + BLACK - color

# Used in the final_status[] array.
const DEAD = 0
const ALIVE = 1
const SEKI = 2
const WHITE_TERRITORY = 3
const BLACK_TERRITORY = 4
const UNKNOWN = 5

type XorRand
  state::UInt32
end

function xor_srand(rand::XorRand, seed::UInt32)
  rand.state = seed
end

function xor_randn(rand::XorRand, n::UInt32)
  rand.state $= rand.state << 13
  rand.state $= rand.state >> 17
  rand.state $= rand.state << 5
  rand.state % n
end

xor_randn(rand::XorRand, n::Int) = convert(Int, xor_randn(rand, convert(UInt32, n)))

# Offsets for the four directly adjacent neighbors. Used for looping.
const deltai = (-1, 1, 0, 0)
const deltaj = (0, 0, -1, 1)
neighbor(i::Int, j::Int, k::Int) = (i + deltai[k], j + deltaj[k])

type Board
  size::Int
  komi::Float64

  # Board represented by a 1D array. The first board_size*board_size
  # elements are used. Vertices are indexed row by row, starting with 0
  # in the upper left corner.
  board::Matrix{Int}

  # Stones are linked together in a circular list for each string.
  next_stone::Matrix{Int}

  # Storage for final status computations.
  final_status::Matrix{Int}

  # Point which would be an illegal ko recapture.
  ko_i::Int
  ko_j::Int

  # xor-shift random number generator.
  rand::XorRand

  function Board(n::Int)
    init(new(), n, convert(UInt32, 2463534242))
  end
end

function init(board::Board, n::Int, seed::UInt32)
  board.size = n
  board.komi = 0.0
  board.board = zeros(Int, n, n)
  board.next_stone = zeros(Int, n, n)
  board.final_status = zeros(Int, n, n)
  board.ko_i = 0
  board.ko_j = 0
  board.rand = XorRand(seed)
  board
end

getindex(board::Board, pos::Int) = board.board[pos]
getindex(board::Board, i::Int, j::Int) = board.board[i, j]

# Functions to convert between 1D and 2D coordinates. The 2D coordinate
# (i, j) points to row i and column j, starting with (1,1) in the
# upper left corner.
POS(board::Board, i::Int, j::Int) = (j - 1) * board.size + i
IJ(board::Board, pos::Int) = (1 + mod((pos - 1), board.size), 1 + fld(pos - 1, board.size))

function set_komi(board::Board, komi::Float64)
  board.komi = komi
end

function set_random_seed(board::Board, seed::UInt32)
  xor_srand(board.rand, seed)
end

function clear(board::Board)
  board.board[:] = EMPTY
end

is_pass_move(i::Int, j::Int) = i == 0 && j == 0

function on_board(board::Board, i::Int, j::Int)
  i >= 1 && i <= board.size && j >= 1 && j <= board.size
end

function legal_move(board::Board, i::Int, j::Int, color::Int)
  other = other_color(color)

  # Pass is always legal.
  if is_pass_move(i, j)
    return true
  end

  # Already occupied.
  if board[i, j] != EMPTY
    return false
  end

  # Illegal ko recapture. It is not illegal to fill the ko so we must
  # check the color of at least one neighbor.
  if i == board.ko_i && j == board.ko_j && ((on_board(board, i - 1, j) && board[i - 1, j] == other) || (on_board(board, i + 1, j) && board[i + 1, j] == other))
    return false
  end

  true
end

# Does the string at (i, j) have any more liberty than the one at (libi, libj)?
function has_additional_liberty(board::Board, i::Int, j::Int, libi::Int, libj::Int)
  start = POS(board, i, j)
  pos = start
  while true
    (ai, aj) = IJ(board, pos)
    for k = 1:4
      (bi, bj) = neighbor(ai, aj, k)
      if on_board(board, bi, bj) && board[bi, bj] == EMPTY && (bi != libi || bj != libj)
        return true
      end
    end

    pos = board.next_stone[pos]
    if pos == start
      break
    end
  end

  false
end

# Does (ai, aj) provide a liberty for a stone at (i, j)?
function provides_liberty(board::Board, ai::Int, aj::Int, i::Int, j::Int, color::Int)
  # A vertex off the board does not provide a liberty.
  if !on_board(board, ai, aj)
    return false
  end

  # An empty vertex IS a liberty.
  if board[ai, aj] == EMPTY
    return true
  end

  # A friendly string provides a liberty to (i, j) if it currently
  # has more liberties than the one at (i, j).
  if board[ai, aj] == color
    return has_additional_liberty(board, ai, aj, i, j)
  end

  # An unfriendly string provides a liberty if and only if it is
  # captured, i.e. if it currently only has the liberty at (i, j).
  !has_additional_liberty(board, ai, aj, i, j)
end

# Is a move at ij suicide for color?
function suicide(board::Board, i::Int, j::Int, color::Int)
  for k = 1:4
    if provides_liberty(board, neighbor(i, j, k)..., i, j, color)
      return false
    end
  end
  true
end

# Remove a string from the board array. There is no need to modify
# the next_stone array since this only matters where there are
# stones present and the entire string is removed.
function remove_string(board::Board, i::Int, j::Int)
  start = POS(board, i, j)
  pos = start
  removed = 0
  while true
    board.board[pos] = EMPTY
    removed += 1
    pos = board.next_stone[pos]
    if pos == start
      break
    end
  end
  removed
end

# Do two vertices belong to the same string? It is required that both
# pos1 and pos2 point to vertices with stones.
function same_string(board::Board, pos1::Int, pos2::Int)
  pos = pos1
  while true
    if pos == pos2
      return true
    end
    pos = board.next_stone[pos]
    if pos == pos1
      break
    end
  end
  false
end

# Play at (i, j) for color. No legality check is done here. We need
# to properly update the board array, the next_stone array, and the
# ko point.
function play_move(board::Board, i::Int, j::Int, color::Int)
  pos = POS(board, i, j)
  captured_stones = 0

  # Reset the ko point.
  board.ko_i = 0
  board.ko_j = 0

  # Nothing more happens if the move was a pass.
  if is_pass_move(i, j)
    return
  end

  # If the move is a suicide we only need to remove the adjacent
  # friendly stones.
  if suicide(board, i, j, color)
    for k = 1:4
      (ai, aj) = neighbor(i, j, k)
      if on_board(board, ai, aj) && board[ai, aj] == color
        remove_string(board, ai, aj)
      end
    end
    return
  end

  # Not suicide. Remove captured opponent strings.
  for k = 1:4
    (ai, aj) = neighbor(i, j, k)
    if on_board(board, ai, aj) && board[ai, aj] == other_color(color) && !has_additional_liberty(board, ai, aj, i, j)
      captured_stones += remove_string(board, ai, aj)
    end
  end

  # Put down the new stone. Initially build a single stone string by
  # setting next_stone[pos] pointing to itself.
  board.board[pos] = color
  board.next_stone[pos] = pos

  # If we have friendly neighbor strings we need to link the strings
  # together.
  for k = 1:4
    (ai, aj) = neighbor(i, j, k)
    pos2 = POS(board, ai, aj)
    # Make sure that the stones are not already linked together. This
    # may happen if the same string neighbors the new stone in more
    # than one direction.
    if on_board(board, ai, aj) && board[pos2] == color && !same_string(board, pos, pos2)
      # The strings are linked together simply by swapping the the
      # next_stone pointers.
      (board.next_stone[pos], board.next_stone[pos2]) = (board.next_stone[pos2], board.next_stone[pos])
    end
  end

  # If we have captured exactly one stone and the new string is a
  # single stone it may have been a ko capture.
  if captured_stones == 1 && board.next_stone[pos] == pos
    # Check whether the new string has exactly one liberty. If so it
    # would be an illegal ko capture to play there immediately. We
    # know that there must be a liberty immediately adjacent to the
    # new stone since we captured one stone.
    for k = 1:4
      (ai, aj) = neighbor(i, j, k)
      if on_board(board, ai, aj) && board[ai, aj] == EMPTY
        if !has_additional_liberty(board, i, j, ai, aj)
          board.ko_i = ai
          board.ko_j = aj
        end
        break
      end
    end
  end
end

# Generate a move.
function generate_move(board::Board, color::Int)
  moves = zeros(Int, 2, board.size * board.size)
  num_moves = 0

  for ai = 1:board.size, aj = 1:board.size
    # Consider moving at (ai, aj) if it is legal and not suicide.
    if legal_move(board, ai, aj, color) && !suicide(board, ai, aj, color)
      # Further require the move not to be suicide for the opponent...
      if !suicide(board, ai, aj, other_color(color))
        num_moves += 1
        moves[:,num_moves] = [ai, aj]
      else
        # ...however, if the move captures at least one stone,
        # consider it anyway.
        for k = 1:4
          (bi, bj) = neighbor(ai, aj, k)
          if on_board(board, bi, bj) && board[bi, bj] == other_color(color)
            num_moves += 1
            moves[:,num_moves] = [ai, aj]
            break
          end
        end
      end
    end
  end

  # Choose one of the considered moves randomly with uniform
  # distribution. (Strictly speaking the moves with smaller 1D
  # coordinates tend to have a very slightly higher probability to be
  # chosen, but for all practical purposes we get a uniform
  # distribution.)
  if num_moves > 0
    move = moves[:,1 + xor_randn(board.rand, num_moves)]
    return (move[1], move[2])
  else
    # But pass if no move was considered.
    return (0, 0)
  end
end

# Set a final status value for an entire string.
function set_final_status_string(board::Board, pos::Int, status::Int)
  start = pos
  while true
    board.final_status[pos] = status
    pos = board.next_stone[pos]
    if pos == start
      break
    end
  end
end

# Compute final status. This function is only valid to call in a
# position where generate_move() would return pass for at least one
# color.
#
# Due to the nature of the move generation algorithm, the final
# status of stones can be determined by a very simple algorithm:
#
# 1. Stones with two or more liberties are alive with territory.
# 2. Stones in atari are dead.
#
# Moreover alive stones are unconditionally alive even if the
# opponent is allowed an arbitrary number of consecutive moves.
# Similarly dead stones cannot be brought alive even by an arbitrary
# number of consecutive moves.
#
# Seki is not an option. The move generation algorithm would never
# leave a seki on the board.
#
# Comment: This algorithm doesn't work properly if the game ends with
#          an unfilled ko. If three passes are required for game end,
#          that will not happen.
#
function compute_final_status(board::Board)
  board.final_status[:] = UNKNOWN

  for i = 1:board.size, j = 1:board.size
    if board[i, j] == EMPTY
      for k = 1:4
        (ai, aj) = neighbor(i, j, k)
        if !on_board(board, ai, aj)
          continue
        end
        # When the game is finished, we know for sure that (ai, aj)
        # contains a stone. The move generation algorithm would
        # never leave two adjacent empty vertices. Check the number
        # of liberties to decide its status, unless it's known
        # already.
        #
        # If we should be called in a non-final position, just make
        # sure we don't call set_final_status_string() on an empty
        # vertex.
        pos = POS(board, ai, aj)
        if board.final_status[ai, aj] == UNKNOWN
          if board[ai, aj] != EMPTY
            if has_additional_liberty(board, ai, aj, i, j)
              set_final_status_string(board, pos, ALIVE)
            else
              set_final_status_string(board, pos, DEAD)
            end
          end
        end
        # Set the final status of the pos vertex to either black
        # or white territory.
        if board.final_status[i, j] == UNKNOWN
          if (board.final_status[ai, aj] == ALIVE) $ (board[ai, aj] == WHITE)
            board.final_status[i, j] = BLACK_TERRITORY
          else
            board.final_status[i, j] = WHITE_TERRITORY
          end
        end
      end
    end
  end
end

get_final_status(board::Board, i::Int, j::Int) = board.final_status[i, j]

function compute_score(board::Board)
  score = board.komi
  compute_final_status(board)
  for i = 1:board.size, j = 1:board.size
    status = get_final_status(board, i, j)
    if status == BLACK_TERRITORY
      score -= 1.0
    elseif status == WHITE_TERRITORY
      score += 1.0
    elseif (status == ALIVE) $ (board[i, j] == WHITE)
      score -= 1.0
    else
      score += 1.0
    end
  end
  score
end

function benchmark(num_games_per_point::Int)
  random_seed = convert(UInt32, 1)
  board_size = 9
  komi = 0.5

  board = Board(board_size)
  set_komi(board, komi)
  set_random_seed(board, random_seed)

  for i = 1:board.size, j = 1:board.size
    white_wins = 0
    black_wins = 0
    for k = 1:num_games_per_point
      passes = 0
      num_moves = 1
      color = WHITE
      clear(board)
      play_move(board, i, j, BLACK)
      while passes < 3 && num_moves < 600
        (movei, movej) = generate_move(board, color)
        play_move(board, movei, movej, color)
        if is_pass_move(movei, movej)
          passes += 1
        else
          passes = 0
        end
        num_moves += 1
        color = other_color(color)
      end
      if passes == 3
        if compute_score(board) > 0
          white_wins += 1
        else
          black_wins += 1
        end
      end
    end
#    @printf("%d %d %f\n", i - 1, j - 1, black_wins / (black_wins + white_wins))
  end
end

function main(args)
  n = 10
  if length(args) > 0
    n = parseint(args[1])
  end
  @time benchmark(n)
end


#
# This script was used to generate results for the camera-ready AABI submission.
#

println("pseudo_point_synthetic_experiments_irregular_obs_sum_separable")

using Revise

using Pkg
Pkg.activate(".");
Pkg.resolve();
Pkg.instantiate();

using ArgParse
using CSV

# Make it possible to parse strings from the command line into Vector{Int}s.
function ArgParse.parse_item(::Type{Vector{Int}}, x::String)
    @assert x[1] == '[' && x[end] == ']'
    return CSV.read(IOBuffer(x[2:prevind(x, end)]); header=false, transpose=true).Column1
end

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table! s begin
        "--run-experiments"
            help = "true / false - run the experiments or not"
            arg_type = Bool
            default = false
        "--N-space-per-time"
            help = "number of points to consider in space per time point"
            arg_type = Int
            default = 10
        "--N-time"
            help = "number of time points to use -- expects multiple settings"
            arg_type = Vector{Int}
            default = [
                10, 25, 50,
                100, 250, 500,
                1_000, 2_500, 5_000,
                10_000, 25_000, 50_000,
                100_000,
            ]
        "--M-per-time"
            help = "number of pseudo-points per time point -- expects multiple settings. 0 means exact"
            arg_type = Vector{Int}
            default = [0, 5, 10, 20]
        "--BLAS-threads"
            help = "number of threads to let BLAS use"
            arg_type = Int
            default = 1
    end
    return parse_args(s)
end

args = parse_commandline()

using BenchmarkTools
using DrWatson
using LinearAlgebra
using Random
using Stheno
using TemporalGPs

BLAS.set_num_threads(args["BLAS-threads"])

const exp_dir_name = "pseudo_point_scaling_irregular_obs_sum_separable"

#
# Specify some functionality required to run the experiments.
#

function separable_kernel(σ::Real, λ_space::Real, λ_time::Real)
    k_space = σ^2 * stretch(EQ(), λ_space)
    k_time = stretch(Matern52(), λ_time)
    return TemporalGPs.Separable(k_space, k_time)
end

function sum_separable_kernel(σ::Real, λ_space::Real, λ_time::Real)
    k1 = separable_kernel(σ, λ_space, λ_time)
    k2 = separable_kernel(1.1 * σ, 0.3 * λ_space, 1.05 * λ_time)
    return k1 + k2
end

#
# Specify experiments to run.
#

# Ensure that output directory for results exists.
const data_dir = joinpath(datadir(), exp_dir_name)
mkpath(data_dir)

# Ensure that directory for experiment specification exists.
mkpath(joinpath(data_dir, "meta"))

wsave(
    joinpath(data_dir, "meta", "settings.bson"),
    :settings => dict_list(
        Dict(
            # The number of kernels to include in the summation.
            :model => [
                (
                    name = "sum-separable",
                    θ = (
                        σ = 0.92,
                        λ_space = 1.1,
                        λ_time = 0.83,
                        S_min = 0.1,
                    ),
                    ctor = sum_separable_kernel,
                ),
            ],

            # The total number of points in space to consider.
            :N_space_per_time => args["N-space-per-time"],

            # The total number of points in time to consider.
            :Ns_time => args["N-time"],

            # The total number of pseudo-points per time point.
            :Ms_per_time => args["M-per-time"],
        ),
    ),
)

#
# Run experiments.
#

if args["run-experiments"]

println("Running experiments")

let
    settings = load(joinpath(data_dir, "meta", "settings.bson"))[:settings]

    for (n, setting) in enumerate(settings)
        println("Starting experiment $n / $(length(settings))")

        # Set seed for reproducibility.
        rng = MersenneTwister(123456)

        # Extract experiment parameters.
        model_name = setting[:model].name
        @show model_name
        θ = setting[:model].θ
        ctor = setting[:model].ctor
        T = setting[:Ns_time]
        N_space_per_time = setting[:N_space_per_time]
        M_per_time = setting[:Ms_per_time]

        # Specify input locations. We place these irregularly through time.
        t = sort(rand(rng, T)) .* 50
        r = [rand(rng, N_space_per_time) * 10 for _ in 1:T]
        x = TemporalGPs.RegularInTime(t, r)
        x_naive = collect(x)

        # Specify inducing inputs. They are regularly-spaced between 0 and 10.
        z_r = collect(range(0.0, 10.0; length=M_per_time))
        z_naive = collect(TemporalGPs.RectilinearGrid(z_r, t))

        # Specify observation noise.
        S = [Diagonal(rand(rng, N_space_per_time) .+ θ.S_min) for _ in 1:T]
        S_naive = vcat(diag.(S)...)

        # Generate data exactly / naively.
        k = ctor(θ.σ, θ.λ_space, θ.λ_time)
        f = GP(k, GPC())

        # Hack together generation.
        N_per_slice = 5_000
        ys_naive = map(1:div(length(x_naive), N_per_slice)) do j
            idx = N_per_slice * (j - 1) + 1:min(N_per_slice * j, length(x_naive))
            return rand(rng, f(x_naive[idx], S_naive[idx]))
        end
        ys_naive = if rem(length(x_naive), N_per_slice) == 0
            ys_naive
        else
            idx = length(x_naive) - rem(length(x_naive), N_per_slice) + 1:length(x_naive)
            vcat(
                ys_naive,
                rand(rng, f(x_naive[idx], S_naive[idx]))
            )
        end
        y_naive = vcat(ys_naive...)
        y = TemporalGPs.match_to(y_naive, x)

        display((length(x_naive), length(y_naive)))
        println()
        if M_per_time == 0 && length(y_naive) <= 10_000
            fx_naive = f(x_naive, S_naive)
            println("Performing exact inference.")
            wsave(
                joinpath(
                    data_dir,
                    savename(
                        Dict(
                            :k => model_name,
                            :N_space_per_time => N_space_per_time,
                            :N_time => T,
                            :M_per_time => M_per_time,
                            :method => :exact,
                        ),
                        "bson",
                    ),
                ),
                Dict(
                    :setting => setting,
                    :lml_result => logpdf(fx_naive, y_naive),
                    :timing => (@benchmark logpdf($fx_naive, $y_naive)),
                    :method => :exact,
                ),
            )
        end

        # State-Space ELBO
        println("Using the state-space ELBO")
        if M_per_time > 0
            wsave(
                joinpath(
                    data_dir,
                    savename(
                        Dict(
                            :k => model_name,
                            :N_space_per_time => N_space_per_time,
                            :N_time => T,
                            :M_per_time => M_per_time,
                            :method => :state_space_elbo,
                        ),
                        "bson",
                    ),
                ),
                Dict(
                    :setting => setting,
                    :lml_result => elbo(to_sde(f)(x, S_naive), y_naive, z_r),
                    :timing => (@benchmark elbo(to_sde($f)($x, $S_naive), $y_naive, $z_r)),
                    :method => :state_space_elbo,
                ),
            )
        end
    end
end

end # run-experiments

#
# Summarise results. Generate output (plots and tables).
#

using ColorTypes, DataFrames, LaTeXStrings, PGFPlotsX, Plots, TexTables
pgfplotsx();

Plots.default(
    tickfontsize=nothing,
    legendfontsize=nothing,
    guidefontsize=nothing,
);

# Use default document font-size if nothing provided.
function Plots.pgfx_font(fontsize::Nothing, thickness_scaling = 1, font = "\\selectfont")
    return string("{", font, "}")
end

# Ensure that results directory exists.
const results_dir = joinpath(projectdir(), "results", exp_dir_name)
mkpath(results_dir)

let
    # Load the results and augment columns with extra info from `settings` column.
    results_unsorted = DataFrame(
        map(eachrow(collect_results(data_dir))) do row
            merge(
                copy(row),
                (
                    N_time = row.setting[:Ns_time],
                    M_per_time = row.setting[:Ms_per_time],
                ),
            )
        end
    )

    results = sort(results_unsorted, [:method, :M_per_time]; rev=true)

    # Each inference method gets its own marker shape.
    marker_shapes = Dict(
        :exact => :square,
        :naive_elbo => :circle,
        :state_space_elbo => :circle,
    )

    # Each inference method gets its own colour.
    colours = Dict(
        :exact => :black,
        :naive_elbo => :red,
        :state_space_elbo => :blue,
    )

    method_name_map = Dict(
        :exact => "exact",
        :naive_elbo => "elbo (naive)",
        :state_space_elbo => "elbo"
    )

    # Each M gets its own line style. This means that the maximum value of `M` is 5.
    # We really shouldn't be using that many anyway though...
    line_styles = [:solid, :solid, :dash, :dot, :dashdot, :dashdotdot]
    Ms_per_time = sort(unique(results.M_per_time))
    line_styles_map = Dict(
        M => style for (M, style) in zip(Ms_per_time, line_styles[eachindex(Ms_per_time)])
    )

    # Locations to save the results to.
    save_dirs = [
        results_dir,
        "/Users/willtebbutt/Dropbox/University/PhD/first_year_report/chapter_3/pseudo_points"
    ]

    # Ensure that save_dirs exist.
    foreach(mkpath, save_dirs)

    let

        # width_pts = 246.0  # or any other value
        # inches_per_points = 1.0 / 72.27
        # width_inches = width_pts *inches_per_points
        width_inches = 2.8
        width_px = width_inches * Plots.DPI  # or  width_inches*DPI
        height_px = width_px

        plt_lml = plot(
            xlabel="T",
            ylabel="ELBO / T",
            legend=:bottomright,
            legendfontsize=9,
            background_color_legend=RGBA(1, 1, 1, 0.6),
            size=(width_px, height_px),
            grid = (1, 0.4),
        );
        plt_comp_time = plot(
            xlabel="T",
            ylabel="Compute Time (s)",
            legend=:bottomright,
            legendfontsize=9,
            background_color_legend=RGBA(1, 1, 1, 0.6),
            size=(width_px, height_px),
            grid = (1, 0.4),
        )
        for group in groupby(results, [:method, :M_per_time])
            method_name = only(unique(group.method))
            M_per_time = only(unique(group.M_per_time))
            if method_name == :exact && M_per_time != first(Ms_per_time)
                continue
            end
            if method_name == :naive_elbo
                continue
            end

            idx = sortperm(group.N_time)
            label = if method_name == :exact
                method_name_map[method_name]
            else
                L"M_\tau=%$M_per_time"
            end

            Ts = group.N_time[idx]

            plot!(plt_lml, Ts, group.lml_result[idx] ./Ts ;
                markershape=marker_shapes[method_name],
                color=colours[method_name],
                linestyle=line_styles_map[M_per_time],
                markerstrokewidth=0,
                markersize=3,
                label=label,
                xscale=:log10,
            )

            min_time = getfield.(minimum.(group.timing), :time)[idx] ./ 1e9
            plot!(plt_comp_time, Ts, min_time;
                markershape=marker_shapes[method_name],
                color=colours[method_name],
                linestyle=line_styles_map[M_per_time],
                markersize=3,
                markerstrokewidth=0,
                label=label,
                yscale=:log10,
                xscale=:log10,
            )
        end
        foreach(
            save_dir -> begin
                savefig(
                    plt_lml,
                    joinpath(save_dir, "irregular_lml_plot_sum_separable.tikz"),
                )
                savefig(
                    plt_comp_time,
                    joinpath(save_dir, "irregular_timing_plot_sum_separable.tikz"),
                )
                savefig(
                    plot(plt_lml, plt_comp_time; layout=(2, 1)),
                    joinpath(save_dir, "irregular_lml_timing_plot_sum_separable.tikz"),
                )
            end,
            save_dirs,
        )
    end
end


using AutoLandmarking
using Flux
using CuArrays
import ImageView

push!(LOAD_PATH, Base.@__DIR__)

# Import data with custom import functions
image_data, imname = load_imgs("C:/Users/immanueldiener/Desktop/Master/scripts/no_enamel", (128,128,128), false)
landmark_data = read_landmarks("C:/Users/immanueldiener/Desktop/Master/scripts/no_enamel", 10, "@1")

# X and Y coordinates are swapped in avizo files, wo I'll swap them back here.
#this step is essential for all the data augmentation to work.
landmark_data = swap_xy(landmark_data)

#some images have a higher value in the fill voxels than in the actual object voxels. This will fix the issue.
# The highest value must always be the filler value!
image_data[:,:,:,[109,130]] = image_data[:,:,:,[109,130]] .- 1
image_data[:,:,:,[109,130]] = image_data[:,:,:,[109,130]] .* -1

image_data2, imname2 = load_imgs("C:/Users/immanueldiener/Desktop/Master/master_data\\additional_teeth_volumes\\", (128,128,128), true)
landmark_data2 = read_landmarks("C:/Users/immanueldiener/Desktop/Master/master_data\\additional_teeth_volumes\\", 10, "@1")

#Swap X and Y again
landmark_data2 = swap_xy(landmark_data2)

#all these images have reversed fill and object values. Tha highest value must always be the filler value!
image_data2 = (image_data2 .- 1) .* -1

images = cat(image_data, image_data2, dims=4)
lms = cat(landmark_data, landmark_data2, dims=2)

#put all the landmarks that are not yet on the surface of an object onto the landmark
lms_on_surface = landmark_to_surface(images, lms, 10)

#define training and testing sets
X_train = images[:,:,:,setdiff(1:243,86:125)]
X_test = images[:,:,:,86:125]
y_train = lms_on_surface[:, setdiff(1:243,86:125)]
y_test = lms_on_surface[:, 86:125]

#augment the training data for better learning and less overfitting
X_train, y_train = mirror_vol(X_train, y_train)
flip1, lm_flipped1 = flip_volume_front(X_train, y_train)
flip2, lm_flipped2 = flip_volume_side(X_train, y_train)
flip3, lm_flipped3 = flip_volume_front(flip1, lm_flipped1)
flip4, lm_flipped4 = flip_volume_side(flip3, lm_flipped3)

# we'll be using flip2 and flip4 and rotate them by 30/60 and 20/70 degrees respectively
flip2_rot1, lm_flip2_rot1 = rotate_volumes(flip2, lm_flipped2, 30)
flip2_rot2, lm_flip2_rot2 = rotate_volumes(flip2, lm_flipped2, 60)

flip4_rot1, lm_flip4_rot1 = rotate_volumes(flip4, lm_flipped4, 20)
flip4_rot2, lm_flip4_rot2 = rotate_volumes(flip4, lm_flipped4, 70)

# lets also rotate the original data (including the mirror images)
og_rot1, lm_og_rot1 = rotate_volumes(X_train, y_train, 50)
og_rot2, lm_og_rot2 = rotate_volumes(X_train, y_train, 80)

# jittering around of the object inside the volume. Jittering is landmark based and
# a 10 Voxel padding will be added after the minimum/maximum landmark coordinate in
# any dimension.

jit1, lm_jit1 = jitter_3D(flip2_rot1, lm_flip2_rot1, 10)
jit2, lm_jit2 = jitter_3D(flip4_rot2, lm_flip4_rot2, 10)
jit3, lm_jit3 = jitter_3D(og_rot2, lm_og_rot2, 10)

#lets also use aligned versions of the original dataset
aligned, lm_aligend, reco = align_principal(X_train, y_train, 192)
resized, lm_resized = resize_relevant(aligned, lm_aligend, 128)
aligned = nothing

jit4, lm_jit4 = jitter_3D(X_train, y_train, 10)

#generate all the depthmaps for input into the network
complete1 = depth_map_all_sides(flip2_rot1)
complete2 = depth_map_all_sides(flip2_rot2)
complete3 = depth_map_all_sides(flip4_rot1)
complete4 = depth_map_all_sides(flip4_rot2)
complete5 = depth_map_all_sides(og_rot1)
complete6 = depth_map_all_sides(og_rot2)
complete7 = depth_map_all_sides(jit1)
complete8 = depth_map_all_sides(jit2)
complete9 = depth_map_all_sides(jit3)
complete10 = depth_map_all_sides(resized)
complete11 = depth_map_all_sides(flip1)
complete12 = depth_map_all_sides(flip3)
complete13 = depth_map_all_sides(jit4)

X_train = cat(complete1, complete2, complete3, complete4, complete5, complete6, complete7, complete8,
    complete9, complete10, complete11, complete12, complete13, dims=4)
y_train = cat(lm_flip2_rot1, lm_flip2_rot2, lm_flip4_rot1, lm_flip4_rot2, lm_og_rot1, lm_og_rot2,
    lm_jit1, lm_jit2, lm_jit3, lm_resized, lm_flipped1, lm_flipped3, lm_jit4, dims=2)

# make depthmaps of the testing data
X_test = depth_map_all_sides(X_test)

# define the cost function
cost(x, y) = sum((model(x)-y).^2)|>gpu

#define the model
model = Flux.mapleaves(cu, AutoLandmarking.vgg19)

# define the trainingrate and optimiser
opt = Flux.ADAM(0.000015)

# redefinition of the dropout function to work with testmode!() since this
# functionality is not working at the moment. This step only needs to be done
# if you wish to perform uncertainty estimation using the dropout method (AutoLM.response_distribution())
# after training.
using Random
function Flux.dropout(x, p; dims = :)
    q = 1 - p
    y = rand!(similar(x, Flux._dropout_shape(x, dims)))
    y .= Flux._dropout_kernel.(y, p, q)
    x .* y
end

import Zygote
Zygote.@adjoint function Flux.dropout(x, p; dims = :)
   q = 1 - p
   y = rand!(similar(x, Flux._dropout_shape(x, dims)))
   y .= Flux._dropout_kernel.(y, p, q)
   return x .* y, Δ -> (Δ .* y, nothing)
end
Zygote.refresh()

function run_model(modell, X, y)
  train_data = Flux.mapleaves(cu, X)
  train_labels = Flux.mapleaves(cu, y)
  dataset = Flux.Data.DataLoader(train_data, train_labels, batchsize = 4, shuffle=true)
  Flux.train!(cost, params(model), dataset, opt)
  testmode!(model)
  cosima = cost_whole_data_2D(train_data, train_labels, cost)
  testmode!(model, false)
  return cosima
end

# create lists to store metrics
accs = []
costs = []

# run the training for 300 epochs. batches of 128 will be loaded onto the GPU,
# which will be further subdevided into minibatches of 4 (as defined in AutoML.run_model).
for i in 1:300
  costr = 0
  for j in 1:128:size(X_train, 4)
    if j+127 > size(X_train, 4)
      train_data = X_train[:,:,:,j:end]
      train_labels = y_train[:,j:end]
    else
      train_data = X_train[:,:,:,j:j+127]
      train_labels = y_train[:,j:j+127]
    end
    cosima = run_model(model, train_data, train_labels)
    costr = costr + cosima
    train_data = nothing
    train_labels = nothing
    # GC.gc()
  end
  push!(costs, costr)
  println("epoch ", i, " finished")
  println(costr)
  testmode!(model)
  acc2, max1, min1 = AutoLandmarking.avg_accuracy_per_point(model, gpu(X_test[:,:,:,:]), gpu(y_test[:,:]), 3)
  testmode!(model, false)
  println("median deviation per point on testing dataset: ", acc2, "with maximum", max1)
  println("median sum is ", sum(acc2))
  push!(accs, sum(acc2))
  if sum(acc2)[1]<7.5
    break
  end
end

#save model
using BSON
model_cpu = cpu(model)
BSON.@save "C:/Users/immanueldiener/Desktop/Master/master_data/model_demo.bson" model_cpu

#load model back in
BSON.@load "C:/Users/immanueldiener/Desktop/Master/master_data/model_demo.bson" model_cpu
model = gpu(model_cpu)

# ================== After training ===================

#solve issue with arrays inside vector
aks = []
for i in accs
  push!(aks, i[1])
end

# check what the best accuracy is
minimum(aks)

#plot the development of cost and accuracy
import Plots
Plots.plot(aks, legend=:topright, label="sum of deviations per point", color= :red, xlabel="epochs", ylabel = "acuracy")
plt = Plots.plot!(Plots.twinx(), costs, label="training loss", legend=:bottomright, ylabel = "loss")

# check for outliers (greatest distance from the procrustes mean).
# teh idea is to eliminate out all the predicted landmark sets that
# deviate a lot from the procrustes mean (reference), since these are
# most likely bad predictions. do get the procrustes mean, we can use training
# testing data together to get a more accureate mean. This is an ideal approch
# because otherwise we would not realize if all the predictions are biased in some
# way that is congruent over all predictions but still wrong.


module TestStatsModels
using DataFrames
using Base.Test
using Compat

# Tests for statsmodel.jl

# A dummy RegressionModel type
immutable DummyMod <: RegressionModel
    beta::Vector{Float64}
    x::Matrix
    y::Vector
end

## dumb fit method: just copy the x and y input over
StatsBase.fit(::Type{DummyMod}, x::Matrix, y::Vector) =
    DummyMod(collect(1:size(x, 2)), x, y)
StatsBase.model_response(mod::DummyMod) = mod.y
## dumb coeftable: just prints the "beta" values
StatsBase.coeftable(mod::DummyMod) =
    CoefTable(reshape(mod.beta, (size(mod.beta,1), 1)),
              ["'beta' value"],
              ["" for n in 1:size(mod.x,2)],
              0)

## Test fitting
d = DataFrame()
d[:y] = [1:4;]
d[:x1] = [5:8;]
d[:x2] = [9:12;]
d[:x3] = [13:16;]
d[:x4] = [17:20;]

f = @formula(y ~ x1 * x2)
m = fit(DummyMod, f, d)
@test model_response(m) == d[:y]

## test prediction method
## vanilla
StatsBase.predict(mod::DummyMod) = mod.x * mod.beta
@test predict(m) == [ ones(size(d,1)) d[:x1] d[:x2] d[:x1].*d[:x2] ] * collect(1:4)

## new data from matrix
StatsBase.predict(mod::DummyMod, newX::Matrix) = newX * mod.beta
mm = ModelMatrix(ModelFrame(f, d))
@test predict(m, mm.m) == mm.m * collect(1:4)

## new data from DataFrame (via ModelMatrix)
@test predict(m, d) == predict(m, mm.m)

d2 = deepcopy(d)
d2[3, :x1] = NA
@test length(predict(m, d2)) == 4

## test copying of names from Terms to CoefTable
ct = coeftable(m)
@test ct.rownms == ["(Intercept)", "x1", "x2", "x1 & x2"]

## show with coeftable defined
io = IOBuffer()
show(io, m)

## with categorical variables
d[:x1p] = PooledDataArray(d[:x1])
f2 = @formula(y ~ x1p)
m2 = fit(DummyMod, f2, d)

@test coeftable(m2).rownms == ["(Intercept)", "x1p: 6", "x1p: 7", "x1p: 8"]

## predict w/ new data missing levels
@test predict(m2, d[2:4, :]) == predict(m2)[2:4]

## predict w/ new data with _extra_ levels (throws an error)
d3 = deepcopy(d)
d3[1, :x1] = 0
d3[:x1p] = PooledDataArray(d3[:x1])
@test_throws ArgumentError predict(m2, d3)

## fit with contrasts specified
d[:x2p] = PooledDataArray(d[:x2])
f3 = @formula(y ~ x1p + x2p)
m3 = fit(DummyMod, f3, d)
fit(DummyMod, f3, d, contrasts = Dict(:x1p => EffectsCoding()))
fit(DummyMod, f3, d, contrasts = Dict(:x1p => EffectsCoding(),
                                      :x2p => DummyCoding()))
@test_throws Exception fit(DummyMod, f3, d, contrasts = Dict(:x1p => EffectsCoding(),
                                                             :x2p => 1))

## Another dummy model type to test fall-through show method
immutable DummyModTwo <: RegressionModel
    msg::Compat.UTF8String
end

StatsBase.fit(::Type{DummyModTwo}, ::Matrix, ::Vector) = DummyModTwo("hello!")
Base.show(io::IO, m::DummyModTwo) = println(io, m.msg)

m2 = fit(DummyModTwo, f, d)
show(io, m2)

end