module MetaGraphsNext
using LightGraphs
using JLD2

import Base:
    eltype,
    show,
    ==,
    Tuple,
    copy,
    length,
    size,
    issubset,
    zero,
    getindex,
    haskey,
    setindex!,
    delete!

import LightGraphs:
    AbstractGraph,
    edgetype,
    nv,
    ne,
    vertices,
    edges,
    is_directed,
    add_vertex!,
    add_edge!,
    rem_vertex!,
    rem_edge!,
    has_vertex,
    has_edge,
    inneighbors,
    outneighbors,
    weights,
    induced_subgraph,
    loadgraph,
    savegraph,
    AbstractGraphFormat,
    reverse

import LightGraphs.SimpleGraphs: SimpleGraph, SimpleDiGraph, fadj, badj

export MetaGraph, weighttype, defaultweight, weightfunction, MGFormat, DOTFormat, reverse

include("metagraph.jl")

function show(
    io::IO,
    g::MetaGraph{<:Any,Label,<:Any,VertexMeta,EdgeMeta},
) where {Label,VertexMeta,EdgeMeta}
    print(
        io,
        "Meta graph based on a $(g.graph) with vertices indexed by $Label(s), $VertexMeta(s) vertex metadata, $EdgeMeta(s) edge metadata, $(repr(g.gprops)) as graph metadata, and default weight $(g.defaultweight)",
    )
end

@inline fadj(g::MetaGraph, x...) = fadj(g.graph, x...)
@inline badj(g::MetaGraph, x...) = badj(g.graph, x...)

eltype(g::MetaGraph) = eltype(g.graph)
edgetype(g::MetaGraph) = edgetype(g.graph)
nv(g::MetaGraph) = nv(g.graph)
vertices(g::MetaGraph) = vertices(g.graph)

ne(g::MetaGraph) = ne(g.graph)
edges(g::MetaGraph) = edges(g.graph)

has_vertex(g::MetaGraph, x...) = has_vertex(g.graph, x...)
@inline has_edge(g::MetaGraph, x...) = has_edge(g.graph, x...)

inneighbors(g::MetaGraph, v::Integer) = inneighbors(g.graph, v)
outneighbors(g::MetaGraph, v::Integer) = fadj(g.graph, v)

issubset(g::T, h::T) where {T<:MetaGraph} = issubset(g.graph, h.graph)

"""
    add_edge!(g, u, v, val)

Add an edge `(u, v)` to MetaGraph `g` having value `val`.

Return true if the edge has been added, false otherwise.
"""
function add_edge!(g::MetaGraph, u::Integer, v::Integer, val)
    added = add_edge!(g.graph, u, v)
    if added
        metaindex = g.metaindex
        g.eprops[arrange(g, metaindex[u], metaindex[v], u, v)] = val
    end
    added
end

@inline function rem_edge!(g::MetaGraph, u::Integer, v::Integer)
    metaindex = g.metaindex
    removed = rem_edge!(g.graph, u, v)
    if removed
        delete!(g.eprops, arrange(g, metaindex[u], metaindex[v], u, v))
    end
    removed
end

"""
    add_vertex!(g, label, val)

Add a vertex to MetaGraph `g` with label `label` having value `val`.

Return true if the vertex has been added, false otherwise.
"""
function add_vertex!(g::MetaGraph, label, val)
    added = add_vertex!(g.graph)
    if added
        v = nv(g)
        g.vprops[label] = (v, val)
        g.metaindex[v] = label
    end
    added
end

function _rem_vertex!(g, label, v)
    vprops = g.vprops
    eprops = g.eprops
    metaindex = g.metaindex
    lastv = nv(g)
    for n in outneighbors(g, v)
        delete!(eprops, arrange(g, label, metaindex[n], v, n))
    end
    for n in inneighbors(g, v)
        delete!(eprops, arrange(g, metaindex[n], label, n, v))
    end
    removed = rem_vertex!(g.graph, v)
    if removed
        if v != lastv # ignore if we're removing the last vertex.
            lastl = metaindex[lastv]
            _, lastvprops = vprops[lastl]
            vprops[lastl] = v, lastvprops
            metaindex[v] = lastl
        end
        delete!(vprops, label)
        delete!(metaindex, lastv)
    end
    removed
end

function rem_vertex!(g::MetaGraph, v::Integer)
    exists = has_vertex(g, v)
    if exists
        _rem_vertex!(g, g.metaindex[v], v)
    else
        false
    end
end

struct MetaWeights{InnerMetaGraph,U<:Real} <: AbstractMatrix{U}
    meta_graph::InnerMetaGraph
end

show(io::IO, x::MetaWeights) = print(io, "metaweights")
show(io::IO, z::MIME"text/plain", x::MetaWeights) = show(io, x)

MetaWeights(g::MetaGraph) = MetaWeights{typeof(g),weighttype(g)}(g)

function getindex(w::MetaWeights, u::Int, v::Int)
    g = w.meta_graph
    metaindex = g.metaindex
    if has_edge(g, u, v)
        g.weightfunction(g[arrange(g, metaindex[u], metaindex[v], u, v)...])
    else
        g.defaultweight
    end
end

function size(d::MetaWeights)
    vertices = nv(d.meta_graph)
    (vertices, vertices)
end

weights(g::MetaGraph) = MetaWeights(g)

getindex(g::MetaGraph) = g.gprops
function getindex(g::MetaGraph, label)
    _, val = g.vprops[label]
    val
end
getindex(g::MetaGraph, label_1, label_2) = g.eprops[arrange(g, label_1, label_2)]

"""
    haskey(g, :label)

Determine whether a graph `g` contains the vertex `:label`.

```jldoctest; setup = :(using MetaGraphsNext; using LightGraphs: Graph)
julia> colors = MetaGraph(Graph(), VertexMeta = String, EdgeMeta = Symbol, gprops = "special");

julia> colors[:red] = "warm";

julia> haskey(colors, :red)
true

julia> haskey(colors, :blue)
false
```
"""
haskey(g::MetaGraph, label) = haskey(g.vprops, label)

"""
    haskey(g, :v1, :v2)

Determine whether a graph `g` contains an edge from `:v1` to `:v2`. The order of `:v1` and `:v2`
only matters if `g` is a digraph.

```jldoctest; setup = :(using MetaGraphsNext; using LightGraphs: Graph)
julia> colors = MetaGraph(Graph(), VertexMeta = String, EdgeMeta = Symbol, gprops = "special");

julia> colors[:red] = "warm"; colors[:blue] = "cool"; colors[:red, :blue] = :purple
:purple

julia> haskey(colors, :red, :blue) && haskey(colors, :blue, :red)
true

julia> haskey(colors, :red, :yellow)
false
```
"""
haskey(g::MetaGraph, label_1, label_2) = haskey(g, label_1) && haskey(g, label_2) && haskey(g.eprops, arrange(g, label_1, label_2))

function setindex!(g::MetaGraph, val, label)
    vprops = g.vprops
    v = if haskey(vprops, label)
        (v, _) = vprops[label]
        v
    else
        add_vertex!(g.graph)
        v = nv(g)
        g.metaindex[v] = label
        v
    end
    vprops[label] = (v, val)
    return nothing
end

function setindex!(g::MetaGraph, val, label_1, label_2)
    vprops = g.vprops
    u, _ = vprops[label_1]
    v, _ = vprops[label_2]
    add_edge!(g.graph, u, v)
    g.eprops[arrange(g, label_1, label_2, u, v)] = val
    return nothing
end

function delete!(g::MetaGraph, label)
    if haskey(g, label)
        v, _ = g.vprops[label]
        _rem_vertex!(g, label, v)
    end
    return nothing
end
function delete!(g::MetaGraph, label_1, label_2)
    vprops = g.vprops
    u, _ = vprops[label_1]
    v, _ = vprops[label_2]
    rem_edge!(g.graph, u, v)
    delete!(g.eprops, arrange(g, label_1, label_2, u, v))
    return nothing
end

"""
    weighttype(g)

Return the weight type for metagraph `g`.

```jldoctest
julia> using MetaGraphsNext

julia> using LightGraphs: Graph

julia> weighttype(MetaGraph(Graph(), defaultweight = 1.0))
Float64
```
"""
weighttype(g::MetaGraph{<:Any,<:Any,<:Any,<:Any,<:Any,<:Any,<:Any,Weight}) where {Weight} =
    Weight

"""
    weightfunction(g)

Return the weight function for metagraph `g`.

```jldoctest
julia> using MetaGraphsNext

julia> using LightGraphs: Graph

julia> weightfunction(MetaGraph(Graph(), weightfunction = identity))(0)
0
```
"""
weightfunction(g::MetaGraph) = g.weightfunction

"""
    defaultweight(g)

Return the default weight for metagraph `g`.

```jldoctest
julia> using MetaGraphsNext

julia> using LightGraphs: Graph

julia> defaultweight(MetaGraph(Graph(), defaultweight = 2))
2
```
"""
defaultweight(g::MetaGraph) = g.defaultweight

function _copy_props!(oldg::T, newg::T, vmap) where {T<:MetaGraph}
    for (newv, oldv) in enumerate(vmap)
        oldl = oldg.metaindex[oldv]
        _, meta = oldg.vprops[oldl]
        newg.metaindex[newv] = oldl
        newg.vprops[oldl] = (newv, meta)
    end
    for newe in edges(newg.graph)
        metaindex = newg.metaindex
        u, v = Tuple(newe)
        label_1 = metaindex[u]
        label_2 = metaindex[v]
        newg.eprops[arrange(newg, label_1, label_2, u, v)] =
            oldg.eprops[arrange(oldg, label_1, label_2)]
    end
    return nothing
end

function induced_subgraph(
    g::T,
    v::AbstractVector{U},
) where {T<:MetaGraph} where {U<:Integer}
    inducedgraph, vmap = induced_subgraph(g.graph, v)
    newg = MetaGraph(
        inducedgraph,
        empty(g.vprops),
        empty(g.eprops),
        g.gprops,
        g.weightfunction,
        g.defaultweight,
        empty(g.metaindex),
    )
    _copy_props!(g, newg, vmap)
    return newg, vmap
end

# TODO - would be nice to be able to apply a function to properties. Not sure
# how this might work, but if the property is a vector, a generic way to append to
# it would be a good thing.

==(x::MetaGraph, y::MetaGraph) = x.graph == y.graph

copy(g::T) where {T<:MetaGraph} = deepcopy(g)

"""
    code_for(meta::MetaGraph, vertex_label)

Find the code associated with a `vertex_label`. This can be useful to pass
to methods inherited from `LightGraphs`. Note, however, that vertex codes could be
reassigned after vertex deletion.

```jldoctest
julia> using MetaGraphsNext

julia> using LightGraphs: Graph

julia> meta = MetaGraph(Graph());

julia> meta[:a] = nothing

julia> code_for(meta, :a)
1
```
"""
function code_for(meta::MetaGraph, vertex_label)
    code, _ = meta.vprops[vertex_label]
    code
end
export code_for

"""
    label_for(meta::MetaGraph, vertex_code)

Find the label associated with a `vertex_code`. This can be useful to interpret the results of methods
inherited from `LightGraphs`. Note, however, that vertex codes could be reassigned after vertex deletion.

```jldoctest
julia> using MetaGraphsNext

julia> using LightGraphs: Graph

julia> meta = MetaGraph(Graph());

julia> meta[:a] = nothing

julia> label_for(meta, 1)
:a
```
"""
function label_for(meta::MetaGraph, vertex_code)
    meta.metaindex[vertex_code]
end
export label_for

include("metadigraph.jl")
include("overrides.jl")
include("persistence.jl")

end # module


#=
FastAI.jl:

Author: Peter Wolf (opus111@gmail.com)
=#

module FastAI

using Random
using StatsBase
using Statistics
using Flux
using Flux: update!
using Flux.Data
using Zygote
using Infiltrator
using Base: length, getindex
using Random: randperm
using DocStringExtensions

export AbstractLearner
export AbstractCallback
export IterableDataset
export MapDataset

export DataBunch
export train
export valid

export DummyCallback
export ProgressCallback

export Learner
export model
export data_bunch
export loss
export loss!
export opt
export opt!
export fit!
export add_cb!

export Recorder

@template (FUNCTIONS, METHODS) =
    """
    $(TYPEDSIGNATURES)
    $(DOCSTRING)
    """

@template (TYPES) =
    """
    $(TYPEDEF)
    $(DOCSTRING)
    """

include("dataset.jl")
include("databunch.jl")
include("learner.jl")
include("callback.jl")
include("recorder.jl")

end


using jInv.Mesh;
using MAT;
using jInv
using jInv.InverseSolve
using ShapeReconstructionPaLS
using ParamLevelSet;
using ShapeReconstructionPaLS.Utils;
using ShapeReconstructionPaLS.PointCloud;
using DelimitedFiles
using LinearAlgebra
using Statistics
using Distributed
using SparseArrays
using Random
ENV["MPLBACKEND"] = "Qt4Agg"

function ind2subv(shape, indices)
    CI = CartesianIndices(shape)
    return getindex.(Ref(CI), indices)
end

plotting = false;

if plotting
	using PyPlot;
	using jInvVis;
	close("all");
end

invertDip = true;
invertVis = false;

locateRBFwithGrads = true;

################################################
### Reading the model and pad ##################
################################################
#model = "skull_160"; ## skull
 model = "hand_160";
file = matopen(string(pwd(),"/models/",model,".mat"));

A = read(file, "B");
close(file);
A = convert(Array{Float32,3},A[1:4:end,1:4:end,1:4:end]);
n_inner = collect(size(A));
pad = ceil.(Int64,0.1*n_inner[1]);
n = n_inner .+ 2*pad;

n_tup = tuple(n...)
volume = prod(n);
n1 = n_tup[1];
n2 = n_tup[2];
n3 = n_tup[3];

