text:
# infrastructure.jl

using CSV
#using DataFrames

include("../api/api.jl")
include("./db.jl")
include("./doc.jl") # database functions

@enum TableName begin
    UNPAID
    PAID
end # enumerator for TableName types

# get last statement number for today
n = 0

process(orders::Array{Order, 1}; path="./test_invoicing.sqlite") = begin
    # connect to db
    db = connect(path)

    # get last order number
    invnbr = 1000 #ToDo

    # create invoices
    invoices = [create(order, "A" * string(invnbr += 1)) for order in orders]

    # archive invoices
    archive(db, string(UNPAID), invoices)

    # close db
    disconnect(db)

    # create journal entries from invoices
    return entries = [conv2entry(inv, 1300, 8000) for inv in invoices]
end # process(path, orders::Array{Order, 1})

#process(bankstm::Array(Bankstatement, 1) = begin
process(invoices::Array{UnpaidInvoice, 1}, stms::Array{BankStatement, 1}; path="./test_invoicing.sqlite") = begin
    # connect to db
    db = connect(path)

    # create array with potential paid invoices based on received bank statements
    potential_paid_invoices = []
    for unpaid_invoice in invoices
      for s in stms # get potential paid invoices
        if occursin(id(unpaid_invoice), descr(s)) # description contains invoice number
          push!(potential_paid_invoices, create(unpaid_invoice, s))
        end
      end
    end

    # convert to an array with PaidInvoice's
    paid_invoices = convert(Array{PaidInvoice, 1}, potential_paid_invoices)

    # archive PaidInvoice's
    archive(db, string(PAID), paid_invoices)

    # close db
    disconnect(db)

    # return array with JournalEntry's
    return entries = [conv2entry(inv, 1150, 1300) for inv in paid_invoices]
end # process(path, invoices::Array{UnpaidInvoice, 1}, stms::Array{BankStatement, 1})

read_bank_statements(path::String) = begin
    # read the CSV file containing bank statements
    df = CSV.read(path) # returns a DataFrame

    # return an array with BankStatement's
    # row[1] is the first value of row, row[2] the second value, etc.
    return [BankStatement(row[1], row[2], row[3], row[4]) for row in eachrow(df)]
end # read_bank_statements

retrieve_unpaid_invoices(;path="./test_invoicing.sqlite")::Array{UnpaidInvoice, 1} = begin
    # connect to db
    db = connect(path)

    # retrieve unpaid invoices as dataframe
    unpaid_records = retrieve(db, string(UNPAID))

    # convert the dataframe to an array with UnpaidInvoice's.
    # row is an array with one element, which is an array.
    # row[1] is the the content of the element, the UnpaidInvoice.
    unpaid_invoices = [row[1] for row in eachrow(unpaid_records.item)]

    # close db
    disconnect(db)

    # return the array with UnpaidInvoice's that
    return unpaid_invoices
end # retrieve_unpaid_invoices

retrieve_paid_invoices(;path="./test_invoicing.sqlite")::Array{PaidInvoice, 1} = begin
    # connect to db
    db = connect(path)

    # retrieve unpaid invoices as dataframe
    paid_records = retrieve(db, string(PAID))

    # convert the dataframe to an array with UnpaidInvoice's.
    # row is an array with one element, which is an array.
    # row[1] is the the content of the element, the UnpaidInvoice.
    paid_invoices = [row[1] for row in eachrow(paid_records.item)]

    # close db
    disconnect(db)

    # return the array with UnpaidInvoice's that
    return paid_invoices
end # retrieve_unpaid_invoices


text:
module BetaTest

using Test
using ForneyLab
using ForneyLab: outboundType, isApplicable, prod!, unsafeMean, unsafeLogMean, unsafeMirroredLogMean, unsafeVar, vague, dims, logPdf, naturalParams, standardDistribution
using ForneyLab: SPBetaOutNPP, SPBetaAMNM, SPBetaBMMN, SPBetaOutNMM, VBBetaOut, VBBetaA, VBBetaB
using SpecialFunctions: digamma

@testset "Beta ProbabilityDistribution and Message construction" begin
    @test ProbabilityDistribution(Univariate, Beta, a=2.0, b=3.0) == ProbabilityDistribution{Univariate, Beta}(Dict(:a=>2.0, :b=>3.0))
    @test_throws Exception ProbabilityDistribution(Multivariate, Beta)
    @test ProbabilityDistribution(Beta, a=2.0, b=3.0) == ProbabilityDistribution{Univariate, Beta}(Dict(:a=>2.0, :b=>3.0))
    @test ProbabilityDistribution(Beta) == ProbabilityDistribution{Univariate, Beta}(Dict(:a=>1.0, :b=>1.0))
    @test Message(Beta) == Message{Beta, Univariate}(ProbabilityDistribution{Univariate, Beta}(Dict(:a=>1.0, :b=>1.0)))
    @test Message(Univariate, Beta) == Message{Beta, Univariate}(ProbabilityDistribution{Univariate, Beta}(Dict(:a=>1.0, :b=>1.0)))
    @test_throws Exception Message(Multivariate, Beta)
end

@testset "dims" begin
    @test dims(ProbabilityDistribution(Beta, a=2.0, b=2.0)) == ()
end

@testset "vague" begin
    @test vague(Beta) == ProbabilityDistribution(Beta, a=1.0, b=1.0)
end

@testset "unsafe mean and variance" begin
    @test unsafeMean(ProbabilityDistribution(Beta, a=2.0, b=2.0)) == 0.5
    @test unsafeLogMean(ProbabilityDistribution(Beta, a=2.0, b=3.0)) == digamma(2.0) - digamma(5.0)
    @test unsafeMirroredLogMean(ProbabilityDistribution(Beta, a=2.0, b=3.0)) == digamma(3.0) - digamma(5.0)
    @test unsafeVar(ProbabilityDistribution(Beta, a=2.0, b=2.0)) == 0.05
end

@testset "log pdf" begin
    @test isapprox(logPdf(ProbabilityDistribution(Beta, a=2.0, b=2.0),0.3), 0.2311117209633866)
end

@testset "prod!" begin
    @test ProbabilityDistribution(Beta, a=2.0, b=2.0) * ProbabilityDistribution(Beta, a=2.0, b=3.0) == ProbabilityDistribution(Beta, a=3.0, b=4.0)
    @test ProbabilityDistribution(Univariate, Beta, a=1.0, b=2.0) * ProbabilityDistribution(Univariate, PointMass, m=0.2) == ProbabilityDistribution(Univariate, PointMass, m=0.2)
    @test ProbabilityDistribution(Univariate, PointMass, m=0.2) * ProbabilityDistribution(Univariate, Beta, a=1.0, b=2.0) == ProbabilityDistribution(Univariate, PointMass, m=0.2)
    @test_throws Exception ProbabilityDistribution(Univariate, PointMass, m=-1.0) * ProbabilityDistribution(Univariate, Beta, a=1.0, b=2.0)
end

@testset "natural parameters" begin
    d = ProbabilityDistribution(Univariate, Beta, a=2.0, b=3.0)
    η = naturalParams(d)
    s = standardDistribution(Univariate, Beta, η=η)
    @test d.params[:a] == s.params[:a] # Test conversion consistency
    @test d.params[:b] == s.params[:b]

    x = [0.1, 0.6, 0.9]
    d_x = logPdf.([d], x)
    η_x = logPdf.(Univariate, Beta, x; η=η)
    @test isapprox(d_x, η_x) # Test pdf consistency
end

#-------------
# Update rules
#-------------

@testset "SPBetaOutNPP" begin
    @test SPBetaOutNPP <: SumProductRule{Beta}
    @test outboundType(SPBetaOutNPP) == Message{Beta}
    @test isApplicable(SPBetaOutNPP, [Nothing, Message{PointMass}, Message{PointMass}])
    @test !isApplicable(SPBetaOutNPP, [Nothing, Message{Gaussian}, Message{PointMass}])

    @test ruleSPBetaOutNPP(nothing, Message(Univariate, PointMass, m=2.0), Message(Univariate, PointMass, m=3.0)) == Message(Univariate, Beta, a=2.0, b=3.0)
end

@testset "SPBetaOutNMM" begin
    @test SPBetaOutNMM <: SumProductRule{Beta}
    @test outboundType(SPBetaOutNMM) == Message{SampleList}
    @test !isApplicable(SPBetaOutNMM, [Nothing, Message{PointMass}, Message{PointMass}])
    @test isApplicable(SPBetaOutNMM, [Nothing, Message{FactorNode}, Message{PointMass}])
    @test isApplicable(SPBetaOutNMM, [Nothing, Message{PointMass}, Message{FactorNode}])

    msg = ruleSPBetaOutNMM(nothing, Message(Univariate, PointMass, m=2.0), Message(Univariate, Gamma, a=300.0, b=100.0))
    @test length(msg.dist.params[:s]) == length(msg.dist.params[:w])
end

@testset "SPBetaAMNM" begin
    @test SPBetaAMNM <: SumProductRule{Beta}
    @test outboundType(SPBetaAMNM) == Message{Function}
    @test isApplicable(SPBetaAMNM, [Message, Nothing, Message])
end

@testset "SPBetaBMMN" begin
    @test SPBetaBMMN <: SumProductRule{Beta}
    @test outboundType(SPBetaBMMN) == Message{Function}
    @test !isApplicable(SPBetaBMMN, [Message, Nothing, Message])
    @test isApplicable(SPBetaBMMN, [Message, Message, Nothing])
end

@testset "VBBetaOut" begin
    @test VBBetaOut <: NaiveVariationalRule{Beta}
    @test outboundType(VBBetaOut) == Message{Beta}
    @test isApplicable(VBBetaOut, [Nothing, ProbabilityDistribution, ProbabilityDistribution])

    @test ruleVBBetaOut(nothing, ProbabilityDistribution(Univariate, PointMass, m=2.0), ProbabilityDistribution(Univariate, PointMass, m=3.0)) == Message(Univariate, Beta, a=2.0, b=3.0)
end

@testset "VBBetaA" begin
    @test VBBetaA <: NaiveVariationalRule{Beta}
    @test outboundType(VBBetaA) == Message{Function}
    @test isApplicable(VBBetaA, [ProbabilityDistribution, Nothing, ProbabilityDistribution])
end

@testset "VBBetaB" begin
    @test VBBetaB <: NaiveVariationalRule{Beta}
    @test outboundType(VBBetaB) == Message{Function}
    @test isApplicable(VBBetaB, [ProbabilityDistribution, ProbabilityDistribution, Nothing])
end

@testset "averageEnergy and differentialEntropy" begin
    @test isapprox(differentialEntropy(ProbabilityDistribution(Univariate, Beta, a=2.0, b=3.0)), averageEnergy(Beta, ProbabilityDistribution(Univariate, Beta, a=2.0, b=3.0), ProbabilityDistribution(Univariate, PointMass, m=2.0), ProbabilityDistribution(Univariate, PointMass, m=3.0)))
end

end # module

ext: jl


text:
function populationrate(spikes::Array{Spike,1}, windowdur, hopsize)
	ts = [s.t for s ∈ spikes]
	t_min = minimum(ts)
	t_max = maximum(ts)
	ts_bases = t_min:hopsize:t_max-windowdur
	rates = zeros(length(ts_bases))
	for (k, t_base) ∈ enumerate(ts_bases)
		# i1 = searchsortedfirst(ts, t_base)
		# i2 = searchsortedlast(ts, t_base+windowdur)
		# rates[k] = length(ts[i1:i2])/windowdur
		spikes = filter(t -> t_base <= t < t_base + windowdur, ts)
		rates[k] = length(spikes)/windowdur
	end
	(ts_bases, rates)
end

gaussian(x, x0, σ) = 1.0/(σ*√(2*π)) * exp(-0.5 * (x - x0)^2 / σ^2 )

function imagefromspikes(spikes::Array{Spike,1}, t1, t2, n1, n2, fs=1000.0, σ=0.002)
	spikes = filter( s -> t1 <= s.t < t2, spikes )
	image = zeros(n2-n1+1, Int( round((t2-t1)*fs )) )
	if length(spikes) == 0
		return image
	end
	@assert minimum(s.n+1 for s ∈ spikes) >= n1 && maximum(s.n for s ∈ spikes) <= n2
	sigmalen = Int(round(3σ*fs))
	ts = range(-3σ, 3σ, length=2*sigmalen+1)
	bump = gaussian.(ts, 0.0, σ)
	bumplen = length(bump)
	halfbumplen = Int(floor(bumplen/2))
	centeridx = Int(ceil(bumplen/2))
	for s ∈ spikes
		n_idx = s.n - n1 + 1 # + 2 to also account for 0 indexing in C
		t_idx = Int(round((s.t - t1)*fs)) + 1
		t_lo = max(1, t_idx - halfbumplen)
		t_hi = min(size(image, 2), t_idx + halfbumplen)
		offset_lo = t_idx - t_lo
		offset_hi = t_hi - t_idx
		image[n_idx, t_lo:t_hi] .+= bump[centeridx-offset_lo:centeridx+offset_hi]
	end
	image
end

