using Mill: NGramMatrix

"""
	struct ExtractString{T} <: AbstractExtractor
		n::Int
		b::Int
		m::Int
		uniontypes::Bool
	end

Represents `String` as `n-`grams (`NGramMatrix` from `Mill.jl`) with base `b` and modulo `m`.

The `uniontypes` field determines whether extractor may or may not accept `missing`.
If `uniontypes` is false, it does not accept missing values. If `uniontypes` is true, it accepts missing values,
and always returns Mill structure of type Union{Missing, T} due to type stability reasons.

# Example
```jldoctest
julia> using Mill: catobs

julia> ExtractString(true)("hello")
2053×1 Mill.ArrayNode{Mill.NGramMatrix{Union{Missing, String}, Vector{Union{Missing, String}}, Union{Missing, Int64}}, Nothing}:
 "hello"

julia> mapreduce(ExtractString(true), catobs, (["hello", "world"]))
2053×2 Mill.ArrayNode{Mill.NGramMatrix{Union{Missing, String}, Vector{Union{Missing, String}}, Union{Missing, Int64}}, Nothing}:
 "hello"
 "world"

julia> mapreduce(ExtractString(true), catobs, ["hello", missing])
2053×2 Mill.ArrayNode{Mill.NGramMatrix{Union{Missing, String}, Vector{Union{Missing, String}}, Union{Missing, Int64}}, Nothing}:
 "hello"
 missing

julia> ExtractString(true)(missing)
2053×1 Mill.ArrayNode{Mill.NGramMatrix{Union{Missing, String}, Vector{Union{Missing, String}}, Union{Missing, Int64}}, Nothing}:
 missing

julia> ExtractString(false)("hello")
2053×1 Mill.ArrayNode{Mill.NGramMatrix{String, Vector{String}, Int64}, Nothing}:
 "hello"

julia> mapreduce(ExtractString(false), catobs, (["hello", "world"]))
2053×2 Mill.ArrayNode{Mill.NGramMatrix{String, Vector{String}, Int64}, Nothing}:
 "hello"
 "world"

julia> ExtractString(false)(["hello", "world"])
ERROR: This extractor does not support missing values

```
"""
struct ExtractString <: AbstractExtractor
	n::Int
	b::Int
	m::Int
	uniontypes::Bool
end

const MaybeString = Union{Missing, String}

ExtractString(n::Int, b::Int, m::Int) = ExtractString(n, b, m, true)
ExtractString(uniontypes::Bool) = ExtractString(3, 256, 2053, uniontypes)
ExtractString() = ExtractString(3, 256, 2053, true)
make_missing_string(s::ExtractString, v, store_input) =
	s.uniontypes ?
	_make_array_node(NGramMatrix(MaybeString[missing], s.n, s.b, s.m), [v], store_input) :
	error("This extractor does not support missing values")
stabilize_types_string(s::ExtractString, x) = s.uniontypes ? Vector{MaybeString}(x) : x
make_empty_string(s::ExtractString, store_input) where {T} =
	ArrayNode(NGramMatrix(s.uniontypes ? Vector{MaybeString}() : Vector{String}(), s.n, s.b, s.m))

(s::ExtractString)(v::String; store_input=false) =
	_make_array_node(NGramMatrix(stabilize_types_string(s, [v]), s.n, s.b, s.m), [v], store_input)
(s::ExtractString)(v::AbstractString; store_input=false) = s(String(v); store_input)
(s::ExtractString)(v::MissingOrNothing; store_input=false) = make_missing_string(s, v, store_input)
(s::ExtractString)(v::ExtractEmpty; store_input=false) = make_empty_string(s, store_input)
(s::ExtractString)(v; store_input=false) = make_missing_string(s, v, store_input)
(s::ExtractString)(v::Symbol; store_input=false) = s(String(v); store_input)

"""
	extractscalar(Type{String}, n = 3, b = 256, m = 2053)

represents strings as ngrams with
- `n` (the degree of ngram),
- `b` base of string,
- `m` modulo on index of the token to reduce dimension

	extractscalar(Type{Number}, m = 0, s = 1)

extracts number subtracting `m` and multiplying by `s`

# Example
```jldoctest
julia> JsonGrinder.extractscalar(String, 3, 256, 2053, true)("5")
2053×1 Mill.ArrayNode{Mill.NGramMatrix{Union{Missing, String}, Vector{Union{Missing, String}}, Union{Missing, Int64}}, Nothing}:
 "5"

julia> JsonGrinder.extractscalar(Int32, 3, 256, true)("5")
1×1 Mill.ArrayNode{Matrix{Union{Missing, Int32}}, Nothing}:
 512

julia> JsonGrinder.extractscalar(String, 3, 256, 2053, false)("5")
2053×1 Mill.ArrayNode{Mill.NGramMatrix{String, Vector{String}, Int64}, Nothing}:
 "5"

julia> JsonGrinder.extractscalar(Int32, 3, 256, false)("5")
1×1 Mill.ArrayNode{Matrix{Int32}, Nothing}:
 512
```
"""
extractscalar(::Type{<:AbstractString}, n = 3, b = 256, m = 2053, uniontypes = true) = ExtractString(n, b, m, uniontypes)

Base.hash(e::ExtractString, h::UInt) = hash((e.n, e.b, e.m, e.uniontypes), h)
Base.:(==)(e1::ExtractString, e2::ExtractString) =
	e1.n === e2.n && e1.b === e2.b && e1.m === e2.m && e1.uniontypes === e2.uniontypes

# to be compatible with the number version
extractscalar(::Type{<:AbstractString}, e::Entry, uniontypes::Bool = true) = ExtractString(uniontypes)


using MeshCatMechanisms
using RigidBodyDynamics
using MeshCat
# using Blink
# AtomShell.isinstalled() || AtomShell.install()
# traj_folder = joinpath(dirname(pathof(TrajectoryOptimization)),"..")
# urdf_folder = joinpath(traj_folder, "dynamics/urdf")
# urdf_kuka_orig = joinpath(urdf_folder, "kuka_iiwa.urdf")
# urdf_kuka = joinpath(urdf_folder, "kuka.urdf")
#
# mech = parse_urdf(urdf_kuka)
# visuals = URDFVisuals(urdf_kuka)
# vis = MechanismVisualizer(mech, visuals)
# open(vis)
# set_configuration!(vis, rand(6))
#
# IJuliaCell(vis)

function animate_trajectory(vis, X::AbstractMatrix{Float64}, dt=0.1)
    animate_trajectory(vis, TrajectoryOptimization.to_dvecs(X), dt)
end

function animate_trajectory(vis, X::Trajectory, dt=0.1)
    for x in X
        set_configuration!(vis, x[1:7])
        sleep(dt)
    end
end

green_ = MeshPhongMaterial(color=RGBA(0, 1, 0, 1.0))
red_ = MeshPhongMaterial(color=RGBA(1, 0, 0, 1.0))
body_collision = MeshPhongMaterial(color=RGBA(1, 0, 0, 0.5))
blue_ = MeshPhongMaterial(color=RGBA(0, 0, 1, 0.5))
orange_ = MeshPhongMaterial(color=RGBA(233/255, 164/255, 16/255, 1.0))
black_ = MeshPhongMaterial(color=RGBA(0, 0, 0, 1.0))
black_transparent = MeshPhongMaterial(color=RGBA(0, 0, 0, 0.1))


#struct Vertex { union { struct { _float x, y, z; }; _float pos[3]; }; };

type ShapeCoordinates
	#std::string& name
	positions::AVec3f
	normals::AVec3f
	texcoords::AVec2f
	origin::Vec3f
	scale::Vec3f

	ShapeCoordinates() = new(AVec3f(),AVec3f(),AVec1f(),Vec3f(),Vec3f())
	ShapeCoordinates(positions,normals,texcoords,origin,scale) = new(positions,normals,texcoords,origin,scale)
end

type VertexIndex
	v::Integer
	vt::Integer
	vn::Integer

	VertexIndex() = new(0,0,0)
	VertexIndex(a::Integer) = new(a,a,a)
	VertexIndex(v::Integer, vn::Integer, vt::Integer) = new(v,vt,vn)
end

const AVertexIndex = Array{VertexIndex,1}
const AAVertexIndex = Array{AVertexIndex,1}

function clear(shape::ShapeCoordinates)
	shape.positions = AVec3f()
	shape.normals = AVec3f()
	shape.texcoords = AVec1f()
	shape.origin = Vec3f()
	shape.scale = Vec3f()
end

# for std::map
function operator<(a::VertexIndex, b::VertexIndex)
	if a.v != b.v return (a.v < b.v) end
	if a.vn != b.vn return (a.vn < b.vn) end
	if a.vt != b.vt return (a.vt < b.vt) end
	false
end

isSpace(_char_ c) = (c == ' ') || (c == '\t')
isNewLine(_char_ c) = (c == '\r') || (c == '\n') || (c == '\0')

function getOrigin(min_BBOX::Vec3f, max_BBOX::Vec3f, scale::Float32)
	(max_BBOX + min_BBOX) * Vec3f(scale * 0.5f)
end

function getScale(min_BBOX::Vec3f, max_BBOX::Vec3f)
	size = max_BBOX - min_BBOX
	scale_value = size.x > size.z ? size.x : size.z
	scale_value = scale_value > size.y ? scale_value : size.y
	if scale_value > 0 scale_value = 1 / scale_value end
	scale_value
end

function parseFloat(token::String)
	token += strspn(token, " \t")
	f = (Float32)atof(token)
	token += strcspn(token, " \t\r")
	f
end

function parsevec(token::String, count::Number)
	a = zeros(count)
	for i=1:count a[i] = parseFloat(token) end
	Vec{count, Float32}(a)
end

# Make index zero-base, and also support relative index.
function fixIndex(idx::Integer, n::Integer)
	if idx > 0 i = idx - 1
	else if idx == 0  i = 0
	else i = n + idx # negative value = relative
	end
	i
end

# Parse triples: i, i/j/k, i//k, i/j
function parseTriple(token::String, vsize::Integer, vnsize::Integer, vtsize::Integer)
	vi = VertexIndex(-1)

	vi.v = fixIndex(atoi(token), vsize)
	token += strcspn(token, "/ \t\r");
	if token[0] != '/' return vi end
	++token;

	# i//k
	if token[0] == '/'
		++token
		vi.vn = fixIndex(atoi(token), vnsize)
		token += strcspn(token, "/ \t\r")
		return vi
	end

	# i/j/k or i/j
	vi.vt = fixIndex(atoi(token), vtsize)
	token += strcspn(token, "/ \t\r")
	if (token[0] != '/') return vi; else ++token;  # skip '/'
	vi.vn = fixIndex(atoi(token), vnsize)
	token += strcspn(token, "/ \t\r")
	vi
end

function updateVertex(vi::VertexIndex, vCache::Dict{VertexIndex, _uint},
	coord::ShapeCoordinates, shape::Mesh)

	idx = 0

	it = vCache.find(vi)

	if it != vCache.end() return it->second end

	if vi.v >= length(coord.position)
		error("_DEBUG_BREAK_IF(((_uint)(vi.v) >= (_uint)(coord.positions.size())))")
		return
	end

	shape->addPosition(coord.positions[vi.v] * coord.scale - coord.origin)

	if vi.vn >= 0
		shape->addNormal(coord.normals[vi.vn])
	end

	if vi.vt >= 0
		shape->addTexcoord(coord.texcoords[vi.vt])
	end

	idx = shape->getPositionSize() - 1
	vCache[vi] = idx

	idx;
end