m = zeros(Float64,n_tup);
mask = zeros(Int8,n_tup);
m[pad+1:end-pad,pad+1:end-pad,pad+1:end-pad] = A; A = 0;
mask[pad+1:end-pad,pad+1:end-pad,pad+1:end-pad] .= 1;
if plotting
	plotModel(m);
	savefig("TrueModel.png");
end
origm = copy(m);
Mesh = getRegularMesh([0.0 5.0 0.0 5.0 0.0 5.0],n);

noiseAnglesDeg = 0.0;
#original was mean(2*Mesh.h);
noiseTrans = mean(0.0*Mesh.h);
noiseAngles = deg2rad(noiseAnglesDeg);
#########################################################################################################
### Dipping Data ########################################################################################
#########################################################################################################
samplingBinning = 1;

Mesh = getRegularMesh([0.0 5.0 0.0 5.0 0.0 5.0],n);
pad = div(pad,samplingBinning);

method = RBF10BasedSimple; methodName="RBF10BasedSimple";

println("Inv mesh:")
println(Mesh.n)

n_tup = tuple(Mesh.n...);
nWorkers = nworkers();
println("nworkers=",nWorkers)
misfun = PCFun; ## least squares

	dipDataFilename = string("pointCloudData",model,"_dataRes",n,1);
	b = [0.0 0 0] #; 0.0 0 0 ]
	theta_phi_dip = [0.0 0.0] #; 0.0 0.0];
	noiseTrans = 0.0*prod(Mesh.h);
	noiseAngles = deg2rad(0.0);
	npc = 1;
	prepareSyntheticPointCloudData(m,Mesh,npc,theta_phi_dip,b,noiseTrans,noiseAngles,dipDataFilename);

	#########################################################################################################
	### Read the dipping data ###############################################################################
	#########################################################################################################
	n,domain,Data,P,Normals,Margin,b,theta_phi_dip,npc = readPointCloudDataFile(dipDataFilename);
	println("initial obtained translation:",b);
        P = Array{Array{Int64,1}}(P);
	writedlm(string("simPCpart2",".txt"),convert(Array{Int64},P[1]));
#       writedlm(string("simPCpart1",".txt"),convert(Array{Int64},P[2]));
#	quit()
	### Set up the dip inversion
	println("type of p:",typeof(P));
	pForDip = getPointCloudParam(Mesh,P,Normals,Margin,theta_phi_dip,b,samplingBinning,nWorkers,method);
	ndips = npc;
	if(locateRBFwithGrads)
		pForDip_MATFree = getPointCloudParam(Mesh,P,Normals,Margin,theta_phi_dip,b,samplingBinning,nWorkers,MATFree);
	end

	### Create Dip param if we use the GRAD of MATFREE to locate new RBFs
	dobsDirect = divideDataToWorkersData(nWorkers,Data);
	println("size of dobs",size(dobsDirect))
	#nnz = filter(x -> x != 0,dobsDirect[1]);
	#println("length of dobsdirect: ",length(dobsDirect[1]),"_",length(nnz));
	Wd_Dip = Array{Array{Float32}}(undef,nWorkers);
	for k=1:nWorkers
		Wd_Dip[k] = ones(Float32,size(dobsDirect[k]));
	end
	println("Wd dip:",size(Wd_Dip))

n_Moves_all = npc;
println("npc = ",npc);

isSimple = 0;
isRBF10 = 0;

###############################################################################################
############## USE RBF dictionary #############################################################
###############################################################################################

	nRBF = 20;

	isRBF10 = (method == RBF10BasedSimple1 || method==RBF10BasedSimple2 || method==RBF10Based);
	isSimple = (method==RBFBasedSimple1 || method==RBF10BasedSimple1 || method==RBFBasedSimple2 || method==RBF10BasedSimple2);

	if isRBF10
		numParamOfRBF = 10;
	else
		numParamOfRBF = 5;
	end
	isSimple = true;
	n_m_simple = numParamOfRBF*nRBF;
	nAll = isSimple ? numParamOfRBF*nRBF : numParamOfRBF*nRBF + n_Moves_all*5;

	modfun = identityMod;

	if isSimple
		modfunForPlotting = (m)->ParamLevelSetModelFunc(Mesh,m;computeJacobian = 0,
			sigma = getDefaultHeavySide() ,bf = 1, numParamOfRBF = numParamOfRBF);
	else
		modfunForPlotting = (m)->ParamLevelSetModelFunc(Mesh,m[1:(numParamOfRBF*nRBF)];computeJacobian = 0,
			sigma = getDefaultHeavySide() ,bf = 1, numParamOfRBF = numParamOfRBF);
	end

	Iact = 1.0;
	IactPlot = 1.0;
	sback = zero(Float32);
	a = -100000000.0;
	b = 10000000.0;
	boundsHigh = zeros(Float32,nAll) .+ b;
	boundsLow = zeros(Float32,nAll) .+ a;
	# m0 - initial guess. mref - reference for regularization.
	mref = zeros(Float32,nAll);
	m0   = zeros(Float64,n_m_simple);
	mCenters = zeros(Float32,n_tup);
	mref[1:numParamOfRBF:n_m_simple] .= ParamLevelSet.centerHeavySide*2.0*rand(nRBF);
	if isRBF10
		# here we initialize B as an identity matrix
		mref[2:10:n_m_simple] .= 0.5;
		mref[5:10:n_m_simple] .= 0.5;
		mref[7:10:n_m_simple] .= 0.5;
		boundsLow[2:10:n_m_simple] .= 0.05;
		boundsLow[5:10:n_m_simple] .= 0.05;
		boundsLow[7:10:n_m_simple] .= 0.05;
	else
		mref[2:5:n_m_simple] .= 1.0;
		boundsLow[2:5:n_m_simple] .= 0.05;
	end
	mref[(numParamOfRBF-2):numParamOfRBF:n_m_simple] .= (Mesh.domain[1] + Mesh.domain[2])/2;
	mref[(numParamOfRBF-1):numParamOfRBF:n_m_simple] .= (Mesh.domain[3] + Mesh.domain[4])/2;
	mref[numParamOfRBF:numParamOfRBF:n_m_simple] .= (Mesh.domain[5] + Mesh.domain[6])/2;

	m0[:] = mref[1:n_m_simple];
	if !isSimple
		m0 = wrapRBFparamAndRotationsTranslations(m0,theta_phi_dip,zeros(size(theta_phi_dip,1),3));
	end

	mref[:] = m0;
	m0[(numParamOfRBF-2):numParamOfRBF:n_m_simple] .+= 0.2*(Mesh.domain[2] - Mesh.domain[1])*randn(nRBF);
	m0[(numParamOfRBF-1):numParamOfRBF:n_m_simple] .+= 0.2*(Mesh.domain[4] - Mesh.domain[3])*randn(nRBF);
	m0[numParamOfRBF:numParamOfRBF:n_m_simple]     .+= 0.2*(Mesh.domain[6] - Mesh.domain[5])*randn(nRBF);

	alpha = 5e-0;

	II = (sparse(1.0I, nAll,nAll));
	regfun = (m, mref, M)->TikhonovReg(m,mref,M,II);
	HesPrec = getEmptyRegularizationPreconditioner();

	if isRBF10
		spd_reg = (m,mref,M) -> RBF_SPD_regularization(m,mref,nRBF);
		regfun 	 		= [regfun;spd_reg];
		alpha 		 	= [alpha;1e-1*alpha]; #was 1e-4 for cube\prism !!!!
		mref 			= [mref zeros(size(mref))];
	end

pMisRFs = getMisfitParam(pForDip, Wd_Dip, dobsDirect, misfun, Iact,sback);

locateRBFwithGrads = true;
if(locateRBFwithGrads)
		mask=0;
		Iact_free = 1.0;
		#### USE A BOUND MODEL or bounds
		a_free = 0.0;
		b_free= 1.0;

		pMis_Free = getMisfitParam(pForDip_MATFree, Wd_Dip, dobsDirect, misfun, Iact_free,sback);
end

cgit  = 50; # not used.
maxit = 10;
pcgTol = 1e-2; #not used

@everywhere function updateAngles(pFor::Array{RemoteChannel},theta_phi)
@sync begin
	@async begin
		for k=1:length(pFor)
			pFor[k] = remotecall_fetch(updateAngles,pFor[k].where,pFor[k],theta_phi);
		end
	end
end
return pFor;
end

@everywhere function updateAngles(pFor::RemoteChannel,theta_phi)
	pForMATFREE  = take!(pFor)
	worker_theta_phi = theta_phi[pForMATFREE.workerSubIdxs,:]
	pForMATFREE.theta_phi_rad=copy(worker_theta_phi);
	put!(pFor,pForMATFREE)
	return pFor;
end

function myDump(mc::Vector,Dc,iter,pInv,pMis,method="",ndips=0,noiseAngle=0,noiseTranslation=0,useVisual=false,plotFun=modfunForPlotting)
		ntup = tuple((pInv.MInv.n)...);
		fullMc = plotFun(mc)[1];
		fullMc = reshape(IactPlot*fullMc[:],ntup);
		writedlm(string("twoHand1PCParams_iter_",iter,".dat"),convert(Array{Float64},mc));
		writedlm(string("twoHand1PCVolume_iter_",iter,"_",ntup,".dat"),convert(Array{Float16},fullMc));
end

pInv = getInverseParam(Mesh,modfun,regfun,alpha,mref,boundsLow,boundsHigh,
                     maxStep = 2e-1,pcgMaxIter=cgit,pcgTol=pcgTol,
					 minUpdate = 0.01, maxIter = maxit,HesPrec=HesPrec);

#### Projected Gauss Newton
mc = m0;
pInv.maxIter = 100;
Dc = nothing;
mc,Dc,flag,his = projGN(mc,pInv,pMisRFs,solveGN=projGNexplicit);

pInv.maxIter = 10;
myDump(mc,0,1,pInv,0,"",0,noiseAnglesDeg,noiseTrans,invertVis);

nRBF=nRBF;
times=0;

if !(method == MATBased || method == MATFree)
	outerIter = 300;
	grad_u = nothing;
	new_RBF_location = [];
	for iterNum=2:outerIter
		println("Inversion using ", nRBF," basis functions")
		new_nRBF =10;
		global times = times + 1;

		# Dc - the data using our method.
		global Dc = Dc;
		u = modfunForPlotting(mc)[1];
		#I = find((u.>0.3) .& (u.<0.6));
		#Ii = P[randperm(length(P))[1:new_nRBF]];
		#Ii = findall(x -> x >0.3 && x<0.7,u);
		#Lets compute the gradient of the misfit function:
		locateRBFwithGrads = true;
		if(locateRBFwithGrads)
			#(m11,theta_phi1,b1) 	= splitRBFparamAndRotationsTranslations(mc,nRBF,ndips,numParamOfRBF)
			#updateAngles(pForDip_MATFree,theta_phi1);
			pMis_Free = getMisfitParam(pForDip_MATFree, Wd_Dip, dobsDirect, misfun, Iact_free,sback);
			computeMisfit(u,pMis_Free);
			gradMis = computeGradMisfit(u,Dc,pMis_Free)
			gradMis = gradMis.^2;
			println("size gradmis:",size(gradMis))
			amax = argmax(gradMis); vmax = maximum(gradMis);
			sp = sortperm(gradMis,rev = true);
			Ii = sp[1:new_nRBF];
			writedlm(string("gradMisfit",".txt"),convert(Array{Float64},gradMis[:]))
		end

		if length(Ii) > new_nRBF
		Ii = Ii[randperm(length(Ii))[1:new_nRBF]];
		end

   println("Ii size:",size(Ii))
    while(Ii[1] in new_RBF_location )#|| Ii.-1 in new_RBF_location || Ii.+1 in new_RBF_location)
      #Ii = findall(x -> x .>0.3 && x.<0.7,u);
	  #Ii = P[randperm(length(P))[1:new_nRBF]];
	  #Ii = Ii[randperm(length(Ii))[1:new_nRBF]];
       tmp = sp[1:outerIter*new_nRBF];
		Ii = tmp[randperm(outerIter*new_nRBF)[1:new_nRBF]];

    end
		global new_RBF_location = append!(new_RBF_location,Ii);
		indices = getindex.(ind2subv(n_tup,Ii),[1 2 3 ]);
		I1 = indices[:,1]; I2= indices[:,2]; I3 = indices[:,3];

		X1 = I1*Mesh.h[1] .- 0.5*Mesh.h[1]; .+ Mesh.domain[1];
		X2 = I2*Mesh.h[2] .- 0.5*Mesh.h[2]; .+ Mesh.domain[3];
		X3 = I3*Mesh.h[3] .- 0.5*Mesh.h[3]; .+ Mesh.domain[5];

		println("X1:",X1);println("X2:",X2);println("X3:",X3);

		tmp  = mc;
		n_mc_old = numParamOfRBF*nRBF;
		n_mc_new = numParamOfRBF*nRBF + length(Ii)*numParamOfRBF;
		mc_new = zeros(length(mc) + length(Ii)*numParamOfRBF);
		if isRBF10
			mref = pInv.mref[:,1];
		else
			mref = pInv.mref;
		end
		mref_new = zeros(Float32,length(mc) + length(Ii)*numParamOfRBF);

		mc_new[1:n_mc_old] = mc[1:n_mc_old];

		mc_new[(n_mc_new+1):end] = mc[(n_mc_old+1):end];
		mref_new[(n_mc_new+1):end] = mref[(n_mc_old+1):end];

		mc_new[(n_mc_old+1):numParamOfRBF:n_mc_new] .= 0.0; #ParamLevelSet.centerHeavySide - ParamLevelSet.deltaHeavySide;
		mc_new[(n_mc_old+numParamOfRBF-2):numParamOfRBF:n_mc_new] .= X1;
		mc_new[(n_mc_old+numParamOfRBF-1):numParamOfRBF:n_mc_new] .= X2;
		mc_new[(n_mc_old+numParamOfRBF):numParamOfRBF:n_mc_new] .= X3;

		pInv.boundsHigh = zeros(Float32,length(mc_new)) .+ b;
		boundsLow = zeros(Float32,length(mc_new)) .+ a;
		boundsLow[1:n_mc_old] = pInv.boundsLow[1:n_mc_old];
		if isRBF10
			mc_new[(n_mc_old+2):10:n_mc_new] .= 3.0;
			mc_new[(n_mc_old+5):10:n_mc_new] .= 3.0;
			mc_new[(n_mc_old+7):10:n_mc_new] .= 3.0;
			boundsLow[(n_mc_old+2):10:n_mc_new] .= 0.05;
			boundsLow[(n_mc_old+5):10:n_mc_new] .= 0.05;
			boundsLow[(n_mc_old+7):10:n_mc_new] .= 0.05;
		else
			mc_new[(n_mc_old+2):5:n_mc_new] .= 1.0;
			boundsLow[(n_mc_old+2):5:n_mc_new] .= 0.05;
		end
		pInv.boundsLow = boundsLow;
		mref_new[1:n_mc_new] = mc_new[1:n_mc_new];

		global mc = mc_new;
		#mref_new[end-9:end] .= 0.0;
		#II = speye(Float32,length(mc_new));
		len = Int64(length(mc_new));
		IIs = sparse(1.0I,len,len);
		pInv.regularizer = (m, mref, M)->TikhonovReg(m,mref,M,IIs);
		global nRBF += new_nRBF;
		if isRBF10
			spd_reg = (m,mref,M) -> RBF_SPD_regularization(m,mref,nRBF);
			pInv.regularizer  = [pInv.regularizer;spd_reg];
			pInv.mref 			  = [mref_new zeros(length(mref_new))];
			#pInv.alpha .*= 0.8;
			pInv.alpha[1] *= 0.8;
			pInv.alpha[2] *= 0.8;
		else
			pInv.mref = mref_new;
			pInv.alpha *= 0.8;
		end
		#mc[end-9:end] .= 0.0;
		mc,Dc,flag,his = projGN(mc,pInv,pMisRFs,solveGN=projGNexplicit);
		myDump(mc,0,iterNum,pInv,0,methodName,0,noiseAnglesDeg,noiseTrans,invertVis);

		if iterNum==outerIter
			pInv.alpha *= 0.1;
			pInv.maxIter = 200;
			mc, = projGN(mc,pInv,pMisRFs,solveGN=projGNexplicit);
			myDump(mc,0,iterNum+1,pInv,0,methodName,0,noiseAnglesDeg,noiseTrans,invertVis);
		end

	end

		meshForPlotting = getRegularMesh(Mesh.domain,Mesh.n*1);
		ntup = tuple((meshForPlotting.n)...);
		if isSimple
			modfunForPlotting = (m)->ParamLevelSetModelFunc(meshForPlotting,m;computeJacobian = 0,
				sigma = getDefaultHeavySide() ,bf = 1, numParamOfRBF = numParamOfRBF);
		else
			modfunForPlotting = (m)->ParamLevelSetModelFunc(meshForPlotting,m[1:(numParamOfRBF*nRBF)];computeJacobian = 0,
				sigma = getDefaultHeavySide() ,bf = 1, numParamOfRBF = numParamOfRBF);
		end
		fullMc = modfunForPlotting(mc)[1];
		writedlm(string("oneHand160_noise0_Volume_final_",ntup,"_from_",tuple((Mesh.n)...),".dat"),convert(Array{Float16},fullMc));
		#m = convert(Array{Float16},m);
		#writedlm(string("Fandisk200Direct_Volume_orig.dat"),m);
		#writedlm(string("Fandisk200Direct_RBF_centers.dat"),mCenters);

