module Plasm

	#export centroid, cuboidGrid, mkpol, view, hpc_exploded, lar2hpc

	using LinearAlgebraicRepresentation
	using PyCall

	@pyimport pyplasm as p

	import Base.view

	"""
		Points = Matrix

	Alias declation of LAR-specific data structure.
	Dense `Matrix` ``M times N`` to store the position of *vertices* (0-cells)
	of a *cellular complex*. The number of rows ``M`` is the dimension
	of the embedding space. The number of columns ``N`` is the number of vertices.
	"""
	const Points = Matrix

	"""
		Cells = Array{Array{Int,1},1}

	Alias declation of LAR-specific data structure.
	Dense `Array` to store the indices of vertices of `P-cells`
	of a cellular complex.
	The linear space of `P-chains` is generated by `Cells` as a basis.
	Simplicial `P-chains` have `P+1` vertex indices for `cell` element in `Cells` array.
	Cuboidal `P-chains` have ``2^P`` vertex indices for `cell` element in `Cells` array.
	Other types of chain spaces may have different numbers of vertex indices for `cell`
	element in `Cells` array.
	"""
	const Cells = Array{Array{Int,1},1}

	"""
		Chain = SparseVector{Int8,Int}

	Alias declation of LAR-specific data structure.
	Binary `SparseVector` to store the coordinates of a `chain` of `N-cells`. It is
	`nnz=1` with `value=1` for the coordinates of an *elementary N-chain*, constituted by
	a single *N-chain*.
	"""
	const Chain = SparseVector{Int8,Int}

	"""
		ChainOp = SparseMatrixCSC{Int8,Int}

	Alias declation of LAR-specific data structure.
	`SparseMatrix` in *Compressed Sparse Column* format, contains the coordinate
	representation of an operator between linear spaces of `P-chains`.
	Operators ``P-Boundary : P-Chain -> (P-1)-Chain``
	and ``P-Coboundary : P-Chain -> (P+1)-Chain`` are typically stored as
	`ChainOp` with elements in `` {-1,0,1 }`` or in `` {0,1 }``, for
	*signed* and *unsigned* operators, respectively.
	"""
	const ChainOp = SparseMatrixCSC{Int8,Int}

	"""
		ChainComplex = Array{ChainOp,1}

	Alias declation of LAR-specific data structure. It is a
	1-dimensional `Array` of `ChainOp` that provides storage for either the
	*chain of boundaries* (from `D` to `0`) or the transposed *chain of coboundaries*
	(from `0` to `D`), with `D` the dimension of the embedding space, which may be either
	`` R^2`` or `` R^3``.
	"""
	const ChainComplex = Array{ChainOp,1}

	"""
		LARmodel = Tuple{Points,Array{Cells,1}}

	Alias declation of LAR-specific data structure.
	`LARmodel` is a pair (*Geometry*, *Topology*), where *Geometry* is stored as
	`Points`, and *Topology* is stored as `Array` of `Cells`. The number of `Cells`
	values may vary from `1` to `N+1`.
	"""
	const LARmodel = Tuple{Points,Array{Cells,1}}

	"""
		LAR = Tuple{Points,Cells}

	Alias declation of LAR-specific data structure.
	`LAR` is a pair (*Geometry*, *Topology*), where *Geometry* is stored as
	`Points`, and *Topology* is stored as `Cells`.
	"""
	const LAR = Tuple{Points,Cells}

	"""
		Hpc = PyCall.PyObject

	Alias declation of LAR-specific data structure.
	`Hpc` stands for *Hierarchical Polyhedral Complex* and is the geometric data structure
	used by `PLaSM` (Programming LAnguage for Solid Modeling). See the Wiley's book
	[*Geometric Programming for Computer-Aided Design*]
	(https://onlinelibrary.wiley.com/doi/book/10.1002/0470013885) and its
	current `Python` library [*https://github.com/plasm-language/pyplasm*]
	(https://github.com/plasm-language/pyplasm).
	"""
	const Hpc = PyCall.PyObject

	"""
		cuboidGrid(shape::Array{Int64,1}[, full=false])::Union{LAR,LARmodel}

	compute a *cellular complex* (mesh) with *cuboidal cells* of either `LARmodel`
	or `LAR` type, depending of the value of optional `full` parameter. The default is
	for returning a `LAR` value, i.e. a pair `(Points, Cells)`.
	The *dimension* of `Cells` is the one of the number `M` of rows of
	cell `Points`. The dimensions of `Array{Cells,1}` in `LARmodel` run
	from ``1`` to ``M``.
	"""
	cuboidGrid = LinearAlgebraicRepresentation.larCuboids

	"""
		centroid( V::Points )::Array{Float64,1}

	*Geometric center* of a `Points` 2-array of `size` ``(M,N)``. Each of the
	``M`` coordinates of *barycenter* of the dense array of ``N`` `points` is the *mean*
	of the corresponding `Points` coordinates.
	"""
	function centroid( V::Points )::Array{Float64,2}
		return sum(V,2)/size(V,2)
	end

	"""
		centroid(V::Array{Float64,2})::Array{Float64,1}

	*Geometric center* of a `Points` 2-array of `size` ``(M,N)``. Each of the
	``M`` coordinates of *barycenter* of the dense array of ``N`` `points` is the *mean*
	of the corresponding `Array` coordinates.
	"""
	function centroid(V::Array{Float64,2})::Array{Float64,2}
	    return sum(V,2)/size(V,2)
    end

	"""
		cells2py(cells::Cells)::PyObject

	Return a `Cells` object in a *Python* source text format. The returned `PyObject` is
	 a list of lists of integers.

	# Example
	``` julia
	julia> FV = LinearAlgebraicRepresentation.cuboid([1,1,1],true)[2][3]
	6-element Array{Array{Int64,1},1}:
	 [1, 2, 3, 4]
	 [5, 6, 7, 8]
	 [1, 2, 5, 6]
	 [3, 4, 7, 8]
	 [1, 3, 5, 7]
	 [2, 4, 6, 8]

	julia> Plasm.cells2py(FV)
	PyObject [[1, 2, 3, 4], [5, 6, 7, 8], [1, 2, 5, 6], [3, 4, 7, 8],
	[1, 3, 5, 7], [2, 4, 6, 8]]
	```
	"""
	function cells2py(cells::Cells)::PyObject
		return PyObject([Any[cell[h] for h=1:length(cell)] for cell in cells])
	end

	"""
		points2py(V::Points)::PyObject

	Return a `Points` object in a *Python* source text format. The returned `PyObject` is
	 a `list of lists of float`.

	# Example
	``` julia
	julia> V = LinearAlgebraicRepresentation.cuboid([1,1,1])[1]
	3×8 Array{Float64,2}:
	 0.0  0.0  0.0  0.0  1.0  1.0  1.0  1.0
	 0.0  0.0  1.0  1.0  0.0  0.0  1.0  1.0
	 0.0  1.0  0.0  1.0  0.0  1.0  0.0  1.0

	julia> Plasm.points2py(V)
	PyObject [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [0.0, 1.0, 1.0],
	[1.0, 0.0, 0.0], [1.0, 0.0, 1.0], [1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]
	```
	"""
	function points2py(V::Points)::PyObject
		return PyObject([Any[V[h,k] for h=1:size(V,1)] for k=1:size(V,2)])
	end

	"""
		mkpol(verts::Points, cells::Cells)::Hpc

	Return an `Hpc` object starting from a `Points` and a `Cells` object. *HPC =
	Hierarchical Polyhedral Complex* is the geometric data structure
	used by `PLaSM` (Programming LAnguage for Solid Modeling). See the Wiley's book
	[*Geometric Programming for Computer-Aided Design*]
	(https://onlinelibrary.wiley.com/doi/book/10.1002/0470013885) and its
	current `Python` library [*https://github.com/plasm-language/pyplasm*]
	(https://github.com/plasm-language/pyplasm).

	```julia
	julia> V,(VV,EV,FV,CV) = LinearAlgebraicRepresentation.cuboid([1,1,1],true);

	julia> Plasm.mkpol(V,EV)
	PyObject <pyplasm.xgepy.Hpc; proxy of <Swig Object of type
	'std::shared_ptr< Hpc > *' at 0x12cf45d50> >

	julia> Plasm.view(Plasm.mkpol(V,EV))
	[...]
	```
	"""
	function mkpol(verts::Points, cells::Cells)::Hpc
		verts = points2py(verts)
		cells = cells2py(cells)
		return p.MKPOL([verts,cells,[]])
	end

	"""
		view(hpc::Hpc)

	Base.view extension.
	Display a *Python*  `HPC` (Hierarchica Polyhedral Complex) `object` using
	the *`PyPlasm` viewer*, written in C++ with `OpenGL` and acceleration algorithms
	for *big geometric data* structures.

	# Example
	``` julia
	julia> m = LinearAlgebraicRepresentation.cuboidGrid([2,2],true)
	([0.0 0.0 … 2.0 2.0; 0.0 1.0 … 1.0 2.0], Array{Array{Int64,1},1}[Array{Int64,1}[[1],
	[2], [3], [4], [5], [6], [7], [8], [9]], Array{Int64,1}[[1, 2], [2, 3], [4, 5], [5,
	6], [7, 8], [8, 9], [1, 4], [2, 5], [3, 6], [4, 7], [5, 8], [6, 9]],
	Array{Int64,1}[[1, 2, 4, 5], [2, 3, 5, 6], [4, 5, 7, 8], [5, 6, 8, 9]]])

	julia> hpc = Plasm.mkpol(m[1],m[2][2])
	PyObject <pyplasm.xgepy.Hpc; proxy of <Swig Object of type 'std::shared_ptr< Hpc > *'
	at 0x140d6c780> >

	julia> Plasm.view(hpc)
	```
	"""
	function view(hpc::Hpc)
		p.VIEW(hpc)
	end

	"""
		view(V::Points, CV::Cells)

	Base.view extension.
	Display a *Python*  `HPC` (Hierarchica Polyhedral Complex) `object` using
	the *`PyPlasm` viewer*, written in C++ with `OpenGL` and acceleration algorithms
	for *big geometric data* structures. Input parameters are of `Points` and `Cells`
	type.

	# Example

	```julia
	julia> V,(VV,EV,FV,CV) = LinearAlgebraicRepresentation.cuboid([1,1,1],true);

	julia> Plasm.mkpol(V,CV)
	PyObject <pyplasm.xgepy.Hpc; proxy of <Swig Object of type
	'std::shared_ptr< Hpc > *' at 0x12cf45d50> >

	julia>
	Plasm.view(Plasm.mkpol(V,CV))
	[...]
	```
	"""
	function view(V::Points, CV::Cells)
		hpc = lar2hpc(V::Points, CV::Cells)
		p.VIEW(hpc)
	end

	"""
		view(model::LARmodel)

	Base.view extension.
	Display a *Python*  `HPC` (Hierarchica Polyhedral Complex) `object` using
	the *`PyPlasm` viewer*, written in C++ with `OpenGL` and acceleration algorithms
	for *big geometric data* structures. The input is a `LARmodel` object.

	# Example

	```julia
	julia> typeof( LinearAlgebraicRepresentation.cuboid([1,1,1], true) )
	Tuple{Array{Float64,2},Array{Array{Int64,1},1}}

	julia> V,(VV,EV,FV,CV) = LinearAlgebraicRepresentation.cuboid([.5,.5,.5], true);

	julia> Plasm.view( (V,[VV,EV,FV,CV]) )
	```
	"""
	function view(model::LARmodel)
		hpc = hpc_exploded(model::LARmodel)(1,1,1)
		p.VIEW(hpc)
	end

	"""
		view(pair::Tuple{Points,Cells})

	Base.view extension.
	Display a *Python*  `HPC` (Hierarchica Polyhedral Complex) `object` using
	the *`PyPlasm` viewer*, written in C++ with `OpenGL` and acceleration algorithms
	for *big geometric data* structures. The input is a `pair` of type
	`Tuple{Points,Cells}`.

	# Example
	```julia
	julia> typeof(LinearAlgebraicRepresentation.cuboid([1,1,1])::LAR)
	Tuple{Array{Float64,2},Array{Array{Int64,1},1}}

	julia> V,(VV,EV,FV,CV) = LinearAlgebraicRepresentation.cuboid([1,1,1], true);

	julia> Plasm.view( (V,FV) );
	```
	"""
	function view(pair::Tuple{Points,Cells})
		V,CV = pair
		hpc = lar2hpc(V::Points, CV::Cells)
		p.VIEW(hpc)
	end

	"""
		view(obj::LinearAlgebraicRepresentation.Struct)

	Display a geometric value of `Struct` type, via conversion to `LAR`
	and then to `Hpc` values.

	# Example
	```julia
	cube = LinearAlgebraicRepresentation.apply( LinearAlgebraicRepresentation.t(-.5,-.5,0), LinearAlgebraicRepresentation.cuboid([1,1,1]));
	tableTop = LinearAlgebraicRepresentation.Struct([ LinearAlgebraicRepresentation.t(0,0,.85), LinearAlgebraicRepresentation.s(1,1,.05), cube ]);
	tableLeg = LinearAlgebraicRepresentation.Struct([ LinearAlgebraicRepresentation.t(-.475,-.475,0), LinearAlgebraicRepresentation.s(.1,.1,.89), cube ]);
	tablelegs = LinearAlgebraicRepresentation.Struct( repeat([ tableLeg, LinearAlgebraicRepresentation.r(0,0,pi/2) ],outer=4) );
	table = LinearAlgebraicRepresentation.Struct([ tableTop, tablelegs ]);

	Plasm.view(table)
	```
	"""
	function view(obj::LinearAlgebraicRepresentation.Struct)
		lar = LinearAlgebraicRepresentation.struct2lar(obj)
		view(lar)
	end

	"""
		view(scene::Array{Any,1})

	Display a geometric `scene`, defined as `Array{Any,1}` of geometric objects
	defined in the *same* coordinate system, i.e. in *World Coordinates*.

	# Example

	A hierarchical `scene` defined in *Local Coordinates* as value of `Struct` type,
	must be converted to `Array{Any,1}` by the expression

		`evalStruct(scene::Struct)::Array{Any,1}`

	```
	cube = LinearAlgebraicRepresentation.apply( LinearAlgebraicRepresentation.t(-.5,-.5,0), LinearAlgebraicRepresentation.cuboid([1,1,1]));
	tableTop = LinearAlgebraicRepresentation.Struct([ LinearAlgebraicRepresentation.t(0,0,.85), LinearAlgebraicRepresentation.s(1,1,.05), cube ]);
	tableLeg = LinearAlgebraicRepresentation.Struct([ LinearAlgebraicRepresentation.t(-.475,-.475,0), LinearAlgebraicRepresentation.s(.1,.1,.89), cube ]);
	tablelegs = LinearAlgebraicRepresentation.Struct( repeat([ tableLeg, LinearAlgebraicRepresentation.r(0,0,pi/2) ],outer=4) );
	table = LinearAlgebraicRepresentation.Struct([ tableTop, tablelegs ]);

	scene = LinearAlgebraicRepresentation.evalStruct(table);
	# output
	# 5-element Array{Any,1}

	Plasm.view(scene)
	```
	"""
	function view(scene::Array{Any,1})
		if prod([isa(item[1:2],LinearAlgebraicRepresentation.LAR) for item in scene])
			p.VIEW(p.STRUCT([Plasm.lar2hpc(item[1],item[2]) for item in scene]))
		end
	end

	"""
		hpc_exploded( model::LARmodel )( sx=1.2, sy=1.2, sz=1.2 )::Hpc

	Convert a `LARmodel` into a `Hpc` object, after exploding all-dimensional cells with
	scale `sx,sy,sz` parameters. Every cell is *translated* by the vector difference
	between its *scaled centroid* and its *centroid*. Every cell is transformed in a
	single `LAR` object before explosion.

	# Example
	```julia
	julia> hpc = Plasm.hpc_exploded(LinearAlgebraicRepresentation.cuboidGrid([3,2,1], true))(1.5,1.5,1.5)

	julia> view(hpc)
	```
	"""
	function hpc_exploded( model )
		function hpc_exploded0( sx=1.2, sy=1.2, sz=1.2 )
			verts,cells = model
			out = []
			for skeleton in cells
				for cell in skeleton
					vcell = hcat([[verts[h,k] for h=1:size(verts,1)] for k in cell]...)

					center = sum([verts[:,v] for v in cell])/length(cell)
					scaled_center = length(center)==2 ? center.*[sx,sy] :
														center.*[sx,sy,sz]
					translation_vector = scaled_center-center
					vcell = vcell .+ translation_vector

					py_verts = Plasm.points2py(vcell)
					py_cells = Plasm.cells2py( [collect(1:size(vcell,2))] )

					hpc = p.MKPOL([ py_verts, py_cells, [] ])
					push!(out, hpc)
				end
			end
			hpc = p.STRUCT(out)
			return hpc
		end
		return hpc_exploded0
	end

	"""
		lar2hpc(V::Points, CV::Cells)::Hpc

	Return an `Hpc` object starting from a `Points` and a `Cells` object. *HPC =
	Hierarchical Polyhedral Complex* is the geometric data structure
	used by `PLaSM` (Programming LAnguage for Solid Modeling). See the Wiley's book
	[*Geometric Programming for Computer-Aided Design*]
	(https://onlinelibrary.wiley.com/doi/book/10.1002/0470013885) and its
	current `Python` library [*https://github.com/plasm-language/pyplasm*]
	(https://github.com/plasm-language/pyplasm).
	# Example

	```julia
	julia> V,(VV,EV,FV,CV) = LinearAlgebraicRepresentation.cuboid([1,1,1],true);

	julia> hpc = lar2hpc( (V, CV)::LAR ... )::Hpc
	PyObject <pyplasm.xgepy.Hpc; proxy of <Swig Object of type
	'std::shared_ptr< Hpc > *' at 0x12cf45d50> >

	julia> view(hpc)
	[...]
	```
	"""
	function lar2hpc(V::Points, CV::Cells)::Hpc
		hpc = mkpol(V,CV)
	end

	"""
		lar2hpc(model::LARmodel)::Hpc

	Return an `Hpc` object starting from a `LARmodel` object. *HPC =
	Hierarchical Polyhedral Complex* is the geometric data structure
	used by `PLaSM` (Programming LAnguage for Solid Modeling). See the Wiley's book
	[*Geometric Programming for Computer-Aided Design*]
	(https://onlinelibrary.wiley.com/doi/book/10.1002/0470013885) and its
	current `Python` library [*https://github.com/plasm-language/pyplasm*]
	(https://github.com/plasm-language/pyplasm).

	# Example
	```julia
	julia> model = LinearAlgebraicRepresentation.cuboid([1,1,1],true);

	julia> view( Plasm.lar2hpc(model) )

	julia> view( Plasm.hpc_exploded(model)(1.5,1.5,1.5) )
	```
	"""
	function lar2hpc(model::LARmodel)::Hpc
		verts = model[1]
		cells = Array{Int,1}[]
		for item in model[2]
			append!(cells, item)
		end
		hpc = mkpol(verts,cells)
	end

	"""
		lar2hpc(scene::Array{Any,1})::Hpc

	# Example
	```julia
	square = ([[0; 0] [0; 1] [1; 0] [1; 1]], [[1, 2, 3, 4]],
	[[1,2], [1,3], [2,4], [3,4]])
	V,FV,EV  =  square
	model  =  V,([[1],[2],[3],[4]],EV,FV)
	table = LinearAlgebraicRepresentation.apply(LinearAlgebraicRepresentation.t(-0.5,-0.5), square)
	chair = LinearAlgebraicRepresentation.Struct([LinearAlgebraicRepresentation.t(0.75,0),LinearAlgebraicRepresentation.s(0.35,0.35),table])
	structo = LinearAlgebraicRepresentation.Struct([LinearAlgebraicRepresentation.t(2,1),table,repeat([LinearAlgebraicRepresentation.r(pi/2),chair],
			outer = 4)...])
	structo1 = LinearAlgebraicRepresentation.Struct(repeat([structo,LinearAlgebraicRepresentation.t(0,2.5)],outer = 10));
	structo2 = LinearAlgebraicRepresentation.Struct(repeat([structo1,LinearAlgebraicRepresentation.t(3,0)],outer = 10));
	scene = LinearAlgebraicRepresentation.evalStruct(structo2);

	Plasm.view(Plasm.lar2hpc(scene))
	```

	"""
	function lar2hpc(scene::Array{Any,1})::Hpc
		hpc = p.STRUCT([ mkpol(item[1],item[2]) for item in scene ])
	end

	"""
		lar2exploded_hpc(V::LinearAlgebraicRepresentation.Points,CV::LinearAlgebraicRepresentation.Cells)::Hpc

	Input `V::Points` and `CV::Cells`. Output an *exploded* `Hpc`
	object,  exploding cells in `CV` with scale `sx,sy,sz` parameters.

	Every cell is *translated* by the vector difference
	between its *scaled centroid* and its *centroid*.

	*HPC = Hierarchical Polyhedral Complex* is the geometric data structure
	used by `PLaSM` (Programming LAnguage for Solid Modeling).

	#	Example

	```
	julia> V,cells = LinearAlgebraicRepresentation.cuboidGrid([3,3,1], true)

	julia> hpc = Plasm.lar2exploded_hpc(V::LinearAlgebraicRepresentation.Points, cells[4]::LinearAlgebraicRepresentation.Cells)()

	julia> Plasm.view(hpc)
	```
	"""
	function lar2exploded_hpc(V::LinearAlgebraicRepresentation.Points, cells::LinearAlgebraicRepresentation.Cells)
		function lar2exploded_hpc0(sx=1.2, sy=1.2, sz=1.2)
			hpc = Plasm.hpc_exploded( (V,[cells]) )(sx,sy,sz)
		end
		return lar2exploded_hpc0
	end

	function viewexploded(V::LinearAlgebraicRepresentation.Points, cells::LinearAlgebraicRepresentation.Cells)
		function lar2exploded_hpc0(sx=1.2, sy=1.2, sz=1.2)
			hpc = Plasm.hpc_exploded( (V,[cells]) )(sx,sy,sz)
			Plasm.view(hpc)
		end
		return lar2exploded_hpc0
	end

   include("./graphic_text.jl")