function parseIndicies(
	std::map<VertexIndex, _uint> vCache,
	const std::vector<std::vector<VertexIndex> >& group,
	#const int material_id,
	#const std::string &name,
	#bool clearCache,
	ShapeCoordinates in,
	shape::Mesh)

	clear(shape) # remove old

	if empty(group) return false end

	# Flatten vertices and indices
	for i=0:(group.size()-1)
		face = group[i]

		i0 = face[0]
		i1 = VertexIndex(-1)
		i2 = face[1]

		npolys = length(face)

		# Polygon -> triangle fan conversion
		for k=2:(npolys-1)
			i1 = i2
			i2 = face[k]

			v0 = updateVertex(i0, vCache, in, shape)
			v1 = updateVertex(i1, vCache, in, shape)
			v2 = updateVertex(i2, vCache, in, shape)

			addIndex(shape,v0)
			addIndex(shape,v1)
			addIndex(shape,v2)

			#shape.material_ids.push_back(material_id);
		end
	end

	#shape.name = name;
	#if (clearCache)	vCache.clear();

	true;

end

function parseIndicies(tmp_vi::Array{Array{VertexIndex,1},1},tmp::ShapeCoordinates,shape::Mesh)

	clear(shape) # remove old

	ig = 0
	iCount = length(tmp_vi)
	vCount = length(tmp.positions)
	vnCount = length(tmp.normals)
	vtCount = length(tmp.texcoords)

	if !iCount
		error("No index coordinates! (tmp)")
		return false
	end

	if !vCount
		error("No vertex coordinates! (tmp)")
		return false
	end

	# For each vertex of each triangle
	for (AAVertexIndex::iterator it = tmp_vi.begin(); it < tmp_vi.end(); ++it, ++ig)
		for (AVertexIndex::iterator jt = (*it).begin(); jt < (*it).end(); ++jt)
			vi = jt
			iv = vi.v
			ivn = vi.vn
			ivt = vi.vt

			if iv == -1
				error("Invalid index for vertex coordinates! %u (%d = -1).", ig, iv)
				return false
			else
				if iv < vCount
					addPosition(shape,*(tmp.positions.begin() + iv))
				else
					error("Missmatch count of vertex coordinates! %u (%d >= %u).", ig, iv, vCount)
					return false
				end
			end

			if ivn > -1
				if ivn < vnCount
					addNormal(shape,*(tmp.normals.begin() + ivn))
				else
					error("Missmatch count of normal coordinates! %u (%d = -1 or >= %u).", ig, ivn, vnCount)
					return false
				end
			end

			if ivt > -1
				if ivt < vtCount
					addTexcoord(shape,*(tmp.texcoords.begin() + ivt))
				else
					error("Missmatch count of texture coordinates! %u (%d = -1 or >= %u).", ig, ivt, vtCount)
					return false
				end
			end
		end
	end

	vCount = length(shape.positions)
	vnCount = length(shape.normals)
	vtCount = length(shape.texcoords)

	if !vCount
		error("No vertex coordinates!")
		return false
	end

	if vnCount > 0 && vCount != vnCount
		error("Missmatch count of vertex and normal coordinates! ($vCount != $vnCount).")
		return false
	end

	if vtCount > 0 && vCount != vtCount
		error("Missmatch count of vertex and texture coordinates! ($vCount != $vtCount).")
		return false
	end

	true
end

function loadOBJ(path::String, shape::Mesh)
	println("Load OBJ file $path...");

	tmp_v = AVec3f()
	tmp_vn = AVec3f()
	tmp_vt = AVec2f()
	tmp_g = AAVertexIndex()

	std::filebuf fb
	std::filebuf* fb_result = fb.open(path, std::ios::in)

	if fb_result == C_NULL
		error("Cannot open file $path!")
		return false
	end

	std::istream is(&fb)
	maxchars = 8192
	buf = zeros(Arrray{Char,1},maxchars)
	vertexCache = Dict{VertexIndex,UInt32}()

	min_BBOX = zeros(Vec3f)
	max_BBOX = zeros(Vec3f)

	while is.peek() != -1
		is.getline(&buf[0], maxchars)
		std::string linebuf(&buf[0])
		size = linebuf.size()
		id = size - 1

		# trim '\r\n', '\n'
		if size > 0 && linebuf[id] == '\n' linebuf.erase(id)
			size = linebuf.size()
			id = size - 1
		end

		if size > 0 && linebuf[id] == '\r' linebuf.erase(id)
			size = linebuf.size()
			id = size - 1
		end

		if empty(linebuf) continue end

		token = linebuf.c_str()
		token += strspn(token, " \t")
		if !token error("!token") end

		first = linebuf[0]
		second = linebuf[1]
		third = linebuf[2]

		if first == '\0' continue # empty line
		elseif first == '#' continue  # comment line

		elseif first == 'v' && isSpace(second)
			token += 2
			v = parsevec3(token)

			# set bounding box
			if v.x < min_BBOX.x min_BBOX.x = v.x end
			if v.y < min_BBOX.y min_BBOX.y = v.y end
			if v.z < min_BBOX.z min_BBOX.z = v.z end
			if v.x > max_BBOX.x max_BBOX.x = v.x end
			if v.y > max_BBOX.y max_BBOX.y = v.y end
			if v.z > max_BBOX.z max_BBOX.z = v.z end

			push!(tmp_v,v)
			continue

		elseif first == 'v' && second == 'n' && isSpace(third)
			token += 3
			push!(tmp_vn,parsevec3(token))
			continue

		elseif first == 'v' && second == 't' && isSpace(third)
			token += 3
			push!(tmp_vt,parsevec2(token))
			continue

		elseif first == 'f' && isSpace(second)
			token += 2
			token += strspn(token, " \t")

			group = Array{VertexIndex,1}()
			while !isNewLine(token[0])
				vi = parseTriple(token, tmp_v.size() / 3, tmp_vn.size() / 3, tmp_vt.size() / 2)
				push!(group,vi)
				n = strspn(token, " \t\r")
				token += n
			end

			push!(tmp_g,group)

			continue
		end
	end
	close(fb)

	scale_value = getScale(min_BBOX, max_BBOX)
	coord = ShapeCoordinates( tmp_v, tmp_vn, tmp_vt, getOrigin(min_BBOX, max_BBOX, scale_value), Vec3f(scale_value) )

	result = parseIndicies(vertexCache, tmp_g, coord, shape) # parseIndicies(tmp_g, in, shape);
	clear(tmp_g)

	if result println("Finsihed loading OBJ file $path.") end

	result
end


export generate_noise_mpi

import CorrNoise
using Random
using FITSIO

try
    import MPI
catch
end

"""
This structure holds informations about a certain chunk of data produced by a certain MPI process to send to another MPI process.

Field              | Type           | Meaning
:----------------- |:-------------- |:------------------------------------------------------------
`pol_number`       | Int            | ID number of the polarimeter
`first_idx`        | Int            | Index of the first element ofthe chunk
`last_idx`         | Int            | Index of the last element of the chunk
`dest_rank`        | Int            | Number of destination rank

# Example
Suppose that the following array of data: [1, 10, 20, 3, 4, 5, 7]
is measured by the polarimeter number 2.
If the chunk [20, 3, 4] needs to be sent to rank number 3, the corresponding structure will be:

chunk_to_send(2, 3, 5, 3)
"""
struct chunk_to_send
    pol_number ::Int
    first_idx::Int
    last_idx::Int
    dest_rank::Int
end

function generate_noisechunks_to_send(chunks, baseline_length_s, fsamp_hz, total_number_of_polarimeters)

    commsize = length(chunks)
    noisechunks_to_send = []

    for cur_rank in 0:(commsize-1) #loop on ranks
        pol_cur_rank  = cur_rank+1 #polarimeter assigned to cur_rank for noise production
        chunk_cur_rank = []

        #Only the first "total_number_of_polarimeters" ranks will generate noise, so they will have something to send.
        if  0 <= cur_rank < total_number_of_polarimeters
            for i in 1:length(chunks) #loop on chunks
                rank_it = i-1
                for j in 1:length(chunks[i]) #loop on detectors per rank
                    if chunks[i][j].pol_number == pol_cur_rank #if cur_rank is going to produce noise for the "correct" polarimeter
                        first_idx = (chunks[i][j].first_idx-1)*baseline_length_s*fsamp_hz+1
                        last_idx = first_idx + chunks[i][j].num_of_elements*baseline_length_s*fsamp_hz -1
                        dest_rank = rank_it
                        chunk = chunk_to_send(pol_cur_rank, first_idx, last_idx, dest_rank)
                        chunk_cur_rank = append!(chunk_cur_rank, [chunk])
                    end
                end
            end
            noisechunks_to_send= append!(noisechunks_to_send, [chunk_cur_rank])
        end
    end

    return noisechunks_to_send

end