end


struct RPCond
    density
    pressure
    velocity
    sound_speed
    energy
end


isinteractive() && include("common.jl")

####
#### test directions mechanism
####

@testset "directions" begin
    directions = Directions(0b110101)
    is_forwards = Vector{Bool}()
    for i in 1:6
        is_forward, directions = next_direction(directions)
        push!(is_forwards, is_forward)
    end
    @test collect(is_forwards) == [true, false, true, false, true, true]
    @test rand(RNG, Directions).flags isa UInt32
end

####
#### dummy trajectory for unit testing
####

"""
Trajectory type that is easy to inspect and reason about, for unit testing.

The field `visited` keeps track of visited nodes, can be reset with `empty!`.
"""
struct DummyTrajectory{L,T,D}
    "Positions that trigger turning."
    turning::T
    "Positions that trigger/mimic divergence."
    divergent::D
    "Log density."
    ℓ::L
    "Visited positions."
    visited::Vector{Int}
end

function DummyTrajectory(ℓ; turning = Set(Int[]), divergent = Set(Int[]))
    DummyTrajectory(turning, divergent, ℓ, Int[])
end

Base.empty!(trajectory) = empty(trajectory.visited)

move(::DummyTrajectory, z, is_forward) = z + (is_forward ? 1 : -1)

function is_turning(::DummyTrajectory, τ)
    turn_flag, positions = τ
    @test length(positions) > 1 # not called on a leaf
    turn_flag
end

function combine_turn_statistics(::DummyTrajectory, τ₁, τ₂)
    turn_flag1, positions1 = τ₁
    turn_flag2, positions2 = τ₂
    @test last(positions1) + 1 == first(positions2) # adjacency and order
    (turn_flag1 && turn_flag2, first(positions1):last(positions2))
end

function combine_visited_statistics(::DummyTrajectory, v₁, v₂)
    a1, s1 = v₁
    a2, s2 = v₂
    (a1 + a2, s1 + s2)
end

# copyied from StatsFuns.jl
log1mexp(x::Real) = x < log(0.5) ? log1p(-exp(x)) : log(-expm1(x))

function combine_proposals(_, ::DummyTrajectory, zeta1, zeta2, logprob2, is_forward)
    lp2 = logprob2 > 0 ? 0.0 : logprob2
    lp1 = logprob2 > 0 ? oftype(lp2, -Inf) : log1mexp(lp2)
    if !is_forward
        # exchange so that we can test for adjacency, and join as a UnitRange
        zeta1, zeta2 = zeta2, zeta1
        lp1, lp2 = lp2, lp1
    end
    z1, p1 = zeta1
    z2, p2 = zeta2
    @test last(z1) + 1 == first(z2) # adjacency and order
    (first(z1):last(z2), vcat(p1 .+ lp1, p2 .+ lp2))
end

function calculate_logprob2(::DummyTrajectory, is_doubling, ω₁, ω₂, ω)
    biased_progressive_logprob2(is_doubling, ω₁, ω₂, ω)
end

function leaf(trajectory::DummyTrajectory, z, is_initial)
    @unpack turning, divergent, ℓ, visited = trajectory
    d = z ∈ divergent
    is_initial && @argcheck !d                              # don't start with divergent
    Δ = ℓ(z)
    v = is_initial ? (0.0, 0) : (min(exp(Δ), 1), 1)
    !is_initial && push!(visited, z)                        # save position
    if d
        nothing, v
    else
        (((z:z, [0.0]),          # ζ = nodes, log prob. within tree
          Δ,                     # ω = Δ for leaf
          (z ∈ turning, z:z)),   # τ
         v)
    end
end

"A log density for testing."
testℓ(z) = -abs2(z - 3) * 0.1

"Total acceptance rate of `ℓ` over `z`"
testA(ℓ, z) = sum(min.(exp.(ℓ.(z)), 1))

"sum of acceptance rates for trajectory."
testA(trajectory::DummyTrajectory) = testA(trajectory.ℓ, trajectory.visited)

@testset "dummy adjacent tree full" begin
    trajectory = DummyTrajectory(testℓ)
    (ζ, ω, τ, z′, i′), v = adjacent_tree(nothing, trajectory, 0, 0, 2, true)
    @test first(ζ) == 1:4
    @test sum(exp, last(ζ)) ≈ 1
    @test trajectory.visited == 1:4
    @test !is_turning(trajectory, τ)
    @test v[1] ≈ testA(trajectory)
    @test v[2] == 4
    @test z′ = i′ == 4
end

@testset "dummy adjacent tree turning" begin
    trajectory = DummyTrajectory(testℓ; turning = 5:7)
    t, v = adjacent_tree(nothing, trajectory, 0, 0, 3, true)
    @test trajectory.visited == 1:6 # [5,6] is turning
    @test t == InvalidTree(5, 6)
    @test v[1] == testA(trajectory)
    @test v[2] == 6
end

@testset "dummy adjacent tree divergent" begin
    trajectory = DummyTrajectory(testℓ; divergent = 5:7)
    t, v = adjacent_tree(nothing, trajectory, 0, 0, 3, true)
    @test trajectory.visited == 1:5 # 5 is divergent
    @test t == InvalidTree(5)
    @test v[1] ≈ testA(testℓ, 1:5)
    @test v[2] == 5
end

@testset "dummy adjacent tree full backward" begin
    trajectory = DummyTrajectory(testℓ)
    (ζ, ω, τ, z′, i′), v = adjacent_tree(nothing, trajectory, 0, 0, 3, false)
    @test first(ζ) == -8:-1
    @test sum(exp, last(ζ)) ≈ 1
    @test trajectory.visited == -(1:8)
    @test !is_turning(trajectory, τ)
    @test v[1] ≈ testA(testℓ, -(1:8))
    @test v[2] == 8
    @test z′ == i′ == -8
end

@testset "dummy sampled tree" begin
    trajectory = DummyTrajectory(testℓ)
    ζ, v, termination, depth = sample_trajectory(nothing, trajectory, 0, 3, Directions(0b101))
    @test trajectory.visited == [1, -1, -2, 2, 3, 4, 5]
    @test first(ζ) == -2:5
    @test sum(exp, last(ζ)) ≈ 1
    @test termination == DynamicHMC.REACHED_MAX_DEPTH
    @test v[1] ≈ testA(trajectory)
    @test v[2] == 7             # initial node does not participate in acceptance rate
end

####
#### Detailed balance tests
####

"An accumulator for log probabilities associated with a position on a trajectory."
empty_accumulator() = Dict{Int,Float64}()

"Add log probabilities `πs` at positions `zs` into `accumulator`."
function add_log_probabilities!(accumulator, zs, πs)
    for (z, π) in zip(zs, πs)
        accumulator[z] = haskey(accumulator, z) ? logaddexp(accumulator[z], π) : π
    end
    accumulator
end

"Normalize an accumulator by depth."
function normalize_accumulator(accumulator, depth)
    D = log(0.5) * depth
    Dict((k => v + D) for (k, v) in pairs(accumulator))
end

"""
An accumulator with the probability of visiting nodes for all trees with `depth`, strarting
from `z`, on `trajectory`.
"""
function visited_log_probabilities(trajectory, z, depth)
    accumulator = empty_accumulator()
    for flags in 0:(2^depth - 1)
        ζ = first(sample_trajectory(nothing, trajectory, z, depth, Directions(UInt32(flags))))
        add_log_probabilities!(accumulator, ζ...)
    end
    normalize_accumulator(accumulator, depth)
end

"""
The probability of visiting node `z′` for all trees with `depth`, strarting from `z`, on
`trajectory`.
"""
function transition_log_probability(trajectory, z, z′, depth)
    p = -Inf
    for flags in 0:(2^depth - 1)
        zs, πs = first(sample_trajectory(nothing, trajectory, z, depth,
                                         Directions(UInt32(flags))))
        ix = findfirst(isequal(z′), zs)
        if ix ≠ nothing
            p = logaddexp(p, πs[ix])
        end
    end
    p + depth * log(0.5)
end

@testset "transition calculations consistency check" begin
    trajectory = DummyTrajectory(testℓ)
    depth = 5
    z = 9
    for (z′, π) in pairs(visited_log_probabilities(trajectory, z, depth))
        @test π ≈ transition_log_probability(trajectory, z, z′, depth)
    end
end

"""
$(SIGNATURES)

For all transitions from `z`, test the detailed balance condition, ie

``ℙ(z) ℙ(j ∣ z) = ℙ(j) ℙ(z ∣ j)``

where ``ℙ(z) = exp(ℓ(z))`` and the transition probabilities ``ℙ(⋅∣⋅)`` are
calculated using `visited_log_probabilities` and `transition_log_probability`.

(We use logs for more accurate calculation.)
"""
function test_detailed_balance(trajectory, z, depth; atol = √eps())
    @unpack ℓ = trajectory
    ℓz = ℓ(z)
    for (z′, π) in pairs(visited_log_probabilities(trajectory, z, depth))
        π′ = transition_log_probability(trajectory, z′, z, depth)
        @test (π + ℓz) ≈ (π′ + ℓ(z′)) atol = atol
    end
end

@testset "detailed balance" begin
    for max_depth in 1:5
        test_detailed_balance(DummyTrajectory(testℓ), 0, max_depth)
    end
    for max_depth in 1:5
        test_detailed_balance(DummyTrajectory(testℓ; turning = 1:2), 3, max_depth)
    end
    for max_depth in 1:6
        test_detailed_balance(DummyTrajectory(testℓ; divergent = 10:11), 3, max_depth)
    end
    for max_depth in 1:6
        test_detailed_balance(DummyTrajectory(testℓ; divergent = 10:12, turning = -3:-2), 3,
                              max_depth)
    end
end


using VisualDL

# initialize
train_logger = VisualDLLogger("tmp", 1, "train")
test_logger = as_mode(train_logger, "test")

# log scalars
for i in 1:100
    with_logger(train_logger) do
        @log_scalar s0=(i,rand()) s1=(i, rand())
    end

    with_logger(test_logger) do
        @log_scalar s0=(i,rand()) s1=(i, rand())
    end
end

# log histograms
for i in 1:100
    with_logger(train_logger) do
       @log_histogram h0=(i, randn(100))
    end
end

# log texts
for i in 1:100
    with_logger(train_logger) do
       @log_text t0=(i, "This is test " * string(i))
    end
end

# log images
for i in 1:100
    with_logger(train_logger) do
       @log_image i0=([3,3,3], rand(27) * 255)
    end