function populationpst(spikes::Array{Spike,1}, inputtimes, dur, n1, n2,
							fs, σ=0.002)
	cumimage = imagefromspikes(spikes, inputtimes[1], inputtimes[1]+dur,
							   n1, n2, fs, σ)
	count = 0
	for time ∈ inputtimes[2:end]
		image = imagefromspikes(spikes, time, time+dur, n1, n2, fs, σ)
		if size(image) == size(cumimage) # in case at end of data
			cumimage .+= image
			count += 1
		end
	end
	cumimage ./ maximum(cumimage)#count
end

function channelsort(spikes::Array{Spike,1}, n::Int)
	times = [Array{Float64,1}(undef,0) for _ ∈ 1:n]
	for s ∈ spikes
		push!(times[s.n], s.t)
	end
	times
end

function sparseac(spiketimes::Array{Float64,1})
	actimes = zeros(length(spiketimes)^2)
	c = 1
	for τ ∈ spiketimes
		for t ∈ spiketimes
			actimes[c] = t - τ
			c += 1
		end
	end
	return actimes |> sort
end

function sparseac(spikes::Array{Spike,1})
	spiketimes = [s.t for s ∈ spikes]
	sparseac(spiketimes)
end

function densehalfac(ts::Array{Float64,1}, fs; norm=true)
	maxdelay = abs(maximum(ts) - minimum(ts))
	maxsamples = Int(round(maxdelay * fs))
	n = maxsamples
	ac = zeros(Float64, n+1)
	for i ∈ 1:length(ts)
		for j ∈ i:length(ts)
			ac[Int(round((ts[j]-ts[i])*fs))+1] += 1.0
		end
	end
	norm && ac ./= ac[1]
	times = (0:n) .* (1/fs)
	return times, ac
end

function densehalfac(spikes::Array{Spike,1}, fs; norm=true)
	ts = [s.t for s ∈ spikes]
	densehalfac(ts, fs; norm=norm)
end

function denseac(spiketimes::Array{Float64,1}, fs; norm=true)
	maxdelay = abs(maximum(spiketimes) - minimum(spiketimes))
	maxsamples = Int(round(maxdelay * fs))
	ac = zeros(2*maxsamples+1)
	mid = maxsamples+1
	for τ ∈ spiketimes
		for t ∈ spiketimes
			ac[Int(round((t-τ)*fs)) + mid] += 1.0
		end
	end
	norm && ac ./= ac[mid]
	times = (-maxsamples:maxsamples) .* (1/fs)
	return times, ac
end

function denseac(spikes::Array{Spike,1}, fs; norm=true)
	spiketimes = [s.t for s ∈ spikes]
	denseac(spiketimes, fs)
end

function spiketodense(spikes::Array{Spike,1}, t1, t2, fs)
	ts = t1:(1.0/fs):t2
	densetrain = zeros(length(ts))
	for spike ∈ spikes
		densetrain[ Int(round((spike.t -t1)*fs)) + 1 ] = 1.0
	end
	return ts, densetrain
end