"""
    function generate_noise_mpi(chunks, baseline_length_s, baselines_per_process, fsamp_hz, σ_k, fknee_hz, rank, comm)

    This MPI based function can be used to generate white noise and 1/f noise when the simulation is splitted in various MPI processes.

    It requires in input:
    - the data chunks (which can be obtained by using the function `split_tod_mpi`)
    - the array containing the number of 1/f baselines to simulate for each process.
            It can be obtained by using the function `split_into_n` in the following way:

            split_into_n(num_of_polarimeters*baselines_per_pol, num_of_MPI_proc)

            where baselines_per_pol = Int64(total_time/baseline_length_s) is the number of baselines of each polarimeter
    - the length (in seconds) of each 1/f noise baseline
    - the total duration (in seconds) of data acquisition
    - the sampling frequency (in Hz)
    - an array containing the RMS of the temperature (in K) for each polarimeter, σ_k[i]= tsys_k[i] / sqrt(β_hz[i] * τ_s)
    - an array containing the knee frequency (in Hz) for each polarimeter
    - an array containing the oof slope for each polarimeter
    - the MPI communicator
    - the random generator seed, as optional argument.

    It returns the noise TOD for the current rank.

    The function behaves differently according to the number of processes used:

        1. If the number of processes is greater than the number of polarimeters:
           the first "number of polarimeters" processes will simultaneously produce the noise, one process for one polarimeter
           (rank 0 for pol.1, rank 1 for pol.2 and so on).
           After that, the noise is redistributed between the ranks according to the input chunk structure.

        2. If the number of processes is lesser than the number of polarimeters,
           Each partial noise TOD is generated by rank 0 and then sent to the correspondent rank.

    N.B. if you want to use this function without MPI, remember to put rank = 0 and comm = missing

"""
function  generate_noise_mpi(chunks, baselines_per_process, baseline_length_s, total_time, fsamp_hz, σ_k, fknee_hz, slope, rank, comm, input_seed = missing)

    total_number_of_polarimeters = length(σ_k)
    commsize = length(chunks)

    if ismissing(input_seed) father_seed = rand(1:10000) else father_seed = input_seed end
    Random.seed!(father_seed)
    seeds = Int64[rand(1:10000) for i in 1:total_number_of_polarimeters]

    ################ IF MORE PROCESSES THAN POLARIMETERS ################
    if(commsize >= total_number_of_polarimeters)

        noisechunks_to_send = generate_noisechunks_to_send(chunks, baseline_length_s, fsamp_hz, total_number_of_polarimeters)

        ##### NOISE PRODUCTION #####
        pol_number = rank+1
        pol_noise = Float64[]
        noise = Array{Float64}(undef, baselines_per_process[rank+1]*baseline_length_s*fsamp_hz)

        if  0 <= rank < total_number_of_polarimeters  #if first "total_number_of_polarimeters" ranks
            rng = CorrNoise.OofRNG(Random.MersenneTwister(seeds[pol_number]), -abs(slope[pol_number]), 1.15e-5, fknee_hz[pol_number], fsamp_hz)
            samples_per_pol = total_time*fsamp_hz
            pol_noise = Float64[CorrNoise.randoof(rng) * σ_k[pol_number] for i in 1:(samples_per_pol)]
        end

        if(!ismissing(comm))
            MPI.Barrier(comm)
        end

        ##### SEND AND RECEIVE #####

        start_idx = 1
        end_idx = 0

        for i in 1:length(noisechunks_to_send)  #loop on ranks
            cur_rank = i-1

            for j in 1:length(noisechunks_to_send[i]) #loop on the noise chunks produced by cur_rank
                start_idx_noise = noisechunks_to_send[i][j].first_idx
                end_idx_noise = noisechunks_to_send[i][j].last_idx
                dest_rank = noisechunks_to_send[i][j].dest_rank

                cur_num_noise_samples = end_idx_noise-start_idx_noise+1

                #if rank has to send noise (since it produced it)
                if cur_rank == rank

                    #if the destination rank is itself, keep noise chunk
                    if dest_rank == rank
                        end_idx = start_idx + cur_num_noise_samples - 1
                        noise[start_idx:end_idx] = pol_noise[start_idx_noise:end_idx_noise]
                        start_idx += cur_num_noise_samples

                    #otherwise send to the destination rank
                    else
                        cur_noise_send = pol_noise[start_idx_noise:end_idx_noise]
                        MPI.Send(cur_noise_send, dest_rank , 0, comm)
                    end

                #if rank is a destination rank for someone
                elseif dest_rank == rank

                    cur_noise_rec = Array{Float64}(undef, cur_num_noise_samples)
                    end_idx = start_idx + cur_num_noise_samples -1
                    MPI.Recv!(cur_noise_rec, cur_rank, 0, comm)
                    noise[start_idx:end_idx] = cur_noise_rec
                    start_idx += cur_num_noise_samples
                end

            end

        end

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

    ################IF LESS PROCESSES THAN POLARIMETERS################
    else
        if rank == 0
            previous_detector = chunks[1][1].pol_number
            noise = Float64[]  #noise rank 0

            rng = CorrNoise.OofRNG(Random.MersenneTwister(seeds[1]), -abs(slope[1]), 1.15e-5, fknee_hz[1], fsamp_hz)

            for i in 1:length(chunks)   #loop on ranks
                num_noise_samples = baselines_per_process[i]*baseline_length_s*fsamp_hz

                for j in 1:length(chunks[i]) #loop on detectors per rank
                    cur_detector = chunks[i][j].pol_number

                    if cur_detector != previous_detector  #if new detector generate new noise
                        rng = CorrNoise.OofRNG(Random.MersenneTwister(seeds[cur_detector]), -abs(slope[cur_detector]), 1.15e-5, fknee_hz[cur_detector], fsamp_hz)
                    end
                    cur_num_noise_samples = chunks[i][j].num_of_elements*baseline_length_s*fsamp_hz
                    cur_noise = Float64[CorrNoise.randoof(rng) * σ_k[cur_detector] for i in 1:(cur_num_noise_samples)]
                    previous_detector =  cur_detector

                    if (i > 1)  #send to other ranks, not rank 0
                        MPI.Send(cur_noise, i-1, 0, comm)
                    else
                        noise = append!(noise, cur_noise)
                    end
                end
            end
        else
            noise = Float64[]
            for j in 1:length(chunks[rank+1])
                cur_noise = Array{Float64}(undef, chunks[rank+1][j].num_of_elements*baseline_length_s*fsamp_hz)
                MPI.Recv!(cur_noise, 0, 0, comm)
                noise = append!(noise, cur_noise)
            end
        end

    end

    return noise

end


push!(LOAD_PATH,"/home/michel/git/PIPS.jl/src")
using PIPS
using Test
using MPI
using Libdl

# hs071
# min x1 * x4 * (x1 + x2 + x3) + x3
# st  x1 * x2 * x3 * x4 >= 25
#     x1^2 + x2^2 + x3^2 + x4^2 = 40
#     1 <= x1, x2, x3, x4 <= 5
# Start at (1,5,5,1)
# End at (1.000..., 4.743..., 3.821..., 1.379...)

function init_x0_wrapper(x0_ptr::Ptr{Float64}, cbd::Ptr{CallBackData})
    data = unsafe_load(cbd)
    userdata = data.prob
    prob = unsafe_pointer_to_objref(userdata)::PipsNlpProblemStruct
    rowid = Int(Int(data.row_node_id))
    colid = Int(Int(data.col_node_id))
    @assert(rowid == colid)
    # n0 = prob.model.get_num_cols(colid)
    x0 = unsafe_wrap(Array, x0_ptr, 4)
    x0[1] = 1.0
    x0[2] = 5.0
    x0[3] = 5.0
    x0[4] = 1.0
    # prob.model.str_init_x0(colid,x0)
            # @assert(id in getLocalBlocksIds(instance.internalModel))
            # mm = getModel(instance.internalModel,id)
            # nvar = getNumVars(instance.internalModel,id)
            # @assert length(x0) == nvar
            #
            # for i=1:nvar
            #     x0[i] = getvalue(Variable(mm,i))
            #     isnan(x0[i]) ? x0[i]=1.0 : nothing
            # end
    return Int32(1)
end

# prob info (prob_info)
function prob_info_wrapper(n_ptr::Ptr{Cint}, col_lb_ptr::Ptr{Float64}, col_ub_ptr::Ptr{Float64}, m_ptr::Ptr{Cint}, row_lb_ptr::Ptr{Float64}, row_ub_ptr::Ptr{Float64}, cbd::Ptr{CallBackData})
    data = unsafe_load(cbd)
    userdata = data.prob
    prob = unsafe_pointer_to_objref(userdata)::PipsNlpProblemStruct
    rowid = Int(Int(data.row_node_id))
    colid = Int(data.col_node_id)
    flag = Int(data.flag)
    @assert(rowid == colid)

	mode = (col_lb_ptr == C_NULL) ? (:Structure) : (:Values)
	if mode==:Structure
        if flag == 1
            m = 0
            println("Only storing this")
            unsafe_store!(m_ptr,convert(Cint,m)::Cint)
        else
            if colid == 0
                n = prob.n
                m = prob.m
            end
            # n = getNumVars(instance.internalModel,id)
            # m = getNumCons(instance.internalModel,id)
    		# prob.model.set_num_rows(colid, m)
    		# prob.model.set_num_cols(colid, n)
            println("Storing m and n")
            unsafe_store!(n_ptr,convert(Cint,n)::Cint)
            unsafe_store!(m_ptr,convert(Cint,m)::Cint)
        end
	else
        # This is for linking constraints which is not supported
        if flag == 1
            m = unsafe_load(m_ptr)
            @assert m==0
        else
    		n = unsafe_load(n_ptr)
    		m = unsafe_load(m_ptr)

    		col_lb = unsafe_wrap(Array,col_lb_ptr,n)
    		col_ub = unsafe_wrap(Array,col_ub_ptr,n)
    		row_lb = unsafe_wrap(Array,row_lb_ptr,m)
    		row_ub = unsafe_wrap(Array,row_ub_ptr,m)
    		# prob.model.str_prob_info(colid,flag,mode,col_lb,col_ub,row_lb,row_ub)
            if colid == 0
                col_lb[1] = prob.x_L[1]
                col_lb[2] = prob.x_L[2]
                col_lb[3] = prob.x_L[3]
                col_lb[4] = prob.x_L[4]
                col_ub[1] = prob.x_U[1]
                col_ub[2] = prob.x_U[2]
                col_ub[3] = prob.x_U[3]
                col_ub[4] = prob.x_U[4]
                row_lb[1] = prob.g_L[1]
                row_lb[2] = prob.g_L[2]
                row_ub[1] = prob.g_U[1]
                row_ub[2] = prob.g_U[2]
        		neq = 0
        		nineq = 0
        		for i = 1:length(row_lb)
        			if row_lb[i] == row_ub[i]
        				neq += 1
        			else
        				nineq += 1
        			end
        		end
        		@assert(neq+nineq == length(row_lb) == m)
            end
    		# prob.model.set_num_eq_cons(colid,neq)
    		# prob.model.set_num_ineq_cons(colid,nineq)
        end
	end
	return Int32(1)
end
# Objective (eval_f)
function eval_f_wrapper(x0_ptr::Ptr{Float64}, x1_ptr::Ptr{Float64}, obj_ptr::Ptr{Float64}, cbd::Ptr{CallBackData})
    data = unsafe_load(cbd)
    userdata = data.prob
    prob = unsafe_pointer_to_objref(userdata)::PipsNlpProblemStruct
    rowid = Int(Int(data.row_node_id))
    colid = Int(data.col_node_id)
    @assert(rowid == colid)
    # n0 = prob.model.get_num_cols(0)
    # n1 = prob.model.get_num_cols(colid)
    # Calculate the new objective
    n0 = prob.n
    n1 = prob.n
    x0 = unsafe_wrap(Array, x0_ptr, n0)
    x1 = unsafe_wrap(Array, x1_ptr, n1)
    obj = x0[1] * x0[4] * (x0[1] + x0[2] + x0[3]) + x0[3]

    # new_obj = convert(Float64, prob.model.str_eval_f(colid,x0,x1))::Float64
    unsafe_store!(obj_ptr, obj)
    return Int32(1)
end

# Constraints (eval_g)
function eval_g_wrapper(x0_ptr::Ptr{Float64}, x1_ptr::Ptr{Float64}, eq_g_ptr::Ptr{Float64}, ineq_g_ptr::Ptr{Float64}, cbd::Ptr{CallBackData})
    # @show " julia - eval_g_wrapper "
    data = unsafe_load(cbd)
    # @show data
    userdata = data.prob
    prob = unsafe_pointer_to_objref(userdata)::PipsNlpProblemStruct
    rowid = Int(data.row_node_id)
    colid = Int(data.col_node_id)
    @assert(rowid == colid)
    # n0 = probd.model.get_num_cols(0)
    # n1 = prob.model.get_num_cols(colid)
    n0 = prob.n
    n1 = prob.n
    x0 = unsafe_wrap(Array, x0_ptr, n0)
    x1 = unsafe_wrap(Array, x1_ptr, n1)
    # Calculate the new constraint values
    # neq = prob.model.get_num_eq_cons(rowid)
    # nineq = prob.model.get_num_ineq_cons(rowid)
    neq = 1
    nineq = 1
    eq_g = unsafe_wrap(Array, eq_g_ptr, neq)
    ineq_g = unsafe_wrap(Array, ineq_g_ptr, nineq)
    eq_g[1] = x0[1]^2 + x0[2]^2 + x0[3]^2 + x0[4]^2
    ineq_g[1] = x0[1]   * x0[2]   * x0[3]   * x0[4]

    return Int32(1)
end

# Objective gradient (eval_grad_f)
function eval_grad_f_wrapper(x0_ptr::Ptr{Float64}, x1_ptr::Ptr{Float64}, grad_f_ptr::Ptr{Float64}, cbd::Ptr{CallBackData})
    data = unsafe_load(cbd)
    userdata = data.prob
    prob = unsafe_pointer_to_objref(userdata)::PipsNlpProblemStruct
    rowid = Int(data.row_node_id)
    colid = Int(data.col_node_id)
    n0 = prob.n
    n1 = prob.n
    x0 = unsafe_wrap(Array, x0_ptr, n0)
    x1 = unsafe_wrap(Array, x1_ptr, n1)
    # grad_len = prob.model.get_num_cols(colid)
    grad_len = prob.n
    grad_f = unsafe_wrap(Array, grad_f_ptr, grad_len)
    grad_f[1] = x0[1] * x0[4] + x0[4] * (x0[1] + x0[2] + x0[3])
    grad_f[2] = x0[1] * x0[4]
    grad_f[3] = x0[1] * x0[4] + 1.0
    grad_f[4] = x0[1] * (x0[1] + x0[2] + x0[3])

    # TODO: Gotta think about this
    # if prob.model.get_sense() == :Max
    #     new_grad_f *= -1.0
    # end
    return Int32(1)