end

for i in 1:100
    with_logger(test_logger) do
        # Array{Number, 3} is also supported
        @log_image image0=rand(10, 10, 3) * 255
    end
end

save(train_logger)
save(test_logger)


#!/usr/bin/env julia

using FeedbackParticleFilters, Test, SafeTestsets

println("Starting tests")

@safetestset "Basics" begin
    println("")
    println("")
    println("")
    printstyled("Basics", bold=true)
    println("")
    include("Basics/test_Abstractions.jl")
    include("Basics/test_AbstractModel.jl")
    include("Basics/test_AbstractFilteringProblem.jl")
    include("Basics/test_AbstractFilterRepresentation.jl")
    include("Basics/test_AbstractFilteringAlgorithm.jl")
    include("Basics/test_HiddenStateModel.jl")
    include("Basics/test_ObservationModel.jl")
    include("Basics/test_ParametricRepresentation.jl")
    include("Basics/test_ParticleRepresentation.jl")
    include("Basics/test_UnweightedParticleRepresentation.jl")
    include("Basics/test_WeightedParticleRepresentation.jl")
    include("Basics/test_GainEstimation.jl")
    include("Basics/test_FilteringProblem.jl")
end

@safetestset "State models" begin
    println("")
    println("")
    println("")
    printstyled("State models", bold=true)
    println("")
    include("StateModels/test_DiffusionStateModel.jl")
    include("StateModels/test_LinearDiffusionStateModel.jl")
    include("StateModels/test_BrownianMotionTorus.jl")
end

@safetestset "Observation models" begin
    println("")
    println("")
    println("")
    printstyled("Observation models", bold=true)
    println("")
    include("ObservationModels/test_DiffusionObservationModel.jl")
    include("ObservationModels/test_LinearDiffusionObservationModel.jl")
end

@safetestset "Filter representations" begin
    println("")
    println("")
    println("")
    printstyled("Filter representations", bold=true)
    println("")
    include("FilterRepresentations/test_MultivariateGaussian.jl")
    include("FilterRepresentations/test_UnweightedParticleEnsemble.jl")
    include("FilterRepresentations/test_WeightedParticleEnsemble.jl")
end

@safetestset "Gain equations" begin
    println("")
    println("")
    println("")
    printstyled("Gain equations", bold=true)
    println("")
    include("GainEquations/test_PoissonEquation.jl")
end

@safetestset "Gain estimation methods" begin
    println("")
    println("")
    println("")
    printstyled("Gain estimation methods", bold=true)
    println("")
    include("GainEstimationMethods/test_ConstantGainApproximation.jl")
    include("GainEstimationMethods/test_ConstantGainEKSPF.jl")
    include("GainEstimationMethods/test_SemigroupMethod.jl")
    include("GainEstimationMethods/test_DifferentialRKHSMethod.jl")
    include("GainEstimationMethods/test_DifferentialRKHSMethodS1.jl")
end

@safetestset "Filtering algorithms" begin
    println("")
    println("")
    println("")
    printstyled("Filtering algorithms", bold=true)
    println("")
    include("FilteringAlgorithms/test_BPF.jl")
    include("FilteringAlgorithms/test_KBF.jl")
    include("FilteringAlgorithms/test_FPF.jl")
    include("FilteringAlgorithms/test_ppFPF.jl")
end

@safetestset "Simulation" begin
    println("")
    println("")
    println("")
    printstyled("Simulation", bold=true)
    println("")
    include("Simulation/test_Simulation.jl")
    include("Simulation/test_SimulationState.jl")
end


export RandomStartPolicy

Base.@kwdef mutable struct RandomStartPolicy{P,R<:RandomPolicy} <: AbstractPolicy
    policy::P
    random_policy::R
    num_rand_start::Int
end

function (p::RandomStartPolicy)(env)
    p.num_rand_start -= 1
    if p.num_rand_start < 0
        p.policy(env)
    else
        p.random_policy(env)
    end
end

function RLBase.update!(
    p::RandomStartPolicy,
    t::AbstractTrajectory,
    e::AbstractEnv,
    s::AbstractStage,
)
    update!(p.policy, t, e, s)
end

for f in (:prob, :priority)
    @eval function RLBase.$f(p::RandomStartPolicy, args...)
        if p.num_rand_start < 0
            $f(p.policy, args...)
        else
            $f(p.random_policy, args...)
        end
    end
end


SUITE["arithmetic"] = BenchmarkGroup()

SUITE["arithmetic"]["DatePeriod"] = begin
    zdt = ZonedDateTime(2000, 1, 2, tz"America/Winnipeg")
    @benchmarkable $zdt + $(Day(1))
end
SUITE["arithmetic"]["TimePeriod"] = begin
    zdt = ZonedDateTime(2000, 1, 2, tz"America/Winnipeg")
    @benchmarkable $zdt + $(Hour(1))
end

SUITE["arithmetic"]["broadcast"] = BenchmarkGroup()

SUITE["arithmetic"]["broadcast"]["VariableTimeZone/TimePeriod"] = begin
    @benchmarkable $(gen_range(Hour(100), tz"America/Winnipeg")) .+ Hour(1)
end
SUITE["arithmetic"]["broadcast"]["VariableTimeZone/DatePeriod"] = begin
    @benchmarkable $(gen_range(Day(100), tz"America/Winnipeg")) .+ Day(1)
end
SUITE["arithmetic"]["broadcast"]["FixedTimeZone/TimePeriod"] = begin
    @benchmarkable $(gen_range(Hour(100), tz"UTC")) .+ Hour(1)
end
SUITE["arithmetic"]["broadcast"]["FixedTimeZone/DatePeriod"] = begin
    @benchmarkable $(gen_range(Day(100), tz"UTC")) .+ Day(1)
end


# !! auto-generated code, do not edit !!
module stereo_msgs
using RobotOSData.Messages
using RobotOSData.StdMsgs
using RobotOSData.CommonMsgs
struct DisparityImage <: Readable
    header::Header
    image::sensor_msgs.Image
    f::Float32
    T::Float32
    valid_window::sensor_msgs.RegionOfInterest
    min_disparity::Float32
    max_disparity::Float32
    delta_d::Float32
end
end


@testset "Gibbs Sampling" begin
    seed!(42)
    L = 3
    D = 10
    N = 20
    nSamples = 10
    b = 5
    i = GibbsSampling(; nBurnin=0)
    @test repr(i) == "Gibbs Sampler"

    @test AGP.ρ(i) == 1
    @test AGP.is_stochastic(i) == false
end


function node_depth(args...;ntip::Int=Int[], nnode::Int=Int[],
    e1::Array{Int64,1}=Int[],e2::Array{Int64,1}=Int[], nedge::Int=Int[],method::Int=Int[],xx::Array{Int64,1}=Int[])

    for i in 1:ntip
        xx[i] = 1
    end

    if method == 1
        for i in 1:nedge
            xx[e1[i]] = xx[e1[i]] + xx[e2[i]]
        end
    elseif method == 2
        for i in 1:nedge
            if (xx[e1[i]] != 0)
                if xx[e1[i]] >=xx[e1[i]] + 1
                    continue
                end
                xx[ei[i]] = xx[e2[i]] + 1
            end
        end
    end
    return xx
end


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

using RadiationDetectorSignals
using Test

using ArraysOfArrays, StaticArrays, TypedTables, Unitful

@testset "detector_hits" begin
    evtno_data = [[1, 1, 1], [2], [3, 4]]
    detno_data = [[2, 4, 4], [3], [2, 5]]
    thit_data = [[0.1, 0.2, 0.4], [0.3], [0.1, 0.7]]
    edep_data = [[45.2, 683.1, 137.4], [64.3], [84.2, 1043.0]]
    pos_data = [
        [SVector(0.4,-5.2,3.3), SVector(2.1,4.7,-5.5), SVector(-7.2,1.1,4.2)],
        [SVector(8.1,-0.9,-4.9)],
        [SVector(2.0,8.7,0.3), SVector(-2.3,9.8,7.4)]
    ]

    @test @inferred(DetectorHits((evtno = evtno_data[1], detno = detno_data[1], thit = thit_data[1], edep = edep_data[1], pos = pos_data[1]))) isa DetectorHits
    x = DetectorHits((evtno = evtno_data[1], detno = detno_data[1], thit = thit_data[1], edep = edep_data[1], pos = pos_data[1]))

    @test @inferred(DetectorHitEvents((evtno = evtno_data, detno = detno_data, thit = thit_data, edep = edep_data, pos = pos_data))) isa DetectorHitEvents
    X = DetectorHitEvents((evtno = evtno_data, detno = detno_data, thit = thit_data, edep = edep_data, pos = pos_data))

    @test @inferred(Table(X[1])) == x
end # testset

@testset "detector_hits with units" begin
    evtno_data = [[1, 1, 1], [2], [3, 4]]
    detno_data = [[2, 4, 4], [3], [2, 5]]
    thit_data = [[0.1, 0.2, 0.4]u"ns", [0.3]u"ns", [0.1, 0.7]u"ns"]
    edep_data = [[45.2, 683.1, 137.4]u"keV", [64.3]u"keV", [84.2, 1043.0]u"keV"]
    pos_data = [
        [SVector(0.4,-5.2,3.3)u"mm", SVector(2.1,4.7,-5.5)u"mm", SVector(-7.2,1.1,4.2)u"mm"],
        [SVector(8.1,-0.9,-4.9)u"mm"],
        [SVector(2.0,8.7,0.3)u"mm", SVector(-2.3,9.8,7.4)u"mm"]
    ]

    @test @inferred(DetectorHits((evtno = evtno_data[1], detno = detno_data[1], thit = thit_data[1], edep = edep_data[1], pos = pos_data[1]))) isa DetectorHits
    x = DetectorHits((evtno = evtno_data[1], detno = detno_data[1], thit = thit_data[1], edep = edep_data[1], pos = pos_data[1]))

    @test @inferred(DetectorHitEvents((evtno = evtno_data, detno = detno_data, thit = thit_data, edep = edep_data, pos = pos_data))) isa DetectorHitEvents
    X = DetectorHitEvents((evtno = evtno_data, detno = detno_data, thit = thit_data, edep = edep_data, pos = pos_data))

    @test @inferred(Table(X[1])) == x
end # testset


using NgSpice
using Test

rootdir = joinpath(@__DIR__, "..")
netpath = joinpath(rootdir, "inputs", "mosfet.cir")

@testset "initialize" begin
    init()
    @test source(netpath) == 0
    @test NgSpice.run() == 0
end

@testset "getvectors" begin
    # as this simulation runs after the one in `ngspice_api.jl`,
    # `tran2` is the current plot
    @test curplot() == "tran2"
    @test listallplots() == ["tran2", "tran1", "const"]
    @test collect(keys(listallvecs())) == collect(sort(listallplots()))
    @test listallvecs()[curplot()] == listcurvecs()
    @test getvec.(listcurvecs()) == get_vector_info.(listcurvecs())
    #listcurvecs()
    for node in split("V(1) emit coll base vcc")
        @test getvec(node)[2] == "voltage"
    end
    for branch in split("vcc#branch vin#branch")
        @test getvec(branch)[2] == "current"
    end
    @test getvec("time")[2] == "time"
end

@testset "passing circuit as an array" begin
    netlist = readlines(netpath)
    @test load_netlist(netlist) == 0
end


using Plots; pyplot()
using Reproduce
using FileIO
using Statistics
using ProgressMeter
using JLD2
using LaTeXStrings
using StatsPlots

include("./plot_utils.jl")
GPU = GeneralPlotUtils
LU = LabelUtils
include("./plotting.jl")

function load_data(data_home)
    experiment_folders = [data_home]
    ic = ItemCollection(joinpath(experiment_folders[1], "data"))
    return ic
end

function load_best(ic)
    algo_divisor_keys = ["behaviour_learner", "demon_learner", "demon_opt", "demon_update"]
    sweep_params = ["alpha_init", "eta", "alpha_init","batch_size"]
    # data_key = :oned_tmaze_dmu_error
    data_key = :oned_tmaze_old_error

    algo_specs_full = GPU.split_algo(ic, algo_divisor_keys)

    all_algos_ics = [search(ic,algo_spec) for algo_spec in algo_specs_full]
    @show length.(all_algos_ics)

    valid_algos_ind = findall(a -> length(a) != 0, all_algos_ics)

    algo_specs = algo_specs_full[valid_algos_ind]
    algo_ics = all_algos_ics[valid_algos_ind]
    best_per_algo_ics = []
    for (i,algo_ic) in enumerate(algo_ics)
        push!(best_per_algo_ics, GPU.get_best_final_perf(algo_ic,sweep_params, data_key, 0.1))
    end
    @show length.(best_per_algo_ics)
    return best_per_algo_ics
end

SR_ER_ic_all = load_data("../data/ER/OneDTMaze_RR_ER/")
SR_ER_ic_all = search(SR_ER_ic_all, Dict("batch_size" => 8))
SR_ER_ic = load_best(SR_ER_ic_all)

Q_ER_ic_all = load_data("../data/ER/OneDTMaze_RR_ER_2/")
Q_ER_ic_all = search(Q_ER_ic_all, Dict("batch_size" => 8))
Q_ER_ic = load_best(Q_ER_ic_all)

SR_ic_all = load_data("../data/OneDTMaze_RR/")
SR_ic_all = search(SR_ic_all, Dict("demon_learner" => "SR"))
SR_ic = load_best(SR_ic_all)

Q_ic_all = load_data("../data/OneDTMaze_RR/")
Q_ic_all = search(Q_ic_all, Dict("demon_learner" => "Q"))
Q_ic = load_best(Q_ic_all)

LSTD_ic_all = load_data("../data/OneDTMaze_RR/")
LSTD_ic_all = search(LSTD_ic_all, Dict("demon_learner" => "LSTD"))
LSTD_ic = load_best(LSTD_ic_all)

@show length(SR_ER_ic)
@show length(Q_ER_ic)
@show length(SR_ic)
@show length(Q_ic)
@show length(LSTD_ic)