end # module


using ForwardDiff
using LinearAlgebra
using Optim
using Pathfinder
using Test

include("test_utils.jl")

function lbfgs_inverse_hessian_explicit(H₀, S, Y)
    B = [H₀ * Y S]
    R = triu(S'Y)
    E = Diagonal(R)
    D = [0*I -inv(R); -inv(R)' R' \ (E + Y' * H₀ * Y)/R]
    return H₀ + B * D * B'
end

@testset "L-BFGS inverse Hessian construction" begin
    @testset "lbfgs_inverse_hessian" begin
        #! format: off
        S0 = [[1.719573, 3.294037, -2.008877, 3.901275, 0.214324, 0.400382, 0.113598, -1.804262, 0.465563, 2.465748], [-0.445476, -0.514915, 1.069617, -1.505506, -0.036535, -0.11386, -0.029737, 0.579662, -0.118694, -1.726258], [-0.013354, 0.1981, 0.081886, -0.194172, 0.010499, -0.007944, -0.001039, 0.061604, -0.002621, 0.11145], [0.00648, 0.255961, -0.011901, -0.059042, 0.013587, -0.002064, 0.000302, 0.023457, 0.00257, -0.002343], [-0.016408, 0.015005, 0.009654, 0.016735, -0.000245, -0.002825, -0.00106, 0.004272, -0.004578, 0.002275]]
        Y0 = -[[-2.357935, -3.343312, 4.659008, -7.065433, -0.228045, -0.584164, -0.156837, 2.863289, -0.631753, -7.152021], [0.610851, 0.522617, -2.480668, 2.726559, 0.038874, 0.166123, 0.041055, -0.9199, 0.161064, 5.007094], [0.018312, -0.201064, -0.189911, 0.351657, -0.011171, 0.011591, 0.001434, -0.097763, 0.003557, -0.323265], [-0.008886, -0.25979, 0.027601, 0.106929, -0.014456, 0.003011, -0.000418, -0.037225, -0.003488, 0.006795], [0.022499, -0.015229, -0.02239, -0.030308, 0.00026, 0.004122, 0.001463, -0.00678, 0.006212, -0.006598]]
        #! format: on
        S = reduce(hcat, S0)
        Y = reduce(hcat, Y0)
        N, history_length = size(S)
        α = rand(N)
        H₀ = Diagonal(α)

        @test @inferred(Pathfinder.lbfgs_inverse_hessian(H₀, S, Y, 0, 0)) ≈ H₀

        H = Pathfinder.lbfgs_inverse_hessian(H₀, S, Y, 3, 3)
        Hexp = lbfgs_inverse_hessian_explicit(H₀, S[:, 1:3], Y[:, 1:3])
        @test H ≈ Hexp

        S2 = S[:, [4:history_length; 1:3]]
        Y2 = Y[:, [4:history_length; 1:3]]
        ilast = argmax(axes(S, 2)[[4:history_length; 1:3]])
        H = Pathfinder.lbfgs_inverse_hessian(H₀, S2, Y2, ilast, history_length)
        Hexp = lbfgs_inverse_hessian_explicit(H₀, S, Y)
        @test H ≈ Hexp

        H = Pathfinder.lbfgs_inverse_hessian(H₀, S, Y, history_length, history_length)
        Hexp = lbfgs_inverse_hessian_explicit(H₀, S, Y)
        @test H ≈ Hexp
    end

    @testset "lbfgs_inverse_hessians" begin
        n = 10
        history_length = 5
        logp(x) = logp_banana(x)
        ∇logp(x) = ForwardDiff.gradient(logp, x)
        nocedal_wright_scaling(α, s, y) = fill!(similar(α), dot(y, s) / sum(abs2, y))
        θ₀ = 10 * randn(n)

        optimizer = Optim.LBFGS(;
            m=history_length, linesearch=Optim.LineSearches.MoreThuente()
        )
        θs, logpθs, ∇logpθs = Pathfinder.maximize_with_trace(logp, ∇logp, θ₀, optimizer)

        # run lbfgs_inverse_hessians with the same initialization as Optim.LBFGS
        Hs = Pathfinder.lbfgs_inverse_hessians(
            θs, ∇logpθs; history_length=history_length, Hinit=nocedal_wright_scaling
        )
        ss = diff(θs)
        ps = (Hs .* ∇logpθs)[1:(end - 1)]
        # check that next direction computed from Hessian is the same as the actual
        # direction that was taken
        @test all(≈(1), dot.(ps, ss) ./ norm.(ss) ./ norm.(ps))
    end
end


mutable struct Topography

    topoz_sT :: Array{Float64, 3}
    Nz_bot_sT :: Array{Int64, 3}

    mask_T      :: AbstractArray{Float64, 3}
    sfcmask_sT  :: AbstractArray{Float64, 3}
    deepmask_T  :: AbstractArray{Float64, 3}

    function Topography(
        Nz_bot        :: Union{AbstractArray{Int64,2}, Nothing},
        Nx            :: Int64,
        Ny            :: Int64,
        z_w           :: AbstractArray{Float64, 1};
        deep_depth    :: Float64 = 500.0
    )

        Nz = length(z_w) - 1

        sT_pts = Nx * Ny

        topoz_sT = zeros(Float64, 1, Nx, Ny)

        if Nz_bot == nothing
            Nz_bot_sT = zeros(Int64, 1, Nx, Ny)
            Nz_bot_sT .= length(z_w) - 1
        else
            if size(Nz_bot) != (Nx, Ny)
                throw(ErrorException(format("Size of `Nz_bot` variable is not the same as given (Nx, Ny) = ({:d}, {:d})", Nx, Ny)))
            end

            Nz_bot_sT = copy(reshape(Nz_bot, 1, Nx, Ny))

        end

        for j=1:Ny, i=1:Nx
            if Nz_bot_sT[1, i, j] > Nz
                Nz_bot_sT[1, i, j] = Nz
            end
            topoz_sT[1, i, j] = z_w[Nz_bot_sT[1, i, j]+1]
        end

        deep_Nz = lookupNz(- deep_depth, z_w)
        println(format("Water depth of {:f} m is considered as deep. The found deep_Nz = {:d}", deep_depth, deep_Nz))
        _lookupNz = (z) -> lookupNz(z, z_w)
        @. Nz_bot_sT = _lookupNz(topoz_sT)

        mask_T = ones(Float64, Nz, Nx, Ny)

        # turn the grids below topography to 0
        for j=1:Ny, i=1:Nx

            if Nz_bot_sT[1, i, j] < Nz
                mask_T[Nz_bot_sT[1, i, j]+1:end, i, j] .= 0.0
            end
        end

        sfcmask_sT = view(mask_T, 1:1, :, :)

        deepmask_T = copy(mask_T)
        # turn the entire column to 0 if it is too shallow
        for j=1:Ny, i=1:Nx
            if Nz_bot_sT[1, i, j] < deep_Nz
                deepmask_T[:, i, j] .= 0.0
            end
        end

        return new(
            topoz_sT,
            Nz_bot_sT,
            mask_T,
            sfcmask_sT,
            deepmask_T,
        )

    end

end

function lookupNz(z, z_w)

    max_Nz = length(z_w) - 1

    if z == 0.0
        return 0
    end

    for i=1:max_Nz-1
        if z_w[i] > z >= z_w[i+1]
            return i
        end
    end

    return max_Nz

end


@testset "Features" begin

for l in 7:5:20
    for offset in 2π .* [0, 2, -1, 10, -50]
        phase_uw = collect(range(-2π; stop=2π, length=l))
        phase = rem2pi.(phase_uw, RoundNearest) .+ offset
        for mag in [nothing, collect(range(10; stop=1, length=l)), collect(range(1; stop=10, length=l))]
            unwrapped = unwrap(phase; mag=mag, correctglobal=true)
            @test unwrapped ≈ phase_uw
        end
    end
end

end


using SparseArrays
import SparseArrays: dropzeros!

"""
dropzeros!(sparseMatrix, drop_criteria)
# Description
- Drops machine zeros in sparse matrix
# Arguments
- `A`: a sparse matrix
- `drop_criteria`: criteria for dropping entries
# return
- nothing
"""
function dropzeros!(A, drop_criteria)
    i, j = findnz(A)
    for loop in 1:length(i)
        if abs(A[i[loop], j[loop]]) < drop_criteria
            A[i[loop], j[loop]] = 0.0
        end
    end
    dropzeros!(A)
end

"""
droprelativezeros!(A; scale=100)
# Description
- Drops machine zeros in sparse matrix
# Arguments
- `A`: a sparse matrix
- `drop_criteria`: criteria for dropping entries
# return
- nothing
# This function calls
dropzeros!(A, eps(scale * maximum(da)))
"""
function droprelativezeros!(A; scale = 100)
    dropzeros!(A, eps(scale * maximum(abs.(A))))
end


### A Pluto.jl notebook ###
# v0.14.4

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 65780f00-ed6b-11ea-1ecf-8b35523a7ac0
begin
	import Pkg
	Pkg.activate("..")
	Pkg.instantiate()
	using Plots
	plotlyjs()
	using PlutoUI
	using FFTW
	using MRIReco
	using LinearAlgebra
	using IterativeSolvers
end

# ╔═╡ ac8ff080-ed61-11ea-3650-d9df06123e1f
md"""

# **Exercise for the ISMRM Lecture on the _Reconstruction of Non-Cartesian Data_**
`Session: Image Reconstruction`
"""

# ╔═╡ a24b4518-e6c3-4845-8b32-b203af6a2d69
html"""
Live Q&A in your local time zone:

"""

# ╔═╡ aff9b50d-d013-4c8b-8510-9a59e0c9fc2e
md"""
This notebook contains _built-in, live answer checks_! In some exercises you will see a colored box, which runs a test case on your code, and provides feedback based on the result. Simply edit the code, run it (`Shift` + `Enter`), and the check runs again as the notebook _reactive_: whenever you change a cell, all affected cells are evaluated, which will become important in Exercises 3-4.

Join the live Q&A session if you have questions, or reach out via [email](mailto:jakob.asslaender@nyumc.org)!
"""

# ╔═╡ 5f95e01a-ee0a-11ea-030c-9dba276aba92
md"""
#### Load packages

Most functionality in Julia in provided by so-called _packages_. On binder, I already installed all required packages. If you are running this exercise on your own computer, you installed everything with the `Pkg.develop(url="https://github.com/JakobAsslaender/2021_ISMRM_nonCartesianReconstruction_Exercises")` command. The only remaining step is to make those packages available:
"""

# ╔═╡ 540ccfcc-ee0a-11ea-15dc-4f8120063397
md"""
## **Exercise 0** - _Julia 101_

Let's start by creating a numerical phantom and calculate the Cartesian k-space. If you are new to Julia, this code will give a glimpse at its syntax.

Let's choose a (quadratic) matrix size:

"""

# ╔═╡ 26709e4b-1675-4ee6-a436-bf9a85683950
nx = 128

# ╔═╡ 97a4d58f-9c32-4de9-9142-31e34fc4e33b
md"""
If you are running this exercise on binder, I would suggest to stick to $n_x \leq 128$ as computational resources on binder are rather limited. If you are runnings it on your own beefy computer, feel free to increase this value!

Let's create a Shepp-Logan phantom:
"""

# ╔═╡ d72dc575-6d2e-489c-b512-aeae6d017f7f
image_org = reverse(shepp_logan(nx),dims=1)

# ╔═╡ 9391828f-041c-49db-93d8-3483bb672f3d
md"
And let's plot it with the `Plots` package:
"

# ╔═╡ 67bedf95-8ff4-430f-8737-a0d21ff368c8
md"Color scaling:"

# ╔═╡ eeed0b36-37be-4af4-a8c2-68c595d6bb2c
@bind cmax_org Slider(0.5:0.1:1.0, default=1.0)

# ╔═╡ c6be8293-7f7c-48a7-b112-92f7ae4fb7c8
md"
FFT functionality is provided by the `FFTW` package and the syntax is very similar to Matlab:
"

# ╔═╡ 35022e16-d376-426c-a9f8-b71760bcaaf7
k_Cart = fftshift(fft(fftshift(image_org)))

# ╔═╡ 40a94bcd-a0cc-4b1a-b370-6405adc0f92b
md"**Dot syntax**: In comparison to Matlab, Julia distinguishes more explicitly between functions that act on a matrix as a whole or on each element. E.g. `exp(x)` calculates the matrix exponential of `x`, while `exp.(x)` calculates the exponential of each element of `x`.

Use the dot-syntax to calculate the absolute value `abs` of `k_Cart`:
. "

# ╔═╡ c0afaa73-aa46-4f61-a5b4-c49d071f4fb6
abs_k_Cart = missing

# ╔═╡ 515ae3f8-fb28-49b9-a023-38d143b6674a
begin
	function calculate_error(reco, org, k_mask)
		k_reco = fftshift(fft(reco))
		k_org  = fftshift(fft(org))
		return norm((k_reco .- k_org)[k_mask]) / norm(k_org[k_mask])
	end
	function calculate_error_norm(reco, org, k_mask)
		reco = reco ./ norm(reco)
		org  = org  ./ norm(org)
		return calculate_error(reco, org, k_mask)
	end

	k_mask = similar(k_Cart, Bool)
	for i ∈ CartesianIndices(k_mask)
		if (i[1]-nx/2)^2 + (i[2]-nx/2)^2 > (nx - 4)^2/4
			k_mask[i] = false
		else
			k_mask[i] = true
		end
	end
end

# ╔═╡ ad6a33b0-eded-11ea-324c-cfabfd658b56
md"""
## Exercise 1 - _filtered back-projection_

For Exercises 1-4, we will use a radial trajectory. In radial sampling, we need to acquire $\pi/2$-times as many _spokes_ compared to Cartesian k-space lines in order to fulfill the Nyquist theorem at the edge of k-space. This is the price you pay for a higher sampling density in the center:
"""

# ╔═╡ 28668257-321b-41f2-bfe2-d5f35441ec8d
nspokes = ceil(Int, nx * π/2 + 20)

# ╔═╡ 137ed83c-171f-48df-be8e-dc538ac18929
md"
The added 20 improve the numerical accuracy. In practice, this should not be necessary.

In the readout direction, we use the usual 2-fold oversampling:
"

# ╔═╡ 04962733-71dd-4e6f-9e23-5b29ee821d68
nr = 2 * nx

# ╔═╡ 74d30291-506f-4e57-9241-c791f0fb8264
md"We can use the MRIReco.jl package to create a trajectory:"

# ╔═╡ 319872dd-d646-4029-9319-4ffb957179a0
trj = trajectory("Radial", nspokes, nr)

# ╔═╡ 4ded3325-4935-4843-b9c2-b0e29469bf3e
plot(trj.nodes[1,:], trj.nodes[2,:], legend=false, xlabel="kx (rad/voxel/2pi)", ylabel="ky (rad/voxel/2pi)")

# ╔═╡ c0f04553-6031-4984-8c80-1260022eb23f
md"With the same package we can also create a non-uniform FFT (NFFT) operator:"

# ╔═╡ 05dafc16-c68f-47e8-8bdf-cfea8736bc94
html"""

F is a linear operator and the MRIReco.jl package uses Julia's

multiple dispatch
to implement a * method for this operator. Despit it being a much more efficient implemenentation, F essentially acts like a matrix and we can calculate the k-space signal with something that looks like a matrix-vector multiplication:
"""

# ╔═╡ f6d6fe5e-3381-4338-817d-366194b3f7ba
function calculate_k_radial(F, image_org)
	k_radial = F * vec(image_org)
	return k_radial
end

# ╔═╡ 424e1625-e14b-4fcb-bd5b-7e426c97d330
md"Let's try to reconstruct this data with a back-projection. As discussed in the lecture, we need to apply density compensation, othewise the image is blurry (see below). Modify below code line to compute the correct density compensation for radial imaging!

Tip: For radial imaging, the density is only a function of `r`, so you only need to replace `ones(nr)`."

# ╔═╡ 16ad1b20-f3c8-4516-b5d9-6150875c326c
density_compensation = repeat(ones(nr), nspokes)

# ╔═╡ 9aff92eb-8094-4a83-ad27-242277b9aee1
md"With the correct density compensation, we can reconstruct the image:"

# ╔═╡ 28c1a017-8977-40a8-91d2-1cac83a0d3cf
md"Color scaling:"

# ╔═╡ bfa47923-d80b-43ea-a21e-f3dad07b4958
md"""
If you change the color scaling, it should become apparent that the image still has artifacts. Take a look at the k-space trajectory plotted above, and compare it to the Cartesian k-space. You should see that the radial k-space trajectory does not acquire the corners of k-space, which results in the here observed artifacts. For comparison, you can cut out the Cartesian k-space and observe the "original image" and it's Cartesian k-space at the very top.
"""

# ╔═╡ 05d84e43-d768-4680-82eb-d3a5419edd85
@bind b_gt Select(["org" => "show original image", "cut" => "cut out k-space"])

# ╔═╡ 2af0c856-430b-41ec-a804-77b3367d1fc5
md"Throughout this notebook, the RMSE tests are calculated on the k-space data and only within the sampled area of k-space."

# ╔═╡ 94ca0194-811f-472f-8815-1f52c6a1619c
if b_gt == "org"
	image_gt = image_org;
	abs_k_Cart_disp = abs_k_Cart;
else
	image_gt = abs.(ifft(fft(image_org) .* ifftshift(k_mask)));
	abs_k_Cart_disp = abs_k_Cart .* k_mask
end;

# ╔═╡ 1adbae4f-c732-457b-aad8-bc0f3a7be13d
heatmap(image_gt, c=:grays, clim=(0,cmax_org))

# ╔═╡ 8a3dc74a-4b20-4c0f-96e8-1ab0451a7da3
try
	heatmap(log.(abs_k_Cart_disp))
catch
end

# ╔═╡ 9d70e2e6-6e45-4662-9792-b36ab36c35b4
md"""
## Exerise 2 - _iterative reconstruction_
In case you didn't get the density compensation right: don't worry. As discussed in the lecture, we can avoid the density compensation by using iterative algorithms.

Here we will use conjugate gradients (CG), which is usually the method of choice for solving linear problem iteratively. To recall, CG solves

$\hat{x} = \arg \min ||Ax - b||_2^2$

where A is a symmetric, positive-definite matrix. In order to bring our imaging problem in the right form we have to define

$A = F' F$

and

$b = F' s,$

where $F$ is our NFFT-Operator, $F'$ its adjoint, and $s$ is the measured signal.

Great, so let's do that! Replace the `missing` in the following two cells:
"""

# ╔═╡ 32bc54d1-4ea8-4663-ba9a-03d357eb8c02
A = randn(256, 256)

# ╔═╡ 51b2f56b-3880-4ba6-828a-104334bad1ec
b = randn(256)

# ╔═╡ 698d0344-ec5d-4c27-a2b1-b52825e398d6
md"Now we have everything in place to run a reconstruction with conjugate gradients:"

# ╔═╡ 57a2c882-760c-4ff9-9bcb-be8f358d7b8a
reco_cg = reshape(cg(A, b, maxiter=50), (isqrt(size(A,1)),isqrt(size(A,1))));

# ╔═╡ 93f84f84-fdf3-41fa-88f8-2a26b8d6982a
md"Color scaling:"

# ╔═╡ 33b23b5c-85e7-49ba-9e4a-27d67943ee73
@bind cmax_cg Slider(0:.1:1, default=1)

# ╔═╡ f6447d81-6612-428d-b3d7-acd7f8756df5
md"""
## Exerise 3 - _the puzzle_

Great, we have mastered back-projection and interative reconstructions. When you look at the definition of `F`, you might be confused by the definition of `kernelSize` and `oversamplingFactor`. In fact they are defined by these two sliders:
"""

# ╔═╡ 98a1c04c-58a3-44f6-995f-5d4d81095576
md"kernelSize = "

# ╔═╡ e6af6037-7eb6-40be-a1b9-846ab1f50b9c
@bind kernelSize Slider(1:9, show_value=true)

# ╔═╡ 8075e80d-b6df-435a-ac4d-44a919909323
md"oversamplingFactor = "

# ╔═╡ e4d80b4c-faa3-4514-876b-63bce1ece81f
@bind oversamplingFactor Slider(1:0.05:2, show_value=true)

# ╔═╡ beccb301-1c40-4c41-8960-462a82eca3ad
F = NFFTOp(size(image_org), trj; oversamplingFactor=oversamplingFactor, kernelSize=kernelSize)

# ╔═╡ ce933bda-ac45-46ec-b577-c4168b1ce569
heatmap(abs.(reco_cg), c=:grays, clim=(0,cmax_cg), annotations=(size(reco_cg,1) ÷ 1.35,size(reco_cg,2) ÷ 1.1,string("kernelSize = ", kernelSize, "
oversamplingFactor = ", oversamplingFactor), :white))

# ╔═╡ 62934cdb-526d-4f97-afc0-20bcd9f9714e
md"""
So we have initialized the NFFT operator with `oversamplingFactor=1` and `kernelSize=1`, yet we get good results. This is in sharp contrast to theory discussed in the lecture.

$(@bind ic_answer Select(["undefined" => "", "green" => "The image looks great, the box is green, and theory rarely has practical value. Let's enjoy the nice weather!", "teacher" => "The teacher said this can't be right, so I guess it's not...", "correct" => "I contemplated about the theory and it's correct. By logical conclusion, there must be something wrong here!"]))

"""

# ╔═╡ aa209a90-0dfb-41ff-bf0e-a26a6c1b8c8b
function dft(image_org, trj)
	k_radial_dft = zeros(ComplexF64, size(trj,2))
	nx2 = size(image_org,1)/2 + 1
	Threads.@threads for i = 1:length(k_radial_dft)
		for j ∈ CartesianIndices(image_org)
	    	k_radial_dft[i] += image_org[j] *
				exp(-1im * 2π * (trj.nodes[1,i] * (j[1]-nx2) +
  							     trj.nodes[2,i] * (j[2]-nx2)))
		end
	end
	return k_radial_dft
end

# ╔═╡ a0e0fba0-39e1-4873-8d7e-ee5cb793062c
@bind use_dft Select(["no" => "Use NFFT Operator to simulate the signal (inverse crime)", "yes" => "Use the discrete Fourier transform to calculate the signal (slow, but more accurate)"])

# ╔═╡ 86af3791-f3e6-45bc-bc4e-072803491d1f
if use_dft == "yes"
	k_radial = dft(image_org, trj)
else
	Fforward = NFFTOp(size(image_org), trj; oversamplingFactor=1, kernelSize=1)
	k_radial = calculate_k_radial(Fforward, image_org)
end

# ╔═╡ b114e592-0d39-43cb-aee3-c0948b2d2f88
reco_bp = reshape(F' * (k_radial .* density_compensation), (nx, nx));

# ╔═╡ 179e6a9d-3728-4bde-8bae-5b6ac4490a74
@bind cmax_bp Slider(0.5*maximum(abs.(reco_bp)):0.1*maximum(abs.(reco_bp)):maximum(abs.(reco_bp)), default=maximum(abs.(reco_bp)))

# ╔═╡ b62f8883-02cd-4ab0-9c27-61283e4839f6
heatmap(abs.(reco_bp), c=:grays, clim=(minimum(abs.(reco_bp)),cmax_bp))

# ╔═╡ 843dcc64-3015-4682-bae7-ddd86a8615fa
md"When scrolling up, you should see that the RMSE test for the CG reconstruction fails again and the image has artifacts (use the slider to change the color scaling). Change the kernel size and oversampling until the RMSE test is once more passed and the ghosting-artifacts dissapear."

# ╔═╡ ac4504d8-ed24-4972-beff-e518a820dc83
md"""
## Exerise 4 - _computation time and accuracy_

Now that we have a working simulation and reconstruction framework, it is time to play with it a bit! I suggest you modify the following parameters, one at a time:

- `nx`
- `nspokes`
- `kernelSize`
- `oversamplingFactor`

At the bottom right of each cell you should see time it took to execute respective cell. When increasing `nx`, you should be able to observe how the reconstruction time increases moderately, but the DFT time will increase substantially. This effect is exactly what computational complexity refers to. When increasing the other three parameters, you should see modest increases in the reconstruction time.

Feel free to reduce `nspokes` below Nyquist and to see modest undersampling artifacts. Create your own cell to perform Cartesian undersampling for comparison!

"""

# ╔═╡ 9ce17af6-6aff-4615-933e-5e67b56d1038
md"""
## Exerise 5 - _off-resonance correction_

This exercise briefly introduces off-resonance correction. Let's assume we have measured a frequency map, e.g. with a dual-echo gradient echo pulse sequence and it looks like this:
"""

# ╔═╡ c04477fc-ab75-4806-9045-50f71f45c707
md"It is physically not particularily plausible and we magically managed to measure the frequency outside of our object, but it will serve the purpose of this excercise. A smart way to representing a frequency map is to store it as the imaginary part of a complex-valued matrix:"

# ╔═╡ a8339197-b1b5-400f-89f7-0daa87e5bf2c
ω = im * quadraticFieldmap(nx,nx,1e3)[:,:,1]

# ╔═╡ 5a9f1bf5-2f04-4a13-859e-0e09957fe20c
heatmap(imag.(ω), c=:berlin, title="ω (rad/s)")

# ╔═╡ 0368a211-0e20-4841-9e93-ba722b7b42f1
t_acq = 15e-3

# ╔═╡ 84c928bd-4dd4-468c-a62e-7347bcede7fd
md"The advantage of this format is that we joinlty describe signal decay in the real part and off-resonance in the imaginary part. The signal evolution in each voxel is then simply given by

$s(t) \propto exp(- t \cdot (1/T_2^* + i \omega)).$

For now, we will stick to off-resonance and assume $1/T_2^* = 0$, but feel free to experiment with $T_2^*$ correction!

For this exercise we will use a spiral trajectory as one of its primary advantage is the possiblity for long readouts. Here, we assume a $(t_acq*1e3) ms duration:
"

# ╔═╡ c393b9a8-2f4c-40e5-adc8-4605655dff89
n_windings = 10

# ╔═╡ d2936afa-b271-464e-b7ef-37f33e7306dc
md"After choosing the length of each spiral, we can calculate the number of spirals required for Nyquist sampling:"

# ╔═╡ 605ee395-e123-4128-9b56-c73a128572ec
n_spirals = ceil(Int, nx / n_windings)

# ╔═╡ 5a58c91c-1383-4c17-9341-fda706eb0c1a
md"And we set the number of readout points along the spiral:"

# ╔═╡ 882de061-061c-4538-b1b4-5aba0a385569
n_ADC = 2^12

# ╔═╡ 1fbbf883-cf0d-461a-98f0-a934395e276a
md"Now we have everything in place to create and plot the trajectory:"

# ╔═╡ bfd13cd5-87ba-4a4a-b59b-2f06fcb2055c
trj_spiral = trajectory("Spiral", n_spirals, n_ADC; windings=n_windings, AQ=t_acq)

# ╔═╡ 123f7947-71ef-4274-b0fc-25ae6d6b1396
md"Here we plot only the first two of the $(n_spirals) spirals:"

# ╔═╡ cfebce07-d5ee-4df4-97c4-28bc1576ccfe
begin
	plot(trj_spiral.nodes[1,1:2^12], trj_spiral.nodes[2,1:2^12], label="spiral 1/$n_spirals")
	plot!(trj_spiral.nodes[1,2^12+1:2^13], trj_spiral.nodes[2,2^12+1:2^13], label="spiral 2/$n_spirals", xlabel="kx (rad/voxel/2pi)", ylabel="ky (rad/voxel/2pi)")
end

# ╔═╡ 51ef7e32-cf88-405e-8b13-2cae310b93ca
md"To avoid an inverse crime, we perform the forward simulation with a discrete Fourier transform. Compare this function to the one in Exercise 3, it's only one additional phase factor!"

# ╔═╡ 1ea1974a-cebb-4c05-a885-89058a34cf15
function dft_or(image_org, trj, ω)
	k_radial_dft = zeros(ComplexF64, size(trj,2))
	nx2 = size(image_org,1)/2 + 1
	Threads.@threads for i = 1:length(k_radial_dft)
		for j ∈ CartesianIndices(image_org)
	    	k_radial_dft[i] += image_org[j] * exp.(-ω[j] * trj.times[i]) *
				exp(-1im * 2π * (trj.nodes[1,i] * (j[1]-nx2) +
  							     trj.nodes[2,i] * (j[2]-nx2)))
		end
	end
	return k_radial_dft
end

# ╔═╡ 7cf1b6fc-2489-4010-aab0-dc973937a704
k_spiral_or = dft_or(image_org, trj_spiral, ω)

# ╔═╡ 814283c0-99e9-4d9e-929f-6aa721c28aee
md"Replace the normal `NFFTOp`-operator with the `FieldmapNFFTOp`-operator. Once you copy and pasted the constructor call, _Live docs_ a the bottom right will tell you the required and (optional) arguments of the function."

# ╔═╡ a49fa464-56f8-4c4a-8e14-00977cb033a6
F_spiral_or = NFFTOp(size(image_org), trj_spiral)

# ╔═╡ 8b894d85-2dcd-4927-9e0d-91eb93c8f3dd
reco_or = reshape(cg(F_spiral_or' * F_spiral_or, F_spiral_or' * k_spiral_or, maxiter=50), (isqrt(size(F_spiral_or,2)),isqrt(size(F_spiral_or,2))));

# ╔═╡ 220675f2-9ef6-4187-bdd9-1dbebb8af4ec
heatmap(abs.(reco_or), c=:grays, clim=(0,1))

# ╔═╡ 756d150a-b7bf-4bf5-b372-5b0efa80d987
md"## Function library

Just some helper functions used in the notebook."

# ╔═╡ 4bc94bec-da39-4f8a-82ee-9953ed73b6a4
hint(text) = Markdown.MD(Markdown.Admonition("hint", "Hint", [text]))

# ╔═╡ aefd5a3d-9b7d-40a1-9cc8-1a5b555d251d
md"""
You can find documentation by clicking on the Live docs in the bottom right of this Pluto window, and typing a function name in the top.

![image](https://user-images.githubusercontent.com/6933510/107848812-c934df80-6df6-11eb-8a32-663d802f5d11.png)

![image](https://user-images.githubusercontent.com/6933510/107848846-0f8a3e80-6df7-11eb-818a-7271ecb9e127.png)

I recommend that you leave the window open while you work on Julia code. It will continually look up documentation for anything you type!
""" |> hint

# ╔═╡ 7906f103-65db-45c6-9d95-ef1483048c14
md"""
The sampling density of radial trajectories increases as we approach the k-space center, more precisely it is proportional to `1/r`, where `r` is the distance from the k-space center.

When back-projecting, we have too much signal in the center of k-space, hence the blurred image. To compensate, we, thus, have to multiply the k-space data with `r`. Try again!
""" |> hint

# ╔═╡ bfa4b5ac-94b6-4edc-8de8-03b6ac184754
answer(text) = Markdown.MD(Markdown.Admonition("hint", "Answer", [text]))

# ╔═╡ c604ce61-08ac-46af-b231-a99e37f52bd6
md"""
The correct answer is

```
density_compensation = repeat(abs.(-nr/2 : nr/2-1), nspokes)
```

or, equivalently:

```
density_compensation = sqrt.(trj.nodes[1,:].^2 + trj.nodes[2,:].^2)
```

""" |> answer

# ╔═╡ 8233b774-7363-463b-9c27-13c2fecb9e25
md"
The solution is:

`F_spiral_or = FieldmapNFFTOp(size(image_org), trj_spiral, ω)`
" |> answer

# ╔═╡ 8ce6ad06-819c-4af5-bed7-56ecc08c97be
almost(text) = Markdown.MD(Markdown.Admonition("warning", "Almost there!", [text]))

# ╔═╡ dfa40e89-03fc-4a7a-825e-92d67ee217b2
still_missing(text=md"Replace `missing` with your answer.") = Markdown.MD(Markdown.Admonition("warning", "Here we go!", [text]))

# ╔═╡ 086ec1ff-b62d-4566-9973-5b2cc3353409
keep_working(text=md"The answer is not quite right.") = Markdown.MD(Markdown.Admonition("danger", "Keep working on it!", [text]))

# ╔═╡ c22f688b-dc04-4a94-b541-fe06266c5446
correct(text=rand(yays)) = Markdown.MD(Markdown.Admonition("correct", "Got it!", [text]))

# ╔═╡ ab3d1b70-88e8-4118-8d3e-601a8a68f72d
not_defined(variable_name) = Markdown.MD(Markdown.Admonition("danger", "Oopsie!", [md"Make sure that you define a variable called **$(Markdown.Code(string(variable_name)))**"]))

# ╔═╡ 810b38a1-490f-4e22-8edb-0e50a8af2f28
if !@isdefined(abs_k_Cart)
	not_defined(:abs_k_Cart)
elseif ismissing(abs_k_Cart)
	still_missing()
elseif !(abs_k_Cart isa Matrix)
	keep_working(md"`abs_kspace` should be a `Matrix`.")
elseif eltype(abs_k_Cart) != Float64
	mytype = typeof(abs_k_Cart[1]);
	almost(md"""
		You generated a matrix, but the elements are of type *$(mytype)*. Instead, they should be of type `Float64`.
		""")
elseif size(abs_k_Cart) != size(abs_k_Cart)
	keep_working(md"`abs_kspace` has the wrong size.")
elseif abs_k_Cart == abs.(k_Cart)
	correct(md"Well done! A Julia Pro in the making! Here is a plot of the absolute value the k-space data (logarithmically):")
else
	keep_working(md"Still something not quite right...")
end

# ╔═╡ a05cce2b-b1b7-4d8c-a823-137b90473d70
begin
	bp_error = round(calculate_error_norm(reco_bp, image_org, k_mask) * 100)/100
	if !@isdefined(reco_bp)
		not_defined(:reco_bp)
	elseif ismissing(reco_bp)
		still_missing()
	elseif all(density_compensation .== 1)
		still_missing(md"modify `density_compensation` to correct for the sampling density of the radial k-space trajectory.")
	elseif !(reco_bp isa Matrix)
		keep_working(md"`reco` should be a `Matrix`.")
	elseif eltype(reco_bp) != ComplexF64
		type_reco = typeof(reco_bp[1]);
		almost(md"""
			You generated a matrix, but the elements are of type *$(type_reco)*. Instead, they 		should be of type `ComplexF64`.
			""")
	elseif size(reco_bp) != size(image_org)
		keep_working(md"`reco` does not have the correct size. I should be the same as `image_org`")
	elseif bp_error > 0.25
		keep_working(md"You are getting an image, but it's not quite looking like `image_org`: the normalized RMSE $bp_error. It should be below 0.25")
	else
		correct(md"Well done! You master the art of filtered back-projection! The normalized RMSE is $bp_error, that is about as good as it gets with filtered backprojection")
	end
end

# ╔═╡ a491a997-0a95-491a-8ca2-1ff85f5f39d3
begin
	if !@isdefined(A)
		not_defined(:A)
	elseif !@isdefined(b)
		not_defined(:b)
	elseif !@isdefined(reco_cg) || ismissing(A) || ismissing(b) || A isa Matrix{Float64}
		still_missing(md"define `A` and `b` correctly!")
	elseif !(A isa LinearOperator{ComplexF64})
		keep_working(md"`A` should be a `LinearOperator{ComplexF64}`.")
	elseif !(b isa Vector{ComplexF64})
		keep_working(md"`b` should be a `Vector{ComplexF64}`.")
	elseif size(reco_cg) != size(image_org)
		keep_working(md"`reco_cg` does not have the correct size. I should be the same as `image_org`")
	else
		error_cg = round(calculate_error(reco_cg, image_org, k_mask) * 10000) / 10000
		if error_cg > 0.003
			keep_working(md"You are getting an image, but it's not quite looking like `image_org`. The normalized RMSE is $error_cg, but it should be below 0.003.")
		else
			correct(md"Well done! You master the art of iterative reconstructions! The normalized RMSE is $error_cg. ")
		end
	end
end

# ╔═╡ 2b0fb5c3-4f96-4559-a28c-dac0de27a328
begin
	if F_spiral_or isa NFFTOp
		still_missing(md"Replace the `NFFTOp` operator with an `FieldmapNFFTOp` operator!")
	elseif !(F_spiral_or isa FieldmapNFFTOp)
		keep_working(md"`F_spiral_or` should be of type `FieldmapNFFTOp`.")
	else
		error_or = round(calculate_error(reco_or, image_org, k_mask) * 10000) / 10000
		if error_or > 0.04
			keep_working(md"Not quiet. The normalized RMSE is $error_or, but it should be below 0.04.")
		else
			correct(md"Well done! Off-resonance artifacts are history! The normalized RMSE is $error_or. Note that it is larger compared to the reconstruction without any off-resonance, which is $error_cg with the current settings (Excercises 2-3). This is a result of the additional approximations made in the `FieldmapNFFTOp`-operator for speed purposes, as discussed in the lecture. However, spatial off-resonance variations can result in a loss of spatial information that we cannot correct for!")
		end
	end
end

# ╔═╡ ed4ef764-e377-4af3-b3a5-86aa442396e8
if !@isdefined(ic_answer)
	not_defined(ic_answer)
elseif ic_answer == "undefined"
	still_missing(md"Choose an answer!")
elseif ic_answer == "green"
	keep_working(md"Nice weather solves almost everything, but not an inverse crime!")
elseif ic_answer == "teacher"
	almost(md"Science is the belief in the ignorance of experts - Richard P. Feynman")
else
	correct(md"What you have discovered in called an *inverse crime*. With `kernelSize=1` we essentially perform nearest-neighbor interpolation. Thus, we just copy the Cartesian k-space data point onto the non-Cartesian grid without changing it, and we copy the same data point back to the Cartesian grid in the reconstruction. The approximation errors cancel each other out and we get great results that have little to do with reality.

To avoid this problem, we should always simulate the forward problem as accurately as possible and use approximations only in the reconstruction. Here, we can simulate the non-Cartesian k-space data *by foot* with the slow discrete Fourier transform:
		")

end

# ╔═╡ 4dd0670b-d84e-4f7f-bfda-ec0628d9ed51
bigbreak = html"

";

# ╔═╡ 4bccda27-d96b-4234-b5e4-933422b4a157
bigbreak

# ╔═╡ 91f4778e-ee20-11ea-1b7e-2b0892bd3c0f
bigbreak

# ╔═╡ Cell order:
# ╟─ac8ff080-ed61-11ea-3650-d9df06123e1f
# ╟─a24b4518-e6c3-4845-8b32-b203af6a2d69
# ╟─aff9b50d-d013-4c8b-8510-9a59e0c9fc2e
# ╟─5f95e01a-ee0a-11ea-030c-9dba276aba92
# ╠═65780f00-ed6b-11ea-1ecf-8b35523a7ac0
# ╟─540ccfcc-ee0a-11ea-15dc-4f8120063397
# ╠═26709e4b-1675-4ee6-a436-bf9a85683950
# ╟─97a4d58f-9c32-4de9-9142-31e34fc4e33b
# ╠═d72dc575-6d2e-489c-b512-aeae6d017f7f
# ╟─9391828f-041c-49db-93d8-3483bb672f3d
# ╟─67bedf95-8ff4-430f-8737-a0d21ff368c8
# ╟─eeed0b36-37be-4af4-a8c2-68c595d6bb2c
# ╟─1adbae4f-c732-457b-aad8-bc0f3a7be13d
# ╟─c6be8293-7f7c-48a7-b112-92f7ae4fb7c8
# ╠═35022e16-d376-426c-a9f8-b71760bcaaf7
# ╟─40a94bcd-a0cc-4b1a-b370-6405adc0f92b
# ╠═c0afaa73-aa46-4f61-a5b4-c49d071f4fb6
# ╟─810b38a1-490f-4e22-8edb-0e50a8af2f28
# ╟─8a3dc74a-4b20-4c0f-96e8-1ab0451a7da3
# ╟─aefd5a3d-9b7d-40a1-9cc8-1a5b555d251d
# ╟─515ae3f8-fb28-49b9-a023-38d143b6674a
# ╟─ad6a33b0-eded-11ea-324c-cfabfd658b56
# ╠═28668257-321b-41f2-bfe2-d5f35441ec8d
# ╟─137ed83c-171f-48df-be8e-dc538ac18929
# ╠═04962733-71dd-4e6f-9e23-5b29ee821d68
# ╟─74d30291-506f-4e57-9241-c791f0fb8264
# ╠═319872dd-d646-4029-9319-4ffb957179a0
# ╟─4ded3325-4935-4843-b9c2-b0e29469bf3e
# ╟─c0f04553-6031-4984-8c80-1260022eb23f
# ╠═beccb301-1c40-4c41-8960-462a82eca3ad
# ╟─05dafc16-c68f-47e8-8bdf-cfea8736bc94
# ╠═f6d6fe5e-3381-4338-817d-366194b3f7ba
# ╟─424e1625-e14b-4fcb-bd5b-7e426c97d330
# ╠═16ad1b20-f3c8-4516-b5d9-6150875c326c
# ╟─a05cce2b-b1b7-4d8c-a823-137b90473d70
# ╟─7906f103-65db-45c6-9d95-ef1483048c14
# ╟─c604ce61-08ac-46af-b231-a99e37f52bd6
# ╟─9aff92eb-8094-4a83-ad27-242277b9aee1
# ╠═b114e592-0d39-43cb-aee3-c0948b2d2f88
# ╟─28c1a017-8977-40a8-91d2-1cac83a0d3cf
# ╟─179e6a9d-3728-4bde-8bae-5b6ac4490a74
# ╟─b62f8883-02cd-4ab0-9c27-61283e4839f6
# ╟─bfa47923-d80b-43ea-a21e-f3dad07b4958
# ╟─05d84e43-d768-4680-82eb-d3a5419edd85
# ╟─2af0c856-430b-41ec-a804-77b3367d1fc5
# ╟─94ca0194-811f-472f-8815-1f52c6a1619c
# ╟─9d70e2e6-6e45-4662-9792-b36ab36c35b4
# ╠═32bc54d1-4ea8-4663-ba9a-03d357eb8c02
# ╠═51b2f56b-3880-4ba6-828a-104334bad1ec
# ╟─698d0344-ec5d-4c27-a2b1-b52825e398d6
# ╠═57a2c882-760c-4ff9-9bcb-be8f358d7b8a
# ╟─a491a997-0a95-491a-8ca2-1ff85f5f39d3
# ╟─93f84f84-fdf3-41fa-88f8-2a26b8d6982a
# ╟─33b23b5c-85e7-49ba-9e4a-27d67943ee73
# ╟─ce933bda-ac45-46ec-b577-c4168b1ce569
# ╟─f6447d81-6612-428d-b3d7-acd7f8756df5
# ╟─98a1c04c-58a3-44f6-995f-5d4d81095576
# ╟─e6af6037-7eb6-40be-a1b9-846ab1f50b9c
# ╟─8075e80d-b6df-435a-ac4d-44a919909323
# ╟─e4d80b4c-faa3-4514-876b-63bce1ece81f
# ╟─62934cdb-526d-4f97-afc0-20bcd9f9714e
# ╟─ed4ef764-e377-4af3-b3a5-86aa442396e8
# ╠═aa209a90-0dfb-41ff-bf0e-a26a6c1b8c8b
# ╟─a0e0fba0-39e1-4873-8d7e-ee5cb793062c
# ╠═86af3791-f3e6-45bc-bc4e-072803491d1f
# ╟─843dcc64-3015-4682-bae7-ddd86a8615fa
# ╟─ac4504d8-ed24-4972-beff-e518a820dc83
# ╟─9ce17af6-6aff-4615-933e-5e67b56d1038
# ╟─5a9f1bf5-2f04-4a13-859e-0e09957fe20c
# ╟─c04477fc-ab75-4806-9045-50f71f45c707
# ╠═a8339197-b1b5-400f-89f7-0daa87e5bf2c
# ╟─84c928bd-4dd4-468c-a62e-7347bcede7fd
# ╠═0368a211-0e20-4841-9e93-ba722b7b42f1
# ╠═c393b9a8-2f4c-40e5-adc8-4605655dff89
# ╟─d2936afa-b271-464e-b7ef-37f33e7306dc
# ╠═605ee395-e123-4128-9b56-c73a128572ec
# ╟─5a58c91c-1383-4c17-9341-fda706eb0c1a
# ╠═882de061-061c-4538-b1b4-5aba0a385569
# ╟─1fbbf883-cf0d-461a-98f0-a934395e276a
# ╠═bfd13cd5-87ba-4a4a-b59b-2f06fcb2055c
# ╟─123f7947-71ef-4274-b0fc-25ae6d6b1396
# ╟─cfebce07-d5ee-4df4-97c4-28bc1576ccfe
# ╟─51ef7e32-cf88-405e-8b13-2cae310b93ca
# ╠═1ea1974a-cebb-4c05-a885-89058a34cf15
# ╠═7cf1b6fc-2489-4010-aab0-dc973937a704
# ╟─814283c0-99e9-4d9e-929f-6aa721c28aee
# ╠═a49fa464-56f8-4c4a-8e14-00977cb033a6
# ╟─2b0fb5c3-4f96-4559-a28c-dac0de27a328
# ╠═8b894d85-2dcd-4927-9e0d-91eb93c8f3dd
# ╟─220675f2-9ef6-4187-bdd9-1dbebb8af4ec
# ╟─8233b774-7363-463b-9c27-13c2fecb9e25
# ╟─4bccda27-d96b-4234-b5e4-933422b4a157
# ╟─91f4778e-ee20-11ea-1b7e-2b0892bd3c0f
# ╟─756d150a-b7bf-4bf5-b372-5b0efa80d987
# ╟─4bc94bec-da39-4f8a-82ee-9953ed73b6a4
# ╟─bfa4b5ac-94b6-4edc-8de8-03b6ac184754
# ╟─8ce6ad06-819c-4af5-bed7-56ecc08c97be
# ╟─dfa40e89-03fc-4a7a-825e-92d67ee217b2
# ╟─086ec1ff-b62d-4566-9973-5b2cc3353409
# ╟─c22f688b-dc04-4a94-b541-fe06266c5446
# ╟─ab3d1b70-88e8-4118-8d3e-601a8a68f72d
# ╟─4dd0670b-d84e-4f7f-bfda-ec0628d9ed51


using Parameters

function rank_matrix(x::Array{Float64, 2}, nt_args::NamedTuple)
    @unpack n_draws, n_pars, n_chains = nt_args
    xc = zeros(Int, n_draws, n_pars, n_chains)
    for i in 1:n_pars
        ranked = Int.(ordinalrank(x[:, i]))
        for j in 1:n_chains
            start_draw = (j-1) * n_draws + 1
            last_draw = j * n_draws
            xc[:, i, j] = ranked[start_draw:last_draw]
        end
    end
    xc
end

function trankplot(nt::NamedTuple, param::Symbol;
        bins=40, n_draws=4000, n_chains=4, n_eff=0)

   if length(size(nt[param])) > 1
        params = [Symbol("$(param)[$i]") for i in 1:size(nt[param], 1)]
        n_pars = length(params)
    else
        params = [param]
        n_pars = 1
    end

    nt_args = (n_draws=n_draws, n_pars=n_pars, n_chains=n_chains)
    if length(params) > 1
        ranks = rank_matrix(Matrix(nt[param]'), nt_args)
    else
        ranks = rank_matrix(reshape(nt[param], length(nt[param]), 1), nt_args)
    end

    figs = Vector{Plots.Plot{Plots.GRBackend}}(undef, length(params))
    for i in 1:length(params)
        figs[i] = stephist(ranks[:,1,1]; bins, alpha=0.9, lab="Chain 1", leg=:bottomright)
        for i in 2:n_chains
            stephist!(ranks[:,1,i]; bins, alpha=0.9, lab="Chain $i", leg=:bottomright)
        end
        title!("Trankplot for $(params[i])")
        annotate!([(30, 0, Plots.text("n_eff = $(n_eff[i])", 10, :darkred, :left))])
    end
    figs
end

export
    rank_matrix,
    trankplot


include("utils.jl")

@testset "Euclidean" begin
    E = Manifolds.Euclidean(3)
    EM = Manifolds.MetricManifold(E,Manifolds.EuclideanMetric())
    @test is_default_metric(EM) == Val{true}()
    @test is_default_metric(E,Manifolds.EuclideanMetric()) == Val{true}()
    x = zeros(3)
    @test det_local_metric(EM,x) == one(eltype(x))
    @test log_local_metric_density(EM,x) == zero(eltype(x))
    @test project_point!(E,x) == x

    manifolds = [ E, EM ]
    types = [Vector{Float64},
             SizedVector{3, Float64},
             MVector{3, Float64},
             Vector{Float32},
             SizedVector{3, Float32},
             MVector{3, Float32},
             Vector{Double64},
             MVector{3, Double64},
             SizedVector{3, Double64}]
    for M in manifolds
        for T in types
            @testset "$M Type $T" begin
                pts = [convert(T, [1.0, 0.0, 0.0]),
                       convert(T, [0.0, 1.0, 0.0]),
                       convert(T, [0.0, 0.0, 1.0])]
                test_manifold(M,
                              pts,
                              test_reverse_diff = isa(T, Vector),
                              test_project_tangent = true,
                              test_musical_isomorphisms = true,
                              test_vector_transport = true,
                              test_mutating_rand = isa(T, Vector),
                              point_distributions = [Manifolds.projected_distribution(M, Distributions.MvNormal(zero(pts[1]), 1.0))],
                              tvector_distributions = [Manifolds.normal_tvector_distribution(M, pts[1], 1.0)])
            end
        end
    end
end


include("../Tests/Tests.jl")
include("utilities.jl")

using Test, JuMP

@testset "Term" begin
    include("term.jl")
end
@testset "Term fixed" begin
    include("term_fixed.jl")
end
@testset "Quadratic" begin
    include("quadratic.jl")
end
@testset "Choi" begin
    include("choi.jl")
end
@testset "Concave then convex cubic" begin
    include("concave_then_convex_cubic.jl")
end
@testset "Horn" begin
    include("horn.jl")
end
@testset "Lyapunov Switched System" begin
    include("lyapunov_switched_system.jl")
end
@testset "Motzkin" begin
    include("motzkin.jl")
end
@testset "BPT12e399" begin
    include("BPT12e399.jl")
end
@testset "Max Cut" begin
    include("maxcut.jl")
end
@testset "Chebyshev" begin
    include("chebyshev.jl")
end
@testset "Quartic comparison" begin
    include("quartic_comparison.jl")
end
@testset "SOSDEMO9" begin
    include("sosdemo9.jl")
end
@testset "SOSDEMO10" begin
    include("sosdemo10.jl")
end


include("header.jl")

@testset "rosenbrock" begin
    # info("Test rosenbrock with map...")
    b0(x) = sum(map((i, j) -> (1 - j)^2 + 10*(i - j^2)^2, x[2:end], x[1:end-1]))
    b1 = grad(b0)
    b1sum(x)=sum(b1(x))
    @test gradcheck(b0,randn(2))
    @test gradcheck(b0,(randn(2)...,))
    @test gradcheck(b0,randn(3))
    @test gradcheck(b0,(randn(3)...,))
    @test gradcheck(b1sum,randn(2))
    @test gradcheck(b1sum,(randn(2)...,))
    @test gradcheck(b1sum,randn(3))
    @test gradcheck(b1sum,(randn(3)...,))
    @time b1sum(randn(10000))
end

nothing


@testset "$TEST $G" begin

g5 = DG(4)
add_edge!(g5,1,2); add_edge!(g5,2,3); add_edge!(g5,1,3); add_edge!(g5,3,4)

@test degree_centrality(g5) == [0.6666666666666666, 0.6666666666666666, 1.0, 0.3333333333333333]
@test in_degree_centrality(g5, normalize=false) == [0.0, 1.0, 2.0, 1.0]
@test out_degree_centrality(g5; normalize=false) == [2.0, 1.0, 1.0, 0.0]

end # testset


# Utility routines

"""
    angle_difference(α, β, degrees::Bool=true) -> Δ

Return the angular difference `Δ` between two angles `α` and `β`, in the direction
α → β (i.e., Δ = (β - α)).  This means that Δ is positive if the direction β is
clockwise of α, and negative otherwise.

Angles are assumed to be degreees unless `degrees` is `false`.
"""
function angle_difference(a::Real, b::Real, degrees::Bool=true)
    whole_circ = degrees ? 360. : 2π
    half_circ = whole_circ/2.
    mod(b - a + half_circ, whole_circ) - half_circ
end

"""
    flip_component!(s::SACtr) -> s

Flip a component so that it points in the opposite direction.
"""
function flip_component!(s::SACtr)
    s[:cmpaz] = mod(s[:cmpaz] + 180, 360)
    s[:cmpinc] = 180 - s[:cmpinc]
    s[:kcmpnm] = sacstring(round(s[:cmpaz], digits=2, base=10))
    multiply!(s, -1)
end

"""
    linear_regression(x, y)

Perform simple linear regression using Ordinary Least Squares. Returns `a` and `b` such
that `a + b*x` is the closest straight line to the given points `(x, y)`, i.e., such that
the squared error between `y` and `a + b*x` is minimized.
"""
function linear_regression(x::AbstractVector, y::AbstractVector)
    size(x) == size(y) || throw(DimensionMismatch("x and y must be the same size"))
    mx, my = mean(x), mean(y)
    b = covm(x, mx, y, my)/varm(x, mx)
    a = my - b*mx
    return a, b
end

"""
    isundefined(x) -> ::Bool

If the SAC value `x` is undefined, return `true`.

    isundefined(s::SACtr, x::Symbol) -> :: Bool

Test whether header `x` is undefined for trace `s`.
"""
isundefined(x::SACFloat) = x == sac_rnull
isundefined(x::SACInt) = x == sac_inull
isundefined(x::SACChar) = strip(x) == strip(sac_cnull)
isundefined(x::SACBool) = false
isundefined(s::SACtr, x::Symbol) = isundefined(getfield(s, x))

"Directory containing sample data"
const SAMPLE_DATA_DIR = joinpath(dirname(@__FILE__()), "..", "data")

"""
    sample() -> ::SACtr
    sample(kind::Symbol) -> ::Array{SACtr}

Return some sample SAC data.

With no arguments, `sample` gives one trace from a local earthquake recorded in
California.

In the second form, a set of traces is returned according to the table below:

|`kind`|Description|
|:-----|:----------|
|`:local`|Livermore Valley, CA.  9 3-component stations|
|`:regional`|Nevada.  4 3-component stations|
|`:teleseism`|**Mid-period** recording of Eureka, CA event.  4 3-c stations|
|`:teleseisl`|**Long-period** recording of Eureka, CA event.  4 3-c stations|
|`:array`|Deep Fiji event.  60 vertical stations in the UK|
"""
function sample()
    file = joinpath(SAMPLE_DATA_DIR, "seis.sac")
    SAC.read(file)
end
const SAMPLE_DATA_KINDS = [:local, :regional, :teleseism, :teleseisl, :array]
function sample(kind::Symbol)
    file_pattern = "*"
    kind in SAMPLE_DATA_KINDS ||
        throw(ArgumentError("Sample data kind must be one of: $SAMPLE_DATA_KINDS"))
    dir = joinpath(SAMPLE_DATA_DIR, string(kind))
    if kind in [:teleseism, :teleseisl]
        file_pattern = "*" * dir[end:end] * ".*"
        dir = dir[1:end-1]
    end
    A, filenames = read_wild(file_pattern, dir, echo=false)
    A[:kcmpnm] = [split(filename, ".")[end] for filename in filenames]
    A
end

"""
    swap_traces!(s1, s2) -> s2, s1

Swap the contents of two traces, `s1` and `s2`, by an in-place transposition of all
header and trace values.
"""
function swap_traces!(s1::SACtr, s2::SACtr)
    for f in fieldnames(s1)
        s1[f], s2[f] = s2[f], s1[f]
    end
    s1, s2
end

"""
    traces_are_orthogonal(s1::SACtr, s2::SACtr, tol=eps(SACFloat)) -> ::Bool

Return `true` if the two traces `s1` and `s2` have component azimuths
90° apart.
"""
traces_are_orthogonal(s1::SACtr, s2::SACtr, tol=eps(SACFloat)) =
    isapprox(SACFloat(abs(angle_difference(s1[:cmpaz], s2[:cmpaz]))), SACFloat(90), atol=tol)

"""
    time(::SACtr) -> t

Return a FloatRange `t` which contains the times for each sample of the SAC trace.
"""
time(s::SACtr) = s.b .+ (0:s.npts-1)*s.delta


of_mean_type(p::Colorant) = ImageSegmentation.meantype(typeof(p))(p)

@testset "Seeded Region Growing" begin
    # 2-D image
    img = zeros(Gray{N0f8}, 10, 10)
    img[6:10,4:8] .= 0.5
    img[3:7,2:6] .= 0.8
    seeds = [(CartesianIndex(3,9),1), (CartesianIndex(5,2),2), (CartesianIndex(9,7),3)]

    expected = ones(Int, 10, 10)
    expected[6:10,4:8] .= 3
    expected[3:7,2:6] .= 2
    expected_labels = [1,2,3]
    expected_means = Dict(1 => of_mean_type(img[3,9]), 2 => of_mean_type(img[5,2]), 3 => of_mean_type(img[9,7]))
    expected_count = Dict(1 => 56, 2 => 25, 3 => 19)

    seg = seeded_region_growing(img, seeds)
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # Custom neighbourhood using a function
    seg = seeded_region_growing(img, seeds, c->[CartesianIndex(c[1]-1,c[2]), CartesianIndex(c[1]+1,c[2]), CartesianIndex(c[1],c[2]-1), CartesianIndex(c[1],c[2]+1)])
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # Offset image
    img = centered(img)
    seeds = [(CartesianIndex(-2,4),1), (CartesianIndex(0,-3),2), (CartesianIndex(4,2),3)]
    expected = centered(expected)
    seg = seeded_region_growing(img, seeds)
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # Custom neighbourhood using a [3,3] vs [5,5] kernel
    img = zeros(Gray{N0f8}, 5, 5)
    img[2:4,2:4] .= 1
    img[3,3] = 0
    seeds = [(CartesianIndex(3,3),1), (CartesianIndex(2,3),2)]

    expected = fill(2,(5,5))
    expected[3,3] = 1
    expected_labels = [1,2]
    expected_means = Dict(1=>Gray{Float64}(0.0), 2=>Gray{Float64}(1/3))
    expected_count = Dict(1=>1, 2=>24)

    seg = seeded_region_growing(img, seeds, (3,3))
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test all(label->(expected_means[label] ≈ seg.segment_means[label]), seg.segment_labels)
    @test seg.image_indexmap == expected

    expected = ones(Int, 5, 5)
    expected[2:4,2:4] .= 2
    expected[3,3] = 1
    expected_labels = [1,2]
    expected_means = Dict(1=>Gray{N0f8}(0.0), 2=>Gray{N0f8}(1.0))
    expected_count = Dict(1=>17, 2=>8)

    seg = seeded_region_growing(img, seeds, (5,5))
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # element-type and InexactError
    img = zeros(Int, 4, 4)
    img[2:end, 2:end] .= 10
    img[end,end] = 11
    seeds = [(CartesianIndex(1,1), 1), (CartesianIndex(2,2), 2)]

    expected = ones(Int, size(img))
    expected[2:end,2:end] .= 2
    expected_labels = [1,2]
    expected_means = Dict(1=>0.0, 2=>91/9)
    expected_count = Dict(1=>7, 2=>9)

    seg = seeded_region_growing(img, seeds, (3,3), (v1,v2)->(δ = abs(v1-v2); return δ > 1 ? δ : zero(δ)))
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # 3-d image
    img = zeros(RGB{N0f8},(9,9,9))
    img[3:7,3:7,3:7] .= RGB{N0f8}(0.5,0.5,0.5)
    img[2:5,5:9,4:6] .= RGB{N0f8}(0.8,0.8,0.8)
    seeds = [(CartesianIndex(1,1,1),1), (CartesianIndex(6,4,4),2), (CartesianIndex(3,6,5),3)]

    expected = ones(Int, (9,9,9))
    expected[3:7,3:7,3:7] .= 2
    expected[2:5,5:9,4:6] .= 3
    expected_labels = [1,2,3]
    expected_means = Dict([(i, of_mean_type(img[seeds[i][1]])) for i in 1:3])
    expected_count = Dict(1=>571, 2=>98, 3=>60)

    seg = seeded_region_growing(img, seeds)
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # custom diff_fn
    img = zeros(RGB{N0f8},(3,3))
    img[1:3,1] .= RGB{N0f8}(0.4,1,0)
    img[1:3,2] .= RGB{N0f8}(0.2,1,0)
    seeds = [(CartesianIndex(2,1),1), (CartesianIndex(2,3),2)]

    expected = ones(Int, (3,3))
    expected[1:3,3] .= 2
    expected_labels = [1,2]
    expected_means = Dict(1=>RGB{Float64}(0.3,1.0,0.0), 2=>RGB{Float64}(0.0,0.0,0.0))
    expected_count = Dict(1=>6, 2=>3)

    seg = seeded_region_growing(img, seeds)
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    expected = ones(Int, (3,3))
    expected[1:3,2] .= 0
    expected[1:3,3] .= 2
    expected_labels = [0,1,2]
    expected_means = Dict(1=>RGB{Float64}(0.4,1.0,0.0), 2=>RGB{Float64}(0.0,0.0,0.0))
    expected_count = Dict(0=>3, 1=>3, 2=>3)

    seg = seeded_region_growing(img, seeds, [3,3], (c1,c2)->abs(of_mean_type(c1).r - of_mean_type(c2).r))
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected
end

@testset "Unseeded Region Growing" begin
    # 2-D image
    img = zeros(Gray{N0f8}, 10, 10)
    img[6:10,4:8] .= 0.5
    img[3:7,2:6] .= 0.8

    expected = ones(Int, 10, 10)
    expected[6:10,4:8] .= 2
    expected[3:7,2:6] .= 3
    expected_labels = [1,2,3]
    expected_means = Dict(1 => of_mean_type(img[3,9]), 3 => of_mean_type(img[5,2]), 2 => of_mean_type(img[9,7]))
    expected_count = Dict(1 => 56, 3 => 25, 2 => 19)

    seg = unseeded_region_growing(img, 0.2)
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # Custom neighbourhood using a function
    seg = unseeded_region_growing(img, 0.2, c->[CartesianIndex(c[1]-1,c[2]), CartesianIndex(c[1]+1,c[2]), CartesianIndex(c[1],c[2]-1), CartesianIndex(c[1],c[2]+1)])
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # Offset image
    img = centered(img)
    expected = centered(expected)
    seg = unseeded_region_growing(img, 0.2)
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # Custom neighbourhood using a [5,5] kernel and varying threshold
    img = zeros(Gray{N0f8}, 5, 5)
    img[2:4,2:4] .= 1
    img[3,3] = 0.8

    expected = fill(1,(5,5))
    expected[2:4,2:4] .= 2
    expected_labels = [1,2]
    expected_means = Dict(1=>Gray{Float64}(0.0), 2=>Gray{Float64}(8.8/9))
    expected_count = Dict(2=>9, 1=>16)

    seg = unseeded_region_growing(img, 0.2, (5,5))
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test all(label->(expected_means[label] ≈ seg.segment_means[label]), seg.segment_labels)
    @test seg.image_indexmap == expected

    expected = ones(Int, 5, 5)
    expected[2:4,2:4] .= 3
    expected[3,3] = 2
    expected_labels = [1,2,3]
    expected_means = Dict(1=>Gray{N0f8}(0.0), 3=>Gray{N0f8}(1.0), 2=>Gray{N0f8}(0.8))
    expected_count = Dict(1=>16, 2=>1, 3=>8)

    seg = unseeded_region_growing(img, 0.1, (5,5))
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    expected = ones(Int, 5, 5)
    expected[2:4,2:4] .= 2
    expected[3,3] = 3
    expected_labels = [1,2,3]
    expected_means = Dict(1=>Gray{N0f8}(0.0), 2=>Gray{N0f8}(1.0), 3=>Gray{N0f8}(0.8))
    expected_count = Dict(1=>16, 3=>1, 2=>8)

    seg = unseeded_region_growing(img, 0.1, (3,3))
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # 3-d image
    img = zeros(RGB{N0f8},(9,9,9))
    img[3:7,3:7,3:7] .= RGB{N0f8}(0.5,0.5,0.5)
    img[2:5,5:9,4:6] .= RGB{N0f8}(0.8,0.8,0.8)

    expected = ones(Int, (9,9,9))
    expected[3:7,3:7,3:7] .= 2
    expected[2:5,5:9,4:6] .= 3
    expected_labels = [1,2,3]
    expected_means = Dict(1=>of_mean_type(img[1,1,1]), 2=>of_mean_type(img[3,3,3]), 3=>of_mean_type(img[2,5,4]))
    expected_count = Dict(1=>571, 2=>98, 3=>60)

    seg = unseeded_region_growing(img, 0.2)
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

    # custom diff_fn
    img = zeros(RGB{N0f8},(3,3))
    img[1:3,1] .= RGB{N0f8}(0.4,1,0)
    img[1:3,2] .= RGB{N0f8}(0.2,1,0)

    expected = ones(Int, (3,3))
    expected[1:3,2] .= 2
    expected[1:3,3] .= 3
    expected_labels = [1,2,3]
    expected_means = Dict(1=>of_mean_type(img[1,1]), 3=>of_mean_type(img[1,3]), 2=>of_mean_type(img[1,2]))
    expected_count = Dict(1=>3, 2=>3, 3=>3)

    seg = unseeded_region_growing(img, 0.2, [3,3], (c1,c2)->abs(of_mean_type(c1).r - of_mean_type(c2).r))
    @test all(label->(label in expected_labels), seg.segment_labels)
    @test all(label->(label in seg.segment_labels), expected_labels)
    @test expected_count == seg.segment_pixel_count
    @test expected_means == seg.segment_means
    @test seg.image_indexmap == expected

end


@testset "Utils" begin
    orthants = Orthants(3)

    @test length(orthants) == 8
    eltype(Orthants) == Vector{Int}
    @test first(orthants) == [1, 1, 1]
    @test last(orthants) == [-1, -1, -1]
    @test orthants[4] == [-1, -1, 1]
end


## R code 11.4
m11.1 <- quap(
    alist(
        pulled_left ~ dbinom( 1 , p ) ,
        logit(p) <- a ,
        a ~ dnorm( 0 , 10 )
    ) , data=d )

## R code 11.7
m11.2 <- quap(
    alist(
        pulled_left ~ dbinom( 1 , p ) ,
        logit(p) <- a + b[treatment] ,
        a ~ dnorm( 0 , 1.5 ),
        b[treatment] ~ dnorm( 0 , 10 )
    ) , data=d )
prior <- extract.prior( m11.2 , n=1e4 )
p <- sapply( 1:4 , function(k) inv_logit( prior$a + prior$b[,k] ) )

## R code 11.9
m11.3 <- quap(
    alist(
        pulled_left ~ dbinom( 1 , p ) ,
        logit(p) <- a + b[treatment] ,
        a ~ dnorm( 0 , 1.5 ),
        b[treatment] ~ dnorm( 0 , 0.5 )
    ) , data=d )

## R code 11.11
m11.4 <- ulam(
    alist(
        pulled_left ~ dbinom( 1 , p ) ,
        logit(p) <- a[actor] + b[treatment] ,
        a[actor] ~ dnorm( 0 , 1.5 ),
        b[treatment] ~ dnorm( 0 , 0.5 )
    ) , data=dat_list , chains=4 , log_lik=TRUE )

## R code 11.19
dat_list2 <- list(
    pulled_left = d$pulled_left,
    actor = d$actor,
    side = d$side,
    cond = d$cond )
m11.5 <- ulam(
    alist(
        pulled_left ~ dbinom( 1 , p ) ,
        logit(p) <- a[actor] + bs[side] + bc[cond] ,
        a[actor] ~ dnorm( 0 , 1.5 ),
        bs[side] ~ dnorm( 0 , 0.5 ),
        bc[cond] ~ dnorm( 0 , 0.5 )
    ) , data=dat_list2 , chains=4 , log_lik=TRUE )

## R code 11.25
dat <- with( d_aggregated , list(
    left_pulls = left_pulls,
    treatment = treatment,
    actor = actor,
    side = side,
    cond = cond ) )

m11.6 <- ulam(
    alist(
        left_pulls ~ dbinom( 18 , p ) ,
        logit(p) <- a[actor] + b[treatment] ,
        a[actor] ~ dnorm( 0 , 1.5 ) ,
        b[treatment] ~ dnorm( 0 , 0.5 )
    ) , data=dat , chains=4 , log_lik=TRUE )

## R code 11.29
dat_list <- list(
    admit = d$admit,
    applications = d$applications,
    gid = ifelse( d$applicant.gender=="male" , 1 , 2 )
)
m11.7 <- ulam(
    alist(
        admit ~ dbinom( applications , p ) ,
        logit(p) <- a[gid] ,
        a[gid] ~ dnorm( 0 , 1.5 )
    ) , data=dat_list , chains=4 )

## R code 11.32
dat_list$dept_id <- rep(1:6,each=2)
m11.8 <- ulam(
    alist(
        admit ~ dbinom( applications , p ) ,
        logit(p) <- a[gid] + delta[dept_id] ,
        a[gid] ~ dnorm( 0 , 1.5 ) ,
        delta[dept_id] ~ dnorm( 0 , 1.5 )
    ) , data=dat_list , chains=4 , iter=4000 )

## R code 11.45
dat <- list(
    T = d$total_tools ,
    P = d$P ,
    cid = d$contact_id )

# intercept only
m11.9 <- ulam(
    alist(
        T ~ dpois( lambda ),
        log(lambda) <- a,
        a ~ dnorm(3,0.5)
    ), data=dat , chains=4 , log_lik=TRUE )

# interaction model
m11.10 <- ulam(
    alist(
        T ~ dpois( lambda ),
        log(lambda) <- a[cid] + b[cid]*P,
        a[cid] ~ dnorm( 3 , 0.5 ),
        b[cid] ~ dnorm( 0 , 0.2 )
    ), data=dat , chains=4 , log_lik=TRUE )

## R code 11.49
dat2 <- list( T=d$total_tools, P=d$population, cid=d$contact_id )
m11.11 <- ulam(
    alist(
        T ~ dpois( lambda ),
        lambda <- exp(a[cid])*P^b[cid]/g,
        a[cid] ~ dnorm(1,1),
        b[cid] ~ dexp(1),
        g ~ dexp(1)
    ), data=dat2 , chains=4 , log_lik=TRUE )

## R code 11.53
# compute the offset
d$log_days <- log( d$days )

# fit the model
m11.12 <- quap(
    alist(
        y ~ dpois( lambda ),
        log(lambda) <- log_days + a + b*monastery,
        a ~ dnorm( 0 , 1 ),
        b ~ dnorm( 0 , 1 )
    ), data=d )

## R code 11.56
code_m11.13 <- "
data{
    int N; // number of individuals
    int K; // number of possible careers
    int career[N]; // outcome
    vector[K] career_income;
}
parameters{
    vector[K-1] a; // intercepts
    real<lower=0> b; // association of income with choice
}
model{
    vector[K] p;
    vector[K] s;
    a ~ normal( 0 , 1 );
    b ~ normal( 0 , 0.5 );
    s[1] = a[1] + b*career_income[1];
    s[2] = a[2] + b*career_income[2];
    s[3] = 0; // pivot
    p = softmax( s );
    career ~ categorical( p );
}
"

## R code 11.59
code_m11.14 <- "
data{
    int N; // number of observations
    int K; // number of outcome values
    int career[N]; // outcome
    real family_income[N];
}
parameters{
    vector[K-1] a; // intercepts
    vector[K-1] b; // coefficients on family income
}
model{
    vector[K] p;
    vector[K] s;
    a ~ normal(0,1.5);
    b ~ normal(0,1);
    for ( i in 1:N ) {
        for ( j in 1:(K-1) ) s[j] = a[j] + b[j]*family_income[i];
        s[K] = 0; // the pivot
        p = softmax( s );
        career[i] ~ categorical( p );
    }
}
"

dat_list <- list( N=N , K=3 , career=career , family_income=family_income )
m11.14 <- stan( model_code=code_m11.14 , data=dat_list , chains=4 )

## R code 11.61
# binomial model of overall admission probability
m_binom <- quap(
    alist(
        admit ~ dbinom(applications,p),
        logit(p) <- a,
        a ~ dnorm( 0 , 1.5 )
    ), data=d )

# Poisson model of overall admission rate and rejection rate
# 'reject' is a reserved word in Stan, cannot use as variable name
dat <- list( admit=d$admit , rej=d$reject )
m_pois <- ulam(
    alist(
        admit ~ dpois(lambda1),
        rej ~ dpois(lambda2),
        log(lambda1) <- a1,
        log(lambda2) <- a2,
        c(a1,a2) ~ dnorm(0,1.5)
    ), data=dat , chains=3 , cores=3 )

m11.15 <- ulam(
    alist(
        days_to_event|adopted==1 ~ exponential( lambda ),
        days_to_event|adopted==0 ~ custom(exponential_lccdf( !Y | lambda )),
        lambda <- 1.0/mu,
        log(mu) <- a[color_id],
        a[color_id] ~ normal(0,1)
    ), data=dat , chains=4 , cores=4 )


module MultivariateCalculusTest

using VarianceComponentModels, Base.Test

srand(123)

# test vech
@testset "vech" begin
  a = reshape(1:9, 3, 3)
  vecha = vech(a)
  @test vecha[1] == a[1, 1]
  @test vecha[2] == a[2, 1]
  @test vecha[3] == a[3, 1]
  @test vecha[4] == a[2, 2]
  @test vecha[5] == a[3, 2]
  @test vecha[6] == a[3, 3]
  a = 1.0
  @test vech(a) == a
end

# test trilind
@testset "trilind" begin
  n = 3
  A = randn(n, n)
  @test vecnorm(vech(A) - A[trilind(A)]) ≈ 0.0
end

# test triuind
@testset "triuind" begin
  n = 2
  idx = triuind(n, n)
  @test length(idx) == binomial(n + 1, 2)
  @test idx[1] == 1
  @test idx[2] == 3
  @test idx[3] == 4
  idx = triuind(n)
  @test length(idx) == binomial(n + 1, 2)
  @test idx[1] == 1
  @test idx[2] == 3
  @test idx[3] == 4
  idx = triuind(randn(n, n))
  @test length(idx) == binomial(n + 1, 2)
  @test idx[1] == 1
  @test idx[2] == 3
  @test idx[3] == 4
  idx = triuind(randn(n, n), 1)
  @test length(idx) == binomial(n, 2)
  @test idx[1] == 3
end

# test commutation
@testset "commutation" begin
  m, n = 3, 2
  A = randn(m, n)
  @test vecnorm(commutation(m, n) * vec(A) - vec(A')) ≈ 0.0
  @test vecnorm(commutation(A) * vec(A) - vec(A')) ≈ 0.0
  @test vecnorm(spcommutation(m, n) * vec(A) - vec(A')) ≈ 0.0
  @test vecnorm(spcommutation(A) * vec(A) - vec(A')) ≈ 0.0
  B = randn(m, m)
  @test vecnorm(commutation(m) * vec(B) - vec(B')) ≈ 0.0
  @test vecnorm(commutation(eltype(B), m) * vec(B) - vec(B')) ≈ 0.0
  @test vecnorm(spcommutation(eltype(B), m) * vec(B) - vec(B')) ≈ 0.0
end

# test duplication
@testset "duplication" begin
  n = 3
  A = randn(n, n)
  A = 0.5(A + A') # symmetrize
  @test vecnorm(duplication(n) * vech(A) - vec(A)) ≈ 0.0
  @test vecnorm(duplication(A) * vech(A) - vec(A)) ≈ 0.0
  @test vecnorm(spduplication(n) * vech(A) - vec(A)) ≈ 0.0
  @test vecnorm(spduplication(A) * vech(A) - vec(A)) ≈ 0.0
end

# test chol_gradient
@testset "chol_gradient" begin
  n = 3
  A = randn(n, n)
  A = 0.5(A + A') # symmetrize
  L = tril(randn(n, n))
  # calculate gradient wrt L using chol_gradient
  dL1 = chol_gradient(vec(A), L)
  # alternative way to calculate gradient wrt L
  dL2 = 2.0vech((A * L) + L' * A' - diagm(diag(A * L)))
  @test vecnorm(dL1 - dL2) ≈ 0.0
end

# test kron_gradient
@testset "kron_gradient" begin
  n, q, p, r = 2, 3, 4, 5
  X = randn(n, q)
  Y = randn(p, r)
  M = kron(X, Y)
  # calculate gradient wrt X using kron_gradient
  dX1 = kron_gradient(vec(M), Y, n, q)
  # alternative way to calculate gradient wrt X
  dX2 = zeros(X)
  for j = 1:q
    for i = 1:n
      dX2[i, j] = vecdot(Y, M[(i-1)*p+1:i*p, (j-1)*r+1:j*r])
    end
  end
  @test vecnorm(dX1 - vec(dX2)) < 1e-8
end

# test duplication
@testset "kronaxpy" begin
  m, n, p, q = 3, 4, 5, 6
  A = randn(m, n)
  X = randn(p, q)
  Y = zeros(m * p, n * q)
  kronaxpy!(A, X, Y)
  @test vecnorm(Y - kron(A, X)) ≈ 0.0
end

# test bump_diagonal
@testset "bump_diag" begin
  A = randn(3, 3)
  B = copy(A)
  bump_diagonal!(B, 1e-3)
  @test B[1, 1] - A[1, 1] ≈ 1e-3
  @test B[2, 2] - A[2, 2] ≈ 1e-3
  @test B[3, 3] - A[3, 3] ≈ 1e-3
end

# test clamp_diagonal
@testset "clamp_diag" begin
  n = 3
  A = randn(n, n)
  B = copy(A)
  clamp_diagonal!(B, 0.0, 1.0)
  @test all(0.0 .≤ diag(B) .≤ 1.0)
end

end # module MultivariateCalculusTest


## Generators for some "classic" graphs ##

"""
    graph_path(n::Integer)

Generate a [path graph](https://en.wikipedia.org/wiki/Path_graph) with `n` nodes.

# Examples
```julia-repl
julia> graph_path(10)
10-node 9-edge directed multigraph
```
"""
function graph_path(n::Integer; kwargs...)
    g = Graph(n; kwargs...)
    for i = 1:n - 1
        addedge!(g, i, i + 1)
    end
    return g
end

"""
    graph_cycle(n::Integer)

Generate a [cycle graph](https://en.wikipedia.org/wiki/Cycle_graph) with `n` nodes.

# Examples
```julia-repl
julia> graph_cycle(10)
10-node 10-edge directed multigraph
```
"""
function graph_cycle(n::Integer; kwargs...)
    g = Graph(n; kwargs...)
    for i = 1:n
        addedge!(g, i, i % n + 1)
    end
    return g
end

"""
    graph_star(n::Integer[; out=true])

Generate a [star graph](https://en.wikipedia.org/wiki/Star_graph) with `n` nodes.
If `out == true`, connections go outwards from the central node.
Otherwise, they go inwards.

# Examples
```julia-repl
julia> graph_star(10)
10-node 9-edge directed multigraph
```
"""
function graph_star(n::Integer; out::Bool=true, kwargs...)
    g = Graph(n; kwargs...)
    for i = 2:n
        if out
            addedge!(g, 1, i)   # outwards
        else
            addedge!(g, i, 1)   # inwards
        end
    end
    return g
end

"""
    graph_wheel(n::Integer[; out=true])

Generate a [wheel graph](https://en.wikipedia.org/wiki/Wheel_graph) with `n` nodes.
If `out == true`, connections go outwards from the central node.
Otherwise, they go inwards.

# Examples
```julia-repl
julia> graph_wheel(10)
10-node 18-edge directed multigraph
```
"""
function graph_wheel(n::Integer; out::Bool=true, kwargs...)
    g = Graph(n; kwargs...)
    for i = 2:n
        out ? addedge!(g, 1, i) : addedge!(g, i, 1)
        addedge!(g, i, i < n ? i + 1 : 2)
    end
    return g
end

"""
    graph_complete(n::Integer[; ordered=false])

Generate a [complete graph](https://en.wikipedia.org/wiki/Complete_graph) with `n` nodes.

If you are generating a directed graph, by default it will have 2 edges between
each pair of nodes. To leave only edges `i → j` where `i < j`, set `ordered=true`.

# Examples
```julia-repl
julia> graph_complete(10)
10-node 90-edge directed multigraph

julia> graph_complete(10, ordered=true)
10-node 45-edge directed multigraph
```
"""
function graph_complete(n::Integer; ordered::Bool=false, kwargs...)
    g = Graph(n; kwargs...)
    if isdirected(g) && !ordered
        for i = 1:n, j = 1:n
            i != j && addedge!(g, i, j)
        end
    else
        for i = 1:n, j = i + 1:n
            addedge!(g, i, j)
        end
    end
    return g
end

"""
    graph_grid(a::Integer[, b=a])

Generate a rectangular grid of size `a`×`b`.
If `b` is ommitted, it is taken to be equal `a`.

# Examples
```julia-repl
julia> graph_grid(10, 5)
50-node 85-edge directed multigraph
```
"""
function graph_grid(a::Integer, b::Integer=a; kwargs...)
    g = Graph(a * b; kwargs...)
    for i = 1:a, j = 1:b
        cur = (i - 1) * b + j
        if i < a
            addedge!(g, cur, cur + b)
        end
        if j < b
            addedge!(g, cur, cur + 1)
        end
    end
    return g
end

"""
    graph_web(c::Integer, r::Integer[; out=true])

Generate a spiderweb-shaped graph with `c` nodes on each of `r` levels.
Overall, it has `c * r + 1` nodes including the center node.

# Examples
```julia-repl
julia> graph_web(20, 4)
81-node 160-edge directed multigraph
```
"""
function graph_web(c::Integer, r::Integer; out=true, kwargs...)
    g = Graph(r * c + 1; kwargs...)
    for i = 1:r, j = 1:c
        cur = 1 + (i - 1) * c + j
        if out                                              # radial
            addedge!(g, i == 1 ? 1 : cur - c, cur)
        else
            addedge!(g, cur, i == 1 ? 1 : cur - c)
        end
        addedge!(g, cur, j < c ? cur + 1 : cur - c + 1)     # lateral
    end
    return g
end

function branch!(g::Graph, root, depth, c, out)
    for j = 1:c
        k = addnode!(g)
        out ? addedge!(g, root, k) : addedge!(g, k, root)
        if depth > 0
            branch!(g, k, depth - 1, c, out)
        end
    end
end

"""
    graph_tree(c::Integer, r::Integer)

Generate a tree with branching factor `c` and `r` levels.
Mind that it has exponential number of nodes in `r`,
so it is not advisable to create graphs with `r > 4`.

# Examples
```julia-repl
julia> graph_tree(3, 3)
40-node 39-edge directed multigraph
```
"""
function graph_tree(c::Integer, r::Integer; out=true, kwargs...)
    g = Graph(; kwargs...)
    addnode!(g)
    branch!(g, 1, r - 1, c, out)
    return g
end


# Computes the inner products between each component and the given response

# STEP component
#correct here to remove U an L
function xdy_step(IT::iterator, y::Vector{Float64}, d::dataCD)
    N = IT.obs
    Nel = IT.nelements[STEP]
    xdy = zeros(Nel)
    y_s = sum(y)

    @inbounds for i in 1:Nel
        for j in i+1:N
            xdy[i] += y[j]
        end
        xdy[i]=(xdy[i] -d.μ[STEP][i]*y_s)
    end

    return xdy
end

# SLOPE component
function xdy_slope(IT::iterator, y::Vector{Float64}, d::dataCD)
    N = IT.obs
    Nel = IT.nelements[SLOPE]
    xdy = zeros(Nel)
    y_s = sum(y)

    @inbounds for i in 1:Nel
        for j in i:N
            xdy[i] += y[j] * (j-i)
        end
        xdy[i] =(xdy[i] -d.μ[SLOPE][i]*y_s)
    end
    return xdy
end

# SPIKE component
function xdy_spike(IT::iterator, y::Vector{Float64}, d::dataCD)
    N = IT.obs
    Nel = IT.nelements[SPIKE]
    xdy = zeros(Nel)
    y_s = sum(y)

    @inbounds for i in 1:Nel
        xdy[i] =(y[i] -d.μ[SPIKE][i]*y_s)
    end
    return xdy
end

# SINE component
function xdy_sin(IT::iterator, y::Vector{Float64}, d::dataCD)
    N = IT.obs
    nf = IT.nelements[SIN]
    xdy0 = zeros(nf)
    y_s = sum(y)

    @inbounds for i in 1:nf
        for j in 1:N
            xdy0[i] += y[j] * sin(d.fs[i]*j)
        end
        xdy0[i] =(xdy0[i] -d.μ[SIN][i]*y_s)
    end
    return xdy0
end

# COSINE component
function xdy_cos(IT::iterator, y::Vector{Float64}, d::dataCD)
    N = IT.obs
    nf = IT.nelements[COS]
    xdy = zeros(nf)
    y_s = sum(y)

    @inbounds for i in 1:nf
        for j in 1:N
            xdy[i] += y[j] * cos(d.fc[i]*j)
        end
        xdy[i] = (xdy[i]-d.μ[COS][i]*y_s)
    end
    return xdy
end

xdy = Vector{Function}(TOTALCOMPONENTS)
xdy[STEP] = xdy_step
xdy[SLOPE] = xdy_slope
xdy[SPIKE] = xdy_spike
xdy[SIN] = xdy_sin
xdy[COS] = xdy_cos


using LinearAlgebra: LinearAlgebra, Symmetric, SymTridiagonal, I

using NQCBase: NQCBase, PeriodicCell, InfiniteCell

struct RingPolymerParameters{T<:AbstractFloat}
    n_beads::Int
    ω_n::T
    ω_n²::T
    springs::Symmetric{T,Matrix{T}}
    normal_mode_springs::Vector{T}
    U::Matrix{T}
    quantum_atoms::Vector{Int}
    tmp::Vector{T}
    function RingPolymerParameters{T}(
        n_beads::Integer, temperature::Real,
        quantum_atoms::Vector{<:Integer}) where {T<:AbstractFloat}

        ω_n = n_beads * temperature
        new(n_beads, ω_n, ω_n^2,
            get_spring_matrix(n_beads, ω_n),
            get_normal_mode_springs(n_beads, ω_n),
            get_normal_mode_transformation(n_beads),
            quantum_atoms,
            zeros(T, n_beads))
    end
end

Base.broadcastable(beads::RingPolymerParameters) = Ref(beads)

"""Constructor for choosing specific elements to be quantum."""
function RingPolymerParameters{T}(n_beads::Integer, temperature::Real, atom_types::AbstractVector{Symbol}, quantum_nuclei::Vector{Symbol}) where {T}
    quantum_atoms = findall(in(quantum_nuclei), atom_types)
    RingPolymerParameters{T}(n_beads, temperature, quantum_atoms)
end

"""Constructor for the case where all nuclei are quantum."""
function RingPolymerParameters{T}(n_beads::Integer, temperature::Real, n_atoms::Integer) where {T}
    RingPolymerParameters{T}(n_beads, temperature, collect(1:n_atoms))
end

"""
    get_L(n_beads, mass, ω_n)

Return the Circulant symmetric matrix for the ring polymer springs.
"""
function get_spring_matrix(n_beads::Integer, ω_n::Real)::Symmetric
    if n_beads == 1
        spring = zeros(1, 1)
    elseif n_beads == 2
        spring = [2 -2; -2 2]
    else
        spring = SymTridiagonal(fill(2, n_beads), fill(-1, n_beads-1))
        spring = convert(Matrix, spring)
        spring[end,1] = spring[1, end] = -1
    end
    Symmetric(spring .*  ω_n^2 / 2)
end

get_normal_mode_springs(n_beads::Integer, ω_n::Real) = get_matsubara_frequencies(n_beads, ω_n) .^2 / 2
get_matsubara_frequencies(n::Integer, ω_n::Real) = 2ω_n*sin.((0:n-1)*π/n)

# """
#     get_normal_mode_transformation(n::Int)::Matrix

# Get the transformation matrix that converts to normal mode coordinates.
# """
# function get_normal_mode_transformation(n::Int)::Matrix
#     a = ([exp(2π*im/n * j * k) for j=0:n-1, k=0:n-1])
#     (real(a) + imag(a)) / sqrt(n)
# end

"""
Creates normal mode transformation for `n` beads.
"""
function get_normal_mode_transformation(n::Integer)::Matrix
    # Real and imaginary parts of the discrete Fourier transform matrix.
    U = sqrt(2/n) .* hcat([cos(2π * j * k / n) for j=0:n-1, k=0:n÷2],
                        [sin(2π * j * k / n) for j=0:n-1, k=n÷2+1:n-1])

    # Normalisation
    U[:, 1] ./= sqrt(2)
    iseven(n) && (U[:, n÷2+1] ./= sqrt(2))

    U
end

function transform_to_normal_modes!(p::RingPolymerParameters, R::AbstractArray{T,3}) where {T}
    @views for i in p.quantum_atoms
        for j in axes(R, 1)
            LinearAlgebra.mul!(p.tmp, p.U', R[j,i,:])
            R[j,i,:] .= p.tmp
        end
    end
end

function transform_from_normal_modes!(p::RingPolymerParameters, R::AbstractArray{T,3}) where {T}
    @views for i in p.quantum_atoms
        for j in axes(R, 1)
            LinearAlgebra.mul!(p.tmp, p.U, R[j,i,:])
            R[j,i,:] .= p.tmp
        end
    end
end

function transform!(p::RingPolymerParameters, A::RingPolymerArray)
    if A.normal
        transform_from_normal_modes!(p, A)
    else
        transform_to_normal_modes!(p, A)
    end
    A.normal = !A.normal
end

nbeads(beads::RingPolymerParameters) = beads.n_beads
Base.length(beads::RingPolymerParameters) = nbeads(beads)
Base.range(beads::RingPolymerParameters) = range(1; length=length(beads))

"""
    cayley_propagator(beads::RingPolymerParameters{T}, dt::Real; half::Bool=true) where {T}

J. Chem. Phys. 151, 124103 (2019); doi: 10.1063/1.5120282
"""
function cayley_propagator(beads::RingPolymerParameters{T}, dt::Real; half::Bool=true) where {T}

    cay(dtA::Matrix)::Matrix = LinearAlgebra.inv(I - dtA/2) * (I + dtA/2)

    ω_k = get_matsubara_frequencies(length(beads), beads.ω_n)
    prop = [Array{T}(undef, 2, 2) for i=1:length(beads)]
    for (i, ω) in enumerate(ω_k)
        A = [0 1; -ω^2 0]
        prop[i] .= half ? real.(sqrt(cay(dt.*A))) : cay(dt.*A)
    end
    prop
end

"""
    get_spring_energy(beads::RingPolymerParameters, masses, R)

Calculate the ring polymer spring potential.
"""
function get_spring_energy(beads::RingPolymerParameters, masses, R)
    E = zero(eltype(R))
    for bead=1:nbeads(beads)-1
        for i in beads.quantum_atoms # Only for quantum nuclei
            for j in axes(R, 1)
                E += masses[i] * (R[j, i, bead] - R[j, i, bead+1])^2
            end
        end
    end
    for i in beads.quantum_atoms
        for j in axes(R, 1)
            E += masses[i] * (R[j, i, end] - R[j, i, 1])^2
        end
    end
    E * beads.ω_n^2/2
end

function NQCBase.apply_cell_boundaries!(cell::PeriodicCell, R::AbstractArray{T,3}, beads::RingPolymerParameters) where {T}
    transform_to_normal_modes!(beads, R)
    R[:,beads.quantum_atoms,1] ./= sqrt(length(beads))
    @views NQCBase.apply_cell_boundaries!(cell, R[:,:,1])
    R[:,beads.quantum_atoms,1] .*= sqrt(length(beads))
    transform_from_normal_modes!(beads, R)
end
NQCBase.apply_cell_boundaries!(::InfiniteCell, ::AbstractArray{T,3}, ::RingPolymerParameters) where {T} = nothing