end

# Jacobian (eval_jac_g)
function eval_jac_g_wrapper(x0_ptr::Ptr{Float64}, x1_ptr::Ptr{Float64},
	e_nz_ptr::Ptr{Cint}, e_values_ptr::Ptr{Float64}, e_row_ptr::Ptr{Cint}, e_col_ptr::Ptr{Cint},
	i_nz_ptr::Ptr{Cint}, i_values_ptr::Ptr{Float64}, i_row_ptr::Ptr{Cint}, i_col_ptr::Ptr{Cint},
	cbd::Ptr{CallBackData}
	)
    data = unsafe_load(cbd)
    userdata = data.prob
    prob = unsafe_pointer_to_objref(userdata)::PipsNlpProblemStruct
    rowid = Int(Int(data.row_node_id))
    colid = Int(Int(data.col_node_id))
    flag = Int(data.flag)
    # n0 = prob.model.get_num_cols(0)
    # n1 = prob.model.get_num_cols(rowid) #we can do this because of 2-level and no linking constraint
    n0 = prob.n
    n1 = prob.n
    x0 = unsafe_wrap(Array, x0_ptr, n0)
    x1 = unsafe_wrap(Array, x1_ptr, n1)
    # nrow = prob.model.get_num_rows(rowid)
    # ncol = prob.model.get_num_cols(colid)
    ncol = prob.n
    mode = (e_col_ptr == C_NULL && i_col_ptr == C_NULL) ? (:Structure) : (:Values)
    if flag != 1
        if mode == :Structure
            if colid == 0 && rowid == 0
                e_values = unsafe_wrap(Array,e_values_ptr,0)
                e_colptr = unsafe_wrap(Array,e_col_ptr,0)
                e_rowidx = unsafe_wrap(Array,e_row_ptr,0)
                i_values = unsafe_wrap(Array,i_values_ptr,0)
                i_colptr = unsafe_wrap(Array,i_col_ptr,0)
                i_rowidx = unsafe_wrap(Array,i_row_ptr,0)

                # (e_nz,i_nz) = prob.model.str_eval_jac_g(rowid,colid,flag,x0,x1,mode,e_rowidx,e_colptr,e_values,i_rowidx,i_colptr,i_values)
                e_nz::Cint = prob.n
                i_nz::Cint = prob.n
                # else
                #     e_nz = 0
                #     i_nz = 0
                # end
                unsafe_store!(e_nz_ptr,convert(Cint,e_nz)::Cint)
                unsafe_store!(i_nz_ptr,convert(Cint,i_nz)::Cint)
            end
        else
            if colid == 0 && rowid == 0
                e_nz = unsafe_load(e_nz_ptr)
                i_nz = unsafe_load(i_nz_ptr)
                if e_nz > 0
                    e_rowidx = unsafe_wrap(Array,e_row_ptr,e_nz)
                    e_colptr = unsafe_wrap(Array,e_col_ptr,ncol+1)
                    e_values = unsafe_wrap(Array,e_values_ptr,e_nz)
                    e_rowidx[1] = 0
                    e_rowidx[2] = 0
                    e_rowidx[3] = 0
                    e_rowidx[4] = 0
                    e_colptr[1] = 0
                    e_colptr[2] = 1
                    e_colptr[3] = 2
                    e_colptr[4] = 3
                    e_colptr[5] = 4
                    e_values[1] = 2*x0[1]  # 2,1
                    e_values[2] = 2*x0[2]  # 2,2
                    e_values[3] = 2*x0[3]  # 2,3
                    e_values[4] = 2*x0[4]  # 2,4
                end
                if i_nz > 0
                    # Constraint (row) 2
                    i_rowidx = unsafe_wrap(Array,i_row_ptr,i_nz)
                    i_colptr = unsafe_wrap(Array,i_col_ptr,ncol+1)
                    i_values = unsafe_wrap(Array,i_values_ptr,i_nz)
                    i_rowidx[1] = 0
                    i_rowidx[2] = 0
                    i_rowidx[3] = 0
                    i_rowidx[4] = 0
                    i_colptr[1] = 0
                    i_colptr[2] = 1
                    i_colptr[3] = 2
                    i_colptr[4] = 3
                    i_colptr[5] = 4
                    # prob.model.str_eval_jac_g(rowid,colid,flag,x0,x1,mode,e_rowidx,e_colptr,e_values,i_rowidx,i_colptr,i_values)
                    # Constraint (row) 2
                    i_values[1] = x0[2]*x0[3]*x0[4]  # 1,1
                    i_values[2] = x0[1]*x0[3]*x0[4]  # 1,2
                    i_values[3] = x0[1]*x0[2]*x0[4]  # 1,3
                    i_values[4] = x0[1]*x0[2]*x0[3]  # 1,4
                end
                return Int32(1)
            end
        end
    else
        @assert flag == 1
        if mode == :Structure
            e_nz = 0
            i_nz = 0
            unsafe_store!(e_nz_ptr,convert(Cint,e_nz)::Cint)
            unsafe_store!(i_nz_ptr,convert(Cint,i_nz)::Cint)
        else
            e_nz = unsafe_load(e_nz_ptr)
            i_nz = unsafe_load(i_nz_ptr)
            @assert e_nz == i_nz == 0
        end
    end
    return Int32(1)
end

# Hessian
function eval_h_wrapper(x0_ptr::Ptr{Float64}, x1_ptr::Ptr{Float64}, lambda_ptr::Ptr{Float64}, nz_ptr::Ptr{Cint}, values_ptr::Ptr{Float64}, row_ptr::Ptr{Cint}, col_ptr::Ptr{Cint}, cbd::Ptr{CallBackData})
    data = unsafe_load(cbd)
    userdata = data.prob
    prob = unsafe_pointer_to_objref(userdata)::PipsNlpProblemStruct
    rowid = Int(data.row_node_id)
    colid = Int(data.col_node_id)
    flag = Int(data.flag)

    high = max(rowid,colid)
    low  = min(rowid,colid)
    n0 = prob.n
    # n0 = prob.model.get_num_cols(0)
    # n1 = prob.model.get_num_cols(high)
    x0 = unsafe_wrap(Array,x0_ptr,n0)
    ncol = prob.n
    # x1 = unsafe_wrap(Array,x1_ptr,n1)
    # ncol = prob.model.get_num_cols(low)
    # g0 = prob.model.get_num_rows(high)
    lambda = unsafe_wrap(Array,lambda_ptr,2)
    obj_factor = 1.0
    mode = (col_ptr == C_NULL) ? (:Structure) : (:Values)
    if mode == :Structure
		if (rowid == 0 && colid == 0)
            nz = prob.nele_hess
        else
            nz = 0
        end
		unsafe_store!(nz_ptr,convert(Cint,nz)::Cint)
    else
    	nz = unsafe_load(nz_ptr)
        if nz > 0
            if colid == 0 && rowid == 0
                values = unsafe_wrap(Array, values_ptr, nz)
                rowidx = unsafe_wrap(Array, row_ptr, nz)
                colptr = unsafe_wrap(Array, col_ptr, ncol+1)
                rowidx[1] = 0

                rowidx[2] = 0
                rowidx[3] = 1

                rowidx[4] = 0
                rowidx[5] = 1
                rowidx[6] = 2

                rowidx[7] = 0
                rowidx[8] = 1
                rowidx[9] = 2
                rowidx[10] = 3

                colptr[1] = 0
                colptr[2] = 1
                colptr[3] = 3
                colptr[4] = 6
                colptr[5] = 10
                # return Int32(1)

                values[1] = obj_factor * (2*x0[4])  # 1,1
                values[2] = obj_factor * (  x0[4])  # 2,1
                values[3] = 0                      # 2,2
                values[4] = obj_factor * (  x0[4])  # 3,1
                values[5] = 0                      # 3,2
                values[6] = 0                      # 3,3
                values[7] = obj_factor * (2*x0[1] + x0[2] + x0[3])  # 4,1
                values[8] = obj_factor * (  x0[1])  # 4,2
                values[9] = obj_factor * (  x0[1])  # 4,3
                values[10] = 0                     # 4,4

                # First constraint
                values[2] += lambda[2] * (x0[3] * x0[4])  # 2,1
                values[4] += lambda[2] * (x0[2] * x0[4])  # 3,1
                values[5] += lambda[2] * (x0[1] * x0[4])  # 4,1
                values[7] += lambda[2] * (x0[2] * x0[3])  # 3,2
                values[8] += lambda[2] * (x0[1] * x0[3])  # 4,2
                values[9] += lambda[2] * (x0[1] * x0[2])  # 4,3

                # Second constraint
                values[1]  += lambda[1] * 2  # 1,1
                values[3]  += lambda[1] * 2  # 2,2
                values[6]  += lambda[1] * 2  # 3,3
                values[10] += lambda[1] * 2  # 4,4
                # prob.model.str_eval_h(rowid,colid,flag,x0,x1,obj_factor,lambda,mode,rowidx,colptr,values)
            end
        end
    end

    return Int32(1)
end

#write solution
function write_solution_wrapper(x_ptr::Ptr{Float64}, y_eq_ptr::Ptr{Float64}, y_ieq_ptr::Ptr{Float64}, cbd::Ptr{CallBackData})
    x = unsafe_wrap(Array,x_ptr,4)
    @test x[1] ≈ 1.0000000000000000 atol=1e-5
    @test x[2] ≈ 4.7429996418092970 atol=1e-5
    @test x[3] ≈ 3.8211499817883077 atol=1e-5
    @test x[4] ≈ 1.3794082897556983 atol=1e-5
    return Int32(1)
end

function PIPScreateProblem(comm::MPI.Comm, model, prof::Bool, numscen::Int, n::Int, x_L::Vector{Float64}, x_U::Vector{Float64},
    m::Int, g_L::Vector{Float64}, g_U::Vector{Float64},
    nele_jac::Int, nele_hess::Int)
	# println(" createProblemStruct  -- julia")
    # TODO: Some callbacks to do not much
	str_init_x0_cb = @cfunction(init_x0_wrapper, Cint, (Ptr{Float64}, Ptr{CallBackData}) )
    str_prob_info_cb = @cfunction(prob_info_wrapper, Cint, (Ptr{Cint}, Ptr{Float64}, Ptr{Float64}, Ptr{Cint}, Ptr{Float64}, Ptr{Float64}, Ptr{CallBackData}) )
    str_eval_f_cb = @cfunction(eval_f_wrapper,Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{CallBackData}) )
    str_eval_g_cb = @cfunction(eval_g_wrapper,Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{CallBackData}) )
    str_eval_grad_f_cb = @cfunction(eval_grad_f_wrapper, Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{CallBackData}) )
    str_eval_jac_g_cb = @cfunction(eval_jac_g_wrapper, Cint, (Ptr{Float64}, Ptr{Float64},
    	Ptr{Cint}, Ptr{Float64}, Ptr{Cint}, Ptr{Cint},
    	Ptr{Cint}, Ptr{Float64}, Ptr{Cint}, Ptr{Cint},
    	Ptr{CallBackData}))
    str_eval_h_cb = @cfunction(eval_h_wrapper, Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Cint}, Ptr{Float64}, Ptr{Cint}, Ptr{Cint}, Ptr{CallBackData}))
    str_write_solution_cb = @cfunction(write_solution_wrapper, Cint, (Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{CallBackData}))

    # println(" callback created ")
    prob = PipsNlpProblemStruct(comm, model, prof, n, x_L, x_U, m, g_L, g_U, nele_jac, nele_hess)
    # @show prob
    ret = ccall(Libdl.dlsym(getPipsLib(),:CreatePipsNlpProblemStruct),Ptr{Nothing},
            (MPI.CComm,
            Cint, Ptr{Nothing}, Ptr{Nothing},
	    Ptr{Nothing}, Ptr{Nothing}, Ptr{Nothing},
	    Ptr{Nothing}, Ptr{Nothing}, Ptr{Nothing},Any
            ),
            MPI.CComm(comm),
            numscen,
            str_init_x0_cb,
            str_prob_info_cb,
            str_eval_f_cb,
            str_eval_g_cb,
            str_eval_grad_f_cb,
            str_eval_jac_g_cb,
            str_eval_h_cb,
            str_write_solution_cb,
            prob
            )
    # println(" ccall CreatePipsNlpProblemStruct done ")
    # @show ret

    if ret == C_NULL
        error("PIPS-NLP: Failed to construct problem.")
    else
        prob.ref = ret
    end
    # @show prob
    # println("end createProblemStruct - julia")
    return prob
