text:
#= License
Copyright 2019, 2020, 2021 (c) Yossi Bokor

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
=#

__precompile__()

module Skyler

#### Requirements ####
using CSV
using DataFrames
using LinearAlgebra
using NearestNeighbors

#=
using Distributed
@everywhere using CSV
@everywhere using DataFrames
@everywhere using LinearAlgebra
@everywhere using NearestNeighbors
=#

using Distances
using Optim
using SimpleGraphs
using Random
using Distributions
using QuadGK
using Plots
using JLD
using SpecialFunctions
using SharedArrays
using Calculus
using NNlib
using FiniteDiff
using ProgressMeter

#### Exports ####

export 	skyler,
		unittest,
		Determine_Dimension,
		Plot_PPC,
		Plot_Graph,
		Load_Example,
		PartitionedPointCloud,
		StratifiedSpace,
		EmbeddedStratifiedSpace,
		Dimension_Split,
		Create_Partition_DataFrame

#### Structures ####

mutable struct PartitionedPointCloud
	Points::AbstractArray
	Strata::Dict{Int, Set}
	Dimensions::Dict{Int, Set}
	Boundaries::AbstractArray
end

mutable struct StratifiedSpace
	StrataDimensions::Dict{Int, Set}
	Boundaries::AbstractArray
end

mutable struct EmbeddedStratifiedSpace
	StratSpace::StratifiedSpace
	Equations::Dict{Int, Array}
end

#### Background Functions ####

function Spherical_Shell(ball_tree::BallTree, point::Array, inner::T, outer::U) where {T<: Number, U<: Number} #this obtains all the samples in the spherical shell around a chosen sample 'point'
	ball_1 = inrange(ball_tree, point, inner, true)
	ball_2 = inrange(ball_tree, point, outer, true)
	annulus = setdiff(ball_2, ball_1)
	return annulus
end

function Determine_Dimension(sample::A, ball_tree::BallTree, index::Int, sample_epsilon::P, radius::R) where {P<:Number,R<:Number, A<:Union{AbstractArray, LinearAlgebra.Adjoint}}# determine whether the `index'th point in the sample looks 0 or 1 dimensional given the coordinates of the other points, a pre-generated ball tree, and the correct algorithm parameters of 'angle_conditon', 'inner_radius', and 'outer_radius'
	point = sample[:,index]
	distance_matrix = pairwise(Euclidean(), sample, dims=2)
	node_neighbours = Spherical_Shell(ball_tree, point, 0., radius+sample_epsilon)
	node_neighbours = push!(node_neighbours, index)
	G_point = IntGraph()
	for i in node_neighbours
		add!(G_point,i)
	end
	for i in vlist(G_point)
		for j in vlist(G_point)
			if distance_matrix[i,j] <= 2*sample_epsilon
				add!(G_point, i, j)
			end
		end
	end
	n_p = size(sample,2) #number of samples
	nc_1 = num_components(G_point)
	if nc_1 != 1 #if the graph on the points within a ball of the chosen points is disconnected, then the point is not close to a vertex
		return 1
	else # if the graph on the points in a ball is connected, then we need to check if 0 or 1 dimensional structure
		node_neighbours_1 = Spherical_Shell(ball_tree, point, radius-1sample_epsilon, radius+sample_epsilon)
		G_point_1 = IntGraph() # construct graph on points in a spherical shell
		for i in node_neighbours_1
			add!(G_point_1, i)
		end
		for i in vlist(G_point_1)
			for j in vlist(G_point_1)
				if distance_matrix[i,j] <= 3*sample_epsilon
					add!(G_point_1, i, j)
				end
			end
		end
		nc_2 =num_components(G_point_1)
		if nc_2 == 2 #if the number of connected components is 2, we need to check if they are roughly opposite each other or not
			cc = collect(SimpleGraphs.components(G_point_1))
			mid_point_1 = []
			for l in 1:size(sample,1)
				mid_point_1_l = 0
				cc_1 = collect(cc[1])
				for p in cc_1
					mid_point_1_l += sample[l, p]
				end
			append!(mid_point_1, mid_point_1_l/(length(cc_1)))
			end
			mid_point_2= []
			for l in 1:size(sample,1)
				mid_point_2_l = 0
				cc_2 = collect(cc[2])
				for p in cc_2
					mid_point_2_l += sample[l,p]
				end
			append!(mid_point_2, mid_point_2_l/(length(cc_2)))
			end
			mid_point_1 = mid_point_1 .- sample[:,index]
			mid_point_2 = mid_point_2 .- sample[:,index]
			#cos_of_angle = dot(mid_point_1, mid_point_2)/(norm(mid_point_1)*norm(mid_point_2))
			if dot(mid_point_1, mid_point_2) <= -radius^2 +2*radius*sample_epsilon +7*sample_epsilon^2
				return 1
			else
				return 0
			end
		else #if the number of connected components is not 2, return dimension 0
			return 0
		end
	end
end

function Dimension_Split(sample::A, sample_epsilon::S, radius::U) where {S<: Number, U<:Number, A<:AbstractArray}# for each point in the sample, determine if it looks 0 or 1 dimensional
	#if nprocs() == 1
		n_p = size(sample,2)
		ball_tree = BallTree(sample)
		dimension_split = Dict{Int,Set}(0 => Set(), 1 =>Set())
		println("Calculating local structure for each sample: ")
		@showprogress for i in 1:n_p
			dim = Determine_Dimension(sample, ball_tree, i, sample_epsilon,radius)
			dimension_split[dim]=push!(dimension_split[dim], i)
		end
		return dimension_split
	#else
	#	n_p = size(sample,2)
	#	@everywhere ball_tree=BallTree(sample)
	#end
end

function Find_Groups(distance_matrix::A, dictionary_of_points::Dict{Int, Set}, dimension::T, connection_threshold::W) where {T <: Number, W<: Number, A<:AbstractArray} # obtain the clusters of points which look 'dimensional'  dimensional
	points = collect(dictionary_of_points[dimension]) #create a set of points
	G = SimpleGraph()
	for i in points
		add!(G, i)
	end
	for i in points
		for j in points
			if distance_matrix[i,j] <= connection_threshold #connect points below the connection threshold to cluster
				add!(G, i, j)
			end
		end
	end
	components = []
	vertices = vlist(G)
	n_v = length(vertices)
	checked = Set()
	discovered=[]
	function DFS(Graph, vertex)
		append!(discovered,vertex)
		for u in neighbors(Graph, vertex)
			if u in(discovered)
			else
				DFS(Graph, u)
			end
		end
	end
	for i in vertices
		if i in checked
		else
			push!(checked, i)
			discovered =[]
			DFS(G, i)
			append!(components, [discovered])
			for n in discovered
				push!(checked, n)
			end
		end
	end
	return components
end

function Generate_Strata_Assignments(distance_matrix::Array, dimension_split::Dict{Int, Set}, v_threshold::T, e_threshold::W) where {T <: Number, W<:Number} # generate the clusters of vertices and edges using the above functions
	v_comps = Find_Groups(distance_matrix, dimension_split, 0, v_threshold)
	number_of_vertices = length(v_comps)
	e_comps = Find_Groups(distance_matrix, dimension_split, 1, e_threshold)
	number_of_edges = length(e_comps)
	change = []
	for i in 1:number_of_edges
		if size(e_comps[i]) == 1
			append!(change, e_comps[i][1])
		end
	end
	if length(change) !=0
		d_s = copy(dimension_split)
		for i in 1:length(change)
			d_s[0] = push!(d_s[0], change[i])
			d_s[1] = setdiff(d_s[1], Set([change[i]]))
		end
		v_comps = Find_Groups(distance_matrix, d_s, 0, v_threshold)
		number_of_vertices = length(v_comps)
		e_comps = Find_Groups(distance_matrix, d_s, 1, e_threshold)
		number_of_edges = length(e_comps)
	end
	strata_assignments = Dict{Int, Set}(1=>Set())
	strata_dimensions = Dict{Int,Set}(0=> Set(), 1=>Set())
	#println("Detected ", number_of_vertices," vertices and ", number_of_edges, " edges.")
	for i in 1:number_of_vertices
		strata_assignments[i] =Set(v_comps[i])
		strata_dimensions[0] =push!(strata_dimensions[0],i)
	end
	for i in 1:number_of_edges
		strata_assignments[number_of_vertices+i] =Set(e_comps[i])
		strata_dimensions[1] =push!(strata_dimensions[1],number_of_vertices+i)
	end
	return strata_assignments, strata_dimensions
end

function Generate_Abstract_Structure(points::A, dimension_split::Dict{Int, Set}, sample_epsilon::T, vertex_threshold::T, edge_threshold::U) where {U<:Number, T<:Number, A<:AbstractArray} # obtain the abstract structure of the graph
	dists = pairwise(Euclidean(), points, dims=2)
	corrected = false
	while corrected == false
		strata_assignments, strata_dimensions = Generate_Strata_Assignments(dists, dimension_split, vertex_threshold, edge_threshold)
		max_dim = maximum(keys(strata_dimensions))
		n_s = maximum(keys(strata_assignments))
		boundaries = zeros(n_s, n_s)
		@showprogress for i in 0:max_dim-1
			for j in strata_dimensions[i]
				s_j = points[:,collect(strata_assignments[j])]
				for k in strata_dimensions[i+1]
					s_k = points[:,collect(strata_assignments[k])]
					d_j_k = Distances.pairwise(Euclidean(), s_j, s_k, dims=2)
					if minimum(d_j_k) <= 3*sample_epsilon
						boundaries[j,k]=1
					end
				end
			end
		end
		println("\r The strata are ", strata_dimensions, " and the boundary matrix is  ", boundaries,".")
		change = []
		if isempty(strata_dimensions[1])
			corrected = true
			StratSpace = StratifiedSpace(strata_dimensions, boundaries)
			return StratSpace, strata_assignments
			break
		else
			for i in strata_dimensions[1]
				if sum(boundaries[:,i]) !=2
					append!(change, i)
				end
			end
		end
		if length(change) == 0
			corrected = true
			StratSpace = StratifiedSpace(strata_dimensions, boundaries)
			return StratSpace, strata_assignments
			break
		else
			println("Stray edge detected, addressing it.")
			for k in 1:length(change)
				dimension_split[0] = union(dimension_split[0], strata_assignments[change[k]])
				dimension_split[1] = setdiff(dimension_split[1], strata_assignments[change[k]])
			end
		end
	end
end

function Partition_Point_Cloud(points::A, sample_epsilon::P, radius::W, vertex_threshold::T, edge_threshold::U)::PartitionedPointCloud where {P<:Number, Q<:Number, R<:Number, W<:Number, T <: Number, U<:Number, A<:AbstractArray}
	dimension = size(points,1)
	distances = pairwise(Euclidean(), points, dims=2)
	dim_split = Dimension_Split(points, sample_epsilon, radius) #partition the sample into 0 or 1 dimensional using a dictionary
	#strata_assignments, strata_dimensions = Generate_Strata_Assignments(distances, dim_split, vertex_threshold, edge_threshold) #find the strata for the partioned point cloud
	StratSpace, strata_assignments = Generate_Abstract_Structure(points, dim_split, sample_epsilon, vertex_threshold, edge_threshold) #obtain the boundary relations
	change = []
	ppc = PartitionedPointCloud(points, strata_assignments, StratSpace.StrataDimensions, StratSpace.Boundaries)
	return ppc
end
#=
Deprecated this was used for modelling the embedding, unsure about future uses
function Generate_Partition(v_clusters, e_clusters)
	lengths = []
	partitions = []
	for i in 1:length(v_clusters)
		append!(lengths, length(v_clusters[i]))
		append!(partitions, v_clusters[i])
	end
	for i in 1:length(e_clusters)
		append!(lengths, length(e_clusters[i]))
		append!(partitions, e_clusters[i])
	end
	return partitions, lengths
end
=#