all_ic = [
    SR_ER_ic[1],
    Q_ER_ic[1],
    SR_ic[1],
    Q_ic[1],
    LSTD_ic[1],
]


println("Hello world!");

@everywhere using Printf;

@show workdir = if length(ARGS) > 0
  ARGS[1];
else
  ".";
end

@everywhere fmt(x) = @sprintf("%07d", round(Int, 1e3*x));
@everywhere fmt_int(x) = @sprintf("%03d", x);

@everywhere function prefix(case)
  "E0-$(fmt(case[:E0]))_K1-$(fmt(case[:K1]))_K2-$(fmt(case[:K2]))_kT-$(fmt(case[:kT]))_Fz-$(fmt(case[:Fz]))_Fx-$(fmt(case[:Fx]))_n-$(fmt(case[:n]))_b-$(fmt(case[:b]))_run-$(fmt_int(case[:run]))";
end

cases = Any[];
#E0s = [1.0; 5.0];
#K1s = [1e-1; 1.0];
#kTs = [1.0; 0.1];
#Fzs = [0.0; 0.5; 1.0];
#Fxs = [0.0; 0.5; 1.0];
#ns = Int[100];
#bs = [1.0];
runs = 1:25;
for run in runs
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 5, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 5, :K1 => 0.1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 0.1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 1,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 2,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 1, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 0.1, :Fz => 0,
                    :Fx => 1, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 0.1, :Fz => 0,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 0.1, :Fz => 1,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 0.1, :K2 => 0.0,
                    :kT => 0.1, :Fz => 1,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 0, :n => 200, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 0, :n => 100, :b => 0.5, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 0, :n => 100, :b => 2.0, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 1,
                    :Fx => 0, :n => 100, :b => 0.5, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 1, :n => 100, :b => 0.5, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 1,
                    :Fx => 1, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 3, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 3,
                    :Fx => 3, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 3,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 5,
                    :Fx => 0, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 3, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 1, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 0,
                    :Fx => 5, :n => 100, :b => 1, :run => run));
  push!(cases, Dict(:E0 => 5, :K1 => 1, :K2 => 0.0,
                    :kT => 1, :Fz => 5,
                    :Fx => 5, :n => 100, :b => 1, :run => run));
end

@info "total number of cases to run: $(length(cases))";

mkpath(workdir);
mkpath(joinpath(workdir, "standard"));
mkpath(joinpath(workdir, "umbrella"));
mkpath(joinpath(workdir, "clustering"));

pmap(case -> begin;

  outfile = joinpath(workdir, "standard", "$(prefix(case)).out");
  if !isfile(outfile)
    println("Running case: $case.");
    command = `julia -O 3 mcmc_eap_chain.jl --chain-type dielectric --energy-type Ising -b $(case[:b]) --E0 $(case[:E0]) --K1 $(case[:K1]) --K2 $(case[:K2]) --kT $(case[:kT]) --Fz $(case[:Fz]) --Fx $(case[:Fx]) -n $(case[:n]) --num-steps 1000000 -v 2 --prefix $(joinpath(workdir, "standard", prefix(case)))`;
    output = read(command, String);
    write(outfile, output);
  else
    println("Case: $case has already been run.");
  end

  outfile = joinpath(workdir, "umbrella", "$(prefix(case)).out");
  if !isfile(outfile)
    println("Running case: $case.");
    command = `julia -O 3 mcmc_eap_chain.jl --chain-type dielectric --energy-type Ising --umbrella-sampling -b $(case[:b]) --E0 $(case[:E0]) --K1 $(case[:K1]) --K2 $(case[:K2]) --kT $(case[:kT]) --Fz $(case[:Fz]) --Fx $(case[:Fx]) -n $(case[:n]) --num-steps 1000000 -v 2 --prefix $(joinpath(workdir, "umbrella", prefix(case)))`;
    output = read(command, String);
    write(outfile, output);
  else
    println("Case: $case has already been run.");
  end

  outfile = joinpath(workdir, "clustering", "$(prefix(case)).out");
  if !isfile(outfile)
    println("Running case: $case.");
    command = `julia -O 3 mcmc_clustering_eap_chain.jl --chain-type dielectric --energy-type Ising -b $(case[:b]) --E0 $(case[:E0]) --K1 $(case[:K1]) --K2 $(case[:K2]) --kT $(case[:kT]) --Fz $(case[:Fz]) --Fx $(case[:Fx]) -n $(case[:n]) --num-steps 1000000 -v 2 --prefix $(joinpath(workdir, "clustering", prefix(case)))`;
    output = read(command, String);
    write(outfile, output);
  else
    println("Case: $case has already been run.");
  end

end, cases);


using COVID
using Test

@testset "COVID.jl" begin
    params = ParameterSetting()
    switch_pts = [(0.3, 5), (0.9, 5)]
    res1 = run_model(params, .4, 20)
    res2 = run_model(params, switch_pts)
    res3 = run_model(params, 0.03, 6, 20)
    @test typeof(res1) == Array{Int64,2}
    @test typeof(res2) == Array{Int64,2}
    @test typeof(res3) == Array{Float64,2}
end


using GLVisualize, GeometryTypes, GLAbstraction

window = glscreen()

rectangles = HyperRectangle(Vec2f0(0), Vec2f0(100))
rectangles = [  HyperRectangle(Vec2f0(0), Vec2f0(100)),
                HyperRectangle(Vec2f0(100), Vec2f0(200)) ]

_view(visualize(rectangles), window)
#_view(visualize(rectangles), window)

renderloop(window)

#http://www.glvisualize.com/examples/sprites/#bouncy


import Base: +, -, *, /, ^

+(x::RObject, y::RObject) = rcall(Symbol("+"), x, y)
-(x::RObject, y::RObject) = rcall(Symbol("-"), x, y)
*(x::RObject, y::RObject) = rcall(Symbol("*"), x, y)
/(x::RObject, y::RObject) = rcall(Symbol("/"), x, y)
^(x::RObject, y::RObject) = rcall(Symbol("^"), x, y)


using JLD2
using DataFrames
using StatsPlots

@load "paper/box.jld2" box_df

p = @df box_df[.!ismissing.(box_df[:, :mean]), :] plot(:n, :mean, group=(:Type, :N), marker=4,
  xlabel="T", ylabel="mean coupling time", xticks=[500;1000;2000;4000], size=(400,250))
Plots.savefig(p, "paper/boxDemoPlot.pdf")

@load "paper/boxInd.jld2" boxInd_df

p = @df boxInd_df[.!ismissing.(boxInd_df[:, :mean]), :] plot(:n, :mean, group=(:Type, :N), marker=4,
  xlabel="T", ylabel="mean coupling time", xticks=[500;1000;2000;4000], size=(450,250))
Plots.savefig(p, "paper/boxDemoIndPlot.pdf")


using Metalhead
using Metalhead: classify

using Images

img = load("tile.png")
vgg = VGG19()

classify(vgg,img)


############################################################
###  ASE Neighborlist implementation
############################################################

"""### type ASENeighborList <: AbstractNeighborList

This makes available the functionality of the ASE neighborlist implementation.
The neighborlist itself is actually stored in `ASEAtoms.po`, but
attaching an `ASENeighborList` will indicate this.

#### Keyword arguments:

 * skin = 0.3
 * sorted
 * self_interaction
 * bothways
"""
type ASENeighborList #  <: AbstractNeighborList
    po::PyObject
    cutoffs::Vector{Float64}
end

# # default constructor from a list of cut-offs
# ASENeighborList(cutoffs::Vector{Float64}; kwargs...) =
#     ASENeighborList(ase_neiglist.NeighborList(cutoffs; kwargs...))

# constructor from an ASEAtoms object, with a single cut-off
# this generates a list of cut-offs, then the neighborulist, then
# builds the neighborlist from the ASEAtoms object
ASENeighborList(atm::ASEAtoms, cutoff::Float64; kwargs...) =
    ASENeighborList(atm, cutoff * ones(length(atm)); kwargs...)

# constructor from an ASEAtoms object, with multiple cutoffs
# this also builds the neighborlist
function ASENeighborList(atm::ASEAtoms, cutoffs::Vector{Float64};
                         bothways=true, self_interaction=false, kwargs...)
    po = ase_neiglist.NeighborList(cutoffs * 0.5;
                                   bothways=bothways,
                                   self_interaction=self_interaction,
                                   kwargs...)
    nlist = ASENeighborList(po, cutoffs)
    update!(nlist, atm)
    return nlist
end

# # regain the cutoffs vector
# _get_cutoffs_ref_(nlist::ASENeighborList) =
#     PyArray(nlist.po["cutoffs"])

"""`update!(nlist::ASENeighborList, atm::ASEAtoms)`

checks whether the atom positions have moved by more than the skin
and rebuilds the list if they have done so.
"""
update!(nlist::ASENeighborList, atm::ASEAtoms) =  nlist.po[:update](atm.po)

"""`build!(nlist::ASENeighborList, atm::ASEAtoms)`

force rebuild of the neighborlist
"""
build!(nlist::ASENeighborList, atm::ASEAtoms) = nlist.po[:build](atm.po)

"""`get_neighbors(n::Integer, neiglist::ASENeighborList) -> (indices, offsets)`

Return neighbors and offsets of atom number n. The positions of the neighbor
atoms can then calculated like this: (python code)

```
          indices, offsets = nl.get_neighbors(42)
          for i, offset in zip(indices, offsets):
              print(atoms.positions[i] + dot(offset, atoms.get_cell()))
```

If `get_neighbors(a)` gives atom b as a neighbor,
    then `get_neighbors(b)` will not return a as a neighbor, unless
    `bothways=True` was used.
"""
function get_neighbors(n::Integer, neiglist::ASENeighborList)
    indices, offset = neiglist.po[:get_neighbors](n-1)
    indices .+= 1
    return (indices, offset)
end
# "alias for `get_neighbors`"
# get_neighbors = get_neighbors
# "alias for `get_neighbors`"
# neighbors = get_neighbors
"alias for `get_neighbors`"
neighbors = get_neighbors

"""`get_neighbors(n::Integer, neiglist::ASENeighborList, atm::ASEAtoms)
              -> (indices::Vector{Int}, s::Vector{Float64}, r::Matrix{Float64})`

 * `indices`: indices of atom positions
 * `s`: scalar distances of neighbours
 * `r`: relative positions of neighbours (vectorial dist)

This is a convenience function that does some of the work of constructing the
neighborhood. This is probably fairly inefficient to use since it has
to construct a `PyArray` object for the positions every time it is called.
Instead, use the iterator. Problem is, this is not much faster, because
the PyCall conversion overhead is so horrendous.
"""
function get_neighbors(n::Integer, neiglist::ASENeighborList, atm::ASEAtoms;
                       rcut=Inf)
    (inds, offsets) = ASE.neighbors(n, neiglist)
    X = _get_positions_ref_(atm)
    cell = get_cell(atm)
    r = X[inds, :]' + cell' * offsets' .- slice(X, n, :)
    s = sqrt(sumabs2(r, 1))
    I = find(s .<= neiglist.cutoffs[n])
    return inds[I], s[I], r[:, I]
end

##################### NEIGHBOURHOOD ITERATOR ###################
# this is about twice as fast as `get_neighbours`
# which indicates that, either, the ASE neighbour list is very slow
#   or, the overhead from the python call is horrendous!
#

type ASEAtomIteratorState
    at::ASEAtoms
    neiglist::ASENeighborList
    n::Int         # iteration index
    X::Matrix{Float64}
    cell_t::Matrix{Float64}
end