end

n = 4
x_L = [1.0, 1.0, 1.0, 1.0]
x_U = [5.0, 5.0, 5.0, 5.0]

m = 2
g_L = [40.0, 25.0]
g_U = [40.0, 2.0e19]

comm = MPI.COMM_WORLD
model = nothing
prof = false
numscen = 0

prob = PIPScreateProblem(comm, model, prof, numscen, n, x_L, x_U, m, g_L, g_U, 8, 10)
@test solveProblemStruct(prob) == :SUCCESSFUL_TERMINATION
# @test PIPSRetCode(solveProblemStruct(prob)) == :SUCCESSFUL_TERMINATION
freeProblemStruct(prob)

# prob.x = [1.0, 5.0, 5.0, 1.0]
# solvestat = solveProblem(prob)


module DataAugmentation

using ColorBlendModes
using CoordinateTransformations
using Distributions: Sampleable, Uniform, Categorical
using IndirectArrays: IndirectArray
using ImageDraw
using Images
using Images: Colorant, permuteddimsview
using ImageTransformations
using ImageTransformations: center, _center, box_extrapolation, warp!
using Interpolations
using MosaicViews: mosaicview
using OffsetArrays: OffsetArray
using LinearAlgebra: I
using Parameters
using Random
using Rotations
using Setfield
using StaticArrays
using Statistics
using Test: @test, @test_nowarn

include("./base.jl")
include("./wrapper.jl")
include("./buffered.jl")
include("./sequence.jl")
include("./items/arrayitem.jl")
include("./projective/base.jl")
include("./items/image.jl")
include("./items/table.jl")
include("./items/keypoints.jl")
include("./items/mask.jl")
include("./projective/compose.jl")
include("./projective/crop.jl")
include("./projective/affine.jl")
include("./projective/warp.jl")
include("./oneof.jl")
include("./preprocessing.jl")
include("./rowtransforms.jl")
include("./colortransforms.jl")
include("testing.jl")
include("./visualization.jl")

export Item,
    Transform,
    ArrayItem,
    MapElem,
    Identity,
    Sequence,
    Project,
    Image,
    TabularItem,
    Keypoints,
    Polygon,
    ToEltype,
    ImageToTensor,
    Normalize,
    NormalizeIntensity,
    MaskMulti,
    MaskBinary,
    BoundingBox,
    ScaleKeepAspect,
    ScaleRatio,
    itemdata,
    Crop,
    CenterCrop,
    RandomCrop,
    ScaleFixed,
    Rotate,
    RandomResizeCrop,
    CenterResizeCrop,
    Buffered,
    BufferedThreadsafe,
    OneHot,
    Zoom,
    OneOf,
    Maybe,
    apply,
    Reflect,
    WarpAffine,
    FlipX,
    FlipY,
    PinOrigin,
    AdjustBrightness,
    AdjustContrast,
    apply!,
    PadDivisible,
    ResizePadDivisible,
    onehot,
    showitems,
    showgrid,
    Bounds,
    getcategorypools

end # module


"""
    LensAssembly{T<:Real}

Structure which contains the elements of the optical system, these can be [`CSGTree`](@ref) or [`Surface`](@ref) objects.

In order to prevent type ambiguities bespoke structs are created for each possible number of elements e.g. `LensAssembly3`.
These are parameterized by the types of the elements to prevent ambiguities.
Basic surface types such as [`Rectangle`](@ref) (which can occur in large numbers) are stored independently in `Vector`s, so type paramters are only needed for CSG objects.

Each struct looks like this:
```julia
struct LensAssemblyN{T,T1,T2,...,TN} <: LensAssembly{T}
    axis::SVector{3,T}
    rectangles::Vector{Rectangle{T}}
    ellipses::Vector{Ellipse{T}}
    hexagons::Vector{Hexagon{T}}
    paraxials::Vector{ParaxialLens{T}}
    E1::T1
    E2::T2
    ...
    EN::TN
end
```

Where `Ti <: Union{Surface{T},CSGTree{T}}`.

To create a LensAssembly object the following functions can be used:
```julia
LensAssembly(elements::Vararg{Union{Surface{T},CSGTree{T},LensAssembly{T}}}; axis = SVector(0.0, 0.0, 1.0)) where {T<:Real}
```
"""
abstract type LensAssembly{T<:Real} end
export LensAssembly

Base.show(io::IO, ::Type{<:LensAssembly}) = print(io, "LensAssembly")
Base.show(io::IO, a::LensAssembly{T}) where {T<:Real} = print(io, "LensAssembly{$T}($(length(a.rectangles) + length(a.ellipses) + length(a.hexagons) + length(typed_elements(a))) elements)")

# in order to prevent type ambiguities lens assembly must be paramaterised by its elements, but we want lens assembly to have an arbitrary number of elements
# as such we have to use macros to create lens assembly structs of a given size, and methods to handle them in a type unambiguous way
# we may often end up with lots of simple surfaces (Rectangle, Ellipse) in the assembly, and we can store these more efficiently in Vectors
macro lensassembly_constructor(N)
    # generate the type definition for a lens assembly of size N
    typesig = [:($(Symbol("T$(i)")) <: Union{Surface{T},CSGTree{T}}) for i in 1:N]
    fields = [:($(Symbol("E$(i)"))::$(Symbol("T$(i)"))) for i in 1:N]
    name = Symbol("LensAssembly$N")
    return esc(quote
        struct $(name){T,$(typesig...)} <: LensAssembly{T}
            axis::SVector{3,T}
            rectangles::Vector{Rectangle{T}}
            ellipses::Vector{Ellipse{T}}
            hexagons::Vector{Hexagon{T}}
            paraxials::Vector{ParaxialLens{T}}
            $(fields...)
        end
    end)
end

# This is the basic idea of what each closestintersection implementation does:
# function closestintersection(ass::LensAssembly{T}, r::AbstractRay{T,N})::Union{Nothing,Intersection{T,N}} where {T<:Real,N}
#     αmin = typemax(T)
#     closest = nothing
#     for element in elements(ass)
#         allintscts = surfaceintersection(element, r)
#         if !(allintscts isa EmptyInterval)
#             intsct = closestintersection(allintscts)
#             if intsct !== nothing
#                 αcurr = α(intsct)
#                 if αcurr < αmin
#                     αmin = αcurr
#                     closest = intsct
#                 end
#             end
#         end
#     end
#     return closest
# end

macro lensassembly_intersection(N)
    # generates the _closestintersection() function for a lens assembly of size N, this should never be called directly, instead call closestintersection() in all cases
    type = :($(Symbol("LensAssembly$(N)")))
    fieldnames = [Symbol("E$(i)") for i in 1:N]
    checks = [quote
        intvl = surfaceintersection(obj.$(fieldnames[i]), r)
        if !(intvl isa EmptyInterval)
            intsct = closestintersection(intvl::Union{Interval{T},DisjointUnion{T}})
            if intsct !== nothing
                αcurr = α(intsct)
                if αcurr < αmin
                    αmin = αcurr
                    closest = intsct
                end
            end
        end
    end for i in 1:N]
    return esc(quote
        function _closestintersection(obj::$(type){T}, r::AbstractRay{T,N})::Union{Nothing,Intersection{T,N}} where {T<:Real,N}
            αmin = typemax(T)
            closest = nothing
            for rect in obj.rectangles
                intvl = surfaceintersection(rect, r)
                if !(intvl isa EmptyInterval)
                    intsct = closestintersection(intvl::Union{Interval{T},DisjointUnion{T}})
                    if intsct !== nothing
                        αcurr = α(intsct)
                        if αcurr < αmin
                            αmin = αcurr
                            closest = intsct
                        end
                    end
                end
            end
            for ell in obj.ellipses
                intvl = surfaceintersection(ell, r)
                if !(intvl isa EmptyInterval)
                    intsct = closestintersection(intvl::Union{Interval{T},DisjointUnion{T}})
                    if intsct !== nothing
                        αcurr = α(intsct)
                        if αcurr < αmin
                            αmin = αcurr
                            closest = intsct
                        end
                    end
                end
            end
            for hex in obj.hexagons
                intvl = surfaceintersection(hex, r)
                if !(intvl isa EmptyInterval)
                    intsct = closestintersection(intvl::Union{Interval{T},DisjointUnion{T}})
                    if intsct !== nothing
                        αcurr = α(intsct)
                        if αcurr < αmin
                            αmin = αcurr
                            closest = intsct
                        end
                    end
                end
            end
            for par in obj.paraxials
                intvl = surfaceintersection(par, r)
                if !(intvl isa EmptyInterval)
                    intsct = closestintersection(intvl::Union{Interval{T},DisjointUnion{T}})
                    if intsct !== nothing
                        αcurr = α(intsct)
                        if αcurr < αmin
                            αmin = αcurr
                            closest = intsct
                        end
                    end
                end
            end
            $(checks...)
            return closest
        end
    end)
end

macro lensassembly_elements(N)
    # generates the _elements() function for a lens assembly of size N, this should never be called directly, instead call elements() in all cases
    type = :($(Symbol("LensAssembly$(N)")))
    fieldnames = [Symbol("E$(i)") for i in 1:N]
    tuple = [:(obj.$(fieldnames[i])) for i in 1:N]
    return esc(quote
        function _typed_elements(obj::$(type){T}) where {T<:Real}
            return tuple($(tuple...))
        end
    end)
end

# pregenerate a small number at compile time
for N in 1:PREGENERATED_LENS_ASSEMBLY_SIZE
    type = :($(Symbol("LensAssembly$(N)")))
    if !isdefined(OpticSim, type)
        @eval @lensassembly_constructor($N)
        @eval @lensassembly_intersection($N)
        @eval @lensassembly_elements($N)
    end
end

function LensAssembly(elements::Vararg{Union{Surface{T},CSGTree{T},LensAssembly{T}}}; axis::SVector{3,T} = SVector{3,T}(0.0, 0.0, 1.0)) where {T<:Real}
    # make the actual object
    actual_elements = []
    rectangles = Vector{Rectangle{T}}(undef, 0)
    ellipses = Vector{Ellipse{T}}(undef, 0)
    hexagons = Vector{Hexagon{T}}(undef, 0)
    paraxials = Vector{ParaxialLens{T}}(undef, 0)
    for e in elements
        if e isa Rectangle{T}
            push!(rectangles, e)
        elseif e isa Ellipse{T}
            push!(ellipses, e)
        elseif e isa Hexagon{T}
            push!(hexagons, e)
        elseif e isa ParaxialLens{T}
            push!(paraxials, e)
        elseif e isa Surface{T} || e isa CSGTree{T}
            push!(actual_elements, e)
        else
            # unpack the lens assembly
            rectangles = vcat(rectangles, e.rectangles)
            ellipses = vcat(ellipses, e.ellipses)
            hexagons = vcat(hexagons, e.hexagons)
            for lae in OpticSim.typed_elements(e)
                push!(actual_elements, lae)
            end
        end
    end
    N = length(actual_elements)
    # make the methods for this N
    type = :($(Symbol("LensAssembly$(N)")))
    if !isdefined(OpticSim, type)
        @eval @lensassembly_constructor($N)
        @eval @lensassembly_intersection($N)
        @eval @lensassembly_elements($N)
    end
    eltypes = typeof.(actual_elements)
    naxis = normalize(axis)
    return eval(:($(type){$T,$(eltypes...)}($naxis, $rectangles, $ellipses, $hexagons, $paraxials, $(actual_elements...))))