function Create_Partition_DataFrame(ppc::PartitionedPointCloud) #create a dataframe we can save the partitioned point cloud in
	dimension = size(ppc.Points,1)
	df = DataFrame(ppc.Points[:,1]')
	insert!(df, 1, [1], :index)
	insert!(df, 2, [1], :dimension)
	deleterows!(df, 1)
	for i in 0:maximum(keys(ppc.Dimensions))
		for j in ppc.Dimensions[i]
			for k in ppc.Strata[j]
				l = [j i]
				for n in 1:dimension
					l=hcat(l, ppc.Points[n,k])
				end
				push!(df, l)
			end
		end
	end
	return sort(df)
end

function Plot_PPC(ppc::PartitionedPointCloud) #plot the partiton with nice colours
	Plots.pyplot()
	if size(ppc.Points,1) == 2
		df = DataFrame(index=Int[], dimension=Int[], c1=[], c2=[])
		for i in 0:maximum(keys(ppc.Dimensions))
			for j in ppc.Dimensions[i]
				for k in ppc.Strata[j]
					push!(df, [j,i,ppc.Points[1,k], ppc.Points[2,k]])
				end
			end
		end
		Plots.scatter( df[!,:c1],df[!,:c2], group=df[!,:index])
	else
		println("Points live in 3 or more dimensions, projecting onto the first 3.")
		df = DataFrame(index=Int[], dimension=Int[], c1=[], c2=[], c3=[])
		for i in 0:maximum(keys(ppc.Dimensions))
			for j in ppc.Dimensions[i]
				for k in ppc.Strata[j]
					push!(df, [j,i,ppc.Points[1,k], ppc.Points[2,k], ppc.Points[3,k]])
				end
			end
		end
		Plots.scatter3d( df[!,:c1],df[!,:c2], df[!,:c3], group=df[!,:index])
	end
end

function Plot_PPC(partition::DataFrame)
	Plots.pyplot()
	dimension = size(partition[!,:point][1],1)
	if dimension == 2
		points = Array{Float64}(undef, 0,2)
		for i in 1:size(partition,2)
			points = hcat(points, partition[i,:point])
		end
		Plots.scatter( points[:,1], points[:,2], group=partition[!,:index])
	else
		println("Points live in 3 or more dimensions, projecting onto the first 3.")
		points = Array{Float64}(undef, 0,3)
		for i in 1:size(partition,1)
			points = vcat(points, partition[i,:point][1:3]')
		end
		Plots.scatter3d( points[:,1], points[:,2], points[:,3], group=partition[!,:index])
	end
end

function Plot_PPC(partition::DataFrame, c_1::Int, c_2::Int)
	Plots.pyplot()
	println("Projecting on to the following coordinates: $c_1 and $c_2.")
	points = Array{Float64}(undef, 0,2)
	for i in 1:size(partition,1)
		points = vcat(points, [partition[i,:point][c_1] partition[i,:point][c_2]])
	end
	Plots.scatter( points[:,1], points[:,2], group=partition[!,:index])
end

function Plot_PPC(partition::DataFrame, c_1::Int, c_2::Int, c_3::Int)
	Plots.pyplot()
	println("Projecting on to the following coordinates: $c_1, $c_2 and $c_3.")
	points = Array{Float64}(undef, 0,3)
	for i in 1:size(partition,1)
		points = vcat(points, [partition[i,:point][c_1] partition[i,:point][c_2] partition[i,:point][c_3]])
	end
	Plots.scatter3d( points[:,1], points[:,2], points[:,3], group=partition[!,:index])
end

function Plot_Graph(vertex_locations, edges, sample) #plot a graph (either exact or modelled)
	Plots.pyplot()
	if length(sample) == 2
		scatter(vertex_locations[sample[1], :], vertex_locations[sample[2], :], s=1)
		for i in 1:(size(edges,1)-1)
			plot!([vertex_locations[sample[1], edges[i,1]], vertex_locations[sample[1],edges[i,2]]], [vertex_locations[sample[2], edges[i,1]], vertex_locations[sample[2], edges[i,2]]], linestyle = :solid)
		end
		i=size(edges,1)
		plot!([vertex_locations[sample[1], edges[i,1]], vertex_locations[sample[1],edges[i,2]]], [vertex_locations[sample[2], edges[i,1]], vertex_locations[sample[2], edges[i,2]]], linestyle = :solid)
	elseif length(sample) == 3
		scatter3d(vertex_locations[sample[1], :], vertex_locations[sample[2], :], vertex_locations[sample[3], :],s=1)
		for i in 1:(size(edges,1)-1)
			plot!([vertex_locations[sample[1], edges[i,1]], vertex_locations[sample[1], edges[i,2]]], [vertex_locations[sample[2], edges[i,1]], vertex_locations[sample[2], edges[i,2]]], [vertex_locations[sample[3], edges[i,1]], vertex_locations[sample[3], edges[i,2]]],  linestyle =:solid)
		end
		i=size(edges,1)
		plot!([vertex_locations[sample[1], edges[i,1]], vertex_locations[sample[1], edges[i,2]]], [vertex_locations[sample[2], edges[i,1]], vertex_locations[sample[2], edges[i,2]]], [vertex_locations[sample[3], edges[i,1]], vertex_locations[sample[3], edges[i,2]]],  linestyle =:solid)
	else
		println("Please specify 2 or 3 sample to project onto.")
		return []
	end
end

function Plot_Graph(vertex_locations, edges) #plot a graph (either exact or modelled)
	Plots.pyplot()
	if size(vertex_locations,1) == 2
		scatter(vertex_locations[sample[1], :], vertex_locations[sample[2], :], s=1)
		for i in 1:(size(edges,1)-1)
			plot!([vertex_locations[1, edges[i,1]], vertex_locations[1,edges[i,2]]], [vertex_locations[2, edges[i,1]], vertex_locations[2, edges[i,2]]], linestyle = :solid)
		end
		i=size(edges,1)
		plot!([vertex_locations[1, edges[i,1]], vertex_locations[1,edges[i,2]]], [vertex_locations[2, edges[i,1]], vertex_locations[2, edges[i,2]]], linestyle = :solid)
	elseif size(vertex_locations,1) == 3
		scatter3d(vertex_locations[1, :], vertex_locations[2, :], vertex_locations[3, :],s=1)
		for i in 1:(size(edges,1)-1)
			plot!([vertex_locations[1, edges[i,1]], vertex_locations[1, edges[i,2]]], [vertex_locations[2, edges[i,1]], vertex_locations[2, edges[i,2]]], [vertex_locations[3, edges[i,1]], vertex_locations[3, edges[i,2]]],  linestyle =:solid,  c=:red)
		end
		i=size(edges,1)
		plot!([vertex_locations[1, edges[i,1]], vertex_locations[1, edges[i,2]]], [vertex_locations[2, edges[i,1]], vertex_locations[2, edges[i,2]]], [vertex_locations[3, edges[i,1]], vertex_locations[3, edges[i,2]]],  linestyle =:solid,  c=:red)
	else
		println("Projecting onto the first 3 coordinates.")
		scatter3d(vertex_locations[1, :], vertex_locations[2, :], vertex_locations[3, :],s=1)
		for i in 1:(size(edges,1)-1)
			plot!([vertex_locations[1, edges[i,1]], vertex_locations[1, edges[i,2]]], [vertex_locations[2, edges[i,1]], vertex_locations[2, edges[i,2]]], [vertex_locations[3, edges[i,1]], vertex_locations[3, edges[i,2]]],  linestyle =:solid)
		end
		i=size(edges,1)
		plot!([vertex_locations[1, edges[i,1]], vertex_locations[1, edges[i,2]]], [vertex_locations[2, edges[i,1]], vertex_locations[2, edges[i,2]]], [vertex_locations[3, edges[i,1]], vertex_locations[3, edges[i,2]]],  linestyle =:solid)
	end
end

function model_fit(data, muT, sigma, E, S; EM_it = 3)

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

	#### Functions for modelling the underlying structure ####

	#=
	To update:
	-keeps sigma fixed at the moment, this can be made optional
	=#

	#=
	MAKE SURE S AND MUT ARE IN THE SAME ORDER
	=#

	#=
	notns:
	n		---	Number of data points
	d		---	Dimension of data
	N		---	Number of strata pieces
	N_0		---	Number of 0-dim strata pieces
	N_1		---	Number of 1-dim strata pieces
	data	---	n dim array of d arrays		---	Sample from embedded graph
	mu		---	N_0 by d dim array	---	Iniitial vertex sample locations
	sigma	---	N dim array			---	Initial error on each strata piece
	E		---	n by N array		---	Initial assignment for each data point
	S		---	N dim str array		---	Abstract graph vertex/edge names in that order
	=#

	# data is a n dim array of d dimensional vectors

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

	####pre-processing

	#get size of different dimensional strata

	#println("S starts as $S")

	#println("mu starts as $muT")

	N_0 	= size(muT,1)

	mu 		= [muT[i,:] for i in 1:N_0]

	N_1 	= size(S ,1) - N_0

	N 		= N_0 + N_1

	n  		= size(data,1)

	d 		= size(data[1],1)

	#load parameters

	E = [E[i,:] for i = 1:n]

	Pi      = 	sum(E)./n

	#println("sigma in EM is", sigma)

	MU		=	Dict{Any,Array}(S[i] => mu[i] for i = 1:N_0)

	SIGMA	=	Dict{Any,Float64}(S[i] => sigma[i] for i = 1:N)

	PI		=	Dict{Any,Float64}(S[i] => Pi[i] for i =1:N)

	#println("MU: ", MU)

	#println("SIGMA: ", SIGMA)

	#println("PI: ", PI)

	function reload_param(mu_t,sigma_t,Pi_t; S_t = S, N_0_t = N_0, N_t = N)

		#updates dictionary values for:
		#mu_t 		N_0 by d float array
		#SIGMA_t 	N float array
		#Pi 		N_float array

		MU_t	=  	Dict{Any,Array}(S_t[i] => mu_t[i] for i = 1:N_0_t)

		SIGMA_t	=  	Dict{Any,Float64}(S_t[i] => sigma_t[i] for i = 1:N_t)

		PI_t	=  	Dict{Any,Float64}(S_t[i] => Pi_t[i] for i = 1:N_t)

		return MU_t, SIGMA_t, PI_t

	end

	####density functions

	#0-dim strata

	function gaussian(x,v,s)

		linear = x - v

		return exp(-dot(linear, linear)/(2*s))

	end

	function p_e(x,v,s)

		d = size(x)[1]

		den = gaussian(x,v,s)/((2*pi*s)^(d/2))

		return den

	end

	function log_p_e(x,v,s)

		d = size(x)[1]

		linear = x - v

		den= -(d/2)*log(2*pi*s^2) - dot(linear,linear)/(2*s^2)

		return den

	end

	#1-dim strata

	#with clipping

	function p_l(xj, v1, v2, s)

		d = size(xj,1)

		xj = - xj

		den = (exp((4*dot(v1 - v2, (v1 + v2)/2 + xj)^2 - 4*norm(v1 - v2)^2*norm((v1 + v2)/2 + xj)^2)/(8*s^2*norm(v1 - v2)^2))*pi^(1/2 - d/2)*(-erf((2*dot(v1 - v2, (v1 + v2)/2 + xj) - norm(v1 - v2)^2)/(2*sqrt(2)*s*norm(v1 - v2))) + erf((2*dot(v1 - v2, (v1 + v2)/2 + xj) + norm(v1 - v2)^2)/(2*sqrt(2)*s*norm(v1 - v2)))))/norm(v1 - v2)

		if den <= -6

			den = -6

		end

  		return den

  	end

  	function log_p_l(Xj, V1, V2, s)

		d = size(Xj,1)

		Xj = - Xj

		den = log(2^((-1 - d)/2)*pi^(1/2 - d/2)*s^(1 - d)) +
				log(-erf((2*dot(V1 - V2, (V1 + V2)/2 + Xj) - norm(V1 - V2)^2)/(2*sqrt(2)*s*norm(V1 - V2))) +
			   	erf((2*dot(V1 - V2, (V1 + V2)/2 + Xj) + norm(V1 - V2)^2)/(2*sqrt(2)*s*norm(V1 - V2)))) -
			 	log(norm(V1 - V2)) + (4*dot(V1 - V2, (V1 + V2)/2 + Xj)^2 - 4*norm(V1 - V2)^2*norm((V1 + V2)/2 + Xj)^2)/
			  	(8*s^2*norm(V1 - V2)^2)

		if den <= -6

			den = -6

		end

  		return den

  	end

  	#density evaluation matrix

	function densities(MU_t,SIGMA_t; S_t = S, N_0 = N_0 , N_1 = N_1, N = N)

		zero_dim 	= [x -> p_e(x,MU_t[strat],SIGMA_t[strat]) for strat = S_t[1:N_0]]

		one_dim  	= [x -> p_l(x,MU_t[strat[2]],MU_t[strat[1]],SIGMA_t[strat]) for strat = S_t[N_0+1:N]]

		return one_dim # vcat(zero_dim,one_dim)

	end

	function log_densities(mu_t,SIGMA_t; S_t = S, N_0 = N_0 , N_1 = N_1, N = N)

		MU_t    	=   Dict{Any,Array}(S_t[i] => mu_t[i] for i = 1:N_0)

		zero_dim 	=   [x -> log_p_e(x,MU_t[strat],SIGMA_t[strat]) for strat = S_t[1:N_0]]

		one_dim  	=   [x -> log_p_l(x,MU_t[strat[1]],MU_t[strat[2]],SIGMA_t[strat]) for strat = S_t[N_0+1:N]]

		return  vcat(zero_dim,one_dim)

	end

	function density_matrix(data,density_functions)

		dm 			= 	[ (broadcast(rho, data))  for rho in density_functions ]

		dm 			=  	hcat(dm...)

		dm 			= 	[dm[i,:] for i = 1:n]

		return dm

	end

	####cost construction

	#expectation matrix for current parameters

	function expectation(data,MU_t,SIGMA_t,Pi_t)

		DM = density_matrix(data,log_densities(MU_t,SIGMA_t))

		logpi = log.(Pi_t)

		E = broadcast(t -> logpi .+ t, DM)

		E = broadcast(t -> softmax(t), E)

		return E

	end

	#cost for current vertex values

	function Cost(data,mu0,SIGMA_t, Pi, E)

		mu_t = [mu0[i,:] for i in 1:N_0]

		log_DM = density_matrix(data,log_densities(mu_t,SIGMA_t))

		#not needed in optimisation but can be checked for EM
		#log_pi = log.(Pi)
		#log_sum = broadcast(t -> t + log_pi, log_DM)

		w_sum = broadcast((e,t) -> sum(e.*t), E, log_DM)

		return sum(w_sum)/n

	end

	####analytic gradient construction

	#gradient values for cost function

	function grad_log_p_l(k, Xj, V1, V2, s)

		xj = Xj[k]

		v1 = V1[k]

		v2 = V2[k]

		if dot(V2 - V1, Xj) > dot(V1, V2 - V1)

			xj = xj - ((dot(V2 - V1, Xj) - dot(V1, V2 - V1) )/(norm(V2 - V1)^2))*(v2 - v1)

			Xj = Xj - ((dot(V2 - V1, Xj) - dot(V1, V2 - V1) )/(norm(V2 - V1)^2)).*(V2 - V1)

		elseif dot(V2 - V1, Xj) < dot(V2, V1 - V2)

			xj = xj - ((dot(V2 - V1, Xj) -  dot(V2, V1 - V2) )/(norm(V2 - V1)^2))*(v2 - v1)

			Xj = Xj - ((dot(V2 - V1, Xj) -  dot(V2, V1 - V2) )/(norm(V2 - V1)^2)).*(V2 - V1)

		end

		g_v1 = -(v1 + v2 + 2*xj)/(4*s^2) + (2*(-1 + exp(dot(V1 - V2, (V1 + V2)/2 + Xj)/s^2))*(v1 - v2)*
			dot(V1 - V2, (V1 + V2)/2 + Xj) + (3*v1 - v2 + 2*xj - exp(dot(V1 - V2, (V1 + V2)/2 + Xj)/s^2)*
		  	(v1 + v2 + 2*xj))*norm(V1 - V2)^2)/
			(exp((2*dot(V1 - V2, (V1 + V2)/2 + Xj) + norm(V1 - V2)^2)^2/(8*s^2*norm(V1 - V2)^2))*sqrt(2*pi)*s*
			(erf((-2*dot(V1 - V2, (V1 + V2)/2 + Xj) + norm(V1 - V2)^2)/(2*sqrt(2)*s*norm(V1 - V2))) +
			erf((2*dot(V1 - V2, (V1 + V2)/2 + Xj) + norm(V1 - V2)^2)/(2*sqrt(2)*s*norm(V1 - V2))))*
			norm(V1 - V2)^3) + ((-v1 + v2)*dot(V1 - V2, (V1 + V2)/2 + Xj)^2 +
			(s^2*(-v1 + v2) + (v1 + xj)*dot(V1 - V2, (V1 + V2)/2 + Xj))*norm(V1 - V2)^2)/(s^2*norm(V1 - V2)^4)

		g_v2 = -(v1 + v2 + 2*xj)/(4*s^2) + (-2*(-1 + exp(dot(V1 - V2, (V1 + V2)/2 + Xj)/s^2))*(v1 - v2)*
			dot(V1 - V2, (V1 + V2)/2 + Xj) + (-v1 - v2 - 2*xj + exp(dot(V1 - V2, (V1 + V2)/2 + Xj)/s^2)*
	 		(-v1 + 3*v2 + 2*xj))*norm(V1 - V2)^2)/
			(exp((2*dot(V1 - V2, (V1 + V2)/2 + Xj) + norm(V1 - V2)^2)^2/(8*s^2*norm(V1 - V2)^2))*sqrt(2*pi)*s*
			(erf((-2*dot(V1 - V2, (V1 + V2)/2 + Xj) + norm(V1 - V2)^2)/(2*sqrt(2)*s*norm(V1 - V2))) +
			erf((2*dot(V1 - V2, (V1 + V2)/2 + Xj) + norm(V1 - V2)^2)/(2*sqrt(2)*s*norm(V1 - V2))))*
			norm(V1 - V2)^3) + ((v1 - v2)*dot(V1 - V2, (V1 + V2)/2 + Xj)^2 -
			(s^2*(-v1 + v2) + (v2 + xj)*dot(V1 - V2, (V1 + V2)/2 + Xj))*norm(V1 - V2)^2)/(s^2*norm(V1 - V2)^4)

	   	return g_v1, g_v2

	end

 	function grad_log_p_l_n(k, Xj, V1, V2, s)

		xj = Xj[k]

		v1 = V1[k]

		v2 = V2[k]

                if dot(V2 - V1, Xj) > dot(V1, V2 - V1)

                        xj = xj - ((dot(V2 - V1, Xj) - dot(V1, V2 - V1) )/(norm(V2 - V1)^2))*(v2 - v1)

                        Xj = Xj - ((dot(V2 - V1, Xj) - dot(V1, V2 - V1) )/(norm(V2 - V1)^2)).*(V2 - V1)

                elseif dot(V2 - V1, Xj) < dot(V2, V1 - V2)

                        xj = xj - ((dot(V2 - V1, Xj) -  dot(V2, V1 - V2) )/(norm(V2 - V1)^2))*(v2 - v1)

                        Xj = Xj - ((dot(V2 - V1, Xj) -  dot(V2, V1 - V2) )/(norm(V2 - V1)^2)).*(V2 - V1)

                end

        function V_r(v_r, V)

        	V[k] = v_r

        	return V

        end

		log_p_l_proj = v_r -> log_p_l(Xj, V_r(v_r[1], V1), V_r(v_r[2], V2), s)

		g = Calculus.gradient(log_p_l_proj)

		return g([v1, v2])

	end

	function grad_log_cost!(Grad, data, mu_t, S_t, E_list, SIGMA_t)

		n  = size(data,1)

		d  = size(data[1],1)

		grad = Dict{Any,Array}(S_t[i] => zeros(Float64, 1 , d) for i = 1:N_0)

		Mu_t = Dict{Any,Array}(S_t[i] => mu_t[i,:] for i = 1:N_0)

		E_array = vcat(transpose(E_list)...)

		E_t = Dict{Any,Array}(S_t[i] => E_array[:,i] for i = 1:N)

		for s in S_t

			if s isa Integer

				grad[s] += transpose(sum(E_t[s].*broadcast( x -> ((MU[s] - x)./SIGMA[s]), data))/n)

			elseif s isa Tuple

				for k in 1:d

						for j in 1:n

							g_v1, g_v2 = grad_log_p_l_n(k, data[j], Mu_t[s[1]], Mu_t[s[2]], SIGMA_t[s])

							grad[s[1]][k] += E_t[s[1]][j]*g_v1/n

							grad[s[2]][k] += E_t[s[2]][j]*g_v2/n

							#println(grad[s[1]][k],grad[s[2]][k])
						end

				end

			end

		end

		for i in 1:N_0

			Grad[i,:] = - grad[S_t[i]]

		end

	end

	####EM algorithm

	mu00 = vcat([m for m in mu]'...)

	for it = 1:EM_it
		#E step

		E = expectation(data,MU,SIGMA,Pi)

		Pi      = 	sum(E)./n

		#M step

		mu0 = vcat([m for m in mu]'...)

		C = t -> -Cost(data,t,SIGMA, Pi ,E)

		#swap if don't want gradient clipping or set tol = Inf
		#G! = (g,t) -> FiniteDiff.finite_difference_gradient!(g,C,t)

		function G!(storage,t; tol = 0.1)

			FiniteDiff.finite_difference_gradient!(storage,C,t)

			storage[abs.(storage) .> tol] = sign.(storage[abs.(storage) .> tol])*tol

			return

		end

		results = Optim.optimize(C,G!,mu0,GradientDescent(),Optim.Options(g_tol = 1e-2))

		mu1 = Optim.minimizer(results)

		mu = [mu1[i,:] for i in 1:N_0]

		#println("Current Pi:")
		#println(Pi)

		#println("Current Sigma:")
		#println(sigma)

		#println("Current vertices:")
		#println(mu)

		print("For iteration: ", it," cost is: ", -C(mu1) + sum(log.(Pi)) )

		#Update parameters

		MU, SIGMA, PI = reload_param(mu, sigma, Pi ; S_t = S, N_0_t = N_0, N_t = N)

		#println(MU)

	end

	return MU, SIGMA, PI, E

end

function Save_PPC(ppc::PartitionedPointCloud, path::String)
	jldopen(path, "w") do file
		addrequire(file, Skyler)
		addrequire(file, LinearAlgebra)
		write(file, "PartitionedPointCloud", ppc)
	end
end # Save_PPC

#### Main functions and wrappers ####
function skyler(points::A, sample_epsilon::P, radius::R; out = "model", EM_it = 3, sig = sample_epsilon/2) where {P<:Number, Q<:Number, R<:Number, A<:Union{LinearAlgebra.Adjoint,AbstractArray}} #wrapper for obtaining the partition

	dimension = size(points,1)
	ppc = Partition_Point_Cloud(points, sample_epsilon, radius, (3*radius)/2 + 2*sample_epsilon, 3*sample_epsilon)
	if out == "PPC"
		return ppc
	elseif out == "DF"
			return Create_Partition_DataFrame(ppc)
	elseif out == "model"
		vertices = size(collect(ppc.Dimensions[0]),1)
		vertex_guess = Array{Float64}(undef, dimension,0)
		for i in 1:vertices
			v_i = zeros(dimension, 1)
			for j in collect(ppc.Strata[i])
				v_i .+= points[:, j]
			end
			v_i = v_i/(length(collect(ppc.Strata[i])))
			vertex_guess = hcat(vertex_guess, v_i)
		end
		#println("Vertex guess is: ",vertex_guess)
		E = zeros(size(points,2), size(ppc.Boundaries,1))

		for i in 1:vertices
			for j in collect(ppc.Strata[i])
				E[j,i] =1
			end
		end
		for i in collect(ppc.Dimensions[1])
			for j in collect(ppc.Strata[i])
				E[j,i] =1
			end
		end

		sigma = [sig for i in 1:size(ppc.Boundaries,1)]

		S = []
		for i in collect(ppc.Dimensions[0])
			push!(S, i)
		end
		for k in 1:maximum(keys(ppc.Dimensions))
			for i in collect(ppc.Dimensions[k])
				t =tuple(findall(x->x==1, ppc.Boundaries[:,i])...,)
				push!(S, t)
			end
		end

		data = vcat([points[:,i] for i in 1:size(points,2)])

		v_guess = Array{Float64}(undef, size(vertex_guess,2), size(vertex_guess,1))

		for i in 1:size(vertex_guess,2)
			for j in 1:size(vertex_guess,1)
				v_guess[i,j] = vertex_guess[j,i]
			end
		end
		#println("S is $S")
		MU, SIGMA, PI, E = model_fit(data, v_guess, sigma, E, S, EM_it= EM_it)
		mu_ans = collect(values(MU))
		mu_ans = vcat([datum' for datum in mu_ans]...)
		return ppc, mu_ans', vertex_guess
	else
		println("Error: output option 'out' selected is not valid, please choose from 'model' to obtain a model of the underlying structure, or 'struct' to receive a DataFrame indicating which strata a point as been associated with and the dimension of this strata, and the boundary relations.")
		return []
	end
end

function skyler(ppc::PartitionedPointCloud, sample_epsilon::P; EM_it=3, sig=sample_epsilon/2) where {P<:Number} #wrapper for modelling a point cloud
	points=ppc.Points
	vertices = size(collect(ppc.Dimensions[0]),1)
	vertex_guess = Array{Float64}(undef, dimension,0)
	for i in 1:vertices
		for j in collect(ppc.Strata[i])
			v_i .+= points[:, j]
		end
		v_i = v_i/(length(collect(ppc.Strata[i])))
		vertex_guess = hcat(vertex_guess, v_i)
	end
	println("Vertex guess is: ",vertex_guess)
	E = zeros(size(points,2), size(ppc.Boundaries,1))

	for i in 1:vertices
		for j in collect(ppc.Strata[i])
			E[j,i] =1
		end
	end
	for i in collect(ppc.Dimensions[1])
		for j in collect(ppc.Strata[i])
			E[j,i] =1
		end
	end

	sigma = [sig for i in 1:size(ppc.Boundaries,1)]

	S = []
	for i in collect(ppc.Dimensions[0])
		push!(S, i)
	end
	for k in 1:maximum(keys(ppc.Dimensions))
		for i in collect(ppc.Dimensions[k])
			t =tuple(findall(x->x==1, ppc.Boundaries[:,i])...,)
			push!(S, t)
		end
	end

	data = vcat([points[:,i] for i in 1:size(points,2)])

	v_guess = Array{Float64}(undef, size(vertex_guess,2),size(vertex_guess,1))

	for i in 1:size(vertex_guess,2)
		for j in 1:size(vertex_guess,1)
			v_guess[i,j] = vertex_guess[j,i]
		end
	end
	MU, SIGMA, PI, E = model_fit(data, v_guess, sigma, E, S, EM_it)
	mu_ans = collect(values(MU))
	mu_ans = vcat([datum' for datum in mu_ans]...)
	return mu_ans', vertex_guess
end #skyler for modelling a partioned point cloud

#### Functions for examples ####
function Load_Example(i)
	dir = pwd()
	cd(@__DIR__)
	cd("..")
	cd("Examples")
	points = CSV.read("Sample-$i.csv", header=false)
	cd(dir)
	return convert(Array,points)
end

#### unittest ####

function test_1()
dir = pwd()
	cd(@__DIR__)
	cd("..")
	cd("Examples")

	points= convert(Array,CSV.read("Sample-1.csv", header=false))
	ppc = skyler(points, 0.01, 0.12, out="PPC")
	par = Create_Partition_DataFrame(ppc)
	saved_partition = CSV.read("Partition-1.csv")

	if par == saved_partition
		return []
	else
		println("Error: test_1 not passed.")
		return par
	end
end

function test_2()
	dir = pwd()
	cd(@__DIR__)
	cd("..")
	cd("Examples")

	points= convert(Array,CSV.read("Sample-2.csv",header=false))'
	ppc = skyler(points, 0.1, 1.2, out="PPC")
	par = Create_Partition_DataFrame(ppc)
	saved_partition = CSV.read("Partition-2.csv")

	if par == saved_partition
		return []
	else
		println("Error: test_2 not passed.")
		return par
	end
end

function unittest()
	x = Array{Any}(undef, 2)

	x[1] = test_1() #expected answer empty
	x[2] = test_2() #expected answer empty

	for p 	= 	1:length(x)
		if !isempty(x[p])
			println(p)
			return x
		end
	end
	return []
end

end #module


text:
# ---
# title: "`opening`"
# id: op_opening
# cover: assets/opening.png
# ---

# `opening` operator is defined as `dilate(erode(img))`. Intuitively, opening operation
# fills the white holes in the image.

using ImageMorphology
using TestImages
using ImageBase
using ImageShow

img = restrict(testimage_dip3e("fig0940")) # rice
img01 = @. Gray(img > 0.5)
mosaic(img, img01; nrow=1)

#-

out1 = opening(img) # default: all spatial dimensions, r=1, a box-shape SE
out2 = opening(img; dims=(2,)) # only apply to the second dimension
out3 = opening(img; r=5) # half-size r=5

## also to the binary version
out1_01 = opening(img01)
out2_01 = opening(img01; dims=(2,))
out3_01 = opening(img01; r=5)

mosaic(out1, out2, out3, out1_01, out2_01, out3_01; nrow=2, rowmajor=true)

# ## See also

# [`closing`](@ref op_closing) is the dual operator of `opening` in the following sense:

complement.(opening(img)) == closing(complement.(img))

# For a comprehensive and more accurate documentation, please check the [`opening`](@ref)
# reference page.

# save cover image #src
using FileIO #src
mkpath("assets") #src
img = Gray.(restrict(testimage("blobs"))) #src
save("assets/opening.png", opening(img; r=2)) #src

ext: jl


text:
# indexing

## utilities

using Base.Cartesian

## logical indexing

Base.getindex(xs::CuArray, bools::AbstractArray{Bool}) = getindex(xs, CuArray(bools))

function Base.getindex(xs::CuArray{T}, bools::CuArray{Bool}) where {T}
  bools = reshape(bools, prod(size(bools)))
  indices = cumsum(bools)  # unique indices for elements that are true

  n = @allowscalar indices[end]  # number that are true
  ys = CuArray{T}(undef, n)

  if n > 0
    function kernel(ys::CuDeviceArray{T}, xs::CuDeviceArray{T}, bools, indices)
        i = threadIdx().x + (blockIdx().x - 1) * blockDim().x

        @inbounds if i <= length(xs) && bools[i]
            b = indices[i]   # new position
            ys[b] = xs[i]
        end

        return
    end

    function configurator(kernel)
        config = launch_configuration(kernel.fun)

        threads = min(length(indices), config.threads)
        blocks = cld(length(indices), threads)

        return (threads=threads, blocks=blocks)
    end

    @cuda name="logical_getindex" config=configurator kernel(ys, xs, bools, indices)
  end

  unsafe_free!(indices)

  return ys
end

## find*

function Base.findall(bools::CuArray{Bool})
    I = keytype(bools)
    indices = cumsum(reshape(bools, prod(size(bools))))

    n = @allowscalar indices[end]
    ys = CuArray{I}(undef, n)

    if n > 0
        function kernel(ys::CuDeviceArray, bools, indices)
            i = threadIdx().x + (blockIdx().x - 1) * blockDim().x

            @inbounds if i <= length(bools) && bools[i]
                i′ = CartesianIndices(bools)[i]
                b = indices[i]   # new position
                ys[b] = i′
            end

            return
        end

        function configurator(kernel)
            config = launch_configuration(kernel.fun)

            threads = min(length(indices), config.threads)
            blocks = cld(length(indices), threads)

            return (threads=threads, blocks=blocks)
        end

        @cuda name="findall" config=configurator kernel(ys, bools, indices)
    end

    unsafe_free!(indices)

    return ys
end

function Base.findall(f::Function, A::CuArray)
    bools = map(f, A)
    ys = findall(bools)
    unsafe_free!(bools)
    return ys
end

function Base.findfirst(testf::Function, xs::CuArray)
    I = keytype(xs)

    y = CuArray([typemax(Int)])

    function kernel(y::CuDeviceArray, xs::CuDeviceArray)
        i = threadIdx().x + (blockIdx().x - 1) * blockDim().x

        @inbounds if i <= length(xs) && testf(xs[i])
            CUDAnative.@atomic y[1] = min(y[1], i)
        end

        return
    end

    function configurator(kernel)
        config = launch_configuration(kernel.fun)

        threads = min(length(xs), config.threads)
        blocks = cld(length(xs), threads)

        return (threads=threads, blocks=blocks)
    end

    @cuda name="findfirst" config=configurator kernel(y, xs)

    first_i = @allowscalar y[1]
    return keys(xs)[first_i]
end

Base.findfirst(xs::CuArray{Bool}) = findfirst(identity, xs)

function Base.findfirst(vals::CuArray, xs::CuArray)
    # figure out which dimension was reduced
    @assert ndims(vals) == ndims(xs)
    dims = [i for i in 1:ndims(xs) if size(xs,i)!=1 && size(vals,i)==1]
    @assert length(dims) == 1
    dim = dims[1]

    ## find the first matching element

    indices = fill(typemax(Int), size(vals))

    # iteration domain across the main dimension
    Rdim = CartesianIndices((size(xs, dim),))

    # iteration domain for the other dimensions
    Rpre = CartesianIndices(size(xs)[1:dim-1])
    Rpost = CartesianIndices(size(xs)[dim+1:end])
    Rother = CartesianIndices((length(Rpre), length(Rpost)))

    function kernel(xs, vals, indices, Rdim, Rpre, Rpost, Rother)
        # iterate the main dimension using threads and the first block dimension
        i = (blockIdx().x-1) * blockDim().x + threadIdx().x
        # iterate the other dimensions using the remaining block dimensions
        j = (blockIdx().z-1) * gridDim().y + blockIdx().y

        if i <= length(Rdim) && j <= length(Rother)
            I = Rother[j]
            Ipre = Rpre[I[1]]
            Ipost = Rpost[I[2]]

            @inbounds if xs[Ipre, i, Ipost] == vals[Ipre, Ipost]
                full_index = LinearIndices(xs)[Ipre, i, Ipost]   # atomic_min only works with integers
                reduced_index = LinearIndices(indices)[Ipre, Ipost] # FIXME: @atomic doesn't handle array ref with CartesianIndices
                CUDAnative.@atomic indices[reduced_index] = min(indices[reduced_index], full_index)
            end
        end

        return
    end

    function configurator(kernel)
        # what's a good launch configuration for this kernel?
        config = launch_configuration(kernel.fun)

        # blocks to cover the main dimension
        threads = min(length(Rdim), config.threads)
        blocks_dim = cld(length(Rdim), threads)
        # NOTE: the grid X dimension is virtually unconstrained

        # blocks to cover the remaining dimensions
        dev = CUDAdrv.device(kernel.fun.mod.ctx)
        max_other_blocks = attribute(dev, CUDAdrv.DEVICE_ATTRIBUTE_MAX_GRID_DIM_Y)
        blocks_other = (min(length(Rother), max_other_blocks),
                        cld(length(Rother), max_other_blocks))

        return (threads=threads, blocks=(blocks_dim, blocks_other...))
    end

    @cuda config=configurator kernel(xs, vals, indices, Rdim, Rpre, Rpost, Rother)

    ## convert the linear indices to an appropriate type

    kt = keytype(xs)

    if kt == Int
        return indices
    else
        indices′ = CuArray{kt}(undef, size(indices))
        broadcast!(indices′, indices, Ref(keys(xs))) do index, keys
            keys[index]
        end

        return indices′
    end
end

function Base.findmin(a::CuArray; dims=:)
    if dims == Colon()
        m = minimum(a)
        i = findfirst(x->x==m, a)
        return m,i
    else
        minima = minimum(a; dims=dims)
        i = findfirst(minima, a)
        return minima,i
    end
end

function Base.findmax(a::CuArray; dims=:)
    if dims == Colon()
        m = maximum(a)
        i = findfirst(x->x==m, a)
        return m,i
    else
        maxima = maximum(a; dims=dims)
        i = findfirst(maxima, a)
        return maxima,i
    end
end


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

using BAT
using Test

using Random, Random123

@testset "rng_init" begin
    @testset "BAT.RNGPartition" begin
        @test @inferred(BAT.RNGPartition(Philox4x(), Base.OneTo(42))) isa BAT.RNGPartition

        _getfields(x::Any) = map(n -> getfield(x, n), fieldnames(typeof(x)))

        rng1 = Philox4x()
        @test BAT.rngpart_getpartctrs(rng1) == (partctrs = (0, 0, 0, 0, 0, 0), depth = 1)
        @test BAT.rngpart_getpartctrs(rng1) == (partctrs = (0, 0, 0, 0, 0, 0), depth = 1)
        prng1 = @inferred BAT.RNGPartition(rng1, 3:3207)
        @test prng1 isa BAT.RNGPartition{<:typeof(rng1)}
        @test _getfields(prng1) == (BAT.rngpart_getseed(rng1), (1, 0, 0, 0, 0, 0), 1, 3:3207)
        @test BAT.rngpart_getpartctrs(rng1) == (partctrs = (3206, 0, 0, 0, 0, 0), depth = 1)

        @test BAT.rngpart_getpartctrs(@inferred AbstractRNG(prng1, 3)) == (partctrs = (1, 0, 0, 0, 0, 0), depth = 2)
        @test BAT.rngpart_getpartctrs(@inferred AbstractRNG(prng1, 3207)) == (partctrs = (3205, 0, 0, 0, 0, 0), depth = 2)
        @test_throws ArgumentError AbstractRNG(prng1, 2)
        @test_throws ArgumentError AbstractRNG(prng1, 3208)
        rng2 = AbstractRNG(prng1, 6)
        @test BAT.rngpart_getpartctrs(rng2) == (partctrs = (4, 0, 0, 0, 0, 0), depth = 2)
        @test all(x -> x != 0, rand(rng1, UInt64, 16) .- rand(rng2, UInt64, 16))
        @test BAT.rngpart_getpartctrs(rng2) == (partctrs = (4, 0, 0, 0, 0, 0), depth = 2)
        @test_throws InexactError BAT.RNGPartition(deepcopy(rng2), -5:800000000000000000)
        @test_throws ArgumentError BAT.RNGPartition(deepcopy(rng2), -5:2147483641)
        prng2 = BAT.RNGPartition(rng2, -5:2147483640)

        rng3 = AbstractRNG(prng2, 2147483640)
        @test BAT.rngpart_getpartctrs(rng3) == (partctrs = (4, 2147483646, 0, 0, 0, 0), depth = 3)
        prng3a = BAT.RNGPartition(rng3, 1:1073741823)
        @test BAT.rngpart_getpartctrs(rng3) == (partctrs = (4, 2147483646, 1073741823+1, 0, 0, 0), depth = 3)
        prng3b = BAT.RNGPartition(rng3, 1:1073741822)
        @test _getfields(prng3b) == (BAT.rngpart_getseed(rng1), (4, 2147483646, 1073741823+2, 0, 0, 0), 3, 1:1073741822)
        @test BAT.rngpart_getpartctrs(rng3) == (partctrs = (4, 2147483646, 1073741823+1073741822+2, 0, 0, 0), depth = 3)
        @test_throws ArgumentError BAT.RNGPartition(deepcopy(rng3), 1:1)

        rng4 = AbstractRNG(prng3b, 1073741822)
        BAT.rngpart_getpartctrs(rng4) == (partctrs = (4, 2147483646, 1073741823+1073741822+1, 0, 0, 0), depth = 4)
        prng4 = BAT.RNGPartition(rng4, 1:1000)

        rng5 = AbstractRNG(prng4, 500)
        prng5 = BAT.RNGPartition(rng5, 1:1000)

        rng6 = AbstractRNG(prng5, 500)
        @test BAT.rngpart_getpartctrs(rng6) == (partctrs = (4, 2147483646, 1073741823+1073741822+1, 500, 500, 0), depth = 6)
        prng6 = BAT.RNGPartition(rng6, 1:1000)
        @test BAT.rngpart_getpartctrs(rng6) == (partctrs = (4, 2147483646, 1073741823+1073741822+1, 500, 500, 1001), depth = 6)

        @test_throws ArgumentError AbstractRNG(prng6, 500)
    end
end


text:
using Monadic
using Test

@test isempty(detect_ambiguities(Monadic))

myflatmap(f, x) = Iterators.flatten(map(f, x))
iteratorresult = @monadic map myflatmap begin
  x = 1:3
  y = [1, 6]
  @pure x + y
end
@test collect(iteratorresult) == [2, 7, 3, 8, 4, 9]

mywrapper(n::Int) = 1:n
mywrapper(any) = any
iteratorresult = @monadic map myflatmap mywrapper begin
  x = 3
  y = [1, 6]
  @pure x + y
end
@test collect(iteratorresult)  == [2, 7, 3, 8, 4, 9]

iteratorresult = @monadic map myflatmap begin
  1:3
  [1, 6]
end
@test collect(iteratorresult) == [1, 6, 1, 6, 1, 6]


text:
using BasisFunctions: DictionarySolverOperator
using BasisFunctions, DomainSets, ..ExtensionFrames, LinearAlgebra
import BasisFunctions: operator, apply!

"""
A fast FE solver based on a low-rank approximation of the plunge region. The plunge region
is isolated using a projection operator.
For more details, see the paper 'Fast algorithms for the computation of Fourier extensions of arbitrary length'
http://arxiv.org/abs/1509.00206
"""
struct FE_TridiagonalSolver{ELT} <: DictionarySolverOperator{ELT}
    Ut      :: Array{ELT,2}
    V       :: Array{ELT,2}
    Sinv    :: Array{ELT,1}
    A       ::  DictionaryOperator
    b
    blinear ::  Array{ELT,1}
    x2
    x1
    y
    y2
    scaling
    function FE_TridiagonalSolver{ELT}(A::DictionaryOperator, scaling; threshold = default_threshold(A), irange = estimate_plunge_range(A,scaling,estimate_plunge_rank(A)), verbose=false,options...) where ELT
        R = estimate_plunge_rank(A)
        finished = false
        V=[]
        U=[]
        S=[]
        while !finished
            (V,U)=mideigs(size(A,2),size(A,1),scaling,irange)
            S = zeros(ELT,size(V,2))
            for i=1:size(V,2)
                S[i]=U[:,i]'*apply(A,V[:,i])
            end
            if all(minimum(abs.(S)).< threshold) || R>size(A,1)
                finished=true
            else
                R = 2*R
                irange=estimate_plunge_range(A,scaling,R)
            end
        end
        I = abs.(S).>(threshold*maximum(abs.(S)))
        b = zeros(dest(A))
        blinear = zeros(ELT, length(dest(A)))
        x1 = zeros(src(A))
        x2 = zeros(src(A))
        y = zeros(ELT,sum(I))
        y2 = zeros(ELT,sum(I))
    new(U[:,I]',V[:,I],S[I].^(-1), A, b,blinear,x1,x2,y,y2,scaling)
    end
end

FE_TridiagonalSolver(A::DictionaryOperator; scaling = nothing, options...) =
    FE_TridiagonalSolver(A, scaling; options...)

# TODO: remove the scaling factor here
FE_TridiagonalSolver(A::DictionaryOperator, scaling; options...) =
    FE_TridiagonalSolver{eltype(A)}(A, scaling; options...)

operator(op::FE_TridiagonalSolver) = op.A

# Estimate for the rank of (A*Zt-I)*A when computing the low rank decomposition. If check fails, rank estimate is steadily increased.
# TODO: remove this estimate and use adaptivity
estimate_plunge_rank(A::DictionaryOperator) =
    estimate_plunge_rank(src(A), dest(A))

estimate_plunge_rank(src::ExtensionFrame, dest::Dictionary) =
    estimate_plunge_rank(superdict(src), support(src), dest)

estimate_plunge_rank(src::Dictionary, dest::Dictionary) =
    default_estimate_plunge_rank(src, dest)

estimate_plunge_rank(src::Dictionary, domain::Domain, dest::Dictionary) =
    default_estimate_plunge_rank(src, dest)

function default_estimate_plunge_rank(src::Dictionary, dest::Dictionary)
    nml=length(src)^2/length(dest)
    N = dimension(src)
    if N==1
        return max(1,min(round(Int, 9*log(nml)),length(src)))
    else
        return max(1,min(round(Int, 9*log(nml)*nml^((N-1)/N)),length(src)))
    end
end

function estimate_plunge_range(A,L,C)
    mid = size(A,2)-round(Int,size(A,1)*size(A,2)/L)
    mid .+ (max(1-mid,-round(Int,C)):min(size(A,2)-mid,round(Int,C)))
end

apply!(s::FE_TridiagonalSolver, coef_dest, coef_src) =
apply!(s, coef_dest, coef_src, s.A, s.A', s.x1, s.x2,s.y,s.y2)

function apply!(s::FE_TridiagonalSolver, coef_dest, coef_src, A, At, x1, x2,y,y2)
    # Applying plunge to the right hand side
    BasisFunctions.linearize_coefficients!(dest(A), s.blinear, coef_src)
    mul!(y,s.Ut,s.blinear)
    for i =1:length(y)
        y[i]*=s.Sinv[i]
        if abs(y[i])>10
            y[i]=0
        end
    end
    mul!(x2,s.V,y)
    # x2 solves the middle guys
    apply!(A, s.b, x2)
    apply!(At, x1, coef_src-s.b)
    # x1 solves the remainder through one application of the Aerator
    for i in eachindex(x1)
        coef_dest[i] = x1[i]/s.scaling+x2[i]
    end
end

function mideigs(N,M,L,irange)
    J2 = Chamzas(L,M-1,N)
    J1 = Chamzas(L,N-1,M)

    E,V1 = eigen(J1,irange)
    for i = 1:size(V1,1)
        for j = 1:size(V1,2)
            if i%2 == 0
                V1[i,j]*=-1
            end
        end
    end
    V1b = circshift(V1,round(Int,(N-1)/2)+1)
    E,V2 = eigen(J2,irange .+ (M-N))
    (V1b,V2)
end

function Chamzas(N,M,F)
    bs = sin.((pi/N)*(1:M)).*sin.((pi/N)*(M .- (0:(M-1))))
    cs = -cos.((pi/N)*(2*(0:M) .- M))*cos.((pi/N)*(F))
    T = SymTridiagonal(cs,bs)
end


text:
module Arrangement
    using LinearAlgebraicRepresentation
    using IntervalTrees
    using NearestNeighbors
    using Triangle

    include("./minimal_cycles.jl")
    include("./dimension_travel.jl")
    include("./planar_arrangement.jl")
    include("./spatial_arrangement.jl")

    export planar_arrangement, spatial_arrangement

end


text:
"""
    combined.jl

Description:
    Example usage of the the DB, PS, and XB CVIs/ICVIs.

Author:
    Sasha Petrenko <sap625@mst.edu>

Date:
    3/24/2021
"""

# --------------------------------------------------------------------------- #
# PACKAGES
# --------------------------------------------------------------------------- #

# ICVIs pulled from the ClusterValidityIndices package
using ClusterValidityIndices

# Quality of life packages for editing and fancy logging
using Revise
using ProgressBars
using Logging
using Plots

# --------------------------------------------------------------------------- #
# USER CONFIGURATION
#   The user config, such as data paths and plotting parameters
# --------------------------------------------------------------------------- #

# Location of the data
data_paths = readdir("data", join=true)

# Plotting dots-per-inch
dpi = 300

# Plotting style
theme(:dark)

# Plotting backend
try
    pyplot()        # PyPlot backend
catch
    gr()            # GR backend (default for Plots.jl)
end

# --------------------------------------------------------------------------- #
# SCRIPT CONFIGURATION
# --------------------------------------------------------------------------- #

# Set the log level
LogLevel(Logging.Info)

# Load the examples helper functions
include("../common.jl")
include("../experiments.jl")

# Construct the cvis
cvis = [
    XB,
    DB,
    PS
]
n_cvis = length(cvis)

# Iterate across all the data
for data_path in data_paths
    # Get the data_name
    data_name = splitext(basename(data_path))[1]

    @info "------- Testing data: $data_name -------"
    # Load the training data
    data, labels = get_cvi_data(data_path)
    # Relabel it into sequential order
    labels = relabel_cvi_data(labels)
    # Iterate across all CVIs
    for cvi in cvis
        @info "------- Testing CVI: $cvi -------"
        test_cvi(cvi(), data, labels, data_name)
    end
end


text:
export SGD

"""
    SGD

Stochastic Gradient Descent Optimizer.

# Arguments
* rate: learning rate
* [momentum=0.0]: momentum coefficient
* [nesterov=false]: use nesterov acceleration or not
"""
mutable struct SGD
    rate::Float64
    momentum::Float64
    nesterov::Bool
    states::IdDict
end

function SGD(rate=0.0; momentum=0.0, nesterov=false)
    SGD(rate, momentum, nesterov, IdDict())
end

(opt::SGD)(x::Var) = opt(x.data, x.grad)

function (opt::SGD)(x::AbstractArray{T,N}, gx::AbstractArray{T,N}) where {T,N}
    if opt.momentum > 0.0
        if haskey(opt.states, x)
            v = opt.states[x]
        else
            v = zeros(x)
            opt.states[x] = v
        end
        m = T(opt.momentum)
        rate = T(opt.rate)
        v .= m .* v - rate * gx
        if opt.nesterov
            v = copy(v)
            BLAS.scal!(length(v), m, v, 1)
            BLAS.axpy!(-rate, gx, v)
        end
        axpy!(T(1), v, x)
    else
        # axpy!(T(0.0005), x, gx)
        axpy!(T(-opt.rate), gx, x)
    end
    fill!(gx, T(0))
end

function (opt::SGD)(x::CuArray{T,N}, gx::CuArray{T,N}) where {T,N}
    if opt.momentum > 0.0
        throw("Not implemented yet.")
        if haskey(opt.states, x)
            v = opt.states[x]
        else
            v = zeros(x)
            opt.states[x] = v
        end
        m = T(opt.momentum)
        rate = T(opt.rate)
        v .= m .* v - rate * gx
        if opt.nesterov
            v = copy(v)
            scal!(length(v), m, v, 1)
            axpy!(-rate, gx, v)
        end
        axpy!(T(1), v, x)
    else
        axpy!(T(-opt.rate), gx, x)
    end
    fill!(gx, T(0))
end


text:
typealias HeaderBinary Tuple{Array{UInt8, 1}, Array{UInt8, 1}}
typealias Header Tuple{AbstractString, AbstractString}

type DynamicTable
    table::Array{HeaderBinary, 1}
    size::Int
    max_size::Int
end

function new_dynamic_table()
    DynamicTable(Array{HeaderBinary, 1}(), 0, 4096)
end

function consolidate_table!(table::DynamicTable)
    while table.size > table.max_size
        last_header = pop!(table.table)
        table.size -= length(last_header[1]) + length(last_header[2]) + 32
    end
end

function add_header!(table::DynamicTable, header::HeaderBinary)
    name = header[1]
    value = header[2]
    table.size += length(name) + length(value) + 32
    unshift!(table.table, (name, value))
    consolidate_table!(table)
end

function add_header!(table::DynamicTable, header::Header)
    add_header!(table, (convert(Array{UInt8, 1}, header[1]),
                        convert(Array{UInt8, 1}, header[2])))
end

function set_max_table_size!(table::DynamicTable, size::Int)
    table.max_size = size
    consolidate_table!(table)
end

STATIC_TABLE =
    [
     (b":authority", b"")
     (b":method", b"GET")
     (b":method", b"POST")
     (b":path", b"/")
     (b":path", b"/index.html")
     (b":scheme", b"http")
     (b":scheme", b"https")
     (b":status", b"200")
     (b":status", b"204")
     (b":status", b"206")
     (b":status", b"304")
     (b":status", b"400")
     (b":status", b"404")
     (b":status", b"500")
     (b"accept-", b"")
     (b"accept-encoding", b"gzip, deflate")
     (b"accept-language", b"")
     (b"accept-ranges", b"")
     (b"accept", b"")
     (b"access-control-allow-origin", b"")
     (b"age", b"")
     (b"allow", b"")
     (b"authorization", b"")
     (b"cache-control", b"")
     (b"content-disposition", b"")
     (b"content-encoding", b"")
     (b"content-language", b"")
     (b"content-length", b"")
     (b"content-location", b"")
     (b"content-range", b"")
     (b"content-type", b"")
     (b"cookie", b"")
     (b"date", b"")
     (b"etag", b"")
     (b"expect", b"")
     (b"expires", b"")
     (b"from", b"")
     (b"host", b"")
     (b"if-match", b"")
     (b"if-modified-since", b"")
     (b"if-none-match", b"")
     (b"if-range", b"")
     (b"if-unmodified-since", b"")
     (b"last-modified", b"")
     (b"link", b"")
     (b"location", b"")
     (b"max-forwards", b"")
     (b"proxy-authenticate", b"")
     (b"proxy-authorization", b"")
     (b"range", b"")
     (b"referer", b"")
     (b"refresh", b"")
     (b"retry-after", b"")
     (b"server", b"")
     (b"set-cookie", b"")
     (b"strict-transport-security", b"")
     (b"transfer-encoding", b"")
     (b"user-agent", b"")
     (b"vary", b"")
     (b"via", b"")
     (b"www-authenticate", b"")
     ]

function get_header(table::DynamicTable, index::Int)
    # IETF's table indexing is 1-based.
    if index <= length(STATIC_TABLE)
        return STATIC_TABLE[index]
    else
        if index > length(STATIC_TABLE) + table.max_size
            throw(DecodeError("Index greater than sum of both static and dynamic tables."))
        else
            if index > length(STATIC_TABLE) + length(table.table)
                throw(DecodeError("Index out of bound."))
            else
                return table.table[index - length(STATIC_TABLE)]
            end
        end
    end
end


text:
using MultivariateStats
using LinearAlgebra
using Test
import Statistics: mean, cov
import Random

@testset "CCA" begin

    Random.seed!(34568)

    dx = 5
    dy = 6
    p = 3

    # CCA with zero means

    X = rand(dx, 100)
    Y = rand(dy, 100)

    Px = qr(randn(dx, dx)).Q[:, 1:p]
    Py = qr(randn(dy, dy)).Q[:, 1:p]

    M = CCA(Float64[], Float64[], Px, Py, [0.8, 0.6, 0.4])

    @test xindim(M) == dx
    @test yindim(M) == dy
    @test xmean(M) == zeros(dx)
    @test ymean(M) == zeros(dy)
    @test xprojection(M) == Px
    @test yprojection(M) == Py
    @test correlations(M) == [0.8, 0.6, 0.4]

    @test xtransform(M, X) ≈ Px'X
    @test ytransform(M, Y) ≈ Py'Y

    ## CCA with nonzero means

    ux = randn(dx)
    uy = randn(dy)

    M = CCA(ux, uy, Px, Py, [0.8, 0.6, 0.4])

    @test xindim(M) == dx
    @test yindim(M) == dy
    @test xmean(M) == ux
    @test ymean(M) == uy
    @test xprojection(M) == Px
    @test yprojection(M) == Py
    @test correlations(M) == [0.8, 0.6, 0.4]

    @test xtransform(M, X) ≈ Px' * (X .- ux)
    @test ytransform(M, Y) ≈ Py' * (Y .- uy)

    ## prepare data

    n = 1000
    dg = 10
    G = randn(dg, n)

    X = randn(dx, dg) * G + 0.2 * randn(dx, n)
    Y = randn(dy, dg) * G + 0.2 * randn(dy, n)
    xm = vec(mean(X, dims=2))
    ym = vec(mean(Y, dims=2))
    Zx = X .- xm
    Zy = Y .- ym

    Cxx = cov(X, dims=2)
    Cyy = cov(Y, dims=2)
    Cxy = cov(X, Y, dims=2)
    Cyx = Cxy'

    ## ccacov

    # X ~ Y
    M = fit(CCA, X, Y; method=:cov, outdim=p)
    Px = xprojection(M)
    Py = yprojection(M)
    rho = correlations(M)
    @test xindim(M) == dx
    @test yindim(M) == dy
    @test outdim(M) == p
    @test xmean(M) == xm
    @test ymean(M) == ym
    @test issorted(rho; rev=true)

    @test Px' * Cxx * Px ≈ Matrix(I, p, p)
    @test Py' * Cyy * Py ≈ Matrix(I, p, p)
    @test Cxy * (Cyy \ Cyx) * Px ≈ Cxx * Px * Diagonal(rho.^2)
    @test Cyx * (Cxx \ Cxy) * Py ≈ Cyy * Py * Diagonal(rho.^2)
    @test Px ≈ MultivariateStats.qnormalize!(Cxx \ (Cxy * Py), Cxx)
    @test Py ≈ MultivariateStats.qnormalize!(Cyy \ (Cyx * Px), Cyy)

    # Y ~ X
    M = fit(CCA, Y, X; method=:cov, outdim=p)
    Py = xprojection(M)
    Px = yprojection(M)
    rho = correlations(M)
    @test xindim(M) == dy
    @test yindim(M) == dx
    @test outdim(M) == p
    @test xmean(M) == ym
    @test ymean(M) == xm
    @test issorted(rho; rev=true)

    @test Px' * Cxx * Px ≈ Matrix(I, p, p)
    @test Py' * Cyy * Py ≈ Matrix(I, p, p)
    @test Cxy * (Cyy \ Cyx) * Px ≈ Cxx * Px * Diagonal(rho.^2)
    @test Cyx * (Cxx \ Cxy) * Py ≈ Cyy * Py * Diagonal(rho.^2)
    @test Px ≈ MultivariateStats.qnormalize!(Cxx \ (Cxy * Py), Cxx)
    @test Py ≈ MultivariateStats.qnormalize!(Cyy \ (Cyx * Px), Cyy)

    ## ccasvd

    # n > d
    M = fit(CCA, X, Y; method=:svd, outdim=p)
    Px = xprojection(M)
    Py = yprojection(M)
    rho = correlations(M)
    @test xindim(M) == dx
    @test yindim(M) == dy
    @test outdim(M) == p
    @test xmean(M) == xm
    @test ymean(M) == ym
    @test issorted(rho; rev=true)

    @test Px' * Cxx * Px ≈ Matrix(I, p, p)
    @test Py' * Cyy * Py ≈ Matrix(I, p, p)
    @test Cxy * (Cyy \ Cyx) * Px ≈ Cxx * Px * Diagonal(rho.^2)
    @test Cyx * (Cxx \ Cxy) * Py ≈ Cyy * Py * Diagonal(rho.^2)
    @test Px ≈ MultivariateStats.qnormalize!(Cxx \ (Cxy * Py), Cxx)
    @test Py ≈ MultivariateStats.qnormalize!(Cyy \ (Cyx * Px), Cyy)

    # different input types
    XX = convert.(Float32, X)
    YY = convert.(Float32, Y)

    MM = fit(CCA, view(XX, :, 1:400), view(YY, :, 1:400); method=:svd, outdim=p)

    # test that mixing types doesn't error
    xtransform(M, XX)
    ytransform(M, YY)
    xtransform(MM, XX)
    ytransform(MM, YY)

    # type stability
    for func in (xmean, ymean, xprojection, yprojection, correlations)
        @test eltype(func(M)) == Float64
        @test eltype(func(MM)) == Float32
    end
end


text:
I_func() = I_DATA
X_func() = X_DATA
Y_func() = Y_DATA
Z_func() = Z_DATA
H_func() = H_DATA
S_func() = S_DATA
SDAG_func() = SDAG_DATA
T_func() = T_DATA
TDAG_func() = TDAG_DATA
SWAP_func() = SWAP_DATA
CNOT_func() = CNOT_DATA
TOFFOLI_func() = TOFFOLI_DATA
FREDKIN_func() = FREDKIN_DATA
CZ_func() = CZ_DATA

Rx_func(θ::Real) = cos(θ/2)I_DATA - sin(θ/2)im* X_DATA
Ry_func(θ::Real) = cos(θ/2)I_DATA - sin(θ/2)im* Y_DATA
Rz_func(θ::Real) = cos(θ/2)I_DATA - sin(θ/2)im* Z_DATA

function Rϕ_func(θ::Real)
	phaseGate = zeros(ComplexF64,2,2)
	phaseGate[1,1] = 1
	phaseGate[2,2] = exp(θ* 1im)
	return phaseGate
end

function ROTATE_func(θ::Real, nx::Real, ny::Real, nz::Real)
	nx, ny, nz = [nx, ny, nz]/norm([nx, ny, nz])
	return cos(θ/2)*I_DATA-sin(θ/2)im* (nx*X_DATA+ny*Y_DATA+nz*Z_DATA)
end

const I_kernel = GateKernel(I_func, "I")
const X_kernel = GateKernel(X_func, "X")
const Y_kernel = GateKernel(Y_func, "Y")
const Z_kernel = GateKernel(Z_func, "Z")
const H_kernel = GateKernel(H_func, "H")
const S_kernel = GateKernel(S_func, "S")
const SDAG_kernel = GateKernel(SDAG_func, "SDAG")
const T_kernel = GateKernel(T_func, "T")
const TDAG_kernel = GateKernel(TDAG_func, "TDAG")
const SWAP_kernel = GateKernel(SWAP_func, "SWAP")
const CNOT_kernel = GateKernel(CNOT_func, "CNOT")
const TOFFOLI_kernel = GateKernel(TOFFOLI_func, "TOFFOLI")
const FREDKIN_kernel = GateKernel(FREDKIN_func, "FREDKIN")
const CZ_kernel = GateKernel(CZ_func, "CZ")
const Rx = GateKernel(Rx_func, 1, "Rx")
const Ry = GateKernel(Ry_func, 1, "Ry")
const Rz = GateKernel(Rz_func, 1, "Rz")
const Rϕ = GateKernel(Rϕ_func, 1, "Rϕ")
const ROTATE = GateKernel(ROTATE_func, 4, "General rotation")


text:
using Statistics
function demere(num::Int)
    numrolls = 4
    numsixes = 0
    for run in 1:num
        roll = rand(1:6, numrolls, 1)
        if any( roll .== 6)
            numsixes += 1
        end
    end
    probsix = numsixes/num
    return probsix
end

function roulette(k, t, num::Int)
    p = 18/37
    numruin = 0
    numbet = 0
    for run in 1:num
        money = k
        while money > 0 && money < t
            if rand(1)[1] < p
                money += 1
            else
                money -= 1
            end
            numbet += 1
        end
        if money == 0
            numruin += 1
        end
    end
    probruin = numruin / num
    averagebet = numbet / num
    return probruin, averagebet
end

function hitmiss(f::Function, box = [0; 1], num=10000; fmin=0, fmax=1)
   x = (box[2]-box[1]).*rand(num,1) .+ box[1]
   y = (fmax-fmin).*rand(num,1) .+ fmin
   rectarea = (fmax-fmin)*(box[2]-box[1])
   prop = sum(y.<f.(x))/num
   return prop*rectarea
end

function crudemc(f::Function, box = [0;1], num=10000)
   x = (box[2]-box[1]).*rand(num,1) .+ box[1]
   fmean = mean(f.(x))
   return (box[2]-box[1])*fmean
end


text:
using Test
using RLCore
using Random

include("test_env.jl")
include("test_agent.jl")
include("test_episode.jl")

function runtests()
    test_env()
    test_agent()
    test_episode()
end

runtests()


text:
using Tables: materializer

"""
    compact_table_ROT(table; empty_values::Vector=[])

Returns a new table of the original `table` type where all rows
and columns filled entirely by `missing` and `empty_values` have been removed.
"""
function compact_table_ROT(table; empty_values::Vector=[])
    return materializer(table)(compact_table(table; empty_values))
end

"""
    compact_columns_ROT(table; empty_values::Vector=[])

Returns a new table of the original `table` type where all columns filled entirely by
`missing` and `empty_values` have been removed.
"""
function compact_columns_ROT(table; empty_values::Vector=[])
    return materializer(table)(compact_columns(table; empty_values))
end

"""
    compact_rows_ROT(table; empty_values::Vector=[])

Returns a new table of the original `table` type where all rows filled entirely by
`missing` and `empty_values` have been removed.
"""
function compact_rows_ROT(table; empty_values::Vector=[])
    return materializer(table)(compact_rows(table; empty_values))
end

"""
    delete_const_columns_ROT(table)

Returns a new table of the original `table` type where all columns filled with just a
constant value have been removed.
"""
function delete_const_columns_ROT(table)
    return materializer(table)(delete_const_columns(table))
end

"""
    polish_names_ROT(table; style::Symbol=:snake_case)

Returns a new table of the original `table` type where column names have been replaced to
be unique and formated using the `style` selected.

# Styles
- snake_case
- camelCase
"""
function polish_names_ROT(table; style::Symbol=:snake_case)
    return materializer(table)(polish_names(table; style=style))
end

"""
    reinfer_schema_ROT(table; max_types::Int=3)

Returns a new table of the original `table` type where it has been tried to minimize the
amount of element types for each column without making the column type `Any`.

For this, will try to make the column of type `Union` with up to max_types and
internally use `Base.promote_typejoin` on all numeric types.
If not possible, leaves the column as-is.
"""
function reinfer_schema_ROT(table; max_types::Int=3)
    return materializer(table)(reinfer_schema(table; max_types=max_types))
end

"""
    row_as_names_ROT(table, i::Int; remove::Bool=true)

Returns a new table of the original `table` type that has been renamed using row `i` as
new names and removes in-place all the rows above row `i` if `remove=true`.
"""
function row_as_names_ROT(table, i::Int; remove::Bool=true)
    return materializer(table)(row_as_names(table, i; remove=remove))
end

"""
    rename_ROT(table, names::Vector{Symbol})

Returns a new table of the original `table` type where its column names have been changed
to be `names`.
"""
function rename_ROT(table, names::Vector{Symbol})
    return materializer(table)(rename(table, names))
end

"""
    drop_missing_ROT(table; missing_values::Vector=[])

Returns a new table of the original `table` type where all rows where `missing` or
`missing_values` have been found were removed.
"""
function drop_missing_ROT(table; missing_values::Vector=[])
    return materializer(table)(drop_missing(table; missing_values=missing_values))
end

"""
    add_index_ROT(table)
Returns a new table of the original `table` type where a new column being the row index
for the table passed have been added.
"""
function add_index_ROT(table)
    return materializer(table)(add_index(table))
end

max_stars_repo_name: TheRoniOne/Cleaner

max_issues_repo_name: TheRoniOne/Cleaner

max_forks_repo_name: TheRoniOne/Cleaner


text:
export SummedInput

"""
    SummedInput{T, N, L, NI} <: AbstractInput{T, N, L}

An AbstractInput iterator which simply sums the results
of several component AbstractInput iterators.
"""
struct SummedInput{T, N, L, NI} <: AbstractInput{T, N, L}
    iters::NTuple{NI, AbstractInput{T}}
end

function SummedInput(iters::Vararg{AbstractInput{<:T, N}, NI}) where {T<:Real, N, NI}
    L = maximum(length.(iters))
    return SummedInput{T, N, L, NI}(iters)
end

function Base.iterate(iter::SummedInput{T, N, L, NI}, i=1) where {T, N, L, NI}
    if i > L
        return nothing
    end

    out = zeros(T, N)
    for itr in iter.iters
        arr, _ = iterate(itr, i)
        if arr != nothing
            out .+= arr
        end
    end
    return out, i+1
end


text:
"""
    corrplot(data)

Creates an `EChart` showing pairwise correlations between columns of data.

## Methods
```julia
corrplot(m::Matrix)
corrplot(df::AbstractDataFrame)
```

## Arguments
* `bubblesize::Int = 45` : size of bubbles
* `layout::String = "lower"` : one of {"lower", "upper", nothing}
* `kwargs` : varargs to set any field of resulting `EChart` struct

## Notes

`corrplot` does not currently support missing values. Users should call `completecases!()` when using a DataFrame, or manually fill missing values as appropriate.

## Examples
```julia
using ECharts, DataFrames, RDatasets
df = dataset("datasets", "mtcars")
cplot = corrplot(df)
```
"""
function corrplot(m::Matrix;
                  bubblesize::Int = 45,
                  layout::String = "lower",
                  labels::Bool = true,
                  ec_height::Real = 650,
                  ec_width::Real = 650,
                  kwargs...)

    #Validate input, since cor in StatsBase returns plain Matrix
    size(m)[1] != size(m)[2] ? error("Matrix needs to be square for a corrplot") : nothing

    #Copy to avoid mutating
    m_ = m
    if layout == "lower"
        tril!(m_)
    elseif layout == "upper"
        triu!(m_)
    else
        nothing
    end

    #convert back to df, melt to pass to arrayofarray
    cor_df = DataFrame(m_)
    original_names = names(cor_df)
    names!(cor_df, [Symbol(x) for x in 1:length(names(cor_df))])
    cor_df[!, :rownames] = names(cor_df)
    df_melt = melt(cor_df, [:rownames])

    #convert to numeric, fixes weird echarts api requiring numeric series in order to rename
    df_melt[!, :variable] = [parse(Int, string(x)) - 1 for x in df_melt[!, :variable]]
    df_melt[!, :rownames] = [parse(Int, string(x)) - 1 for x in df_melt[!, :rownames]]

    #plot
    ec = newplot(kwargs, ec_charttype = "corrplot")
    ec.ec_height = ec_height
    ec.ec_width = ec_width
    ec.xAxis = [Axis(data = 1:length(original_names), _type = "category")]
    ec.yAxis = [Axis(data = 1:length(original_names), _type = "category", inverse = true)]

    ec.series = [XYSeries(data = arrayofarray(df_melt[!, :variable], df_melt[!, :rownames], df_melt[!, :value]),
                        _type = "scatter",
                        symbolSize = JSFunction("""function (data) {return $bubblesize * Math.sqrt(Math.abs(data[2]))}""")
                        )
                ]

    ec.visualMap = VisualMap(left = "right", top = "middle",
                             min = -1.0,
                             max = 1.0,
                             inRange = Dict("color" => colorpalettes["RdBu"]["9"]),
                             calculable = true,
                             precision = 2
                            )

    #overlay labels. figure out what might be worth a keyword in method definition
    labels ?
    ec.series[1].label = ECharts.Label(show = true,
                                      position = "inside",
                                      formatter = JSFunction("""function (params) {return params.data[2].toFixed(2);}"""),
                                      textStyle = TextStyle(fontWeight = "bold", color = "black", fontSize = 14)

                                  ) : nothing
    ec

end

function corrplot(df::AbstractDataFrame;
                  bubblesize::Int = 45,
                  layout::String = "lower",
                  labels::Bool = true,
                  ec_height::Real = 650,
                  ec_width::Real = 650,
                  kwargs...)

    #get numeric columns from df, call cor from StatsBase
    df_num = df[!, [x <: Union{Missing, Number} for x in eltypes(df)]]
    c = cor(convert(Matrix, df_num))

    ec = corrplot(c, bubblesize = bubblesize, layout = layout, labels = labels, ec_height = ec_height, ec_width = ec_width, kwargs...)

    #add column names
    yaxis!(ec, data = names(df_num))
    xaxis!(ec, data = names(df_num))

    ec

end


text:
function forward(backend::GPUBackend, state::BinaryCrossEntropyLossLayerState, inputs::Vector{Blob})
  pred    = inputs[1]
  label   = inputs[2]
  data_type = eltype(pred)

  num = get_num(pred)
  dim = length(pred)

  x_block = Int(ceil(convert(Float64, dim)/CUDA.THREADS_PER_BLOCK_X))

  loss_blob = make_zero_blob(backend, Float32, 1, 1, 1, 1)

  if data_type == Float32
    kernel = backend.mocha.binary_cross_entropy_loss_forward_float
  elseif data_type == Float64
    kernel = backend.mocha.binary_cross_entropy_loss_forward_double
  else
    error("Unsupported data type $data_type")
  end

  CUDA.launch(kernel, x_block, (CUDA.THREADS_PER_BLOCK_X, 1),
        (pred.ptr.p, label.ptr.p, dim, loss_blob.ptr.p))

  loss = Float32[0]
  copy!(loss, loss_blob)
  state.loss = state.layer.weight * loss[1] / num
  destroy(loss_blob)
end

function backward(backend::GPUBackend, state::BinaryCrossEntropyLossLayerState, inputs::Vector{Blob}, diffs::Vector{Blob})
  if !any([isa(d, CuTensorBlob) for d in diffs])
    return
  end

  pred    = inputs[1]
  label   = inputs[2]
  data_type = eltype(pred)

  num = get_num(pred)
  dim = length(pred)

  x_block = Int(ceil(convert(Float64, dim)/CUDA.THREADS_PER_BLOCK_X))

  if data_type == Float32
    kernel = backend.mocha.binary_cross_entropy_loss_backward_float
  elseif data_type == Float64
    kernel = backend.mocha.binary_cross_entropy_loss_backward_double
  else
    error("Unsupported data type $data_type")
  end

  null_ptr = convert(Ptr{data_type}, 0)
  grad_pred =  isa(diffs[1], CuTensorBlob) ? diffs[1].ptr.p : null_ptr
  grad_label = isa(diffs[2], CuTensorBlob) ? diffs[2].ptr.p : null_ptr

  CUDA.launch(kernel, x_block, (CUDA.THREADS_PER_BLOCK_X, 1),
        (pred.ptr.p, label.ptr.p, dim, grad_pred, grad_label,
         convert(data_type, state.layer.weight/num)))

end


text:
export run_tests

using ProgressMeter

show_progress = true

function set_show_progress(val::Bool)
  global show_progress = val
end

# We use a GUID to make sure that we don't accidentaly collide with other
# values in task_local_storage
const COLLECT_MODULE = :fixtures_runtests_collect_module_75dc91fb0fb74aef840d4c36c0c6ee69

function collect_module(m::Module)
  # If we have collector set then we'll call that if not we will run tests
  # this allows users to run individual modules directly while still allowing
  # defered execution.
  const func = try
                 task_local_storage(COLLECT_MODULE)
               catch
                 run_tests
               end
  func(m)
end

push_test!(tests::Vector{Tuple{AbstractString, Function}}, name::AbstractString, test::Any) = nothing

function push_test!(tests::Vector{Tuple{AbstractString, Function}}, name::AbstractString, test::Function)
  function run_test()
    test()
    return true, :green, "$name PASSED"
  end
  push!(tests, (name, run_test))
end

run_tests(testdir::AbstractString=".") = run_tests(findtestfiles(testdir))

function run_tests(filenames::Vector{AbstractString})
  function collector(m::Module)
    push!(test_modules, m)
  end

  const test_modules = Module[]
  task_local_storage(COLLECT_MODULE, collector) do
    for filename in filenames
      include(filename)
    end
  end
  run_tests(test_modules)
end

function run_tests(modules::Vector{Module})
  const tests = Tuple{AbstractString, Function}[]
  for m in modules, name in Base.names(m, true)
    val = getfield(m, name)
    if startswith(string(name), "test_")
      push_test!(tests, "$m.$name", val)
      if show_progress
        ProgressMeter.printover("found $(length(tests)) tests...")
      end
    end
  end

  const pm = Progress(length(tests), 0, "Running $(length(tests)) tests ", 30)
  const results = Tuple{Bool, Symbol, AbstractString}[]
  const failures = Tuple{AbstractString, AbstractString, AbstractString}[]
  for (i, (name, test)) in enumerate(sort(tests, by=x->x[1]))
    _, output, err = capture_output() do
      push!(results, test())
    end

    if err!=nothing
      push!(results, (false, :red, "$name FAILED"))
      push!(failures, (name, output, stringify_error(err, catch_backtrace())))
    end

    if show_progress
      next!(pm, all(x->x[1], results) ? :green : :red)
    end
  end

  const allok = isempty(results) || all(x->x[1], results)
  if !allok
    println("\n", underline("Tests:"), "\n")
    for (ok, color, result) in results
      print_with_color(color, result, "\n")
    end
  end

  if !isempty(failures)
    println("\n\n",centre("Failures"), "\n")
    for (name, output, err) in failures
      println("\n", centre("$name", '-'))
      println(err)
      if !isempty(output)
        println("\n", underline("Captured Output:"), "\n\n", output)
      end
      println("\n", breakwater('-'))
    end
  end

  if !isempty(results)
    const status_counts = Dict{AbstractString, Integer}()
    for (ok, colour, result) in results
      const status = lowercase(split(result)[end])
      status_counts[status] = get(status_counts, status, 0) + 1
    end
    const status_summary = join(["$count $status" for (status, count) in sort(collect(status_counts))], " ")
    println("\n", centre(status_summary), "\n")
  end

  return allok ? 0 : 1
end

run_tests(m::Module) = run_tests([m])


text:
#------------------------------------------------------------------------------
"""
    ARMAFilter(x,rho=[],theta=[],theta0=1.0,z0=[])

Calculate ARMA(p,q) transformation of an input series x. Uses explicit loop (instead of
DSP.filt).

# Input
- `x::Vector`:          T-vector with an input series
- `rho::Vector`:        (optional) p-vector of autoregression coefficients, could be []
- `theta::Vector`:      (optional) q-vector of moving average coefficients (lag 1 to q), could be []
- `theta0::Number`:     (optional) scalar, coefficient on x[t] in MA part, [1]
- `z0::Vector`:         (optional) p-vector, initial values of z, eg. [z_1;z_0], []

# Output
- `z::Vector`:          T-vector with output from the filter

# Notice
(a) The process is
    z[t] = rho[1]*z[t-1] + ... + rho[p]*z[t-p] + y[t], where
    y[t] = theta0*x[t]   +  theta[1]*x[t-1] + ... + theta[q]*x[t-q]
(b) The initial values of y are assumed to be zero (or z0)
(c) The (time) order of z0 is as for x
(d) The initial values of x are assumed to be zero
(e) To calculate impulse response functions, use x = [1;zeros(T-1,1)]
(f) There are no initial values in a pure MA and the case of q > p
    is handled with padding with zeros (see the code below)

"""
function ARMAFilter(x,rho=Float64[],theta=Float64[],theta0=1,z0=Float64[])

  (T,q,p) = (length(x),length(theta),length(rho))
  r       = max(p,q)                         #r obs are needed for start-up

  x  = [zeros(r);x]                          #pad with zeros, clumsy but easy
  Tb = r + T

  theta_ud = q > 0 ? reverse([theta0;theta]) : theta0  #flip upside down
  rho_ud   = p > 1 ? reverse(rho)            : rho

  Tz0 = size(z0,1)
  if isempty(z0)                             #to store output in
    z = zeros(Tb)
  else
    z = [zeros(r-p);z0;zeros(Tb-(r-p)-Tz0)]  #pad with zeros if q > p
  end                                        #zeros(r-p) is empty if r==p

  for t = r+1:Tb
      if p == 0                              #pure MA
        z[t] = dot(theta_ud,x[t-q:t])
      else                                   #AR or ARMA
        z[t] = dot(rho_ud,z[t-p:t-1]) + dot(theta_ud,x[t-q:t])
    end
  end
  z = z[r+1:Tb]                              #cut padded part

  return z

end
#------------------------------------------------------------------------------

max_stars_repo_name: PaulSoderlind/FinancialEconometrics

max_issues_repo_name: PaulSoderlind/FinancialEconometrics

max_forks_repo_name: PaulSoderlind/FinancialEconometrics


text:
# ---
# title: 1417. Reformat The String
# id: problem1417
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: String
# link: <https://leetcode.com/problems/reformat-the-string/description/>
# hidden: true
# ---
#
# Given alphanumeric string `s`. ( **Alphanumeric string** is a string
# consisting of lowercase English letters and digits).
#
# You have to find a permutation of the string where no letter is followed by
# another letter and no digit is followed by another digit. That is, no two
# adjacent characters have the same type.
#
# Return _the reformatted string_ or return **an empty string** if it is
# impossible to reformat the string.
#
#
#
# **Example 1:**
#
#
#
#     Input: s = "a0b1c2"
#     Output: "0a1b2c"
#     Explanation: No two adjacent characters have the same type in "0a1b2c". "a0b1c2", "0a1b2c", "0c2a1b" are also valid permutations.
#
#
# **Example 2:**
#
#
#
#     Input: s = "leetcode"
#     Output: ""
#     Explanation: "leetcode" has only characters so we cannot separate them by digits.
#
#
# **Example 3:**
#
#
#
#     Input: s = "1229857369"
#     Output: ""
#     Explanation: "1229857369" has only digits so we cannot separate them by characters.
#
#
# **Example 4:**
#
#
#
#     Input: s = "covid2019"
#     Output: "c2o0v1i9d"
#
#
# **Example 5:**
#
#
#
#     Input: s = "ab123"
#     Output: "1a2b3"
#
#
#
#
# **Constraints:**
#
#   * `1 <= s.length <= 500`
#   * `s` consists of only lowercase English letters and/or digits.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
################################################################################
#
#   DEFINITION OF LUTTINGER TISZA STATES AND GROUNDSTATES
#
#   FILE CONTAINS
#   1) Luttinger Tisza states
#       - abstract type definition
#       - interface function definition
#       - concrete type imlementation
#       - interface implementation
#   2) Luttinger Tisza ground states and ground state manifolds
#       - abstract type definition
#       - ground states in real space
#       - ground state manifolds in k space
#
################################################################################

abstract type AbstractLTState end

function eigenvectors(lts) end
function alphas(lts) end
function spinconfiguration(lts, lattice, u=[1,0,0], v=[0,1,0]) :: LTSpinConfiguration{N,L} end

mutable struct LTState <: AbstractLTState
    eigenvectors :: Vector{Vector{Float64}}
    alphas :: Vector{Complex}
end

abstract type AbstractLTGroundstate{LTS <: AbstractLTState} end

function spinconfiguration(ltgs, index=1) end

mutable struct LTGroundstateRealSpace{LTS,L<:AbstractLattice{S,B,U} where {S,B,U},H<:AbstractBondHamiltonian{L,N} where {L,N}} <: AbstractLTGroundstate{LTS}

    # lattice
    lattice :: L
    # bond hamiltonian
    h_bond :: H

    # the minimal energy
    E_min :: Float64

    # the degenerate LT States
    lt_states :: Vector{LTS}

    # their respective constraint values
    constraint_values :: Vector{Float64}

end

mutable struct LTGroundstateKSpace{LTS,U<:AbstractUnitcell{S,B} where {S,B},H<:AbstractBondHamiltonian{L,N} where {L,N}} <: AbstractLTGroundstate{LTS}

    # unitcell
    unitcell :: U
    # bond hamiltonian
    h_bond :: H

    # the minimal energy
    E_min :: Float64

    # the k_vectors with minimal energy, columns correspond to components
    k_vectors :: Vector{Vector{Float64}}
    # the corresponding LT States
    lt_states :: Vector{LTS}

    # their respective constraint values
    constraint_values :: Vector{Float64}

end