function corranalysis(spikes::Array{Spike,1}, input::Array{Spike,1},
					  intimes, windur, steps, fs)
	ts, in_ac = densehalfac(input, fs)
	corrvalues = [Array{Float64,1}(undef, length(steps)) for _ ∈ 1:length(intimes)]
	for (i, intime) ∈ enumerate(intimes)
		for (j, step) ∈ enumerate(steps)
			ss = filter(s -> intime+step <= s.t < intime+step+windur, spikes)
			ts, out_ac = densehalfac(ss, fs)
			len = min(length(out_ac), length(in_ac)) # calculate outside loop
			# corrvalues[i][j] = sum(in_ac[1:len] .* out_ac[1:len])
			corrvalues[i][j] = cor([in_ac[1:len]'; out_ac[1:len]']')[1,2]
			println(corrvalues)
		end
	end
	return corrvalues
end

function pstspikes(spikes::Array{Spike,1}, intimes, windur; offset=0.0)
	spikeblocks = [ Array{Spike,1}(undef,0) for _ ∈ 1:length(intimes) ]
	for (k, time) ∈ enumerate(intimes)
		spikeblocks[k] = filter(s -> time+offset <= s.t < time+offset+windur,
							    spikes)
	end
	spikeblocks
end

function averagedelaylength(synapses)
	delays = unique([s.delay for s ∈ synapses]) |> sort
	delay_counts = Dict{UInt16,Float64}()
	delay_vals = Dict{UInt16,Float64}()
	for d ∈ delays
		delay_counts[d] = 0.0
		delay_vals[d] = 0.0
	end

	for s ∈ synapses
		delay_counts[s.delay] += 1.0
		delay_vals[s.delay] += s.strength
	end

	avg_vals = zeros(Float64, length(delays))
	for (k,d) ∈ enumerate(delays)
		avg_vals[k] = delay_vals[d] / delay_counts[d]
	end
	return Int64.(delays), avg_vals
end

## PG search prototyping -- now in C/MPI
# struct PolychronousGroup
# 	nodes::Array{Int64,1}
# 	target::Int64
# 	times::Array{Float64,1}
# 	weights::Array{Float64,1}
# end
#
#
# """
# 	Finds all polychronous groups associated with a set of "base" neurons.
# 	Loops should be fused/optimized (annoying but doable and necessary)
# """
# function testbasegroup(group, graph::DenseGraph{Int64}, weights::DenseGraph{Float64}; threshold=20.0)
# 	s = length(group)
# 	pgs = Array{PolychronousGroup,1}(undef,0)
# 	contacts = Array{Array{Int64,1},1}(undef,s)
# 	for i ∈ 1:s
# 		contacts[i] = graph.contacts[graph.startidcs[group[i]]:graph.startidcs[group[i]]+graph.counts[group[i]]-1]
# 	end
# 	commoncontacts = copy(contacts[1])
# 	for i ∈ 2:s
# 		commoncontacts = filter!(x -> x ∈ contacts[i], commoncontacts)
# 	end
# 	idcs = zeros(Int64,s)
# 	for commoncontact ∈ commoncontacts
# 		for i ∈ 1:s
# 			idcs[i] = findfirst(x -> x == commoncontact, contacts[i])
# 		end
# 		ws = [weights.vals[weights.startidcs[group[k]]+idcs[k]-1] for k ∈ 1:s]
# 		if sum(ws) >= threshold
# 			delays  = [graph.vals[graph.startidcs[group[k]]+idcs[k]-1] for k ∈ 1:s]
# 			delays = maximum(delays) .- delays
# 			push!(pgs, PolychronousGroup(copy(group), commoncontact, delays, ws))
# 		end
# 	end
# 	return pgs
# end
#
#
#
# """
#
# Generate appropriate loop structure with a macro...
# """
# function findpgs(graph, synapses, s=3; threshold=18.0, maxpgs = 10000)
# 	n = size(graph)[1]
# 	group = zeros(Int64, s)
# 	graph = densegraph(graph) 	# remove after resolving graph in C code
# 	#weights = densegraph(weights)
# 	weights = densifysynapses(synapses, n)
# 	pgs = Array{PolychronousGroup,1}(undef,maxpgs)
# 	numpgs = 0
# 	if s == 1
# 		for i ∈ 1:n
# 			group[1] = i
# 			pgnew = testbasegroup(group, graph, weights; threshold=threshold)
# 			if length(pgnew) > 0
# 				push!(pgs, pgnew...)
# 			end
# 		end
# 	elseif s == 2
# 		for i ∈ 1:n
# 			for j ∈ i+1:n
# 				group[1], group[2] = i, j
# 				pgnew = pgingroup(group, graph, weights; threshold=threshold)
# 				if length(pgnew) > 0
# 					push!(pgs, pgnew...)
# 				end
# 			end
# 		end
# 	elseif s == 3
# 		for i ∈ 1:n
# 			println("Finished $i")
# 			for j ∈ i+1:n
# 				for k ∈ j+1:n
# 					group[1],group[2],group[3] = i,j,k
# 					pgnew = pgingroup(group, graph, weights; threshold=threshold)
# 					if length(pgnew) > 0
# 						for a ∈ 1:length(pgnew)
# 							numpgs += 1
# 							pgs[numpgs] = pgnew[a]
# 							if numpgs == maxpgs
# 								return pgs
# 							end
# 						end
# 					end
# 				end
# 			end
# 		end
# 	elseif s == 4
# 		for i ∈ 1:n
# 			for j ∈ i+1:n
# 				for k ∈ j+1:n
# 					for l ∈ k+1:n
# 						group[1],group[2],group[3],group[4] = i,j,k,l
# 						pgnew = pgingroup(group, graph, weights; threshold=threshold)
# 						if length(pgnew) > 0
# 							push!(pgs, pgnew...)
# 						end
# 					end
# 				end
# 			end
# 		end
# 	end
# 	resize!(pgs,numpgs)
# end
#
#
# """
# 	pginspikes(pg::Array{Spike,1},spikes::Array{Spike,1})
#
# Spikes must be time-sorted
# """
# # export pginspikes
# function pginspikes(pg::Array{Spike,1},spikes::Array{Spike,1};
# 					timetol=0.001,hitthreshold=1.0)
# 	baseidcs = findall(x -> x.n == pg[1].n, spikes)
# 	groupinstances = Array{Array{Spike,1},1}(undef,0)
# 	for baseidx ∈ baseidcs
# 		numfound = 0
# 		t1 = spikes[baseidx].t
# 		t2 = t1 + pg[end].t + timetol
# 		k = baseidx
# 		while k < length(spikes)
# 			k += 1
# 			if spikes[k].t > t2
# 				break
# 			end
# 		end
# 		endidx = k
# 		subspikes = @view spikes[baseidx:endidx]
# 		# subspikes = @view spikes[:]
# 		# println(length(subspikes))
# 		localinstance = Array{Spike,1}(undef,length(pg))
# 		localinstance[1] = spikes[baseidx]
# 		numhits = 1
# 		for pgspike ∈ pg[2:end]
# 			for spike ∈ subspikes
# 				if pgspike.n == spike.n
# 					if abs((pgspike.t+t1)-spike.t) <= timetol
# 						numhits += 1
# 						localinstance[numhits] = spike
# 					end
# 				end
# 			end
# 		end
# 		if numhits/length(pg) >= hitthreshold
# 			localinstance = resize!(localinstance,numhits)
# 			push!(groupinstances,localinstance)
# 		end
# 	end
#
# 	return groupinstances
# end
#
# """
# 	resultsintochunks(results::ExperimentOutput, dur=0.35)
#
# Optimize this later.
# """
# function resultsintochunks(results::ExperimentOutput, dur=0.35)
#
# end
#
#
# """
# """
# # export showfoundgroup
# function showfoundgroup(foundgroup, spikes)
# 	t1 = foundgroup[1].t-0.1
# 	t2 = foundgroup[end].t+0.1
# 	p = spikeraster(filter(s->t1<=s.t<=t2, spikes); reuse=false)
# 	spikeraster(p, foundgroup; markersize=7, markeralpha=0.3,markercolor=:red)
# end
#
#

max_stars_repo_name: analogouscircuit/DelayNetTools

max_issues_repo_name: analogouscircuit/DelayNetTools

max_forks_repo_name: analogouscircuit/DelayNetTools


text:
using ReinforcementLearning

using Documenter
using Documenter.Writers.HTMLWriter
using Documenter.Utilities.DOM
using Documenter.Utilities.DOM: Tag, @tags

"filter concrete subtypes"
is_concrete_type_of(t) = x -> begin
    println(x)
    x isa Type && x <: t && !isabstracttype(x)
end

const top_nav = """

    <button class="navbar-toggler" type="button" data-toggle="collapse" data-target="#navbarTogglerDemo01" aria-controls="navbarTogglerDemo01" aria-expanded="false" aria-label="Toggle navigation">

    </button>

          JuliaReinforcementLearning

        Get Started

        Guide

          Contribute

        Blog

        Doc

        Github

"""

function HTMLWriter.render_html(
    ctx,
    navnode,
    head,
    sidebar,
    navbar,
    article,
    footer,
    scripts::Vector{DOM.Node} = DOM.Node[],
)
    @tags html body div script
    DOM.HTMLDocument(
        html[:lang=>"en"](
            head,
            body(
                Tag(Symbol("#RAW#"))(top_nav),
                div[".documenter-wrapper#documenter"](
                    sidebar,
                    div[".docs-main"](navbar, article, footer),
                    HTMLWriter.render_settings(ctx),
                ),
            ),
            scripts...,
        ),
    )
end

makedocs(
    modules = [
        ReinforcementLearning,
        ReinforcementLearningBase,
        ReinforcementLearningCore,
        ReinforcementLearningEnvironments,
        ReinforcementLearningZoo,
    ],
    format = Documenter.HTML(
        prettyurls = true,
        analytics = "UA-149861753-1",
        canonical = "https://juliareinforcementlearning.github.io/ReinforcementLearning.jl/latest/",
        assets = [
            "assets/favicon.ico",
            "assets/custom.css",
            asset(
                "https://stackpath.bootstrapcdn.com/bootstrap/4.5.0/css/bootstrap.min.css",
            ),
        ],
    ),
    sitename = "ReinforcementLearning.jl",
    linkcheck = !("skiplinks" in ARGS),
    pages = [
        "Home" => "index.md",
        "Manual" => [
            "RLBase" => "rl_base.md",
            "RLCore" => "rl_core.md",
            "RLEnvs" => "rl_envs.md",
            "RLZoo" => "rl_zoo.md",
        ],
    ],
)

deploydocs(
    repo = "github.com/JuliaReinforcementLearning/ReinforcementLearning.jl.git",
    target = "build",
    deps = nothing,
    make = nothing,
)


text:
"""
```
    UpdatableGivensQR <: AbstractQR <: Factorization
```
Holds storage for Givens rotations and the R matrix for a QR factorization.
Allows for efficient addition and deletion of columns.
See `add_column!` and `remove_column!`.
"""
mutable struct UpdatableGivensQR{T, QFT, RFT<:AbstractMatrix{T}, QXT, PT} <: AbstractQR{T}
    rotations_full::QFT
    rot_index::Int # need to track it as varying numbers of deletions can create an unpredictible number of rotations
    R_full::RFT
    QX::QXT # holds space for Q'*x when adding columns
    perm_full::PT # allows efficient insertion of column at arbitrary indices, while only appending columns in memory and keeping track of the permutations
    n::Int
    m::Int
end
const UGQR = UpdatableGivensQR
const UpdatableQR = UpdatableGivensQR
const UQR = UpdatableQR

function UpdatableGivensQR(A::AbstractMatrix, r::Int = size(A, 1))
    UpdatableGivensQR!(copy(A), r)
end

# memory footprint: O(nr + r^2)
function UpdatableGivensQR!(A::AbstractMatrix, r::Int = size(A, 1))
    n, m = size(A)
    n ≥ r || throw("maximum rank r = $r has to be smaller or equal to n = $n")
    r ≥ m || throw("maximum rank r = $r has to be larger or equal to m = $m")
    T = eltype(A)
    rotations_full = allocate_rotations(T, n, r)
    rot_index = number_of_rotations(n, m)
    rotations = @view rotations_full[1:rot_index]
    R_full = zeros(T, r, r)
    F = GivensQR!(A, rotations)
    @. R_full[1:m, 1:m] = F.R
    QX = zeros(T, n, r) # stores Q'*X for addition of multiple columns X
    perm_full = collect(1:r)
    UpdatableGivensQR(rotations_full, rot_index, R_full, QX, perm_full, n, m)
end
# constructor that pre-allocates memory without starting a qr factorization
UpdatableGivensQR(n::Int, r::Int) = UpdatableGivensQR(Float64, n, r)
function UpdatableGivensQR(T::DataType, n::Int, r::Int)
    n ≥ r || throw("maximum rank r = $r has to be smaller or equal to n = $n")
    rotations_full = allocate_rotations(T, n, r)
    rot_index = 0
    R_full = zeros(T, r, r)
    QX = zeros(T, n, r)
    perm_full = collect(1:r)
    m = 0
    UpdatableGivensQR(rotations_full, rot_index, R_full, QX, perm_full, n, m)
end

# copying underlying data so that future changes in F don't change GivensQR structure
GivensQR(F::UpdatableGivensQR) = GivensQR(copy(F.Q), copy(F.R), F.n, F.m)
function Base.copy(F::UpdatableGivensQR)
    UpdatableGivensQR(copy(F.rotations_full), F.rot_index, copy(F.R_full),
                      copy(F.Qx), copy(F.perm_full), F.n, F.m)
end

function Base.getproperty(F::UpdatableGivensQR, s::Symbol)
    if s == :Q
        GivensQ(F.rotations, F.n, F.m)
    elseif s == :rotations
        @view F.rotations_full[1:F.rot_index]
    elseif s == :R
        R = @view F.R_full[1:F.m, 1:F.m]
        UpperTriangular(R)
    elseif s == :perm
        @view F.perm_full[1:F.m]
    else
        getfield(F, s)
    end
end
Base.eltype(F::UGQR{T}) where T = T
Base.size(F::UGQR) = (F.n, F.m)
Base.size(F::UGQR, i::Int) = i > 2 ? 1 : size(F)[i]
Base.AbstractMatrix(F::UGQR) = Matrix(F)
function Base.Matrix(F::UQR)
    n, m = size(F)
    A = zeros(eltype(F), n, m)
    @. A[1:m, 1:m] = F.R
    lmul!(F.Q, A) # avoids every explicitly creating F.Q
    ip = invperm(F.perm) # this allocates, IDEA: could pre-allocate
    for c in eachrow(A) # A[:, invperm(F.perm)]
        permute!(c, ip)
    end
    return A
end

# overwrites the first m elements of x with the solution
function LinearAlgebra.ldiv!(F::UQR, x::AbstractVecOrMat)
    ldiv!(F.Q, x)
    xm = typeof(x) <: AbstractVector ? view(x, 1:F.m) : view(x, 1:F.m, :)
    ldiv!(F.R, xm)
    p = invperm(F.perm)
    for xi in eachcol(xm)
        permute!(xi, p)
    end
    return xm
end

"""
```
    add_column!(F::UpdatableGivensQR, x::AbstractVector, k::Int = size(F, 2) + 1)
```
Given an existing QR factorization `F` of a matrix, computes the factorization of
the same matrix after a new column `x` has been added as the `k`ᵗʰ column.
Computational complexity: O(nm) where size(F) = (n, m).
WARNING: Overwrites existing factorization.
"""
function add_column!(F::UpdatableGivensQR, x::AbstractVector, k::Int = size(F, 2) + 1)
    Qx = @view F.QX[:, 1]
    @. Qx = x # so that x doesn't get overwritten
    add_column!!(F, Qx, k)
end
function add_column!(F::UpdatableGivensQR, X::AbstractMatrix, k::AbstractArray{Int} = (size(F, 2) + 1) : (size(F, 2) + size(X, 2)))
    QX = @view F.QX[:, 1:length(k)]
    @. QX = X # so that X doesn't get overwritten
    add_column!!(F, QX, k)
end

"""
Same as `add_column!` but WARNING: mutates `x`!
"""
function add_column!!(F::UpdatableGivensQR, x::AbstractVector, k::Int = size(F, 2) + 1)
    length(x) == size(F, 1) || throw(DimensionMismatch("length of input not equal to first dimension of UpdatableQR factorization"))
    add_permutation!(F, k)
    F.rot_index = append_column!(x, F.rotations_full, F.rot_index, F.R_full, F.m)
    F.m += 1
    return F
end

# adding columns in X to factorization, mutates X!
function add_column!!(F::UpdatableGivensQR, X::AbstractMatrix, k::AbstractArray = (size(F, 2) + 1) : (size(F, 2) + size(X, 2)))
    size(X, 1) == size(F, 1) || throw(DimensionMismatch("length of input not equal to first dimension of UpdatableQR factorization"))
    size(X, 2) == length(k) || throw(DimensionMismatch("length of index array k ($(length(k))) not equal to second dimension of data matrix X ($(size(X, 2)))"))
    add_permutation!(F, k)
    F.rot_index = append_column!(X, F.rotations_full, F.rot_index, F.R_full, F.m)
    F.m += length(k)
    return F
end

# X are columns to be appended to the factorization
# rotations is a vector of Givens rotations,
# R is matrix where R[1:m, 1:m] has already been filled with the factorization of the previously added columns
function append_column!(X::AbstractVecOrMat, rotations::AbstractVector{<:Givens},
                        rot_index::Int, R::AbstractMatrix, m::Int)
    n, k = size(X, 1), size(X, 2) # if X is a vector, size(X, 2) = 1
    size(R, 2) > m || throw("cannot add another column to factorization with current rank $m and maximum rank $(size(R, 2))")

    # check if there's enough space in rotations left to carry out this addition
    ensure_space_to_append_column!(rotations, rot_index, n, m, k)

    Q = GivensQ(@view(rotations[1:rot_index]), n, m)
    QX = lmul!(Q', X) # overwriting X
    for j in 1:k
        QXj = @view QX[:, j:end] # do not need to treat columns 1:j-1 further
        # for i in n:-1:m+2 # zero out the "spike"
        for i in n:-1:m+1+j # zero out the "spike"
            G, r = givens(QXj[i-1, 1], QXj[i, 1], i-1, i)
            lmul!(G, QXj)
            rot_index += 1
            rotations[rot_index] = G
        end
    end
    @. R[1:m+k, m+1:m+k] = QX[1:m+k, :] # only necessary if these are not already equal in memory
    return rot_index
end

function add_permutation!(F::UpdatableQR, k::Union{Int, AbstractVector{Int}})
    F.m + length(k) ≤ length(F.perm_full) || throw("cannot add $(length(k)) columns to factorization with current rank $(F.m) and maximum rank $(length(F.perm_full))")
    for (i, p) in enumerate(F.perm)
        for ki in k
            if p ≥ ki # incrementing indices above ki
                F.perm[i] += 1
            end
        end
    end
    F.perm_full[F.m+1:F.m+length(k)] .= k
    return F
end

# current factorization size is (n, m), number of columns to be added is k
function ensure_space_to_append_column!(rotations::AbstractVector{<:Givens},
                        rot_index::Int, n::Int, m::Int, k::Int = 1, verbose::Bool = false)
    nrot = length(rotations) - rot_index # number of rotations left to add in rotations vector
    nrot_to_add = number_of_rotations_to_append_column(n, m, k)
    if nrot < nrot_to_add
        verbose && println("INFO: adding more memory for Givens rotations in append_column!")
        append!(rotations, Vector{eltype(rotations)}(undef, nrot_to_add - nrot))
    end
    return rotations
end

"""
```
remove_column!(F::UpdatableGivensQR, k::Int = size(F, 2))
```
Updates the existing QR factorization `F` of a matrix to the factorization of
the same matrix after its kᵗʰ column has been deleted.
Computational complexity: O(m²) where size(F) = (n, m).
"""
function remove_column!(F::UpdatableGivensQR, k::Int = size(F, 2))
    i = remove_permutation!(F, k)
    F.rot_index = remove_column!(F.R, F.rotations_full, F.rot_index, i)
    F.m -= 1
    return F
end

function remove_column!(R::AbstractMatrix, rotations::AbstractVector{<:Givens},
                        rot_index::Int, k::Int = size(F, 2))
    m = size(R, 2)
    1 <= k <= m || throw(DimensionMismatch("index $k not in range [1, $m]"))

    # check if there's enough space in rotations left to carry out this removal
    ensure_space_to_remove_column!(rotations, rot_index, m, k)

    for i in k+1:m # zero out subdiagonal of submatrix following kth column
        Ri = @view R[:, i:end]
        G, rho = givens(R[i-1, i], R[i, i], i-1, i)
        lmul!(G, Ri)
        rot_index += 1
        rotations[rot_index] = G
    end
    @inbounds for j in k:m-1 # overwrite deleted column
        @simd for i in 1:j # guarantees we only access upper triangular part
            R[i, j] = R[i, j+1]
        end
    end
    return rot_index # this is generally m-k, important to increment the rot_index of the factorization
end

# returns the index i of the column in R that is referenced by the kth column of the factorization
# those are not equal because of the permutation of the columns given by F.perm
function remove_permutation!(F::UpdatableQR, k::Int)
    i = findfirst(==(k), F.perm)
    isnothing(i) && throw(DimensionMismatch("index k = $k not found in F.perm, can't remove the associated column"))
    perm = F.perm
    for j in i+1:F.m # shifts all indices above k down to keep active indices adjacent in memory
        perm[j-1] = perm[j]
    end
    for (i, p) in enumerate(perm) # decrements indices above k (since we are removing it)
        if p ≥ k
            perm[i] = p-1
        end
    end
    return i
end

function ensure_space_to_remove_column!(rotations::AbstractVector{<:Givens},
                        rot_index::Int, m::Int, k::Int, verbose::Bool = false)
    nrot = length(rotations) - rot_index
    nrot_to_remove = number_of_rotations_to_remove_column(m, k)
    if nrot < nrot_to_remove
        verbose && println("INFO: adding more memory for Givens rotations in remove_column!")
        append!(rotations, Vector{eltype(rotations)}(undef, nrot_to_remove-nrot))
    end
    return rotations
end


text:
##  Simply-supported/clamped beam modal analysis. Corotational beam
#
# Vibration analysis of beam which is simply supported in one plane, and clamped in another.
# The results are compared with analytical expressions.
module beam_modal_examples

using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsFlexBeams.CrossSectionModule: CrossSectionRectangle
using FinEtoolsFlexBeams.MeshFrameMemberModule: frame_member
using FinEtoolsFlexBeams.FEMMCorotBeamModule
using FinEtoolsFlexBeams.FEMMCorotBeamModule: FEMMCorotBeam
stiffness = FEMMCorotBeamModule.stiffness
mass = FEMMCorotBeamModule.mass
using FinEtoolsFlexBeams.RotUtilModule: initial_Rfield
using LinearAlgebra: dot
using Arpack
using LinearAlgebra
using SparseArrays
using PlotlyJS
using Test

function beam_modal()
##
# Parameters:
E=30002.0*1000;#30002ksi
nu=0.0;
b=18; h=18; L=300;# in, cross-section and length
g=32.17*12; # in/second^2
rho=0.2/g;# Pound/g mass density per unit volume

##
# Cross-sectional properties
cs = CrossSectionRectangle(s -> b, s -> h, s -> [1.0, 0.0, 0.0])

##
# Choose the mass formulation:
mass_type=2;

##
# Analytical frequencies
#
# The analytical frequencies were taken from table 8-1 of Formulas for
# natural frequency and mode shape, Robert D. Blevins, Krieger publishing
# company, Malabar Florida, reprint edition 2001.
#
# The beam has cylindrical supports at either end. Therefore in the
# vertical plane it behaves as a simply supported beam, while in the
# horizontal plane it behaves as a clamped beam.
#
# Vertical plane, simply supported:
# Horizontal plane, clamped:
A=b*h;# in^2
I2=b*h^3/12;# cm^4
I3=b^3*h/12;# cm^4
@show analyt_Frequencies = [(1*pi)^2/(2*pi*L^2)*sqrt(E*I2/rho/A),
(4.73004074)^2/(2*pi*L^2)*sqrt(E*I2/rho/A)];
neigvs = length(analyt_Frequencies);

# Select the number of elements per half the length.
xyz = [[0 -L/2 0];[0 L/2 0]]
n=4;
fens, fes = frame_member(xyz, n, cs)
@show fens

# Material properties
material = MatDeforElastIso(DeforModelRed3D, rho, E, nu, 0.0)

##
# Construct the requisite fields, geometry and displacement
# Initialize configuration variables
geom0 = NodalField(fens.xyz)
u0 = NodalField(zeros(size(fens.xyz,1), 3))
Rfield0 = initial_Rfield(fens)
dchi = NodalField(zeros(size(fens.xyz,1), 6))

# Apply EBC's
l1 = selectnode(fens; box = [0 0 -L/2 -L/2 0 0], tolerance = L/10000)
for i in [1,2,3,5,6]
    setebc!(dchi, l1, true, i)
end
l1 = selectnode(fens; box = [0 0 L/2  L/2 0 0], tolerance = L/10000)
for i in [1,2,3,5,6]
    setebc!(dchi, l1, true, i)
end
applyebc!(dchi)
numberdofs!(dchi);

# Assemble the global discrete system
femm = FEMMCorotBeam(IntegDomain(fes, GaussRule(1, 2)), material)
K = stiffness(femm, geom0, u0, Rfield0, dchi);
M = mass(femm, geom0, u0, Rfield0, dchi);

# Solve the eigenvalue problem
d,v,nconv = eigs(K, M; nev=2*neigvs, which=:SM, explicittransform=:none)
fs = real(sqrt.(complex(d)))/(2*pi)
println("Eigenvalues: $fs [Hz]")

# #
# for i=1:neigvs
#     disp(['  Eigenvector ' num2str(i) ' frequency ' num2str(Frequencies(i)) ' (analyt=' num2str(analyt_Frequencies(i)) ') [Hz]']);
# end
# disp(' ')
return true
end # beam_modal

function allrun()
    println("#####################################################")
    println("# beam_modal ")
    beam_modal()
    return true
end # function allrun

end # module


text:
struct NesterovMomentum end

function Base.iterate(::NesterovMomentum, t=1.0)
    t_next = (1 + sqrt(1 + 4 * t^2)) / 2
    return (t - 1) / t_next, t_next
end

struct NesterovIterable{T, F, S, M}
    w0::T
    f::F
    stepsize::S
    momentum::M
end

NesterovIterable(w0, f; stepsize, momentum=NesterovMomentum()) = NesterovIterable(w0, f, to_iterator(stepsize), to_iterator(momentum))

Base.IteratorSize(::Type{<:NesterovIterable}) = Base.IsInfinite()

mutable struct NesterovState{T, G, S, M}
    w::T
    f_w
    grad_f_w::G
    stepsize_iterator::S
    momentum_iterator::M
    z::T
    z_prev::T
end

function Base.iterate(iter::NesterovIterable)
    w = copy(iter.w0)
    grad_f_w, f_w = gradient(iter.f, w)
    state = NesterovState(
        w, f_w, grad_f_w,
        Iterators.Stateful(iter.stepsize), Iterators.Stateful(iter.momentum),
        copy(w), zero(w)
    )
    return IterationOutput(state.w, state.f_w, state.grad_f_w), state
end

function Base.iterate(iter::NesterovIterable, state::NesterovState)
    stepsize = popfirst!(state.stepsize_iterator)
    state.z .= state.w .- stepsize .* state.grad_f_w
    momentum = popfirst!(state.momentum_iterator)
    state.w .= state.z .+ momentum .* (state.z .- state.z_prev)
    state.z_prev, state.z = state.z, state.z_prev
    state.grad_f_w, state.f_w = gradient(iter.f, state.w)
    return IterationOutput(state.w, state.f_w, state.grad_f_w), state
end

struct Nesterov
    kwargs
    Nesterov(; kwargs...) = new(kwargs)
end

(alg::Nesterov)(args...) = NesterovIterable(args...; alg.kwargs...)


text:
function testreflect(func)
    @testset "$func invariance under mirror transforms" begin
        noise = two_phase_noise_3d()
        for phase in 0:1
            f = mean ∘ func
            # Calculate correlation function on the original image
            corr1 = f(noise, phase)
            # Reflect from yOz plane and calculate correlation function
            corr2 = f(noise[end:-1:1,:,:], phase)
            # Reflect from xOz plane and calculate correlation function
            corr3 = f(noise[:,end:-1:1,:], phase)
            # Reflect from xOy plane and calculate correlation function
            corr4 = f(noise[:,:,end:-1:1], phase)
            @test corr1 ≈ corr2 ≈ corr3 ≈ corr4
        end
    end
end

testreflect(Directional.s2)
testreflect(Directional.l2)
testreflect(Directional.c2)
testreflect(Directional.surfsurf)
# TODO: somehow test histograms returned by pore_size and chord_length

function testsurface(func)
    @testset "Check $(func)⁰(a) = $(func)¹(a) for two phase media" begin
        noise = two_phase_noise_3d()
        @test mean(func(noise, 0; len = 50)) ≈
              mean(func(noise, 1; len = 50))
    end
end

testsurface(Directional.surfsurf)
testsurface(Directional.surfvoid)


text:
# show method for html output of question types
function _markdown_to_html(x)
    length(x) == 0 && return("")
    x = Markdown.parse(x)
    x = sprint(io -> Markdown.html(io, x))
    x = replace(x, r"^
"=>"", r"
$"=>"")
    return x
end

function Base.show(io::IO, m::MIME"text/html", x::Numericq)

    ID = randstring()

    FORM = Mustache.render(html_templates["inputq_form"];
                           ID=ID,
                           PLACEHOLDER = "Numeric answer",
                           UNITS=x.units,
                           TYPE="number",
                           HINT = length(x.label) == 0 ? x.hint : ""
                           )

    GRADING_SCRIPT =
        Mustache.render(html_templates["input_grading_script"];
                        ID = ID,
                        CORRECT_ANSWER = "(Math.abs(this.value - $(x.val)) <= $(x.tol))"
                        )

    Mustache.render(io,
                    html_templates["question_tpl"];
                    ID = ID,
                    TYPE = "text",
                    STATUS = "",
                    LABEL=_markdown_to_html(x.label),
                    HINT = length(x.label) > 0 ? x.hint : "",
                    FORM = FORM,
                    GRADING_SCRIPT = GRADING_SCRIPT
                    )

end

function Base.show(io::IO, m::MIME"text/html", x::Stringq)

    ID = randstring()

    FORM = Mustache.render(html_templates["inputq_form"];
                           ID=ID,
                           PLACEHOLDER = "Text answer",
                           TYPE="text",
                           HINT = length(x.label) == 0 ? x.hint : ""
                           )

    GRADING_SCRIPT =
        Mustache.render(html_templates["input_grading_script"];
                        ID = ID,
                        CORRECT_ANSWER = """RegExp('$(x.re.pattern)').test(this.value)"""
                        )

    Mustache.render(io, html_templates["question_tpl"];
                    ID = ID,
                    TYPE = "text",
                    STATUS = "",
                    LABEL=_markdown_to_html(x.label),
                    HINT = length(x.label) > 0 ? x.hint : "",
                    FORM = FORM,
                    GRADING_SCRIPT = GRADING_SCRIPT
                    )

end

function _make_item(i, choice)
    choice′ = sprint(io -> Markdown.html(io, Markdown.parse(choice)))
    # strip
 tag
    choice′ = chomp(choice′)
    choice′ = replace(choice′, r"^
" => "")
    choice′ = replace(choice′, r"
$" => "")

    return (NO=i, LABEL=choice′, VALUE=i)
end

function Base.show(io::IO, m::MIME"text/html", x::Radioq)

    ID = randstring()

    choices = string.(x.choices)
    items = [_make_item(i, choice) for (i,choice) ∈ enumerate(choices)]

    GRADING_SCRIPT = Mustache.render(html_templates["radio_grading_script"];
                             ID = ID,
                             CORRECT_ANSWER = x.answer
                             )
    FORM = Mustache.render(html_templates["Radioq"];
                           ID = ID,
                           ITEMS = items,
                           INLINE = x.inline ? " inline" : ""
                           )

    Mustache.render(io, html_templates["question_tpl"],
                    ID = ID,
                    TYPE = "radio",
                    FORM = FORM,
                    GRADING_SCRIPT = GRADING_SCRIPT,
                    LABEL=_markdown_to_html(x.label),
                    HINT = x.hint # use HINT in question
                    )

end

function Base.show(io::IO, m::MIME"text/html", x::Multiq)

    ID = randstring()

    choices = string.(x.choices)
    items = [_make_item(i, choice) for (i,choice) ∈ enumerate(choices)]

    GRADING_SCRIPT = Mustache.render(html_templates["multi_grading_script"];
                                     ID = ID,
                                     CORRECT_ANSWER = x.answer,
                                     INCORRECT = "Not yet",
                                     CORRECT = "Correct"
                             )
    FORM = Mustache.render(html_templates["Multiq"];
                           ID = ID,
                           ITEMS = items,
                           INLINE = x.inline ? " inline" : ""
                           )

    Mustache.render(io, html_templates["question_tpl"],
                    ID = ID,
                    TYPE = "radio",
                    FORM = FORM,
                    GRADING_SCRIPT = GRADING_SCRIPT,
                    LABEL=_markdown_to_html(x.label),
                    HINT = x.hint # use HINT in question
                    )

end


text:
# sorting functionality

"""
The main quicksort kernel uses dynamic parallelism. Let's call blocksize `M`. The first part
of the kernel bubble sorts `M` elements with maximal stride between `lo` and `hi`. If the
sublist is <= `M` elements, `stride` = 1 and no recursion happens. Otherwise, we pick
element `lo + M ÷ 2 * stride` as a pivot. This is an efficient choice for random lists and
pre-sorted lists.

Partition is done in stages:
1. For batches of M values, cumsum how many > pivot are left of each index. The comparison
   alternates between < and <= with recursion depth. This makes no difference when there are
   many unique values, but when there are many duplicates, this effectively partitions into
   <, =, and >.
2. Consolidate batches. This runs inside the quicksort kernel.

Sublists (ranges of the list being sorted) are denoted by `lo` and one of `L` and `hi`. `lo`
is an exclusive lower bound, `hi` is an inclusive upperboard, `L` is their difference.
`b_sums` is "batch sums", the number of values in a batch which are >= pivot or > pivot
depending on the relevant `parity`

Originally developed by @xaellison (Alex Ellison).
"""
module QuickSortImpl

export quicksort!

using ..CUDA
using ..CUDA: i32

# Comparison

@inline function flex_lt(a, b, eq, lt, by)
    a′ = by(a)
    b′ = by(b)
    (eq && a′ == b′) || lt(a′, b′)
end

# Batch partitioning
"""
Performs in-place cumsum using shared memory. Intended for use with indexes
"""
function cumsum!(sums)
    shift = 1

    while shift < length(sums)
        to_add = 0
        @inbounds if threadIdx().x - shift > 0
            to_add = sums[threadIdx().x - shift]
        end

        sync_threads()
        @inbounds if threadIdx().x - shift > 0
            sums[threadIdx().x] += to_add
        end

        sync_threads()
        shift *= 2
    end
end

"""
Partition the region of `values` after index `lo` up to (inclusive) `hi` with
respect to `pivot`. Computes each value's comparison to pivot, performs a cumsum
of those comparisons, and performs one movement using shmem. Comparison is
affected by `parity`. See `flex_lt`. `swap` is an array for exchanging values
and `sums` is an array of Ints used during the merge sort.
Uses block y index to decide which values to operate on.
"""
@inline function batch_partition(values, pivot, swap, sums, lo, hi, parity, lt::F1, by::F2) where {F1,F2}
    sync_threads()
    blockIdx_yz = (blockIdx().z - 1i32) * gridDim().y + blockIdx().y
    idx0 = lo + (blockIdx_yz - 1i32) * blockDim().x + threadIdx().x
    val = idx0 <= hi ? values[idx0] : one(eltype(values))
    comparison = flex_lt(pivot, val, parity, lt, by)

    @inbounds if idx0 <= hi
         sums[threadIdx().x] = 1 & comparison
    else
         sums[threadIdx().x] = 1
    end
    sync_threads()

    cumsum!(sums)

    dest_idx = @inbounds comparison ? blockDim().x - sums[end] + sums[threadIdx().x] : threadIdx().x - sums[threadIdx().x]
    @inbounds if idx0 <= hi && dest_idx <= length(swap)
        swap[dest_idx] = val
    end
    sync_threads()

    @inbounds if idx0 <= hi
         values[idx0] = swap[threadIdx().x]
    end
    sync_threads()
end

"""
Each block evaluates `batch_partition` on consecutive regions of length blockDim().x
from `lo` to `hi` of `values`.
"""
function partition_batches_kernel(values::AbstractArray{T}, pivot, lo, hi, parity, lt::F1,
                                  by::F2) where {T,F1,F2}
    sums = CuDynamicSharedArray(Int, blockDim().x)
    swap = CuDynamicSharedArray(T, blockDim().x, sizeof(sums))
    batch_partition(values, pivot, swap, sums, lo, hi, parity, lt, by)
    return
end

# Batch consolidation

"""
Finds the index in `array` of the last value <= `pivot` if `parity` = true or the
last value < `pivot` if `parity` = false.
Searches after index `lo` up to (inclusive) index `hi`
"""
function find_partition(array, pivot, lo, hi, parity, lt::F1, by::F2) where {F1,F2}
    low = lo + 1
    high = hi
    @inbounds while low <= high
        mid = (low + high) ÷ 2
        if flex_lt(pivot, array[mid], parity, lt, by)
            high = mid - 1
        else
            low = mid + 1
        end
    end
    return low - 1
end

"""
This assumes the region of `vals` of length `L` starting after `lo`
has been batch partitioned with respect to `pivot`. Further, it assumes that
these batches are of size `blockDim().x`.

Using 1 step per batch, consolidate these partitioned batches such that the
region is fully partitioned. Each step moves at most `blockDim().x` values.

`b_sums`: either shared memory or a global array which serves as scratch space
for storing the partition of each batch.

`parity`: see top docstring

Must only run on 1 SM.
"""
@inline function consolidate_batch_partition(vals::AbstractArray{T}, pivot, lo, L, b_sums,
                                             parity, lt::F1, by::F2) where {T,F1,F2}
    sync_threads()
    @inline N_b() = cld(L , blockDim().x)
    @inline batch(k) = threadIdx().x + k * blockDim().x

    my_iter = 0
    a = 0
    b = 0

    @inbounds for batch_i in 1:N_b()
        if batch_i % blockDim().x == 1
            if batch(my_iter) <= N_b()
                seek_lo = lo + (batch(my_iter) - 1) * blockDim().x
                seek_hi = lo + min(L, batch(my_iter) * blockDim().x)
                b_sums[threadIdx().x] =
                    seek_hi - find_partition(vals, pivot, seek_lo, seek_hi, parity, lt, by)
            end
            my_iter += 1
        end

        function n_eff()
            if batch_i != N_b() || L % blockDim().x == 0
                blockDim().x
            else
                L % blockDim().x
            end
        end

        sync_threads()
        d = b_sums[batch_i - (my_iter - 1) * blockDim().x]
        c = n_eff() - d
        to_move = min(b, c)
        sync_threads()
        swap = if threadIdx().x <= to_move
            vals[lo + a + threadIdx().x]
        else
            zero(eltype(vals))  # unused value
        end
        sync_threads()
        if threadIdx().x <= to_move
            vals[lo + a + threadIdx().x] = vals[lo + a + b + c - to_move + threadIdx().x]
        end
        sync_threads()
        if threadIdx().x <= to_move
            vals[lo + a + b + c - to_move + threadIdx().x] = swap
        end
        sync_threads()
        a += c
        b += d
    end

    sync_threads()
    return lo + a
end

# Sorting
"""
Finds the median of `vals` starting after `lo` and going for `blockDim().x`
elements spaced by `stride`. Performs bitonic sort in shmem, returns middle value.
Faster than bubble sort, but not as flexible. Does not modify `vals`
"""
function bitonic_median(vals :: AbstractArray{T}, swap, lo, L, stride, lt::F1, by::F2) where {T,F1,F2}
    sync_threads()
    bitonic_lt(i1, i2) = @inbounds flex_lt(swap[i1 + 1], swap[i2 + 1], false, lt, by)

    @inbounds swap[threadIdx().x] = vals[lo + threadIdx().x * stride]
    sync_threads()
    old_val = zero(eltype(swap))

    log_blockDim = begin
        out = 0
        k = blockDim().x
        while k > 1
            k = k >> 1
            out += 1
        end
        out
    end

    log_k = 1
    while log_k <= log_blockDim
        k = 1 << log_k
        j = k ÷ 2

        while j > 0
            i = threadIdx().x - 1i32
            l = xor(i, j)
            to_swap = (i & k) == 0 && bitonic_lt(l, i) || (i & k) != 0 && bitonic_lt(i, l)
            to_swap = to_swap == (i < l)

            if to_swap
                @inbounds old_val = swap[l + 1]
            end
            sync_threads()
            if to_swap
                @inbounds swap[i + 1] = old_val
            end
            sync_threads()
            j = j ÷ 2
        end
        log_k += 1
    end
    sync_threads()
    return @inbounds swap[blockDim().x ÷ 2]
end

"""
Performs bubble sort on `vals` starting after `lo` and going for min(`L`, `blockDim().x`)
elements spaced by `stride`. Good for sampling pivot values as well as short sorts.
"""
@inline function bubble_sort(vals, swap, lo, L, stride, lt::F1, by::F2) where {F1,F2}
    sync_threads()
    L = min(blockDim().x, L)
    @inbounds begin
        if threadIdx().x <= L
            swap[threadIdx().x] = vals[lo + threadIdx().x * stride]
        end
        sync_threads()
        for level in 0:L
            # get left/right neighbor depending on even/odd level
            buddy = threadIdx().x - 1i32 + 2i32 * (1i32 & (threadIdx().x % 2i32 != level % 2i32))
            buddy_val = if 1 <= buddy <= L && threadIdx().x <= L
                 swap[buddy]
            else
                zero(eltype(swap)) # unused value
            end
            sync_threads()
            if 1 <= buddy <= L && threadIdx().x <= L
                is_left = threadIdx().x < buddy
                # flex_lt needs to handle equivalence in opposite ways for the
                # two threads in each swap pair. Otherwise, if there are two
                # different values with the same by, one will overwrite the other
                if is_left != flex_lt(swap[threadIdx().x], buddy_val, is_left, lt, by)
                    swap[threadIdx().x] = buddy_val
                end
            end
            sync_threads()
        end
        if threadIdx().x <= L
            vals[lo + threadIdx().x * stride] = swap[threadIdx().x]
        end
    end
    sync_threads()
end

"""
Launch batch partition kernel and sync
"""
@inline function call_batch_partition(vals::AbstractArray{T}, pivot, swap, b_sums, lo, hi,
                                      parity, sync::Val{true}, lt::F1, by::F2) where {T, F1, F2}
    L = hi - lo
    if threadIdx().x == 1
        blocks_y = cld(L, blockDim().x)

        # TODO: add wrappers
        device = Ref{Cint}()
        CUDA.cudaGetDevice(device)
        max_blocks_y = Ref{Cint}()
        CUDA.cudaDeviceGetAttribute(max_blocks_y, CUDA.cudaDevAttrMaxGridDimY, device[])

        blocks_z, blocks_y = fldmod1(blocks_y, max_blocks_y[])

        @cuda(blocks=(1,blocks_y,blocks_z), threads=blockDim().x, dynamic=true,
              shmem=blockDim().x*(sizeof(Int)+sizeof(T)),
              partition_batches_kernel(vals, pivot, lo, hi, parity, lt, by))
        device_synchronize()
    end

    # XXX: this global fence shouldn't be needed, according to the CUDA docs a sync_threads
    #      should be sufficient to propagate the global memory writes from the child grid
    #      to the (single-block) parent. See JuliaGPU/CUDA.jl#955 for more details.
    threadfence()
end

"""
Partition batches in a loop using a single block
"""
@inline function call_batch_partition(vals::AbstractArray{T}, pivot, swap, b_sums, lo, hi,
                                      parity, sync::Val{false}, lt::F1, by::F2) where {T, F1, F2}
    while lo <= hi
        batch_partition(vals, pivot, swap, b_sums, lo, min(hi, lo + blockDim().x), parity, lt, by)
        lo += blockDim().x
    end
end

"""
Quicksort recursion condition
For a full sort, `partial` is nothing so it shouldn't affect whether recursion
happens.
"""
function partial_range_overlap(lo, hi, partial :: Nothing)
    true
end

"""
Quicksort recursion condition
If the domain to sort `lo` to `hi` overlaps with `partial`, then we should
do recursion on it, and this returns true (if not, then false)
"""
function partial_range_overlap(lo, hi, partial_k)
    return !(lo > last(partial_k) || hi < first(partial_k))
end

"""
Perform quicksort on dimension `dims` of `vals` for the region with `lo` as an exclusive
floor and `hi` as an inclusive ceiling. `parity` is a boolean which says whether to
partition by < or <= with respect to the pivot. `sync_depth` is how many (more) levels of
recursion with `qsort_kernel` can be done before reaching `cudaLimitDevRuntimeSyncDepth`.
From the host, this value must not exceed that limit.

`sync` and enclosed type `S` determine how partition occurs: If `sync` is `true`, the kernel
partitions batches in a child kernel, synchronizes, and then consolidates the batches. The
benefit of this kernel is that it distributes the work of partitioning batches across
multiple SMs. If `sync` is `false`, the kernel partitions without launching any child
kernels, then has recursive `qsort_kernel` children for left and right partitions.
`device_synchronize` is never called from this kernel, so there is no practical limit on
recursion.

To detect the scenario of all values in the region being the same, we have two args:
`prev_pivot` and `stuck`. If two consecutive partitions have the same pivot and both failed
to split the region in two, that means all the values are equal. `stuck` is incremented when
the pivot hasn't changed and partition = `lo` or `hi`. If `stuck` reaches 2, recursion ends.
`stuck` is initialized at -1 because `prev_pivot` must be initialized to some value, and
it's possible that the first pivot will be that value, which could lead to an incorrectly
early end to recursion if we started `stuck` at 0.
"""
function qsort_kernel(vals::AbstractArray{T,N}, lo, hi, parity, sync::Val{S}, sync_depth,
                      prev_pivot, lt::F1, by::F2, ::Val{dims}, partial=nothing, stuck=-1) where {T, N, S, F1, F2, dims}
    b_sums = CuDynamicSharedArray(Int, blockDim().x)
    swap = CuDynamicSharedArray(T, blockDim().x, sizeof(b_sums))
    shmem = sizeof(b_sums) + sizeof(swap)
    L = hi - lo

    # extract the dimension that we need to sort (selected from the rest by the block x index)
    slice = if N == 1
        vals
    else
        # dimensions that are not part of the sort; index them using the block index
        otherdims = ntuple(i -> i == dims ? 1 : size(vals, i), N)
        other = CartesianIndices(otherdims)[blockIdx().x]

        # create a view that keeps the sorting dimension but indexes across the others
        slicedims = map(Base.Slice, axes(vals))
        idxs = ntuple(i->i==dims ? slicedims[i] : other[i], N)
        view(vals, idxs...)
    end

    # step 1: single block sort. It'll either finish sorting a subproblem or
    # help select a pivot value

    if L <= blockDim().x
        bubble_sort(slice, swap, lo, L, 1, lt, by)
        return
    end

    pivot = bitonic_median(slice, swap, lo, L, L ÷ blockDim().x, lt, by)

    # step 2: use pivot to partition into batches
    call_batch_partition(slice, pivot, swap, b_sums, lo, hi, parity, sync, lt, by)

    # step 3: consolidate the partitioned batches so that the sublist from [lo, hi) is
    #         partitioned, and the partition is stored in `partition`. Dispatching on P
    #         cleaner and faster than an if statement

    partition = consolidate_batch_partition(slice, pivot, lo, L, b_sums, parity, lt, by)

    # step 4: recursion
    if threadIdx().x == 1
        stuck = (pivot == prev_pivot && partition == lo || partition == hi) ? stuck + 1 : 0

        if stuck < 2 && partition > lo && partial_range_overlap(lo, partition, partial)
            s = CuDeviceStream()
            if S && sync_depth > 1
                @cuda(threads=blockDim().x, dynamic=true, stream=s, shmem=shmem,
                      qsort_kernel(slice, lo, partition, !parity, Val(true), sync_depth - 1,
                      pivot, lt, by, Val(1), partial, stuck))
            else
                @cuda(threads=blockDim().x, dynamic=true, stream=s, shmem=shmem,
                      qsort_kernel(slice, lo, partition, !parity, Val(false), sync_depth - 1,
                      pivot, lt, by, Val(1), partial, stuck))
            end
            CUDA.unsafe_destroy!(s)
        end

        if stuck < 2 && partition < hi && partial_range_overlap(partition, hi, partial)
            s = CuDeviceStream()
            if S && sync_depth > 1
                @cuda(threads=blockDim().x, dynamic=true, stream=s, shmem=shmem,
                      qsort_kernel(slice, partition, hi, !parity, Val(true), sync_depth - 1,
                      pivot, lt, by, Val(1), partial, stuck))
            else
                @cuda(threads=blockDim().x, dynamic=true, stream=s, shmem=shmem,
                      qsort_kernel(slice, partition, hi, !parity, Val(false), sync_depth - 1,
                      pivot, lt, by, Val(1), partial, stuck))
            end
            CUDA.unsafe_destroy!(s)
        end
    end

    return
end

function sort_args(args, partial_k :: Nothing)
    return args
end

function sort_args(args, partial_k)
    return (args..., partial_k)
end

#function sort

function quicksort!(c::AbstractArray{T,N}; lt::F1, by::F2, dims::Int, partial_k=nothing, block_size_shift=0) where {T,N,F1,F2}
    max_depth = CUDA.limit(CUDA.LIMIT_DEV_RUNTIME_SYNC_DEPTH)
    len = size(c, dims)

    1 <= dims <= N || throw(ArgumentError("dimension out of range"))
    otherdims = ntuple(i -> i == dims ? 1 : size(c, i), N)

    my_sort_args = sort_args((c, 0, len, true, Val(N==1 && max_depth > 1),
             max_depth, nothing, lt, by, Val(dims)), partial_k)

    kernel = @cuda launch=false qsort_kernel(my_sort_args...)

    get_shmem(threads) = threads * (sizeof(Int) + sizeof(T))
    config = launch_configuration(kernel.fun, shmem=threads->get_shmem(threads))
    threads = prevpow(2, config.threads)
    threads = threads >> block_size_shift   # for testing purposes

    kernel(my_sort_args...;
           blocks=prod(otherdims), threads=threads, shmem=get_shmem(threads))

    return c
end

end

"""
This is an iterative bitonic sort that mimics a recursive version to support
non-power2 lengths.

Credit for the recursive form of this algorithm goes to:
https://www.inf.hs-flensburg.de/lang/algorithmen/sortieren/bitonic/oddn.htm

CUDA.jl implementation originally by @xaellison

Overview: `comparator_kernel` implements a layer of sorting network comparators
generally. The sort could run just by looping over `comparator`, but
`comparator_small_kernel` copies values into shmem and loops over several
comparators that don't need to access any values outside the range held in
shared memory. It provides a moderate speedup.

Notation:
`k`, `j` denote the level of the sorting network (equivalently, recursion depth).
`vals` is the array of values of type `T` that is either being `sort`-ed or `sortperm`-ed.
`inds` is an array of indices of type `J` that gets permuted in `sortperm!` (standard 1-indexed)
`i1`, `i2` index either `vals` or `inds` depending on the operation.
`lo`, `n`, and `m` are integers of type `I` used to denote/calculate ranges as
    described in the recursive algorithm link above. Note these follow the 0-indexing
    convention from the above source.
"""
module BitonicSortImpl

export bitonic_sort!

using ..CUDA
using ..CUDA: i32

# General functions

@inline two(::Type{Int}) = 2
@inline two(::Type{Int32}) = 2i32

@inline function gp2lt(x::Int)::Int
    x -= 1
    x |= x >> 1
    x |= x >> 2
    x |= x >> 4
    x |= x >> 8
    x |= x >> 16
    x |= x >> 32
    xor(x, x >> 1)
end

@inline function gp2lt(x::Int32)::Int32
    x -= 1i32
    x |= x >> 1i32
    x |= x >> 2i32
    x |= x >> 4i32
    x |= x >> 8i32
    x |= x >> 16i32
    xor(x, x >> 1i32)
end

@inline function bisect_range(index::I, lo::I, n::I) where {I}
    if n <= one(I)
        return -one(I), -one(I)
    end
    m = gp2lt(n)
    if index < lo + m
        n = m
    else
        lo = lo + m
        n = n - m
    end
    return lo, n
end

@inline function rev_lt(a::T, b::T, lt, rev::Val{R}) where {T,R}
    if R
        return lt(b, a)
    else
        return lt(a, b)
    end
end

@inline function rev_lt(a::Tuple{T,J}, b::Tuple{T,J}, lt, rev::Val{R}) where {T,J,R}
    if R
        if a[1] == b[1]
            return a[2] < b[2]
        else
            return lt(b[1], a[1])
        end
    else
        return lt(a, b)
    end
end

# Functions specifically for "large" bitonic steps (those that cannot use shmem)

@inline function compare!(vals::AbstractArray{T}, i1::I, i2::I, dir::Bool, by, lt, rev) where {T,I}
    i1′, i2′ = i1 + one(I), i2 + one(I)
    @inbounds if dir != rev_lt(by(vals[i1′]), by(vals[i2′]), lt, rev)
        vals[i1′], vals[i2′] = vals[i2′], vals[i1′]
    end
end

@inline function compare!(vals_inds::Tuple, i1::I, i2::I, dir::Bool, by, lt, rev) where {I}
    i1′, i2′ = i1 + one(I), i2 + one(I)
    vals, inds = vals_inds
    # comparing tuples of (value, index) guarantees stability of sort
    @inbounds if dir != rev_lt((by(vals[inds[i1′]]), inds[i1′]), (by(vals[inds[i2′]]), inds[i2′]), lt, rev)
        inds[i1′], inds[i2′] = inds[i2′], inds[i1′]
    end
end

@inline function get_range_part1(n::I, index::I, k::I)::Tuple{I,I,Bool} where {I}
    lo = zero(I)
    dir = true
    for iter = one(I):k-one(I)
        if n <= one(I)
            return -one(I), -one(I), false
        end

        if index < lo + n ÷ two(I)
            n = n ÷ two(I)
            dir = !dir
        else
            lo = lo + n ÷ two(I)
            n = n - n ÷ two(I)
        end
    end
    return lo, n, dir
end

@inline function get_range_part2(lo::I, n::I, index::I, j::I)::Tuple{I,I} where {I}
    for iter = one(I):j-one(I)
        lo, n = bisect_range(index, lo, n)
    end
    return lo, n
end

"""
Determines parameters for swapping when the grid index directly maps to an
Array index for swapping
"""
@inline function get_range(n, index, k, j)
    lo, n, dir = get_range_part1(n, index, k)
    lo, n = get_range_part2(lo, n, index, j)
    return lo, n, dir
end

"""
Performs a step of bitonic sort requiring swaps between indices further apart
than the size of block allows (eg, 1 <--> 10000)

The grid index directly maps to the index of `c` that will be used in the swap.

Note that to avoid synchronization issues, only one thread from each pair of
indices being swapped will actually move data. This does mean half of the threads
do nothing, but it works for non-power2 arrays while allowing direct indexing.
"""
function comparator_kernel(vals, length_vals::I, k::I, j::I, by::F1, lt::F2, rev) where {I,F1,F2}
    index = (blockDim().x * (blockIdx().x - one(I))) + threadIdx().x - one(I)

    lo, n, dir = get_range(length_vals, index, k, j)

    if !(lo < zero(I) || n < zero(I)) && !(index >= length_vals)
        m = gp2lt(n)
        if lo <= index < lo + n - m
            i1, i2 = index, index + m
            @inbounds compare!(vals, i1, i2, dir, by, lt, rev)
        end
    end
    return
end

# Functions for "small" bitonic steps (those that can use shmem)

@inline function compare_small!(vals::AbstractArray{T}, i1::I, i2::I, dir::Bool, by, lt, rev) where {T,I}
    i1′, i2′ = i1 + one(I), i2 + one(I)
    @inbounds if dir != rev_lt(by(vals[i1′]), by(vals[i2′]), lt, rev)
        vals[i1′], vals[i2′] = vals[i2′], vals[i1′]
    end
end

@inline function compare_small!(vals_inds::Tuple, i1::I, i2::I, dir::Bool, by, lt, rev) where {I}
    i1′, i2′ = i1 + one(I), i2 + one(I)
    vals, inds = vals_inds
    # comparing tuples of (value, index) guarantees stability of sort
    @inbounds if dir != rev_lt((by(vals[i1′]), inds[i1′]), (by(vals[i2′]), inds[i2′]), lt, rev)
        vals[i1′], vals[i2′] = vals[i2′], vals[i1′]
        inds[i1′], inds[i2′] = inds[i2′], inds[i1′]
    end
end

"""
For each thread in the block, "re-compute" the range which would have been
passed in recursively. This range only depends on the block, and guarantees
all threads perform swaps accessible using shmem.

Various negative exit values just for debugging.
"""
@inline function block_range(n::I, block_index::I, k::I, j::I)::Tuple{I,I,Bool} where {I}
    lo = zero(I)
    dir = true
    tmp = block_index * two(I)

    # analogous to `get_range_part1`
    for iter = one(I):(k-one(I))
        tmp ÷= two(I)
        if n <= one(I)
            return -one(I), -one(I), false
        end

        if tmp % two(I) == zero(I)
            n = n ÷ two(I)
            dir = !dir
        else
            lo = lo + n ÷ two(I)
            n = n - n ÷ two(I)
        end
    end

    # analogous to `get_range_part2`
    for iter = one(I):(j-one(I))
        tmp ÷= two(I)
        if n <= one(I)
            return -one(I), -one(I), false
        end

        m = gp2lt(n)
        if tmp % two(I) == zero(I)
            n = m
        else
            lo = lo + m
            n = n - m
        end

    end
    if zero(I) <= n <= one(I)
        return -one(I), -one(I), false
    end
    return lo, n, dir
end

"""
For sort/sort! `c`, allocate and return shared memory view of `c`
Each view is indexed along block x dim: one view per pseudo-block
`index` is expected to be from a 0-indexing context
"""
@inline function initialize_shmem!(vals::AbstractArray{T}, index::I, in_range, offset = zero(I)) where {T,I}
    swap = CuDynamicSharedArray(T, (blockDim().x, blockDim().y), offset)
    if in_range
        @inbounds swap[threadIdx().x, threadIdx().y] = vals[index+one(I)]
    end
    sync_threads()
    return @view swap[:, threadIdx().y]
end

"""
For sortperm/sortperm!, allocate and return shared memory views of `c` and index
array. Each view is indexed along block x dim: one view per pseudo-block.
`index` is expected to be from a 0-indexing context, but the indices stored in
`val_inds` are expected to be 1-indexed
"""
@inline function initialize_shmem!(vals_inds::Tuple{AbstractArray{T},AbstractArray{J}}, index, in_range) where {T,J}
    offset = prod(blockDim()) * sizeof(J)
    vals, inds = vals_inds
    inds_view = initialize_shmem!(inds, index, in_range)
    vals_view = initialize_shmem!(vals, inds_view[threadIdx().x] - one(J), in_range, offset)
    return vals_view, inds_view
end

"""
For sort/sort!, copy shmem view `swap` back into global array `c`
`index` is expected to be from a 0-indexing context
"""
@inline function finalize_shmem!(vals::AbstractArray, swap::AbstractArray, index::I, in_range::Bool) where {I}
    if in_range
        @inbounds vals[index+one(I)] = swap[threadIdx().x]
    end
end

"""
For sortperm/sortperm!, copy shmem view `swap` back to global index array
`index` is expected to be from a 0-indexing context, but the indices stored in
`val_inds` are expected to be 1-indexed
"""
@inline function finalize_shmem!(vals_inds::Tuple, swap::Tuple, index, in_range::Bool)
    vals, inds = vals_inds
    swap_vals, swap_inds = swap
    finalize_shmem!(inds, swap_inds, index, in_range)
end

"""
Performs consecutive steps of bitonic sort requiring swaps between indices no
further apart than the size of block allows. This effectively moves part of the
inner loop (over j, below) inside of a kernel to minimize launches and do
swaps in shared mem.

Note that the x dimension of a thread block is treated as a comparator,
so when the maximum size of a comparator in this kernel is small, multiple
may be executed along the block y dimension, allowing for higher occupancy.
This is captured by `pseudo_block_idx`.

Note that this moves the array values copied within shmem, but doesn't copy them
back to global the way it does for indices.
"""
function comparator_small_kernel(c, length_c::I, k::I, j_0::I, j_f::I, by::F1, lt::F2, rev) where {I,F1,F2}
    pseudo_block_idx = (blockIdx().x - one(I)) * blockDim().y + threadIdx().y - one(I)
    # immutable info about the range used by this kernel
    _lo, _n, dir = block_range(length_c, pseudo_block_idx, k, j_0)
    index = _lo + threadIdx().x - one(I)
    in_range = (threadIdx().x <= _n && _lo >= zero(I))

    swap = initialize_shmem!(c, index, in_range)

    # mutable copies for pseudo-recursion
    lo, n = _lo, _n

    for j = j_0:j_f
        if !(lo < zero(I) || n < zero(I)) && in_range
            m = gp2lt(n)
            if lo <= index < lo + n - m
                i1, i2 = index - _lo, index - _lo + m
                compare_small!(swap, i1, i2, dir, by, lt, rev)
            end
        end
        lo, n = bisect_range(index, lo, n)
        sync_threads()
    end

    finalize_shmem!(c, swap, index, in_range)
    return
end

# Host side code
function bitonic_shmem(c::AbstractArray{T}, threads) where {T}
    return prod(threads) * sizeof(T)
end

function bitonic_shmem(c, threads)
    return prod(threads) * sum(map(a -> sizeof(eltype(a)), c))
end

"""
Call bitonic sort on `c` which can be a CuArray of values to `sort!` or a tuple
of values and an index array for doing `sortperm!`. Cannot provide a stable
`sort!` although `sortperm!` is properly stable. To reverse, set `rev=true`
rather than `lt=!isless` (otherwise stability of sortperm breaks down).
"""
function bitonic_sort!(c; by = identity, lt = isless, rev = false) where {T}
    c_len = if typeof(c) <: Tuple
        length(c[1])
    else
        length(c)
    end
    k0 = c_len |> log2 |> ceil |> Int

    blocks_per_mp = CUDA.attribute(device(), CUDA.DEVICE_ATTRIBUTE_MAX_BLOCKS_PER_MULTIPROCESSOR)

    # These two outer loops are the same as the serial version outlined here:
    # https://en.wikipedia.org/wiki/Bitonic_sorter#Example_code
    # Notation: our k/j are the base2 logs of k/j in Wikipedia example

    for k = k0:-1:1
        j_final = (1 + k0 - k)
        for j = 1:j_final

            # use Int32 args for indexing when possible --> ~10% faster kernels
            I = c_len <= typemax(Int32) ? Int32 : Int
            args1 = (c, map(I, (c_len, k, j, j_final))..., by, lt, Val(rev))
            kernel1 = @cuda launch=false comparator_small_kernel(args1...)
            config1 = launch_configuration(kernel1.fun, shmem = threads -> bitonic_shmem(c, threads))
            threads1 = prevpow(2, config1.threads)

            args2 = (c, map(I, (c_len, k, j))..., by, lt, Val(rev))
            kernel2 = @cuda launch=false comparator_kernel(args2...)
            config2 = launch_configuration(kernel2.fun, shmem = threads -> bitonic_shmem(c, threads))
            threads2 = prevpow(2, config2.threads)

            threads = min(threads1, threads2)
            min_pseudo_block = threads ÷ blocks_per_mp
            log_threads = threads |> log2 |> Int

            if k0 - k - j + 2 <= log_threads
                pseudo_block_size = 1 << abs(j_final + 1 - j)
                _block_size = if pseudo_block_size <= min_pseudo_block
                    (pseudo_block_size, min_pseudo_block ÷ pseudo_block_size)
                else
                    (pseudo_block_size, 1)
                end
                b = nextpow(2, cld(c_len, prod(_block_size)))
                kernel1(args1...; blocks = b, threads = _block_size, shmem = bitonic_shmem(c, _block_size))
                break
            else
                b = nextpow(2, cld(c_len, threads))
                kernel2(args2...; blocks = b, threads = threads)
            end
        end
    end

    return c
end
end

# Base interface implementation

using .BitonicSortImpl
using .QuickSortImpl

abstract type SortingAlgorithm end
struct QuickSortAlg <: SortingAlgorithm end
struct BitonicSortAlg <: SortingAlgorithm end

const QuickSort = QuickSortAlg()
const BitonicSort = BitonicSortAlg()

function Base.sort!(c::AnyCuVector, alg::QuickSortAlg; lt=isless, by=identity, rev=false)
    # for reverse sorting, invert the less-than function
    if rev
        lt = !lt
    end

    quicksort!(c; lt, by, dims=1)
    return c
end

function Base.sort!(c::AnyCuVector, alg::BitonicSortAlg; kwargs...)
    return bitonic_sort!(c; kwargs...)
end

function Base.sort!(c::AnyCuVector; alg :: SortingAlgorithm = BitonicSort, kwargs...)
    return sort!(c, alg; kwargs...)
end

function Base.sort!(c::AnyCuArray; dims::Integer, lt=isless, by=identity, rev=false)
    # for multi dim sorting, only quicksort is supported so no alg keyword
    if rev
        lt = !lt
    end

    quicksort!(c; lt, by, dims)
    return c
end

function Base.sort(c::AnyCuArray; kwargs...)
    return sort!(copy(c); kwargs...)
end

function Base.partialsort!(c::AnyCuVector, k::Union{Integer, OrdinalRange}; lt=isless, by=identity, rev=false)
    # for reverse sorting, invert the less-than function
    if rev
        lt = !lt
    end

    function out(k :: OrdinalRange)
        return copy(c[k])
    end

    # work around disallowed scalar index
    function out(k :: Integer)
        return Array(c[k:k])[1]
    end

    quicksort!(c; lt, by, dims=1, partial_k=k)
    return out(k)
end

function Base.partialsort(c::AnyCuArray, k::Union{Integer, OrdinalRange}; kwargs...)
    return partialsort!(copy(c), k; kwargs...)
end

function Base.sortperm!(I::AnyCuArray{T}, c::AnyCuArray; initialized=false, kwargs...) where T
    if length(I) != length(c)
        throw(ArgumentError("index vector must have the same length/indices as the source vector"))
    end
    if !initialized
        I .= one(T):T(length(I))
    end
    bitonic_sort!((c, I); kwargs...)
    return I
end

function Base.sortperm(c::AnyCuArray; kwargs...)
    sortperm!(CuArray(1:length(c)), c; initialized=true, kwargs...)
end


text:
"""
    YaoAPI

API definitions for Yao.
"""
module YaoAPI

include("interface.jl")
include("registers.jl")
include("blocks.jl")

end # module


text:
@enum Shape CIRCLE RECTANGLE ROUNDED_RECTANGLE DISTANCEFIELD TRIANGLE
@enum RaymarchAlgorithm IsoValue Absorption MaximumIntensityProjection AbsorptionRGBA IndexedAbsorptionRGBA
@enum CubeSides TOP BOTTOM FRONT BACK RIGHT LEFT

struct Grid{N, T <: Range}
    dims::NTuple{N, T}
end
Base.ndims(::Grid{N,T}) where {N,T} = N

Grid(ranges::Range...) = Grid(ranges)
function Grid(a::Array{T, N}) where {N, T}
    s = Vec{N, Float32}(size(a))
    smax = maximum(s)
    s = s./smax
    Grid(ntuple(Val{N}) do i
        linspace(0, s[i], size(a, i))
    end)
end

Grid(a::AbstractArray, ranges...) = Grid(a, ranges)

"""
This constructor constructs a grid from ranges given as a tuple.
Due to the approach, the tuple `ranges` can consist of NTuple(2, T)
and all kind of range types. The constructor will make sure that all ranges match
the size of the dimension of the array `a`.
"""
function Grid(a::AbstractArray{T, N}, ranges::Tuple) where {T, N}
    length(ranges) =! N && throw(ArgumentError(
        "You need to supply a range for every dimension of the array. Given: $ranges
        given Array: $(typeof(a))"
    ))
    Grid(ntuple(Val{N}) do i
        linspace(first(ranges[i]), last(ranges[i]), size(a, i))
    end)
end

Base.length(p::Grid) = prod(size(p))
Base.size(p::Grid) = map(length, p.dims)
function Base.getindex(p::Grid{N,T}, i) where {N,T}
    inds = ind2sub(size(p), i)
    Point{N, eltype(T)}(ntuple(Val{N}) do i
        p.dims[i][inds[i]]
    end)
end

Base.start(g::Grid) = 1
Base.done(g::Grid, i) = i > length(g)
Base.next(g::Grid, i) = g[i], i+1

GLAbstraction.isa_gl_struct(x::Grid) = true
GLAbstraction.toglsltype_string(t::Grid{N,T}) where {N,T} = "uniform Grid$(N)D"
function GLAbstraction.gl_convert_struct(g::Grid{N, T}, uniform_name::Symbol) where {N,T}
    return Dict{Symbol, Any}(
        Symbol("$uniform_name.ref") => Vec{N, Float32}(map(x-> eltype(x)(x.ref), g.dims)),
        Symbol("$uniform_name.offset") => Vec{N, Float32}(map(x-> eltype(x)(x.offset), g.dims)),
        Symbol("$uniform_name._step") => Vec{N, Float32}(map(x-> step(x), g.dims)),
        Symbol("$uniform_name.dims") => Vec{N, Cint}(map(length, g.dims))
    )
end
function GLAbstraction.gl_convert_struct(g::Grid{1, T}, uniform_name::Symbol) where T
    x = g.dims[1]
    return Dict{Symbol, Any}(
        Symbol("$uniform_name.ref") => Float32(eltype(x)(x.ref)),
        Symbol("$uniform_name.offset") => Float32(eltype(x)(x.offset)),
        Symbol("$uniform_name._step") => Float32(step(x)),
        Symbol("$uniform_name.dims") => Cint(length(x))
    )
end
import Base: getindex, length, next, start, done

to_cpu_mem(x) = x
to_cpu_mem(x::GPUArray) = gpu_data(x)

const ScaleTypes = Union{Vector, Vec, AbstractFloat, Void, Grid}
const PositionTypes = Union{Vector, Point, AbstractFloat, Void, Grid}

mutable struct ScalarRepeat{T}
    scalar::T
end
Base.ndims(::ScalarRepeat) = 1
Base.getindex(s::ScalarRepeat, i...) = s.scalar
#should setindex! really be allowed? It will set the index for the whole row...
Base.setindex!(s::ScalarRepeat{T}, value, i...) where {T} = (s.scalar = T(value))
Base.eltype(::ScalarRepeat{T}) where {T} = T

Base.start(::ScalarRepeat) = 1
Base.next(sr::ScalarRepeat, i) = sr.scalar, i+1
Base.done(sr::ScalarRepeat, i) = false

struct Instances{P,T,S,R}
    primitive::P
    translation::T
    scale::S
    rotation::R
end

function _Instances(position,px,py,pz, scale,sx,sy,sz, rotation, primitive)
    args = (position,px,py,pz, scale,sx,sy,sz, rotation, primitive)
    args = map(to_cpu_mem, args)
    p = const_lift(ArrayOrStructOfArray, Point3f0, args[1:4]...)
    s = const_lift(ArrayOrStructOfArray, Vec3f0, args[5:8]...)
    r = const_lift(ArrayOrStructOfArray, Vec4f0, args[9])
    const_lift(Instances, args[10], p, s, r)
end
function _Instances(position, scale, rotation, primitive)
    p = const_lift(ArrayOrStructOfArray, Point3f0, position)
    s = const_lift(ArrayOrStructOfArray, Vec3f0, scale)
    r = const_lift(ArrayOrStructOfArray, Vec4f0, rotation)
    const_lift(Instances, primitive, p, s, r)
end

struct GridZRepeat{G, T, N} <: AbstractArray{Point{3, T}, N}
    grid::G
    z::Array{T, N}
end
Base.size(g::GridZRepeat) = size(g.z)
Base.size(g::GridZRepeat, i) = size(g.z, i)
Base.IndexStyle(::Type{<:GridZRepeat}) = Base.IndexLinear()

function Base.getindex(g::GridZRepeat{G, T}, i) where {G,T}
    pxy = g.grid[i]
    Point{3, T}(pxy[1], pxy[2], g.z[i])
end

function ArrayOrStructOfArray(::Type{T}, array::Void, a, elements...) where T
    StructOfArrays(T, a, elements...)
end
function ArrayOrStructOfArray(::Type{T}, array::StaticVector, a, elements...) where T
    StructOfArrays(T, a, elements...)
end
function ArrayOrStructOfArray(::Type{T}, scalar::StaticVector, a::Void, elements::Void...) where T
    ScalarRepeat(transform_convert(T, scalar))
end
function ArrayOrStructOfArray(::Type{T1}, array::Array{T2}, a::Void, elements::Void...) where {T1,T2}
    array
end
function ArrayOrStructOfArray(::Type{T1}, grid::Grid, x::Void, y::Void, z::Array) where T1<:Point
    GridZRepeat(grid, z)
end
function ArrayOrStructOfArray(::Type{T1}, array::Grid, a::Void, elements::Void...) where T1<:Point
    array
end
function ArrayOrStructOfArray(::Type{T}, scalar::T) where T
    ScalarRepeat(scalar)
end
function ArrayOrStructOfArray(::Type{T}, array::Array) where T
    array
end

struct TransformationIterator{T,S,R}
    translation::T
    scale::S
    rotation::R
end
function TransformationIterator(instances::Instances)
    TransformationIterator(
        instances.translation,
        instances.scale,
        instances.rotation
    )
end
function start(t::TransformationIterator)
    start(t.translation), start(t.scale), start(t.rotation)
end

function done(t::TransformationIterator, state)
    (done(t.translation, state[1]) ||
    done(t.scale, state[2]) ||
    done(t.rotation, state[3]))::Bool
end

import GeometryTypes: transform_convert

function qmul(quat, vec)
    num = quat[1] * 2f0;
    num2 = quat[2] * 2f0;
    num3 = quat[3] * 2f0;
    num4 = quat[1] * num;
    num5 = quat[2] * num2;
    num6 = quat[3] * num3;
    num7 = quat[1] * num2;
    num8 = quat[1] * num3;
    num9 = quat[2] * num3;
    num10 = quat[4] * num;
    num11 = quat[4] * num2;
    num12 = quat[4] * num3;
    return Point3f0(
        (1f0 - (num5 + num6)) * vec[1] + (num7 - num12) * vec[2] + (num8 + num11) * vec[3],
        (num7 + num12) * vec[1] + (1f0 - (num4 + num6)) * vec[2] + (num9 - num10) * vec[3],
        (num8 - num11) * vec[1] + (num9 + num10) * vec[2] + (1f0 - (num4 + num5)) * vec[3]
    )
end

# For quaternions
function to_rotation_mat(x::Vec{4, T}) where T
    Mat4{T}(Quaternions.Quaternion(x[4], x[1], x[2], x[3], true))
end
# For relative rotations of a vector
function to_rotation_mat(x::StaticVector{2, T}) where T
    to_rotation_mat(Vec3f0(x[1], x[2], 0))
end
function to_rotation_mat(x::StaticVector{3, T}) where T
    rotation = Vec3f0(transform_convert(Point3f0, x))
    v, u = normalize(rotation), Vec3f0(0, 0, 1)
    # Unfortunately, we have to check for when u == -v, as u + v
    # in this case will be (0, 0, 0), which cannot be normalized.
    q = if (u == -v)
        # 180 degree rotation around any orthogonal vector
        other = (abs(dot(u, Vec{3, T}(1,0,0))) < 1.0) ? Vec{3, T}(1,0,0) : Vec{3, T}(0,1,0)
        Quaternions.qrotation(normalize(cross(u, other)), T(180))
    else
        half = normalize(u+v)
        vc = cross(u, half)
        Quaternions.Quaternion(dot(u, half), vc[1], vc[2], vc[3])
    end
    Mat4{T}(q)
end
function next(t::TransformationIterator, state)
    _translation, st = next(t.translation, state[1])
    _scale, ss = next(t.scale, state[2])
    _rotation, sr = next(t.rotation, state[3])

    translation = Point3f0(transform_convert(Vec3f0, _translation))
    scale = Vec3f0(transform_convert(Point3f0, _scale))
    rotation = to_rotation_mat(_rotation)
    ((translation, scale, rotation), (st, ss, sr))
end

struct Intensity{T <: AbstractFloat} <: FieldVector{1, T}
    i::T
end
@inline (I::Type{Intensity{T}})(i::Tuple) where {T <: AbstractFloat} = I(i...)
@inline (I::Type{Intensity{T}})(i::Intensity) where {T <: AbstractFloat} = I(i.i)
Intensity{T}(x::Color{Tc, 1}) where {T <: AbstractFloat, Tc} = Intensity{T}(gray(x))

const GLIntensity = Intensity{Float32}
export Intensity, GLIntensity

NOT(x) = !x

struct GLVisualizeShader <: AbstractLazyShader
    paths::Tuple
    kw_args::Dict{Symbol, Any}
    function GLVisualizeShader(paths::String...; view = Dict{String, String}(), kw_args...)
        # TODO properly check what extensions are available
        @static if !is_apple()
            view["GLSL_EXTENSIONS"] = "#extension GL_ARB_conservative_depth: enable"
            view["SUPPORTED_EXTENSIONS"] = "#define DETPH_LAYOUT"
        end
        args = Dict{Symbol, Any}(kw_args)
        args[:view] = view
        args[:fragdatalocation] = [(0, "fragment_color"), (1, "fragment_groupid")]
        new(map(x-> assetpath("shader", x), paths), args)
    end
end


text:
@safetestset "HEAD requests" begin
  @safetestset "HEAD requests should be by default handled by GET" begin
    using Genie
    using HTTP

    port = nothing
    port = rand(8500:8900)

    route("/") do
      "GET request"
    end

    server = up(port)

    response = try
      HTTP.request("GET", "http://127.0.0.1:$port", ["Content-Type" => "text/html"])
    catch ex
      ex.response
    end

    @test response.status == 200
    @test String(response.body) == "GET request"

    response = try
      HTTP.request("HEAD", "http://127.0.0.1:$port", ["Content-Type" => "text/html"])
    catch ex
      ex.response
    end

    @test response.status == 200
    @test String(response.body) == ""

    down()
    sleep(1)
    server = nothing
    port = nothing
  end;

  @safetestset "HEAD requests have no body" begin
    using Genie
    using HTTP

    port = nothing
    port = rand(8500:8900)

    route("/") do
      "Hello world"
    end

    route("/", method = HEAD) do
      "Hello world"
    end

    server = up(port; open_browser = false)

    response = try
      HTTP.request("GET", "http://127.0.0.1:$port", ["Content-Type" => "text/html"])
    catch ex
      ex.response
    end

    @test response.status == 200
    @test String(response.body) == "Hello world"

    response = try
      HTTP.request("HEAD", "http://127.0.0.1:$port", ["Content-Type" => "text/html"])
    catch ex
      ex.response
    end
    @test response.status == 200
    @test isempty(String(response.body)) == true

    down()
    sleep(1)
    server = nothing
    port = nothing
  end;

  @safetestset "HEAD requests should overwrite GET" begin
    using Genie
    using HTTP

    port = nothing
    port = rand(8500:8900)

    request_method = ""

    route("/", named = :get_root) do
      request_method = "GET"
      "GET request"
    end

    route("/", method = "HEAD", named = :head_root) do
      request_method = "HEAD"
      "HEAD request"
    end

    server = up(port)
    sleep(1)

    response = try
      HTTP.request("GET", "http://127.0.0.1:$port", ["Content-Type" => "text/html"])
    catch ex
      ex.response
    end

    @test response.status == 200
    @test request_method == "GET"

    response = try
      HTTP.request("HEAD", "http://127.0.0.1:$port", ["Content-Type" => "text/html"])
    catch ex
      ex.response
    end

    @test response.status == 200
    @test request_method == "HEAD"

    down()
    sleep(1)
    server = nothing
    port = nothing
  end;

end;


text:
using QuantumStateBase
using QuantumStatePlots
using Documenter

############
# gen logo #
############

ENV["GKSwstype"]="nul"
logo_path = mkpath(joinpath(@__DIR__, "src/assets"))

dim = 50
state = displace!(squeeze!(SinglePhotonState(dim=dim), ξ(0.6, 1π)), α(5., 3π/2))
@time wf = WignerFunction(-10:0.1:10, -10:0.1:10, dim=dim)

plot_wigner(wf(state), Heatmap, size=(600, 550), file_path=joinpath(logo_path, "heatmap.png"))
plot_all(wf(state), state, levels=8, size=(600, 550), file_path=joinpath(logo_path, "all.png"))

##################
# build document #
##################

DocMeta.setdocmeta!(QuantumStatePlots, :DocTestSetup, :(using QuantumStatePlots); recursive=true)

makedocs(;
    modules=[QuantumStatePlots],
    authors="JingYu Ning <foldfelis@gmail.com> and contributors",
    repo="https://github.com/foldfelis-QO/QuantumStatePlots.jl/blob/{commit}{path}#{line}",
    sitename="QuantumStatePlots.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://foldfelis-qo.github.io/QuantumStatePlots.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/foldfelis-QO/QuantumStatePlots.jl",
)


text:
module RDST

using Random

import Base: rand, show, copy

abstract type AbstractRNGStream end
abstract type AbstractStreamableRNG <: AbstractRNG end # object of subtype of AbstractStreamableRNG are StreamableRNG
                                                       #currently just mrg32k3a and xoshiro256plus

###mrg32k3a
include("mrg32k3a/mrg32k3a.jl")
include("mrg32k3a/mrg32k3a_types.jl")
include("mrg32k3a/add.jl")  #just some of my addition to the stolen code

#xoshiro256+

include("xoshiro/xoshiro256plus.jl")
include("xoshiro/xoshiro256plus_types.jl")

export AbstractRNGStream, AbstractStreamableRNG
export copy
export checkseed, MRG32k3a, rand, srand, reset_stream!, reset_substream!, next_substream!, MRG32k3aGen, show, get_state, next_stream
export Xoshiro256p, Xoshiro256plusGen

end # module


text:
using Revise
using ADCME
using PyCall
using Test
using MAT
using NNFEM
using PyPlot
using JLD2
using LinearAlgebra
using Distributions, Random
using ForwardDiff
using DelimitedFiles
np = pyimport("numpy")

"""
Force load on one edge of the plate,
The edge is [0,L], which is discretized to ne elements, with porder-th order polynomials
type: a string of load type, which can be "constant", or "Gaussian"
args: an array,
      for Constant load, it has p1 and p2, in tangential and normal direction
      for "Gaussian" load, it has p, x0, and σ, the force is in the normal direction

"""
function ComputeLoad(L, ne, porder, ngp, type,  args)
    @assert(ngp <= 4)
    dx = L/ne

    xx = np.linspace(0, L, ne*porder + 1)

    Ft, Fn = zeros(Float64, ne*porder + 1), zeros(Float64, ne*porder + 1)

    elem_coords = zeros(Float64, porder + 1, 2)

    # construct pressure load function
    if type == "Constant"
        pt, pn = args
        f_px, f_py = x->pt, x->pn
    elseif type == "Gaussian"
        pn, x0, σ = args
        f_px, f_py = x->0, x-> pn * 1.0/(sqrt(2*pi*σ^2)) * exp.(-0.5*(x .- x0).^2/σ^2)
    else
        error("Force load type is not recogonized ", type)
    end

    # assemble the force

    @show ne
    for e = 1:ne
        if porder == 1
            loc_id = [e, e+1]
        elseif porder == 2
            loc_id = [2*(e-1)+1, 2*(e-1)+3, 2*(e-1)+2]
        else
            error("porder error porder ==", porder)
        end

        loc_xx = xx[loc_id]         #x-coordinates of element nodes
        elem_coords[:,1] .= loc_xx

        #return list
        weights, hs = get1DElemShapeData( elem_coords, ngp)

        for igp = 1:ngp
            gp_xx = loc_xx' * hs[igp]   #x-coordinates of element Gaussian points
            Ft[loc_id] += f_px(gp_xx) * hs[igp] * weights[igp]
            Fn[loc_id] += f_py(gp_xx) * hs[igp] * weights[igp]

        end
    end

    return Ft, Fn
end

function BoundaryCondition(tid, nx, ny, porder=2, Lx = 1.0, Ly = 0.5; force_scale=5.0)
    nnodes, neles = (nx*porder + 1)*(ny*porder + 1), nx*ny

    x = np.linspace(0.0, Lx, nx*porder + 1)
    y = np.linspace(0.0, Ly, ny*porder + 1)

    X, Y = np.meshgrid(x, y)
    nodes = zeros(nnodes,2)
    nodes[:,1], nodes[:,2] = X'[:], Y'[:]
    ndofs = 2

    EBC, g = zeros(Int64, nnodes, ndofs), zeros(nnodes, ndofs)
    FBC, fext = zeros(Int64, nnodes, ndofs), zeros(nnodes, ndofs)
    ft = nothing
    gt = nothing

    #todo only change 203

    # setting EBC:
    # tid = 1XX   |==|-->
    # tid = 2XX
    if div(tid,100)==1
        EBC[collect(1:nx*porder+1), :] .= -1 # fix bottom
        FBC[collect((nx*porder+1)*ny*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -2 # force on the top
    elseif div(tid,100)==2
        EBC[collect(1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -1 # fix left
        FBC[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -2 # force on the right
    elseif div(tid,100)==3
        EBC[collect(1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -1 # fix left
        FBC[collect(2:nx*porder+1), :] .= -2 # force on the bottom
    elseif div(tid,100)==4
        EBC[collect(1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -1 # fix left
        EBC[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -2 #pull right
    end

    # Uniform conservative pressure load
    # F1 = 2e2 /force_scale #gcm/ms^2 compress/pull
    # F2 = 2e1 /force_scale #gcm/ms^2 bend

    P1 = 8 /force_scale #gcm/ms^2 compress/pull
    P2 = 0.8 /force_scale #gcm/ms^2 bend
    P3 = 3/force_scale
    @show P1, P2
    ngp = 3
    #Bending or Pulling
    #
    if tid == 200
        F1, F2 = ComputeLoad(Ly, ny, porder, ngp, "Constant",  [P1, 0])
        fext[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), 1] .= F1
        fext[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), 2] .= F2

    elseif tid == 400
        g[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), 1] .= 0.00508e-3 *10#m/ms
        g[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), 2] .= 0
    else
        error("tid = $tid is not understood")
    end

    dof_to_active = findall(FBC[:].==-2)
    if length(dof_to_active) != 0
        ft = t->fext[:][dof_to_active] #*sin(π*t/(T))
    end

    dof_to_active = findall(EBC[:].==-2)
    if length(dof_to_active) != 0
        gt = t->[g[:][dof_to_active]*t, g[:][dof_to_active], g[:][dof_to_active]*0.0]
    end

    npoints = (nx+1)*(ny+1)
    node_to_point = zeros(Int64, nnodes) .- 1
    pointid = 1
    for j = 1:2:(ny*porder + 1)
        for i = 1:2:(nx*porder + 1)
            node_to_point[(j-1)*(nx*porder + 1) + i] = pointid
            pointid += 1
        end
    end

    return nodes, EBC, g, gt, FBC, fext, ft, npoints, node_to_point
end


text:
module GameUI

include("window.jl")

export window

end # module


text:
using UtilityFunctionsLH
using Test

@testset "UtilityFunctionsLH" begin
    include("helper_test.jl")
    include("crra_test.jl")
end

# ----------