end

function typed_elements(ass::LensAssembly{T}) where {T<:Real}
    # under certain circumstances we run into the world age problem (see https://discourse.julialang.org/t/how-to-bypass-the-world-age-problem/7012)
    # to avoid this we need to use invokelatest(), but this is type ambiguous and ruins performance, and in most cases isn't necessary
    # so we can use try/catch to only call it when we need (and warn appropriately), retaining performance in almost all cases
    # NEVER call _typed_elements directly, always use this method
    try
        _typed_elements(ass)
    catch MethodError
        @warn "First time call to LensAssembly of previously unused size, performance will be worse than normal." maxlog = 1
        Base.invokelatest(_typed_elements, ass)
    end
end

elements(ass::LensAssembly{T}) where {T<:Real} = (ass.rectangles..., ass.ellipses..., ass.hexagons..., ass.paraxials..., typed_elements(ass)...)
export elements

function closestintersection(ass::LensAssembly{T}, r::AbstractRay{T,N})::Union{Nothing,Intersection{T,N}} where {T<:Real,N}
    # under certain circumstances we run into the world age problem (see https://discourse.julialang.org/t/how-to-bypass-the-world-age-problem/7012)
    # to avoid this we need to use invokelatest(), but this is type ambiguous and ruins performance, and in most cases isn't necessary
    # so we can use try/catch to only call it when we need (and warn appropriately), retaining performance in almost all cases
    # NEVER call _closestintersection directly, always use this method
    try
        _closestintersection(ass, r)
    catch MethodError
        @warn "First time call to LensAssembly of previously unused size, performance will be worse than normal." maxlog = 1
        Base.invokelatest(_closestintersection, ass, r)
    end
end

function BoundingBox(la::LensAssembly{T}) where {T<:Real}
    es = elements(la)
    bbox = BoundingBox(es[1])
    for b in es[2:end]
        if b isa ParametricSurface{T} || b isa CSGTree{T}
            bbox = union(bbox, b)
        end
    end
    return bbox
end

axis(a::LensAssembly{T}) where {T<:Real} = a.axis
export axis

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

"""
    LensTrace{T<:Real,N}

Contains an intersection point and the ray segment leading to it from within an optical trace.
The ray carries the path length, power, wavelength, number of intersections and source number, all of which are accessible directly on this class too.

Has the following accessor methods:
```julia
ray(a::LensTrace{T,N}) -> OpticalRay{T,N}
intersection(a::LensTrace{T,N}) -> Intersection{T,N}
power(a::LensTrace{T,N}) -> T
wavelength(a::LensTrace{T,N}) -> T
pathlength(a::LensTrace{T,N}) -> T
point(a::LensTrace{T,N}) -> SVector{N,T}
uv(a::LensTrace{T,N}) -> SVector{2,T}
sourcenum(a::LensTrace{T,N}) -> Int
nhits(a::LensTrace{T,N}) -> Int
```
"""
struct LensTrace{T<:Real,N}
    ray::OpticalRay{T,N}
    intersection::Intersection{T,N}
end
export LensTrace

ray(a::LensTrace{T,N}) where {T<:Real,N} = a.ray
intersection(a::LensTrace{T,N}) where {T<:Real,N} = a.intersection
power(a::LensTrace{T,N}) where {T<:Real,N} = power(ray(a))
wavelength(a::LensTrace{T,N}) where {T<:Real,N} = wavelength(ray(a))
pathlength(a::LensTrace{T,N}) where {T<:Real,N} = pathlength(ray(a))
point(a::LensTrace{T,N}) where {T<:Real,N} = point(intersection(a))
uv(a::LensTrace{T,N}) where {T<:Real,N} = uv(intersection(a))
sourcenum(a::LensTrace{T,N}) where {T<:Real,N} = sourcenum(ray(a))
nhits(a::LensTrace{T,N}) where {T<:Real,N} = nhits(ray(a))
export intersection

function Base.print(io::IO, a::LensTrace{T,N}) where {T,N}
    println(io, "Ray\n$(ray(a))")
    println(io, "Lensintersection\n$(intersection(a))")
end

struct NoPower end
const nopower = NoPower()

# isnopower(::NoPower) = true
# isnopower(::Any) = false

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

"""
    trace(assembly::LensAssembly{T}, r::OpticalRay{T}, temperature::T = 20.0, pressure::T = 1.0; trackrays = nothing, test = false)

Returns the ray as it exits the assembly in the form of a [`LensTrace`](@ref) object if it hits any element in the assembly, otherwise `nothing`.
Recursive rays are offset by a small amount (`RAY_OFFSET`) to prevent it from immediately reintersecting the same lens element.

`trackrays` can be passed an empty vector to accumulate the `LensTrace` objects at each intersection of `ray` with a surface in the assembly.
"""
function trace(assembly::LensAssembly{T}, r::OpticalRay{T,N}, temperature::T = T(OpticSim.GlassCat.TEMP_REF), pressure::T = T(OpticSim.GlassCat.PRESSURE_REF); trackrays::Union{Nothing,Vector{LensTrace{T,N}}} = nothing, test::Bool = false, recursion::Int = 0)::Union{Nothing,NoPower,LensTrace{T,N}} where {T<:Real,N}
    if power(r) < POWER_THRESHOLD || recursion > TRACE_RECURSION_LIMIT
        return nopower
    end
    intsct = closestintersection(assembly, r)
    if intsct === nothing
        return nothing
    else
        surfintsct = point(intsct)
        nml = normal(intsct)
        opticalinterface = interface(intsct)
        λ = wavelength(r)

        if VERSION < v"1.6.0-DEV"
            # TODO REMOVE
            temp = @unionsplit OpticalInterface T opticalinterface processintersection(opticalinterface, surfintsct, nml, r, temperature, pressure, test, recursion == 0)
        else
            temp = processintersection(opticalinterface, surfintsct, nml, r, temperature, pressure, test, recursion == 0)
        end

        if temp === nothing
            return nopower
        else
            raydirection, raypower, raypathlength = temp
            if trackrays !== nothing
                # we want the power that is hitting the surface, i.e. power on incident ray, but the path length at this intersection
                # i.e. with the path length for this intersection added and power modulated by absorption only
                push!(trackrays, LensTrace(OpticalRay(ray(r), raypower, λ, opl = raypathlength, nhits = nhits(r), sourcenum = sourcenum(r), sourcepower = sourcepower(r)), intsct))
            end
            offsetray = OpticalRay(surfintsct + RAY_OFFSET * raydirection, raydirection, raypower, λ, opl = raypathlength, nhits = nhits(r) + 1, sourcenum = sourcenum(r), sourcepower = sourcepower(r))
            res = trace(assembly, offsetray, temperature, pressure, trackrays = trackrays, test = test, recursion = recursion + 1)
            if res === nothing
                return LensTrace(OpticalRay(surfintsct, raydirection, raypower, λ, opl = raypathlength, nhits = nhits(r), sourcenum = sourcenum(r), sourcepower = sourcepower(r)), intsct)
            else
                return res
            end
        end
    end
end

# """approximate focal length of simple double convex lens. For plano convex set one of the radii to Inf"""
focallength(index::Float64, rf::Float64, rb::Float64) = 1 / ((index - 1) * (1 / rf - 1 / rb))


#######
# Demo setup

cool(x) = x + 1
cool(x, y) = x + y + 1

# a rule we define so we can test rules
dummy_identity(x) = x
@scalar_rule(dummy_identity(x), true)

nice(x) = 1
@scalar_rule(nice(x), ZeroTangent())

sum_two(x, y) = x + y
@scalar_rule(sum_two(x, y), (true, true))

complex_times(x) = (1 + 2im) * x
@scalar_rule(complex_times(x), 1 + 2im)

# Tests that aim to ensure that the API for frules doesn't regress and make these things
# hard to implement.

varargs_function(x...) = sum(x)
function ChainRulesCore.frule(dargs, ::typeof(varargs_function), x...)
    Δx = Base.tail(dargs)
    return sum(x), sum(Δx)
end

mixed_vararg(x, y, z...) = x + y + sum(z)
function ChainRulesCore.frule(
    dargs::Tuple{Any, Any, Any, Vararg},
    ::typeof(mixed_vararg), x, y, z...,
)
    Δx = dargs[2]
    Δy = dargs[3]
    Δz = dargs[4:end]
    return mixed_vararg(x, y, z...), Δx + Δy + sum(Δz)
end

type_constraints(x::Int, y::Float64) = x + y
function ChainRulesCore.frule(
    (_, Δx, Δy)::Tuple{Any, Int, Float64},
    ::typeof(type_constraints), x::Int, y::Float64,
)
    return type_constraints(x, y), Δx + Δy
end

mixed_vararg_type_constaint(x::Float64, y::Real, z::Vararg{Float64}) = x + y + sum(z)
function ChainRulesCore.frule(
    dargs::Tuple{Any, Float64, Real, Vararg{Float64}},
    ::typeof(mixed_vararg_type_constaint), x::Float64, y::Real, z::Vararg{Float64},
)
    Δx = dargs[2]
    Δy = dargs[3]
    Δz = dargs[4:end]
    return x + y + sum(z), Δx + Δy + sum(Δz)
end

ChainRulesCore.frule(dargs, ::typeof(Core._apply), f, x...) = frule(dargs[2:end], f, x...)

#######

_second(t) = Base.tuple_type_head(Base.tuple_type_tail(t))