ASEAtomIteratorState(at::ASEAtoms, neiglist::ASENeighborList) =
    ASEAtomIteratorState( at, neiglist, 0,
                          get_positions(at),
                          get_cell(at)' )

import Base.start
start(I::Tuple{ASEAtoms,ASENeighborList}) =
    ASEAtomIteratorState(I...)

import Base.done
done(I::Tuple{ASEAtoms,ASENeighborList}, state::ASEAtomIteratorState) =
    (state.n == length(state.at))

import Base.next
function next(I::Tuple{ASEAtoms,ASENeighborList}, state::ASEAtomIteratorState)
    state.n += 1
    inds, offsets = neighbors(state.n, state.neiglist)
    r = state.X[:,inds] + state.cell_t * offsets' .- state.X[:,state.n]
    s = sqrt(sumabs2(r, 1))
    I = find(s .<= I[2].cutoffs[state.n])
    return (state.n, inds, s, r), state
    # now find the indices that are actually within the cut-off
    #  TODO: resolve the nasty issue of 1/2 cut-off first?!?!?
    # I = find(s .< 3.0)   # state.cutoffs[state.n])  (hard-coded for debugging)
    # return (state.n, inds[I], s[I], r[:,I]), state
end


# Julia wrapper for the sgp4 Python library:
# https://pypi.python.org/pypi/sgp4/

VERSION >= v"0.4.0-dev+6521" && __precompile__()

module SGP4

import Compat.ASCIIString

using PyCall

function __init__()
    try
        global const sgp4io = PyCall.pywrap(pyimport("sgp4.io"))
        global const earth_gravity = PyCall.pywrap(pyimport("sgp4.earth_gravity"))
    catch e
        error("Error loading sgp4 python package - check to make sure it's installed")
    end
end

export GravityModel,
       twoline2rv,
       propagate

immutable GravityModel
    model::PyObject # can be any of {wgs72old, wgs72, wgs84}
end
GravityModel(ref::AbstractString) = earth_gravity.pymember(ref)

# sgp4.io convenience functions
twoline2rv(args...) = sgp4io.twoline2rv(args...)

"""
Propagate the satellite from its epoch to the date/time specified

Returns (position, velocity) at the specified time
"""
function propagate( sat::PyObject,
                    year::Real,
                    month::Real,
                    day::Real,
                    hour::Real,
                    min::Real,
                    sec::Real )
    (pos, vel) = sat[:propagate](year,month,day,hour,min,sec)

    # check for errors
    if sat[:error] != 0
        println(sat[:error_message])
    end

    return ([pos...],[vel...])
end

function propagate( sat::PyObject,
                    t::DateTime )
    propagate(sat,
              Dates.year(t),
              Dates.month(t),
              Dates.day(t),
              Dates.hour(t),
              Dates.minute(t),
              Dates.second(t))
end

"Propagate many satellites to a common time"
function propagate( sats::Vector{PyObject},
                    year::Real,
                    month::Real,
                    day::Real,
                    hour::Real,
                    min::Real,
                    sec::Real )
    pos = zeros(3,length(sats))
    vel = zeros(3,length(sats))
    for i = 1:length(sats)
        (pos[:,i],vel[:,i]) = propagate(sats[i],year,month,day,hour,min,sec)
    end
    return (pos,vel)
end

function propagate( sats::Vector{PyObject},
                    t::DateTime )
    propagate(sats,
              Dates.year(t),
              Dates.month(t),
              Dates.day(t),
              Dates.hour(t),
              Dates.minute(t),
              Dates.second(t))
end

end #module


module Tpetra
using CxxWrap, MPI
import ..libjltrilinos
import ..CxxUnion
import ..Teuchos
import ..Kokkos

import LinearAlgebra
using SparseArrays

@wrapmodule(libjltrilinos, :register_tpetra)

LinearAlgebra.dot(a::CxxUnion{Tpetra.Vector{ST,LT,GT,NT}}, b::CxxUnion{Tpetra.Vector{ST,LT,GT,NT}}) where {ST,LT,GT,NT} = Tpetra.dot(a,b)

Map(num_indices::Integer, index_base::Integer, comm::CxxWrap.SmartPointer{<:Teuchos.Comm}) = Map(num_indices, index_base, comm, Kokkos.default_node_type())
Map(num_indices::Integer, index_list, index_base::Integer, comm::CxxWrap.SmartPointer{<:Teuchos.Comm}) = Map(num_indices, index_list, index_base, comm, Kokkos.default_node_type())
getGlobalElement(map, idx::UInt64) = getGlobalElement(map, convert(Int, idx))

"""
Construct a Trilinos sparse matrix from a SparseMatrixCSC. Note that the resulting matrix is the transpose of the Julia matrix, since Trilinos is row-based.
The resulting matrix has a fixed structure (using a constant CrsGraph in Tpetra) and its rows are distributed over the available MPI processes.
The matrix A is supposed to be completely supplied on each rank.
"""
function CrsMatrix(A::SparseMatrixCSC{ScalarT, GlobalOrdinalT}, rowmap) where {ScalarT,GlobalOrdinalT}
  # Create the sparse matrix structure
  matrix_graph = CrsGraph(rowmap, GlobalOrdinalT(0))
  rows = rowvals(A)
  vals = nonzeros(A)
  num_my_elements = getNodeNumElements(rowmap)
  for local_row = 0:num_my_elements-1
    global_row = getGlobalElement(rowmap, local_row)
    insertGlobalIndices(matrix_graph, global_row, rows[nzrange(A, global_row+1)].-1)
  end
  fillComplete(matrix_graph)

  tpetra_matrix = CrsMatrix(matrix_graph)
  resumeFill(tpetra_matrix)
  for local_row = 0:num_my_elements-1
    global_row = getGlobalElement(rowmap, local_row)
    replaceGlobalValues(tpetra_matrix, global_row, rows[nzrange(A, global_row+1)].-1, vals[nzrange(A, global_row+1)])
  end
  fillComplete(tpetra_matrix)
  return tpetra_matrix
end

"""
CrsMatrix apply function with defauls arguments
"""
apply(A::CxxUnion{CrsMatrix{ST,LT,GT,NT}},
                   X::CxxUnion{MultiVector{ST,LT,GT,NT}},
                   Y::CxxUnion{MultiVector{ST,LT,GT,NT}},
                   mode::Teuchos.ETransp=Teuchos.NO_TRANS,
                   α::Number=1,
                   β::Number=0) where {ST,LT,GT,NT} = _apply(A,X,Y,mode,α,β)

"""
Gets an AbstractArray-compatible device view of the local part of the multivector. Calls getLocalView internally.
"""
device_view(mv::CxxUnion{MultiVector{ST,LT,GT,NT}}) where {ST,LT,GT,NT} = Kokkos.view(getLocalView(device_view_type(mv),mv),Val{2})

"""
Gets an AbstractArray-compatible host view of the local part of the multivector. Calls getLocalView internally.
"""
host_view(mv::CxxUnion{MultiVector{ST,LT,GT,NT}}) where {ST,LT,GT,NT} = Kokkos.view(getLocalView(host_view_type(mv),mv),Val{2})

"""
Gets an AbstractArray-compatible device view of the local part of the vector. Calls getLocalView internally.
"""
device_view(v::CxxUnion{Vector{ST,LT,GT,NT}}) where {ST,LT,GT,NT} = Kokkos.view(getLocalView(device_view_type(v),v),Val{1})

"""
Gets an AbstractArray-compatible host view of the local part of the vector. Calls getLocalView internally.
"""
host_view(v::CxxUnion{Vector{ST,LT,GT,NT}}) where {ST,LT,GT,NT} = Kokkos.view(getLocalView(host_view_type(v),v),Val{1})

function CrsGraph(rowmap::CxxUnion{Map{LT,GT,NT}}, colmap::CxxUnion{Map{LT,GT,NT}}, graph::Kokkos.StaticCrsGraph) where {LT,GT,NT}
  return CrsGraph(rowmap, colmap, Kokkos.to_cpp(graph,local_graph_type(CrsGraph{LT,GT,NT})))
end

# Serial dump to regular dense array
function to_array(A)
  nb_rows = Int(getNodeNumRows(A))
  result = zeros(nb_rows,nb_rows)
  row_indices = zeros(Int32,100)
  row_values = zeros(Float64,100)

  indices_view = Teuchos.ArrayView(row_indices)
  values_view = Teuchos.ArrayView(row_values)
  rowlength_ref = Ref(UInt(0))

  for row in 1:nb_rows
    getLocalRowCopy(A, row-1, indices_view, values_view, rowlength_ref)
    for i in 1:rowlength_ref[]
      result[row,row_indices[i]+1] = row_values[i]
    end
  end
  return result
end

end


const Data{M, T} = NamedTuple{(:poly, :sigtail, :sigratio),
                              Tuple{Polynomial{M, T}, Polynomial{M, T}, M}}
const SigTable{I, M, T} = Dict{Tuple{I, M}, Data{M, T}}
const GenData{I} = NamedTuple{(:position, :att_key, :tag),
                              Tuple{I, I, Symbol}}
gendata(i::I, j::I, tg::Symbol) where I = (position = i, att_key = j, tag = tg)

mutable struct SigPolynomialΓ{I, M, T, MΓ<:Context{M}, TΓ<:Context{T}, PΓ<:PolynomialΓ{M, T, MΓ, TΓ}, S}<:Context{Tuple{I, M}}
    po::PΓ
    tbl::SigTable{I, M, T}
    ord_indices::Dict{I, GenData{I}}
end

pos_type(::SigPolynomialΓ{I}) where {I} = I
mon_type(::SigPolynomialΓ{I, M}) where {I, M} = M
coeff_type(::SigPolynomialΓ{I, M, T}) where {I, M, T} = T
mod_order(::SigPolynomialΓ{I, M, T, MΓ, TΓ, PΓ, S}) where {I, M, T, MΓ, TΓ, PΓ, S} = S

gendata(ctx::SigPolynomialΓ{I}) where I = (position = zero(I), att_key = zero(I), tag = :f)

function pos(ctx::SigPolynomialΓ{I, M},
             p::Tuple{I, M}) where {I, M}

    iszero(p[1]) && return zero(I)
    ctx.ord_indices[p[1]][:position]
end

function gettag(ctx::SigPolynomialΓ{I, M},
                p::Tuple{I, M}) where {I, M}

    iszero(p[1]) && return :nil
    ctx.ord_indices[p[1]][:tag]
end

getattkey(ctx::SigPolynomialΓ{I, M}, p::Tuple{I, M}) where {I, M} = ctx.ord_indices[p[1]][:att_key]

function new_gen!(ctx::SigPolynomialΓ{I, M, T},
                  info_hashmap::Dict{I, Info},
                  posit::I,
                  attached_to::I,
                  tagg::Symbol,
                  pol::Polynomial{M, T}) where {I, M, T}

    # register new generator
    posit_key = register!(ctx, pol, info_hashmap)

    # rebuild ord_indices
    new_pos!(ctx, posit_key, posit, attached_to, tagg)
end

function register!(ctx::SigPolynomialΓ{I, M, T},
                   pol::Polynomial{M, T},
                   info_hashmap::Dict{I, Dict{String, Int}}) where {I, M, T}

    posit_key = maximum(keys(ctx.ord_indices)) + one(I)
    ctx((posit_key, one(ctx.po.mo)), pol)
    info_hashmap[posit_key] = new_info()
    ctx.ord_indices[posit_key] = gendata(ctx)
    return posit_key
end

function new_pos!(ctx::SigPolynomialΓ{I},
                  posit_key::I,
                  posit::I,
                  attached_to::I,
                  tagg::Symbol) where I

    delete!(ctx.ord_indices, posit_key)
    new_dict_arr = [(k, i >= posit ? gendata(i + one(I), j, tg) : gendata(i, j, tg))
                    for (k, (i, j, tg)) in ctx.ord_indices]
    push!(new_dict_arr, (posit_key, gendata(posit, attached_to, tagg)))
    ctx.ord_indices = Dict(new_dict_arr)
end

function new_pos!(ctx::SigPolynomialΓ{I}, posit_key::I, posit::I) where I
    new_pos!(ctx, posit_key, posit, ctx.ord_indices[posit_key][:att_key],
             ctx.ord_indices[posit_key][:tag])
end

# find maximal index
function maxpos(ctx::SigPolynomialΓ{I, M}) where {I, M}

    maximum(v -> v[:position], values(ctx.ord_indices))
end

# return original generator of higher index than pos if it exists
function f_left(ctx::SigPolynomialΓ{I, M}, pos::I) where {I, M}

    result = zero(I)
    for (i, v) in ctx.ord_indices
        if v[:tag] == :f && v[:position] > pos
            if iszero(result) || v[:position] < ctx.ord_indices[result][:position]
                result = i
            end
        end
    end
    return result
end

function idxsigpolynomialctx(coefficients,
                             ngens;
                             monomials=nothing,
                             index_type=UInt32,
                             mask_type=UInt32,
                             pos_type=UInt32,
                             mod_order=:POT,
                             deg_bound = 0,
                             kwargs...)
    if isnothing(monomials)
        moctx = ixmonomialctx(; indices=index_type, mask_type=mask_type, deg_bound=deg_bound, kwargs...)
    end
    po = polynomialctx(coefficients, monomials = moctx)
    tbl = SigTable{pos_type, index_type, eltype(coefficients)}()
    ord_indices = Dict([(pos_type(i), gendata(pos_type(i), zero(pos_type(i)), :f)) for i in 1:ngens])
    SigPolynomialΓ{pos_type, eltype(moctx), eltype(coefficients),
                   typeof(moctx), typeof(coefficients), typeof(po), mod_order}(po, tbl, ord_indices)
end

# registration functions

function (ctx::SigPolynomialΓ{I, M, T})(sig::Tuple{I, M},
                                        pol::Polynomial{M, T},
                                        sigtail::Polynomial{M, T}) where {I, MO, M, T}
    if iszero(pol)
        ratio = one(ctx.po.mo)
    else
        ratio = div(ctx.po.mo, sig[2], leadingmonomial(pol))
    end
    val = Data{M, T}((pol, sigtail, ratio))
    try
        ctx.tbl[sig] = val
    catch
        insert!(ctx.tbl, sig, val)
    end
end

(ctx::SigPolynomialΓ{I, M, T})(sig::Tuple{I, M}, pol::Polynomial{M, T}) where {I, M, T} = ctx(sig, pol, zero(pol))

Base.getindex(ctx::SigPolynomialΓ{I, M}, sig::Tuple{I, M}) where {I, M} = getindex(ctx.tbl, sig)

# get functions

# WARNING: if sig::Tuple{J, MO} where J != I or MO != M then this will convert sig to a Tuple{I, M}

@inline function (ctx::SigPolynomialΓ{I, M})(sig::Tuple{I, M}) where {I, M, MO}
    ctx.tbl[sig]
    # get(ctx.tbl, sig) do
    #     error("Nothing registered under the signature $(sign)")
    # end
end

function (ctx::SigPolynomialΓ{I, M, T})(m::M, sig::Tuple{I, M}; orig_elem = false) where {I, M, T}
    key = (sig[1], mul(ctx.po.mo, m, sig[2]))
    if !(orig_elem)
        get(ctx.tbl, key) do
            val = ctx.tbl[sig]
            Data{M, T}((mul(ctx.po, val[:poly], m), mul(ctx.po, val[:sigtail], m), val[:sigratio]))
        end
    end
    val = ctx.tbl[sig]
    Data{M, T}((mul(ctx.po, val[:poly], m), mul(ctx.po, val[:sigtail], m), val[:sigratio]))
end

# get projection to highest index

function project(sig::Tuple{I, M},
                 sigtail::Polynomial{M, T}) where {I, M, T}

    Polynomial{M, T}(vcat(sig[2], sigtail.mo), vcat(one(T), sigtail.co))
end

function project(ctx::SigPolynomialΓ{I, M},
                 sig::Tuple{I, M}) where {I, M}
    project(sig, ctx(sig)[:sigtail])
end

# forwarding of functions on polynomials/monomials

function mul(ctx::SigPolynomialΓ{I, M}, m::M, sig::Tuple{I, M}) where {I, M}
    (sig[1], mul(ctx.po.mo, m, sig[2]))
end

function divides(ctx::SigPolynomialΓ{I, M}, s1::Tuple{I, M}, s2::Tuple{I, M}) where {I, M}
    s1[1] == s2[1] && divides(ctx.po.mo, s1[2], s2[2])
end

@inline leadingmonomial(ctx::SigPolynomialΓ{I, M}, sig::Tuple{I, M}) where {I, M} = leadingmonomial(ctx(sig)[:poly])

@inline leadingmonomial(ctx::SigPolynomialΓ{I, M}, m::M, sig::Tuple{I, M}) where {I, M} = leadingmonomial(ctx(m, sig)[:poly])

# sorting

@inline @generated function lt(ctx::SigPolynomialΓ{I, M, T, MΓ, TΓ, PΓ, S},
                               a::Tuple{I, M},
                               b::Tuple{I, M}) where {I, M, T, MΓ, TΓ, PΓ, S}

    if S == :POT
        quote
            if a[1] == b[1]
                return lt(ctx.po.mo, a[2], b[2])
            end
            return ctx.ord_indices[a[1]][:position] < ctx.ord_indices[b[1]][:position]
        end
    elseif S == :TOP
        quote
            if a[2] == b[2]
                return ctx.ord_indices[a[1]][:position] < ctx.ord_indices[b[1]][:position]
            end
            return lt(ctx.po.mo, a[2], b[2])
        end
    end
end

# Abstract Algebra

(ctx::SigPolynomialΓ)(i, m::AA.MPolyElem) = (pos_type(ctx)(i), ctx.po.mo(m))
(ctx::SigPolynomialΓ{I, M})(sig::Tuple{I, M}, p::AA.MPolyElem) where {I, M} = ctx(sig, ctx.po(p))

function (R :: AA.MPolyRing)(ctx::SigPolynomialΓ{I, M},
                             sig::Tuple{I, M}) where {I, M}

    R(ctx.po, ctx(sig)[:poly])
end


module ORBIT

using ..ParticleSwarm
using ..Sequencing # Includes Clifford, Waveforms, and Qubit

# Fitness calls here are of the form
# "How well does <waveform> implement <pulse> on <qubit>?"

# The overarching "ORBIT" call will include a qubit and a gate, and define the
# fitness anonymously through the below.
function fitness(params::Vector{UInt16}, q::Qubit, p::Pulse)
  pulselength = div(length(params), 3)
  q[p] = ExactWaveform(params[1:pulselength], params[pulselength:2*pulselength],
                       params[2:pulselength:3*pulselength], true)
  pulseIndex = findfirst(SQPulse, p)

  # Every "1" in the readout detracts from the fitness score.
  1 - mean([sendSequenceAndReadout(q, benchmark1Qubit(
                                   pulseIndex, orbitSequenceLength))
        for _ in 1:orbitAveraging])
end

# Same but for FloatWaveforms, where we are optimizing vectors of Float64
function fitness(params::Vector{Float64}, q::Qubit, p::Pulse)
  q[p] = params
  pulseIndex = findfirst(SQPulse, p)
  1 - mean([sendSequenceAndReadout(q, benchmark1Qubit(
                                   pulseIndex, orbitSequenceLength))
        for _ in 1:orbitAveraging])
end

function optimize(q::Qubit, p::Pulse, numpoints::Integer, init::Waveform)
  if isa(init, ExactWaveform)
    if isa(q, QubitWithZ)
      varCount = 3*numpoints
      guess = [init.XYI; init.XYQ; init.Z]
    else
      varCount = 2*numpoints
      guess = [init.XYI; init.XYQ]
    end
  else
    varCount = numpoints
    guess = init.wavedata
  end

  if p[1] < 7 && isa(q.lineXYI[1], InstrumentControl.InsAWG5014C) # XY on AWG
end

end # End module ORBIT


@testset "Penalty Parameter Interface" begin
    struct Fake <: PH.AbstractPenaltyParameter end
    fake = Fake()
    xhid = PH.XhatID(PH.NodeID(0), PH.Index(0))
    phd = fake_phdata(3, 4; add_leaves=false)
    @test_throws PH.UnimplementedError PH.get_penalty_value(fake)
    @test_throws PH.UnimplementedError PH.get_penalty_value(fake, xhid)
    @test_throws PH.UnimplementedError PH.is_initial_value_dependent(Fake)
    @test_throws PH.UnimplementedError PH.is_subproblem_dependent(Fake)
    @test_throws PH.UnimplementedError PH.is_variable_dependent(Fake)
    @test_throws PH.UnimplementedError PH.penalty_map(fake)
    @test_throws PH.UnimplementedError PH.process_penalty_initial_value(fake, phd)
    @test_throws PH.UnimplementedError PH.process_penalty_subproblem(fake,
                                                                     phd,
                                                                     PH.scid(0),
                                                                     Dict{PH.VariableID,
                                                                          Float64}()
                                                                  )
end

@testset "Scalar Penalty" begin
    @test PH.is_initial_value_dependent(PH.ScalarPenaltyParameter) == false
    @test PH.is_subproblem_dependent(PH.ScalarPenaltyParameter) == false
    @test PH.is_variable_dependent(PH.ScalarPenaltyParameter) == false

    rflt = 10.0*rand()
    r = PH.ScalarPenaltyParameter(rflt)
    @test PH.get_penalty_value(r) == rflt
    @test PH.get_penalty_value(r, PH.XhatID(PH.NodeID(0), PH.Index(0))) == rflt
end

@testset "Proportional Penalty" begin
    @test PH.is_initial_value_dependent(PH.ProportionalPenaltyParameter) == false
    @test PH.is_subproblem_dependent(PH.ProportionalPenaltyParameter) == true
    @test PH.is_variable_dependent(PH.ProportionalPenaltyParameter) == true

    N = 2
    st = PH.two_stage_tree(N)
    var_map = Dict{PH.ScenarioID,Dict{PH.VariableID,PH.VariableInfo}}()
    coef_map = Dict{Int,Dict{PH.VariableID,Float64}}()
    for s in 0:(N-1)
        js = two_stage_model(PH.scid(s))
        svm = PH.report_variable_info(js, st)
        var_map[PH.scid(s)] = svm
        (br_vars, lf_vars) = PH._split_variables(st, collect(keys(svm)))
        coef_map[s] = PH.report_penalty_info(js,
                                             br_vars,
                                             PH.ProportionalPenaltyParameter
                                             )
    end

    c_rho = 10.0 * (rand() + 1e-12) # Avoid the possibility of c_rho == 0.0
    phd = PH.PHData(PH.ProportionalPenaltyParameter(c_rho),
                    st,
                    Dict(1 => Set([PH.scid(0), PH.scid(1)])),
                    var_map,
                    TimerOutputs.TimerOutput()
                    )

    for s in 0:(N-1)
        PH.process_penalty_subproblem(phd.r, phd, PH.scid(s), coef_map[s])
    end

    for (xhid, rho) in pairs(phd.r.penalties)
        @test st.tree_map[xhid.node].stage == PH.stid(1)
        @test isapprox(rho, c_rho) || isapprox(rho, phd.r.constant)
    end
end

@testset "SEP Penalty" begin
    @test PH.is_initial_value_dependent(PH.SEPPenaltyParameter) == true
    @test PH.is_subproblem_dependent(PH.SEPPenaltyParameter) == true
    @test PH.is_variable_dependent(PH.SEPPenaltyParameter) == true

    N = 2
    st = PH.two_stage_tree(N)
    var_map = Dict{PH.ScenarioID,Dict{PH.VariableID,PH.VariableInfo}}()
    var_vals = Dict{Int,Dict{PH.VariableID,Float64}}()
    coef_map = Dict{Int,Dict{PH.VariableID,Float64}}()
    for s in 0:(N-1)
        js = two_stage_model(PH.scid(s))
        svm = PH.report_variable_info(js, st)
        var_map[PH.scid(s)] = svm
        (br_vars, lf_vars) = PH._split_variables(st, collect(keys(svm)))
        PH.solve_subproblem(js)
        var_vals[s] = PH.report_values(js, collect(keys(svm)))
        coef_map[s] = PH.report_penalty_info(js,
                                             br_vars,
                                             PH.SEPPenaltyParameter
                                             )
    end

    phd = PH.PHData(PH.SEPPenaltyParameter(),
                    st,
                    Dict(1 => Set([PH.scid(0), PH.scid(1)])),
                    var_map,
                    TimerOutputs.TimerOutput()
                    )

    for s in 0:(N-1)
        PH._copy_values(phd.scenario_map[PH.scid(s)].branch_vars, var_vals[s])
        PH.process_penalty_subproblem(phd.r, phd, PH.scid(s), coef_map[s])
    end

    PH.compute_and_save_xhat(phd)
    PH.process_penalty_initial_value(phd.r, phd)

    value = 0.5 * (10.0 + 3.0)
    value = 0.5 * (10.0 - value) + 0.5 * (value - 3.0)
    rho_val = 1.0 / value
    for (xhid, rho) in pairs(phd.r.penalties)
        @test st.tree_map[xhid.node].stage == PH.stid(1)
        @test isapprox(rho, rho_val) || isapprox(rho, phd.r.default)
    end

    if isdefined(Main, :Xpress)
        N = 8
        st = PH.two_stage_tree(N)
        var_map = Dict{PH.ScenarioID,Dict{PH.VariableID,PH.VariableInfo}}()
        var_vals = Dict{Int,Dict{PH.VariableID,Float64}}()
        coef_map = Dict{Int,Dict{PH.VariableID,Float64}}()
        for s in 0:(N-1)
            js = two_stage_int(PH.scid(s))
            svm = PH.report_variable_info(js, st)
            var_map[PH.scid(s)] = svm
            (br_vars, lf_vars) = PH._split_variables(st, collect(keys(svm)))
            PH.solve_subproblem(js)
            var_vals[s] = PH.report_values(js, collect(keys(svm)))
            coef_map[s] = PH.report_penalty_info(js,
                                                 br_vars,
                                                 PH.SEPPenaltyParameter
                                                 )
        end

        phd = PH.PHData(PH.SEPPenaltyParameter(),
                        st,
                        Dict(1 => Set([PH.scid(s) for s in 0:(N-1)])),
                        var_map,
                        TimerOutputs.TimerOutput()
                        )

        for s in 0:(N-1)
            PH._copy_values(phd.scenario_map[PH.scid(s)].branch_vars, var_vals[s])
            PH.process_penalty_subproblem(phd.r, phd, PH.scid(s), coef_map[s])
        end

        PH.compute_and_save_xhat(phd)
        PH.process_penalty_initial_value(phd.r, phd)

        value = (10 - 4 + 1)
        rho_val = 2.0 / value
        for (xhid, rho) in pairs(phd.r.penalties)
            @test st.tree_map[xhid.node].stage == PH.stid(1)
            @test isapprox(rho, rho_val)
        end
    else
        println("Skipping SEP integer setting tests.")
    end
end

@testset "Reset penalty parameters" begin
    for r in (PH.ProportionalPenaltyParameter(1.0), PH.SEPPenaltyParameter(1.0))
        @test isempty(r.penalties) == true
        r.penalties[PH.XhatID(PH.NodeID(0), PH.Index(0))] = 1.0
        @test isempty(r.penalties) == false
        PH.reset_penalty_map(r)
        @test isempty(r.penalties) == true
    end
end


#####
#
# Example script for TintiNet.jl
# Model IGBT (InceptGOR-Bert): 3-state SS Regressor
#
# To run training, revise your code and uncomment the last line.
# On slower GPUs, an epoch may take several minutes with the default configuration.
# Adjust model size using the `hyperpar` and `config` dictionaries.
#
#####

#####
# 0. Loading Packages and defining Hyperparameters
#####

using CUDA, Zygote, Flux, Transformers, JSON, BSON, Random, Dates, Printf
using IterTools: ncycle;
using Transformers.Basic, Transformers.Pretrain, Transformers.Datasets, Transformers.Stacks, Transformers.BidirectionalEncoder
using BSON: @load, @save
using Flux: @functor, onehot, onecold, mse, shuffle, sigmoid, gradient, unsqueeze
using Flux.Optimise: update!
using Statistics: mean, std

using TintiNet
dformat = Dates.DateFormat("yyyy-m-d-HH-MM-SS")
Random.seed!(102528)

##
# Algorithm Hyperparameters

hyperpar = Dict(
    "max_seq_len"       => 128,     # Maximum sequence length on the model (default 192)
    "batch_size"        => 128,     # Batch size for training (default 128)
    "number_of_epochs"  => 200,     # Number of epochs for training (10 for the example file)
    "learning_rate"     => 1e-3,    # Learning rate for optimizer (default 0.001)
    "smooth_epsilon"    => 1e-6     # Small value to smooth onehot result before logcrossentropy (default 1e-6)
)

##
# Model Configuration

const config = Dict(
  "input_vocab_size"                => 22,      # Size of the SEQUENCE/FASTA vocabulary, usually 20 aminoacids + "X" + "-"
  "output_size"                     => 5,       # Output size. Usually 5 due to sin/cos phi, sin/cos psi and acc.
  "embed_dimension"                 => 128,     # Size of the individual Embed vectors
  "bert_num_hidden_layers"          => 2,       # Number of consecutive Transformer layers on the BERT encoder
  "bert_transformer_layer_heads"    => 8,       # Number of heads in each Transformer layer
  "bert_transformer_layer_hidden"   => 128,     # Hidden size of the positionwise linear layer in each Transformer layer
  "bert_transformer_dropout_prob"   => 0.1,     # Dropout probability for the Transformer layers
  "bert_transformer_activation"     => gelu,    # Activation function for each Transformer layer (default: gelu)
  "detection_num_filters"           => 32       # Number of filters in each convolutional decoder layer
)

#####
# 2. Obtaining and shaping the data
#####

## 2.1. Training sequences, structures and lengths

remapnan(x) = x == nothing ? missing : x

training_set_datapath = "../data/fold01training_dataset.json"

open(training_set_datapath, "r") do file

    entries = JSON.parse(file)
    num_entries = length(entries)

    parsed_sequences = Vector{Vector{String}}(undef, num_entries)
    parsed_phis = Vector{Vector{Union{Missing, Float64}}}(undef, num_entries)
    parsed_psis = Vector{Vector{Union{Missing, Float64}}}(undef, num_entries)
    parsed_accs = Vector{Vector{Union{Missing, Float64}}}(undef, num_entries)
    parsed_lengths = Vector{Int}(undef, num_entries)

    for (i,j) in enumerate(entries)
        parsed_sequences[i] = string.(j["fasta_seq"])
        parsed_phis[i] = map(remapnan, j["dssp_phi"])
        parsed_psis[i] = map(remapnan, j["dssp_psi"])
        parsed_accs[i] = map(remapnan, j["dssp_acc"])
        parsed_lengths[i] = length(parsed_sequences[i])
    end

    order_idx = sortperm(parsed_lengths)

    global train_sequences_ordered = parsed_sequences[order_idx]
    global train_phis_ordered = parsed_phis[order_idx]
    global train_psis_ordered = parsed_psis[order_idx]
    global train_accs_ordered = parsed_accs[order_idx]
    global train_lengths_ordered = parsed_lengths[order_idx]

end

## 2.2. Testing sequences, structures and lengths

testing_set_datapath = "../data/fold01testing_dataset.json"

open(testing_set_datapath, "r") do file

    entries = JSON.parse(file)
    num_entries = length(entries)

    parsed_sequences = Vector{Vector{String}}(undef, num_entries)
    parsed_phis = Vector{Vector{Union{Missing, Float64}}}(undef, num_entries)
    parsed_psis = Vector{Vector{Union{Missing, Float64}}}(undef, num_entries)
    parsed_accs = Vector{Vector{Union{Missing, Float64}}}(undef, num_entries)
    parsed_lengths = Vector{Int}(undef, num_entries)

    for (i,j) in enumerate(entries)
        parsed_sequences[i] = string.(j["fasta_seq"])
        parsed_phis[i] = map(remapnan, j["dssp_phi"])
        parsed_psis[i] = map(remapnan, j["dssp_psi"])
        parsed_accs[i] = map(remapnan, j["dssp_acc"])
        parsed_lengths[i] = length(parsed_sequences[i])
    end

    order_idx = sortperm(parsed_lengths)

    global test_sequences_ordered = parsed_sequences[order_idx]
    global test_phis_ordered = parsed_phis[order_idx]
    global test_psis_ordered = parsed_psis[order_idx]
    global test_accs_ordered = parsed_accs[order_idx]
    global test_lengths_ordered = parsed_lengths[order_idx]

end

## 2.3. Lazy DataLoaders

training_dataset = Flux.Data.DataLoader((train_sequences_ordered, train_phis_ordered, train_psis_ordered, train_accs_ordered, train_lengths_ordered); batchsize=hyperpar["batch_size"], shuffle=true)
eval_train_dataset = Flux.Data.DataLoader((train_sequences_ordered, train_phis_ordered, train_psis_ordered, train_accs_ordered, train_lengths_ordered); batchsize=64,shuffle=false)
eval_test_dataset = Flux.Data.DataLoader((test_sequences_ordered, test_phis_ordered, test_psis_ordered, test_accs_ordered, test_lengths_ordered); batchsize=64,shuffle=false)

#####
# 3. Network Structure
#####

model = Stack(
        @nntopo(x → enc → enc → enc → inc → inc → inc → inc → inc → inc → inc → inc → inc → inc → pe:(inc, pe) → bert_input → bert_output → bert_output → bert_output → y → y → y → y → y),
         # Embedding Layers
        Embed(config["embed_dimension"], config["input_vocab_size"]),
        unsqueeze(3),
        x -> permutedims(x, [3,2,1,4]),
        # Processing Layers
        Inception8(config["embed_dimension"]),
        Dropout(config["bert_transformer_dropout_prob"]),
        Inception8(config["embed_dimension"]),
        Dropout(config["bert_transformer_dropout_prob"]),
        Inception8(config["embed_dimension"]),
        Dropout(config["bert_transformer_dropout_prob"]),
        Inception8(config["embed_dimension"]),
        Dropout(config["bert_transformer_dropout_prob"]),
        x -> permutedims(x, [3,2,1,4]),
        x -> reshape(x, size(x,1), size(x,2), size(x,4)),
        PositionEmbedding(config["embed_dimension"], hyperpar["max_seq_len"]; trainable = false),
        (e, pe) -> e .+ pe,
        # Processing Layers
        Bert(
            config["embed_dimension"],
            config["bert_transformer_layer_heads"],
            config["bert_transformer_layer_hidden"],
            config["bert_num_hidden_layers"];
            act = config["bert_transformer_activation"],
            pdrop = config["bert_transformer_dropout_prob"],
            attn_pdrop = config["bert_transformer_dropout_prob"]
        ),
        # Detection Layers
        unsqueeze(3),
        bert_output -> permutedims(bert_output, [3,2,1,4]),
        Conv((1,7), config["embed_dimension"] => config["detection_num_filters"], relu; pad = (0,3)),
        Conv((1,7), config["detection_num_filters"] => config["detection_num_filters"], relu; pad = (0,3)),
        Conv((1,7), config["detection_num_filters"] => config["output_size"], tanh; pad = (0,3)),
        x -> permutedims(x, [3,2,1,4]),
        x -> reshape(x, size(x,1), size(x,2), size(x,4)),
    ) |> gpu

#####
# 4. Model training
#####

opt = Flux.Optimise.ADAM(hyperpar["learning_rate"])     # Initializes the ADAM optimizer
ps = params(model)      # Collects model parameters

function regressor_train_loss(model, x, y, mask)
    sum( ( (model(x) .- y) .^ 2 ) .* mask ) / sum(mask)
end

function regressor_eval_loss(predictions, y, mask)
    sum( ( (predictions .- y) .^ 2 ) .* mask ) / sum(mask)
end

function build_regressor_target_array(
    phis::Vector{Vector{Union{Missing, Float64}}},
    psis::Vector{Vector{Union{Missing, Float64}}},
    accs::Vector{Vector{Union{Missing, Float64}}},
    seqlen::Int64)

    y = zeros(5, seqlen, length(phis))

    for seq_i in 1:length(phis)
        for aa_j in 1:length(phis[seq_i])
            if !ismissing(phis[seq_i][aa_j])
                y[1, aa_j, seq_i] = sind(phis[seq_i][aa_j])
                y[2, aa_j, seq_i] = cosd(phis[seq_i][aa_j])
            end

            if !ismissing(psis[seq_i][aa_j])
                y[3, aa_j, seq_i] = sind(psis[seq_i][aa_j])
                y[4, aa_j, seq_i] = cosd(psis[seq_i][aa_j])
            end

            if !ismissing(accs[seq_i][aa_j])
                y[5, aa_j, seq_i] = (accs[seq_i][aa_j] - 100) / 100
            end

        end
    end

    return(y)

end

function prepare_regressor_mask(batchofseqs::Vector{Vector{String}},
                                phis::Vector{Vector{Union{Missing, Float64}}},
                                psis::Vector{Vector{Union{Missing, Float64}}},
                                accs::Vector{Vector{Union{Missing, Float64}}},
                                seqlen::Int64)

    mask = getmask(preprocess(batchofseqs, seqlen)) # Uses getmask() from Transformers.Basic

    for i in 1:length(batchofseqs)
        for j in 1:length(batchofseqs[i])
            if ( ismissing(phis[i][j]) | ismissing(psis[i][j]) | ismissing(accs[i][j]) )
                mask[1, j, i] = 0.0
            end
        end
    end

    return(mask)

end

#define training loop
function train!(hyperpar, training_dataset,
                model, ps, opt,
                train_loss=regressor_train_loss, eval_loss=regressor_eval_loss,
                eval_train_dataset=eval_train_dataset, eval_test_dataset=eval_test_dataset,
                fastaVocab=fastaVocab)
    @info "start training"
    for epoch in 1:hyperpar["number_of_epochs"]
    i = 0
        for batch in training_dataset
            i += 1
            x, phi, psi, acc, len = batch

            train_x_processed = fastaVocab(preprocess(x, hyperpar["max_seq_len"])) |> gpu
            train_y_processed = build_regressor_target_array(phi, psi, acc, hyperpar["max_seq_len"]) |> gpu
            train_mask_processed = prepare_regressor_mask(x, phi, psi, acc, hyperpar["max_seq_len"]) |> gpu

#            grad = gradient(()->train_loss(model, train_x_processed, train_y_processed, train_mask_processed), ps)

            l, back = Flux.pullback(ps) do
                train_loss(model, train_x_processed, train_y_processed, train_mask_processed)
            end
            grad = back(Flux.Zygote.sensitivity(l))

            update!(opt, ps, grad)

            if i%64 == 0

                Flux.testmode!(model)

                let
                batch_loss_ideal_train_set = 0.0
                batch_loss_ideal_test_set = 0.0

                for eval_train_batch in eval_train_dataset

                    eval_x, eval_phi, eval_psi, eval_acc, eval_len = eval_train_batch

                    eval_x_processed = fastaVocab(preprocess(x, hyperpar["max_seq_len"])) |> gpu
                    eval_y_processed = build_regressor_target_array(phi, psi, acc, hyperpar["max_seq_len"]) |> gpu
                    eval_mask_processed = prepare_regressor_mask(x, phi, psi, acc, hyperpar["max_seq_len"]) |> gpu

                    network_outputs = model(eval_x_processed)
                    batch_loss_ideal_train_set += eval_loss(network_outputs, eval_y_processed, eval_mask_processed)

                end #end IDEAL_TRAIN_EVAL loop

                for eval_test_batch in eval_test_dataset

                    eval_x, eval_phi, eval_psi, eval_acc, eval_len = eval_test_batch

                    eval_x_processed = fastaVocab(preprocess(x, hyperpar["max_seq_len"])) |> gpu
                    eval_y_processed = build_regressor_target_array(phi, psi, acc, hyperpar["max_seq_len"]) |> gpu
                    eval_mask_processed = prepare_regressor_mask(x, phi, psi, acc, hyperpar["max_seq_len"]) |> gpu

                    network_outputs = model(eval_x_processed)
                    batch_loss_ideal_test_set += eval_loss(network_outputs, eval_y_processed, eval_mask_processed)

                end #end IDEAL_TEST_EVAL loop

            ideal_train_loss = batch_loss_ideal_train_set
            ideal_test_loss = batch_loss_ideal_test_set

            Flux.trainmode!(model)

            checkpoint_model = cpu(model)
            @save "TintiNet-IGBT-Regressor-Checkpoint-$(Dates.format(Dates.now(), dformat)).bson" checkpoint_model opt
            @printf("Epoch: %d\t Batch: %d\t IdTrLos: %2.6f\t IdTsLoss: %2.6f\t Time: %s\n",
                    epoch, i, ideal_train_loss, ideal_test_loss, Dates.format(now(), "yyyy-mm-dd HH:MM:SS"))

            end # end let

            end # endif

        end # end for batch

    end # end for epoch

end # end function train!()

#####
# Uncomment next line to allow the script to undergo training.
# train!(hyperpar, training_dataset, model, ps, opt)


# some implementations for convenience
# maintained by Zach Sunberg

actions{S}(mdp::MDP{S,Bool}) = (true, false)
actions{S}(mdp::MDP{S,Bool}, s::S, t::Tuple{Bool,Bool}=(true,false)) = (true,false)
actions{S,O}(mdp::POMDP{S,Bool,O}) = (true, false)
actions{S,O}(mdp::POMDP{S,Bool,O}, s::S, t::Tuple{Bool,Bool}=(true,false)) = (true,false)
actions{S,O}(mdp::RPOMDP{S,Bool,O}) = (true, false)
actions{S,O}(mdp::RPOMDP{S,Bool,O}, s::S, t::Tuple{Bool,Bool}=(true,false)) = (true,false)
n_actions{S}(mdp::MDP{S,Bool}) = 2
n_actions{S,O}(mdp::POMDP{S,Bool,O}) = 2
n_actions{S,O}(mdp::RPOMDP{S,Bool,O}) = 2

actions{S}(mdp::MDP{S,Int}, s::S, r::Range) = actions(mdp)

rand(rng::AbstractRNG, t::Tuple{Bool, Bool}) = rand(rng, Bool)
rand(t::Tuple{Bool, Bool}) = rand(Bool)

iterator(s::AbstractVector) = s
iterator(s::Tuple) = s
iterator(r::Range) = r
iterator(g::Base.Generator) = g

states(mdp::MDP{Bool}) = (true, false)
states(mdp::POMDP{Bool}) = (true, false)
states(mdp::RPOMDP{Bool}) = (true, false)
n_states(mdp::MDP{Bool}) = 2
n_states(mdp::POMDP{Bool}) = 2
n_states(mdp::RPOMDP{Bool}) = 2

observations{S,A}(::POMDP{S,A,Bool}) = (true,false)
observations{S,A}(::POMDP{S,A,Bool}, s::S) = (true,false)
observations{S,A}(::RPOMDP{S,A,Bool}) = (true,false)
observations{S,A}(::RPOMDP{S,A,Bool}, s::S) = (true,false)
n_observations{S,A}(::POMDP{S,A,Bool}) = 2
n_observations{S,A}(::RPOMDP{S,A,Bool}) = 2

state_index(mdp::Union{MDP, POMDP, RPOMDP}, s::Int) = s
action_index(mdp::Union{MDP, POMDP, RPOMDP}, a::Int) = a
obs_index(mdp::Union{MDP, POMDP, RPOMDP}, o::Int) = o

state_index(mdp::Union{MDP, POMDP, RPOMDP}, s::Bool) = s+1
action_index(mdp::Union{MDP, POMDP, RPOMDP}, a::Bool) = a+1
obs_index(mdp::Union{MDP, POMDP, RPOMDP}, o::Bool) = o+1