@testset "frule and rrule" begin
    dself = ZeroTangent()
    @test frule((dself, 1), cool, 1) === nothing
    @test frule((dself, 1), cool, 1; iscool=true) === nothing
    @test rrule(cool, 1) === nothing
    @test rrule(cool, 1; iscool=true) === nothing

    # add some methods:
    ChainRulesCore.@scalar_rule(Main.cool(x), one(x))
    @test hasmethod(rrule, Tuple{typeof(cool),Number})
    ChainRulesCore.@scalar_rule(Main.cool(x::String), "wow such dfdx")
    @test hasmethod(rrule, Tuple{typeof(cool),String})
    # Ensure those are the *only* methods that have been defined
    cool_methods = Set(m.sig for m in methods(rrule) if _second(m.sig) == typeof(cool))
    only_methods = Set([Tuple{typeof(rrule),typeof(cool),Number},
                        Tuple{typeof(rrule),typeof(cool),String}])
    @test cool_methods == only_methods

    frx, cool_pushforward = frule((dself, 1), cool, 1)
    @test frx === 2
    @test cool_pushforward === 1
    rrx, cool_pullback = rrule(cool, 1)
    self, rr1 = cool_pullback(1)
    @test self === NoTangent()
    @test rrx === 2
    @test rr1 == 1.0   # projection may make this ===

    frx, nice_pushforward = frule((dself, 1), nice, 1)
    @test nice_pushforward === ZeroTangent()
    rrx, nice_pullback = rrule(nice, 1)
    @test (NoTangent(), ZeroTangent()) === nice_pullback(1)

    # Test that these run. Do not care about numerical correctness.
    @test frule((nothing, 1.0, 1.0, 1.0), varargs_function, 0.5, 0.5, 0.5) == (1.5, 3.0)

    @test frule((nothing, 1.0, 2.0, 3.0, 4.0), mixed_vararg, 1.0, 2.0, 3.0, 4.0) == (10.0, 10.0)

    @test frule((nothing, 3, 2.0), type_constraints, 5, 4.0) == (9.0, 5.0)
    @test frule((nothing, 3.0, 2.0im), type_constraints, 5, 4.0) == nothing

    @test(frule(
        (nothing, 3.0, 2.0, 1.0, 0.0),
        mixed_vararg_type_constaint, 3.0, 2.0, 1.0, 0.0,
    ) == (6.0, 6.0))

    # violates type constraints, thus an frule should not be found.
    @test frule(
        (nothing, 3, 2.0, 1.0, 5.0),
        mixed_vararg_type_constaint, 3, 2.0, 1.0, 0,
    ) == nothing

    @test frule((nothing, nothing, 5.0), Core._apply, dummy_identity, 4.0) == (4.0, 5.0)

    @testset "broadcasting true" begin
        sx = @SVector [1, 2]
        sy = @SVector [3, 4]

        # Test that @scalar_rule and `true` play nice together, w.r.t broadcasting
        @inferred frule((ZeroTangent(), sx, sy), sum_two, 1, 2)
    end

    @testset "complex numbers" begin
        x, ẋ, Ω̄ = randn(ComplexF64, 3)
        Ω = complex_times(x)

        # forwards
        Ω_fwd, Ω̇ = frule((nothing, ẋ), complex_times, x)
        @test Ω_fwd == Ω
        @test Ω̇ ≈ (1 + 2im) * ẋ

        # reverse
        Ω_rev, back = rrule(complex_times, x)
        @test Ω_rev == Ω
        ∂self, ∂x = back(Ω̄)
        @test ∂self == NoTangent()
        @test ∂x ≈ (1 - 2im) * Ω̄

        # real argument, complex output
        xr = rand()
        Ωr = complex_times(xr)
        Ωr_rev, backr = rrule(complex_times, xr)
        ∂selfr, ∂xr = backr(Ω̄)
        @test_skip ∂xr isa Float64  # to be made true with projection
        @test_skip ∂xr ≈ real(∂x)
    end

    @testset "@opt_out" begin
        first_oa(x, y) = x
        @scalar_rule(first_oa(x, y), (1, 0))
        @opt_out ChainRulesCore.rrule(::typeof(first_oa), x::T, y::T) where T<:Float32
        @opt_out(
            ChainRulesCore.frule(::Any, ::typeof(first_oa), x::T, y::T) where T<:Float32
        )

        @testset "rrule" begin
            @test rrule(first_oa, 3.0, 4.0)[2](1) == (NoTangent(), 1, 0)
            @test rrule(first_oa, 3f0, 4f0) === nothing

            @test !isempty(Iterators.filter(methods(ChainRulesCore.no_rrule)) do m
                m.sig <:Tuple{Any, typeof(first_oa), T, T} where T<:Float32
            end)
        end

        @testset "frule" begin
            @test frule((NoTangent(), 1,0), first_oa, 3.0, 4.0) == (3.0, 1)
            @test frule((NoTangent(), 1,0), first_oa, 3f0, 4f0) === nothing

            @test !isempty(Iterators.filter(methods(ChainRulesCore.no_frule)) do m
                m.sig <:Tuple{Any, Any, typeof(first_oa), T, T} where T<:Float32
            end)
        end
    end
end


# surface(::Matrix, ::Matrix, ::Matrix)
function _default(main::Tuple{MatTypes{T}, MatTypes{T}, MatTypes{T}}, s::Style{:surface}, data::Dict) where T <: AbstractFloat
    @gen_defaults! data begin
        position_x = main[1] => (Texture, "x position, must be a `Matrix{Float}`")
        position_y = main[2] => (Texture, "y position, must be a `Matrix{Float}`")
        position_z = main[3] => (Texture, "z position, must be a `Matrix{Float}`")
        scale = Vec3f0(0) => "scale must be 0, for a surfacemesh"
    end
    surface(position_z, s, data)
end

# surface(Vector or Range, Vector or Range, ::Matrix)
function _default(main::Tuple{VectorTypes{T}, VectorTypes{T}, MatTypes{T}}, s::Style{:surface}, data::Dict) where T <: AbstractFloat
    @gen_defaults! data begin
        position_x = main[1] => (Texture, "x position, must be a `Vector{Float}`")
        position_y = main[2] => (Texture, "y position, must be a `Vector{Float}`")
        position_z = main[3] => (Texture, "z position, must be a `Matrix{Float}`")
        scale = Vec3f0(0) => "scale must be 0, for a surfacemesh"
    end
    surface(position_z, s, data)
end

# surface(::Matrix)
function _default(main::MatTypes{T}, s::Style{:surface}, data::Dict) where T <: AbstractFloat
    @gen_defaults! data begin
        ranges = ((-1f0, 1f0), (-1f0,1f0)) => "x, and y position given as `(start, endvalue)` or any `Range`"
    end
    delete!(data, :ranges) # no need to have them in the OpenGL data
    _default((Grid(to_value(main), to_value(ranges)), main), s, data)
end

# surface(::Matrix)
function _default(main::Tuple{G, MatTypes{T}}, s::Style{:surface}, data::Dict) where {G <: Grid{2}, T <: AbstractFloat}
    xrange = main[1].dims[1]; yrange = main[1].dims[2]
    xscale = (maximum(xrange) - minimum(xrange)) / (length(xrange)-1)
    yscale = (maximum(yrange) - minimum(yrange)) / (length(yrange)-1)
    @gen_defaults! data begin
        position    = main[1] =>" Position given as a `Grid{2}`.
        Can be constructed e.g. `Grid(LinRange(0,2,N1), LinRange(0,3, N2))`"
        position_z  = main[2] => (Texture, "height offset for the surface, must be `Matrix{Float}`")
        scale       = Vec3f0(xscale, yscale, 1) => "scale of the grid planes forming the surface. Can be made smaller, to let the grid show"
    end
    surface(position_z, s, data)
end

_extrema(x::FRect3D) = Vec2f0(minimum(x)[3], maximum(x)[3])
nothing_or_vec(x) = x
nothing_or_vec(x::Array) = vec(x)

function normal_calc(x::Bool, invert_normals::Bool = false)
    i = invert_normals ? "-" : ""
    if x
        "$(i)getnormal(position, position_x, position_y, position_z, o_uv);"
        # "getnormal_fast(position_z, ind2sub(dims, index1D));"
    else
        "vec3(0, 0, $(i)1);"
    end
end

function light_calc(x::Bool)
    if x
        """
        vec3 L      = normalize(o_lightdir);
        vec3 N      = normalize(o_normal);
        vec3 light1 = blinnphong(N, o_camdir, L, color.rgb);
        color       = vec4(light1, color.a); // + light2 * 0.4
        """
    else
        ""
    end
end

function native_triangle_mesh(mesh)
    return gl_convert(triangle_mesh(mesh))
end

function surface(main, s::Style{:surface}, data::Dict)
    @gen_defaults! data begin
        primitive = Rect2D(0f0,0f0,1f0,1f0) => native_triangle_mesh
        scale = nothing
        position = nothing
        position_x = nothing => Texture
        position_y = nothing => Texture
        position_z = nothing => Texture
        wireframe = false
        glow_color = RGBA{Float32}(0,0,0,0) => GLBuffer
        stroke_color = RGBA{Float32}(0,0,0,1) => GLBuffer
        stroke_width = wireframe ? 0.03f0 : 0f0
        glow_width = 0f0
        uv_offset_width = Vec4f0(0) => GLBuffer
        shape = RECTANGLE
        image = nothing => Texture
        distancefield = nothing => Texture
        shading = true
        normal = shading
        invert_normals = false
    end
    @gen_defaults! data begin
        color = nothing => Texture
        color_map = nothing => Texture
        color_norm = nothing
        fetch_pixel = false
        matcap = nothing => Texture
        uv_scale = Vec2f0(1)
        instances = const_lift(x->(size(x,1)-1) * (size(x,2)-1), main) => "number of planes used to render the surface"
        shader = GLVisualizeShader(
            "fragment_output.frag", "util.vert", "surface.vert",
            to_value(wireframe) ? "distance_shape.frag" : "standard.frag",
            view = Dict(
                "position_calc" => position_calc(position, position_x, position_y, position_z, Texture),
                "normal_calc" => normal_calc(normal, to_value(invert_normals)),
                "light_calc" => light_calc(shading),
            )
        )
    end
end

function position_calc(x...)
    _position_calc(Iterators.filter(x->!isa(x, Nothing), x)...)
end

function _position_calc(
        position_x::MatTypes{T}, position_y::MatTypes{T}, position_z::MatTypes{T}, target::Type{Texture}
    ) where T<:AbstractFloat
    """
    int index1D = index + offseti.x + offseti.y * dims.x + (index/(dims.x-1));
    ivec2 index2D = ind2sub(dims, index1D);
    vec2 index01 = vec2(index2D) / (vec2(dims)-1.0);
    pos = vec3(
        texelFetch(position_x, index2D, 0).x,
        texelFetch(position_y, index2D, 0).x,
        texelFetch(position_z, index2D, 0).x
    );
    """
end

function _position_calc(
        position_x::VectorTypes{T}, position_y::VectorTypes{T}, position_z::MatTypes{T},
        target::Type{Texture}
    ) where T<:AbstractFloat
    """
    int index1D = index + offseti.x + offseti.y * dims.x + (index/(dims.x-1));
    ivec2 index2D = ind2sub(dims, index1D);
    vec2 index01 = vec2(index2D) / (vec2(dims)-1.0);
    pos = vec3(
        texelFetch(position_x, index2D.x, 0).x,
        texelFetch(position_y, index2D.y, 0).x,
        texelFetch(position_z, index2D, 0).x
    );
    """
end

function _position_calc(
        position_xyz::VectorTypes{T}, target::Type{TextureBuffer}
    ) where T <: StaticVector
    "pos = texelFetch(position, index).xyz;"
end

function _position_calc(
        position_xyz::VectorTypes{T}, target::Type{GLBuffer}
    ) where T <: StaticVector
    len = length(T)
    filler = join(ntuple(x->0, 3-len), ", ")
    needs_comma = len != 3 ? ", " : ""
    "pos = vec3(position $needs_comma $filler);"
end

function _position_calc(
        grid::Grid{2}, position_z::MatTypes{T}, target::Type{Texture}
    ) where T<:AbstractFloat
    """
    int index1D = index + offseti.x + offseti.y * dims.x + (index/(dims.x-1));
    ivec2 index2D = ind2sub(dims, index1D);
    vec2 index01 = vec2(index2D) / (vec2(dims)-1.0);
    float height = texture(position_z, index01).x;
    pos = vec3(grid_pos(position, index01), height);
    """
end


module MadNLPGPU

import LinearAlgebra
# CUDA
import CUDA: CUBLAS, CUSOLVER, CuVector, CuMatrix, CuArray, toolkit_version, R_64F, has_cuda, @allowscalar, runtime_version
# Kernels
import KernelAbstractions: @kernel, @index, wait, Event
import CUDAKernels: CUDADevice

import MadNLP

import MadNLP:
    @kwdef, Logger, @debug, @warn, @error,
    AbstractOptions, AbstractLinearSolver, AbstractNLPModel, set_options!, MadNLPLapackCPU,
    SymbolicException,FactorizationException,SolveException,InertiaException,
    introduce, factorize!, solve!, improve!, is_inertia, inertia, tril_to_full!

include("kernels.jl")

if has_cuda()
    include("lapackgpu.jl")
    export MadNLPLapackGPU
end

end # module


###############################################################################
# Extract designs from Sort-seq run
# -----------------------------------------------------------------------------
# Author:  Thomas E. Gorochowski <thomas.gorochowski@bristol.ac.uk>
# Licence: MIT
###############################################################################

using CSV
using DataFrames
using Query
using Glob
using Printf
using BioSequences
using BioAlignments
using FASTX

###############################################################################
# PARAMETERS – Should be optimised for particular use case
###############################################################################

# Minimum alignment score for start and end of memory register
# Previously 300, 100
global MIN_ALN_SCORE_REGION = 20
global MIN_ALN_SCORE_SPACER = 0

# Alignemnt parameters (these have been tweaked for 25 bp barcodes)
#global ALN_MATCH = 5
#global ALN_MISMATCH = -4
#global ALN_GAP_OPEN = -4
#global ALN_GAP_EXTEND = -2

# Parameters from Matt's paper (need to adjust the thresholds)
global ALN_MATCH = 2
global ALN_MISMATCH = -3
global ALN_GAP_OPEN = -5
global ALN_GAP_EXTEND = -2

global BASE_FRACTION_CALL_INSERT = 0.5

###############################################################################
# FUNCTIONS TO DO THE HEAVY LIFTING
###############################################################################

function load_references(filename)
    ref_seqs = Dict{String,LongDNASeq}()
    reader = FASTA.Reader(open(filename, "r"))
    record = FASTA.Record()
    read_count = 0
    while !eof(reader)
        read!(reader, record)
        ref_seqs[FASTA.identifier(record)] = FASTA.sequence(record)
        read_count = read_count + 1
    end
    close(reader)
    @info string("Processed ", filename, " containing ", read_count, " sequences")
    return ref_seqs
end

refs = load_references("data/references.fasta")

function ref_alignment_range(aln)
    start_pos = -1
    end_pos = -1
    aln_anchors = aln.a.aln.anchors
    first_match = true
    for a in aln_anchors
        if a.op == OP_SEQ_MATCH
            if first_match == true
                start_pos = a.refpos
                first_match = false
                end_pos = a.refpos
            else
                end_pos = a.refpos
            end
        end
    end
    return (start_pos, end_pos)
end

function check_length_fix_orientation(test_seq, start_seq, end_seq)
    # Alignment type used and scoring model
    problem = SemiGlobalAlignment()
    scoremodel = AffineGapScoreModel(
                  match=ALN_MATCH,
                  mismatch=ALN_MISMATCH,
                  gap_open=ALN_GAP_OPEN,
                  gap_extend=ALN_GAP_EXTEND
                )
    test_seq_rc = reverse_complement(test_seq)
    alignment_start = pairalign(problem, start_seq, test_seq, scoremodel)
    alignment_start_rc = pairalign(problem, start_seq, test_seq_rc, scoremodel)
    alignment_end = pairalign(problem, end_seq, test_seq, scoremodel)
    alignment_end_rc = pairalign(problem, end_seq, test_seq_rc, scoremodel)
    score_start = score(alignment_start)
    score_start_rc = score(alignment_start_rc)
    score_end = score(alignment_end)
    score_end_rc = score(alignment_end_rc)
    # Test to see what start and end are most likely and return sequence
    if score_start > score_start_rc && score_end > score_end_rc &&
        score_start > MIN_ALN_SCORE_REGION && score_end > MIN_ALN_SCORE_REGION
        (s_idx0, s_idx1) = ref_alignment_range(alignment(alignment_start))
        (e_idx0, e_idx1) = ref_alignment_range(alignment(alignment_end))
        return test_seq[s_idx0:e_idx1]
    elseif score_start_rc > score_start && score_end_rc > score_end &&
        score_start_rc > MIN_ALN_SCORE_REGION && score_end_rc > MIN_ALN_SCORE_REGION
        (s_idx0, s_idx1) = ref_alignment_range(alignment(alignment_start_rc))
        (e_idx0, e_idx1) = ref_alignment_range(alignment(alignment_end_rc))
        return test_seq_rc[s_idx0:e_idx1]
    else
        return nothing
    end
end

function find_spacer_list(seq, ref_seqs)
    spacers = []
    # Alignment type used and scoring model
    problem = SemiGlobalAlignment()
    scoremodel = AffineGapScoreModel(
                  match=ALN_MATCH,
                  mismatch=ALN_MISMATCH,
                  gap_open=ALN_GAP_OPEN,
                  gap_extend=ALN_GAP_EXTEND
                )
    seq_len = length(seq)
    seq_rc = reverse_complement(seq)
    for (ref_key, ref_seq) in ref_seqs
        # Check reference is for a spacer
        if ref_key != "start" && ref_key != "end"
            spacer_aln = pairalign(problem, ref_seq, seq, scoremodel)
            spacer_aln_rc = pairalign(problem, ref_seq, seq_rc, scoremodel)
            spacer_score = score(spacer_aln)
            spacer_score_rc = score(spacer_aln_rc)
            # Check alignments meet minimum score
            if spacer_score > spacer_score_rc && spacer_score > MIN_ALN_SCORE_SPACER
                (s_idx0, s_idx1) = ref_alignment_range(alignment(spacer_aln))
                push!(spacers, [ref_key, s_idx0])
            elseif spacer_score_rc > spacer_score && spacer_score_rc > MIN_ALN_SCORE_SPACER
                (s_idx0, s_idx1) = ref_alignment_range(alignment(spacer_aln_rc))
                push!(spacers, [string(ref_key, "_r"), seq_len-s_idx1])
            end
        end
    end
    # Sort the found spacers on their position
    sort!(spacers, lt=(x,y)->isless(x[2], y[2]))
    # Generate string for configuration

    #=
    # Only return results where all spacers are present
    if length(spacers) == length(ref_seqs) - 2
        spacer_string = ""
        for el in spacers
            if spacer_string == ""
                spacer_string = el[1]
            else
                spacer_string = string(spacer_string, " ", el[1])
            end
        end
        if spacers == []
            return "None"
        else
            return spacer_string
        end
    else
        return "None"
    end
    =#

    spacer_string = ""
    for el in spacers
        if spacer_string == ""
            spacer_string = el[1]
        else
            spacer_string = string(spacer_string, " ", el[1])
        end
    end
    if spacers == []
        return "None"
    else
        return spacer_string
    end
end

function find_spacers(fastq_filename, ref_seqs)
    @info string("Started processing ", fastq_filename, "...\n")
    results = Dict{String, Int}()
    try
        reader = FASTQ.Reader(open(fastq_filename, "r"))
        record = FASTQ.Record()
        read_count = 0
        processed_read_count = 0
        while !eof(reader)
            # Comment out if processing full data
            if read_count >= 10000
                break
            end
            if mod(read_count, 1000) == 0
                print(".")
            end
            read!(reader, record)
            fixed_seq = check_length_fix_orientation(FASTQ.sequence(record), ref_seqs["start"], ref_seqs["end"])
            if fixed_seq != nothing
                # Search for order and orientation of spacers and add to list
                cur_key = find_spacer_list(fixed_seq, ref_seqs)
                if haskey(results, cur_key)
                    results[cur_key] += 1
                else
                    results[cur_key] = 1
                end
                processed_read_count = processed_read_count + 1
            end
            read_count = read_count + 1
        end
        close(reader)
        @info string("Processed ", fastq_filename, " containing ", read_count, " reads (",
                     processed_read_count, " full length memory registers)\n")
    catch y
        @info string("Error in ", fastq_filename, "\n")
        println(y)
    end
    return results
end

function write_results_to_file(results, filename)
    open(filename, "w") do file
        write(file, "spacer_config,count\n")
        for (key, val) in results
            write(file, string(key, ",", val, "\n"))
        end
    end
end

###############################################################################
# RUN THE ANALYSIS
###############################################################################

# Load the metadata for the run
metadata = DataFrame(CSV.File("metadata.csv", header=1, delim=","))

# Create mapping from barcode to sample name using the metadata for run 1
bc_map1 = @from i in metadata begin
    @where i.run == 1
    @select i.barcode=>i.name
    @collect Dict
end

# Generate the input and output file names
read_filenames = []
output_filenames = []

for (bc, name) in bc_map1
    read_filename = "fastq/barcode$(@sprintf("%02d", bc)).fastq"
    output_filename = "output/$(name).csv"
    push!(read_filenames, read_filename)
    push!(output_filenames, output_filename)
end

# Load the spacer sequences to search for
spacer_refs_filename = "spacers.fasta"
ref_seqs = load_spacers(spacer_refs_filename)

# Extract the spacer configurations from each sample
Threads.@threads for idx in 1:length(read_filenames)
    read_filename = read_filenames[idx]
    output_filename = output_filenames[idx]
    results = find_spacers(read_filename, ref_seqs)
    write_results_to_file(results, output_filename)
end
=#


facts("Timing functions") do

    context("fpswhen") do
        b = Signal(false)
        t = fpswhen(b, 10)
        acc = foldp((x, y) -> x+1, 0, t)
        sleep(0.14)

        @fact queue_size() --> 0
        push!(b, true)

        dt = @elapsed Reactive.run(11) # the first one starts the timer
        push!(b, false)
        Reactive.run(1)

        sleep(0.11) # no more updates
        @fact queue_size() --> 0

        @fact dt --> roughly(1, atol=0.2) # mac OSX needs a lot of tolerence here
        @fact value(acc) --> 10

    end

    context("fps") do
        t = fps(10)
        acc = foldp(push!, Float64[], t)
        Reactive.run(11) # Starts with 0
        log = copy(value(acc))

        @fact log[1] --> roughly(0.1, atol=0.1) # First one's always crappy
        log = log[2:end]

        @fact sum(log) --> roughly(1.0, atol=0.05)
        @fact [1/10 for i=1:10] --> roughly(log, atol=0.03)

        @fact queue_size() --> 0
        sleep(0.11)
        @fact queue_size() --> 1
        sleep(0.22)
        @fact queue_size() --> 1

        close(acc)
        close(t)

        Reactive.run_till_now()
    end

    context("every") do
        t = every(0.1)
        acc = foldp(push!, Float64[], t)
        Reactive.run(11)
        end_t = time()
        log = copy(value(acc))

        @fact log[end-1] --> roughly(end_t, atol=0.001)

        close(acc)
        close(t)
        Reactive.run_till_now()

        @fact [0.1 for i=1:10] --> roughly(diff(log), atol=0.03)

        sleep(0.2)
        # make sure close actually also closed the timer
        @fact queue_size() --> 0
    end

    context("throttle") do
        x = Signal(0)
        y = throttle(0.1, x)
        y′ = throttle(0.2, x, push!, Int[], x->Int[]) # collect intermediate updates
        z = foldp((acc, x) -> acc+1, 0, y)
        z′ = foldp((acc, x) -> acc+1, 0, y′)

        push!(x, 1)
        step()

        push!(x, 2)
        step()

        push!(x, 3)
        t0=time()
        step()

        @fact value(y) --> 0
        @fact value(z) --> 0
        @fact queue_size() --> 0

        sleep(0.1)

        @fact queue_size() --> 1
        step()
        @fact value(y) --> 3
        @fact value(z) --> 1
        sleep(0.1)

        @fact queue_size() --> 1
        step()
        @fact value(z′) --> 1
        @fact value(y′) --> Int[1,2,3]

        push!(x, 3)
        step()

        push!(x, 2)
        step()

        push!(x, 1)
        step()
        sleep(0.2)

        @fact queue_size() --> 2
        step()
        step()
        @fact value(y) --> 1
        @fact value(z′) --> 2
        @fact value(y′) --> Int[3,2,1]

        # type safety
        s1 = Signal(3)
        s2 = Signal(rand(2,2))
        m = merge(s1, s2)
        t = throttle(1/60, m; typ=Any)
        r = rand(3,3)
        push!(s2, r)
        Reactive.run(1)
        sleep(0.05)
        Reactive.run(1)
        @fact value(t) --> r
    end
end