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

# The functions
#
# * _radware_sigcompress_encode_impl!
# * _radware_sigcompress_encode_impl!
#
# are almost literal translations of the functions compress_signal and
# decompress_signal from the radware-sigcompress v1.0 C-code (sigcompress.c
# as of 2017-10-25).
#
# The code was translated to Julia and is released here under MIT license with
# permission of the author of the original C-code
# (Copyright (c) 2018, David C. Radford <radforddc@ornl.gov>).

# C-style post-increment emulation
macro _pincr(x)
    quote
        begin
           tmp = $(esc(x))
           $(esc(x)) = $(esc(x)) + 1
           tmp
       end
    end
end

function _radware_sigcompress_encode_impl!(sig_out::AbstractVector{UInt16}, sig_in::AbstractVector{Int16})
    sig_len_in = length(eachindex(sig_in))

    max_output_len = 2 * sig_len_in
    if (length(eachindex(sig_out)) < max_output_len)
        resize!(sig_out, max_output_len)
    end

    i::Int32 = 0; j::Int32 = 0; max1::Int32 = 0; max2::Int32 = 0; min1::Int32 = 0; min2::Int32 = 0; ds::Int32 = 0; nb1::Int32 = 0; nb2::Int32 = 0
    iso::Int32 = 0; nw::Int32 = 0; bp::Int32 = 0; dd1::Int32 = 0; dd2::Int32 = 0

    db = UInt16[0, 0]
    dd = reinterpret(UInt32, db)

    mask = UInt16[1, 3, 7,15, 31, 63, 127, 255, 511, 1023, 2047, 4095, 8191, 16383, 32767, 65535]

    #static int len[17] = {4096, 2048,512,256,128, 128,128,128,128,
    #                      128,128,128,128, 48,48,48,48};
    #= ------------ do compression of signal ------------ =#
    j = iso = 1; bp = 0

    sig_out[@_pincr(iso)] = sig_len_in     # signal length
    while (j <= sig_len_in)         # j = starting index of section of signal
        # find optimal method and length for compression of next section of signal
        max1 = min1 = sig_in[j]
        max2 = -16000
        min2 = 16000
        nb1 = nb2 = 2
        nw = 1
        i=j+1; while (i <= sig_len_in && i < j+48)  # FIXME; # 48 could be tuned better?
            if (max1 < sig_in[i]) max1 = sig_in[i] end
            if (min1 > sig_in[i]) min1 = sig_in[i] end
            ds = Int32(sig_in[i]) - Int32(sig_in[i-1])
            if (max2 < ds) max2 = ds end
            if (min2 > ds) min2 = ds end
            @_pincr(nw)
            @_pincr(i)
        end
        if (max1-min1 <= max2-min2)  # use absolute values
            nb2 = 99
            while (max1 - min1 > mask[nb1]) @_pincr(nb1) end
            #for (; i <= sig_len_in && i < j+len[nb1]; @_pincr(i)) {
            while (i <= sig_len_in && i < j+128)  # FIXME; # 128 could be tuned better?
                if (max1 < sig_in[i]) max1 = sig_in[i] end
                dd1 = max1 - min1
                if (min1 > sig_in[i]) dd1 = max1 - sig_in[i] end
                if (dd1 > mask[nb1]) break end
                if (min1 > sig_in[i]) min1 = sig_in[i] end
                @_pincr(nw)
                @_pincr(i)
            end
        else                      # use difference values
            nb1 = 99
            while (max2 - min2 > mask[nb2]) @_pincr(nb2) end
            #for (; i <= sig_len_in && i < j+len[nb1]; @_pincr(i)) {
            while (i <= sig_len_in && i < j+128)  # FIXME; # 128 could be tuned better?
                ds = Int32(sig_in[i]) - Int32(sig_in[i-1])
                if (max2 < ds) max2 = ds end
                dd2 = max2 - min2
                if (min2 > ds) dd2 = max2 - ds end
                if (dd2 > mask[nb2]) break end
                if (min2 > ds) min2 = ds end
                @_pincr(nw)
                @_pincr(i)
            end
        end

        if (bp > 0) @_pincr(iso) end
        #=  -----  do actual compression  -----  =#
        sig_out[@_pincr(iso)] = nw  # compressed signal data, first byte = # samples
        bp = 0               # bit pointer
        if (nb1 <= nb2)
            #=  -----  encode absolute values  -----  =#
            sig_out[@_pincr(iso)] = nb1                    # # bits used for encoding
            sig_out[@_pincr(iso)] = min1 % UInt16  # min value used for encoding
            i = iso; while (i <= iso + div(nw*nb1, 16)) sig_out[i] = 0; @_pincr(i) end
            i = j; while (i < j + nw)
                dd[1] = sig_in[i] - min1              # value to encode
                dd[1] = dd[1] << (32 - bp - nb1)
                sig_out[iso] |= db[2]
                bp += nb1
                if (bp > 15)
                    sig_out[iso+=1] = db[1]
                    bp -= 16
                end
                @_pincr(i)
            end
        else
            #=  -----  encode derivative / difference values  -----  =#
            sig_out[@_pincr(iso)] = nb2 + 32  # # bits used for encoding, plus flag
            sig_out[@_pincr(iso)] = sig_in[j] % UInt16  # starting signal value
            sig_out[@_pincr(iso)] = min2 % UInt16       # min value used for encoding
            i = iso; while (i <= iso + div(nw*nb2, 16)) sig_out[i] = 0; @_pincr(i) end
            i = j+1; while (i < j + nw)
                dd[1] = Int32(sig_in[i]) - Int32(sig_in[i-1]) - min2     # value to encode
                dd[1] = dd[1] << (32 - bp - nb2)
                sig_out[iso] |= db[2]
                bp += nb2
                if (bp > 15)
                    sig_out[iso+=1] = db[1]
                    bp -= 16
                end
                @_pincr(i)
            end
        end
        j += nw
    end

    if (bp > 0) @_pincr(iso) end
    if ((iso - 1)%2 > 0) @_pincr(iso) end     # make sure iso is even for 4-byte padding

    resize!(sig_out, iso - 1)   # number of shorts in compressed signal data
    sig_out
end #= compress_signal =#

function _radware_sigcompress_decode_impl!(sig_out::AbstractVector{Int16}, sig_in::AbstractVector{UInt16})
    sig_len_in = length(eachindex(sig_in))

    i::Int32 = 0; j::Int32 = 0; min::Int32 = 0; nb::Int32 = 0; isi::Int32 = 0; iso::Int32 = 0; nw::Int32 = 0; bp::Int32 = 0; siglen::Int32 = 0
    db = UInt16[0, 0]
    dd = reinterpret(UInt32, db)
    mask = UInt16[1, 3, 7,15, 31, 63, 127, 255, 511, 1023, 2047, 4095, 8191, 16383, 32767, 65535]

    #= ------------ do decompression of signal ------------ =#
    j = isi = iso = 1; bp = 0
    siglen = sig_in[@_pincr(isi)] % Int16  # signal length

    if (length(eachindex(sig_out)) < siglen)
        resize!(sig_out, siglen)
    end

    #printf("<<< siglen = %d\n", siglen);
    fill!(sig_out, 0)
    while (isi <= sig_len_in && iso <= siglen)
        if (bp > 0) @_pincr(isi) end
        bp = 0              # bit pointer
        nw = sig_in[@_pincr(isi)]  # number of samples encoded in this chunk
        nb = sig_in[@_pincr(isi)]  # number of bits used in compression

        if (nb < 32)
            #=  -----  decode absolute values  -----  =#
            min = sig_in[@_pincr(isi)] % Int16  # min value used for encoding
            db[1] = sig_in[isi]
            i = 0; while (i < nw && iso <= siglen)
                if (bp+nb > 15)
                    bp -= 16
                    db[2] = sig_in[@_pincr(isi)]
                    if isi <= lastindex(sig_in)
                        db[1] = sig_in[isi]
                    end
                    dd[1] = dd[1] << (bp+nb)
                else
                    dd[1] = dd[1] << nb
                end
                sig_out[@_pincr(iso)] = (db[2] & mask[nb]) + min
                bp += nb
                @_pincr(i)
            end
        else
            nb -= 32
            #=  -----  decode derivative / difference values  -----  =#
            sig_out[@_pincr(iso)] = sig_in[@_pincr(isi)] % Int16  # starting signal value
            min = sig_in[@_pincr(isi)] % Int16    # min value used for encoding
            if isi <= lastindex(sig_in)
                db[1] = sig_in[isi]
            end
            i = 1; while (i < nw && iso <= siglen)
                if (bp+nb > 15)
                    bp -= 16
                    db[2] = sig_in[@_pincr(isi)]
                    if isi <= lastindex(sig_in)
                        db[1] = sig_in[isi]
                    end
                    dd[1] = dd[1] << (bp+nb)
                else
                    dd[1] = dd[1] << nb
                end
                sig_out[iso] = ((db[2] & mask[nb]) + min + sig_out[iso-1]) % Int16; @_pincr(iso)
                bp += nb
                @_pincr(i)
            end
        end
        j += nw
    end

    outlen = iso - 1
    if (siglen != outlen)
        throw(ErrorException("ERROR in decompress_signal: outlen $outlen != siglen $siglen"))
    end

    resize!(sig_out, siglen)
    sig_out
end

"""
    RadwareSigcompress <: AbstractArrayCodec
"""
struct RadwareSigcompress <: AbstractArrayCodec
    shift::Int
end
export RadwareSigcompress

RadwareSigcompress(::Type{T}) where {T<:Signed} = RadwareSigcompress(0)
RadwareSigcompress(::Type{T}) where {T<:Unsigned} = RadwareSigcompress(typemin(typeof(signed(zero(T)))))

function EncodedArrays.encode_data!(encoded::AbstractVector{UInt8}, codec::RadwareSigcompress, data::AbstractVector{T}) where {T}
    shift = convert(T, codec.shift)
    data_conv = Int16.(data .+ shift)
    output = similar(data_conv, UInt16)
    _radware_sigcompress_encode_impl!(output, data_conv)
    resize!(encoded, 0)
    append!(encoded, reinterpret(UInt8, hton.(output)))
    encoded
end

function EncodedArrays.decode_data!(data::AbstractVector{T}, codec::RadwareSigcompress, encoded::AbstractVector{UInt8}) where {T}
    encoded_conv = ntoh.(reinterpret(UInt16, encoded))
    output = similar(encoded_conv, Int16)
    tmp_data = _radware_sigcompress_decode_impl!(output, encoded_conv)
    copyto!(data, tmp_data)
    data .-= convert(T, codec.shift)

    data
end

function read_from_properties(read_property::Function, src::Any, ::Type{RadwareSigcompress})
    shft = read_property(src, :radware_sigcompress_shift, 0)
    RadwareSigcompress(shft)
end

function write_to_properties!(write_property::Function, dest::Any, codec::RadwareSigcompress) where T
    write_property(dest, :radware_sigcompress_shift, codec.shift)
    nothing
end


###################
# DISPLAY METHODS #
###################

# mathjax = "https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/MathJax.js?config=TeX-AMS-MML_SVG"
#COLORS
function theme_chooser(darkmode)
	if darkmode
		bgcolor = "rgba(0,0,0,0)"#"rgb(13,16,17)"
		text_color = "gray"
		plot_bgcolor = "rgb(22,26,29)" #"rgb(33,37,41)"
		grid_color = "rgb(40,52,66)" #rgb(40,40,40)
		sep_color = "white"
	else
		bgcolor = "rgba(0,0,0,0)"#"white"
		text_color = "gray"#"rgb(49,70,101)"
		plot_bgcolor = "rgb(229,236,246)"
		grid_color = "white"
		sep_color = "white"
	end
	bgcolor, text_color, plot_bgcolor, grid_color, sep_color
end
function interp_map(c_map, t_interp)
	idx = [c[1] for c = c_map]
	R = [parse.(Int, split(c[2][5:end-1],", "))[1] for c = c_map]
	G = [parse.(Int, split(c[2][5:end-1],", "))[2] for c = c_map]
	B = [parse.(Int, split(c[2][5:end-1],", "))[3] for c = c_map]
	r = LinearInterpolation(idx,R)(t_interp)
	g = LinearInterpolation(idx,G)(t_interp)
	b = LinearInterpolation(idx,B)(t_interp)
	c_map_interp = ["hsv($r, $g, $b)" for (r,g,b)=zip(r,g,b)]
	c_map_interp
end

plot_seq(seq::Sequence; width=nothing, height=nothing, slider=true, show_seq_blocks=false, darkmode=false, max_rf_samples=100) = begin
	bgcolor, text_color, plot_bgcolor, grid_color, sep_color = theme_chooser(darkmode)
	idx = ["Gx" "Gy" "Gz"]
	N = length(seq)
	O = size(seq.RF,1)
	ΔT = durs(seq)
	T0 = cumsum([0; ΔT],dims=1)
	off_val = Inf #This removes the unnecessary points in the plot
	#GRADS
	t1x = vcat([get_theo_t(seq.GR[1,i]) .+ T0[i] for i=1:N]...)
	t1y = vcat([get_theo_t(seq.GR[2,i]) .+ T0[i] for i=1:N]...)
	t1z = vcat([get_theo_t(seq.GR[3,i]) .+ T0[i] for i=1:N]...)
	Gx =  vcat([get_theo_A(seq.GR[1,i];off_val) for i=1:N]...)
	Gy =  vcat([get_theo_A(seq.GR[2,i];off_val) for i=1:N]...)
	Gz =  vcat([get_theo_A(seq.GR[3,i];off_val) for i=1:N]...)
	#RFS
	t2 =  vcat([get_theo_t(seq.RF[1,i];max_rf_samples) .+ T0[i] for i=1:N]...)
	R =   vcat([get_theo_A(r;off_val,max_rf_samples) for r = seq.RF]...)
	#ADC
	t3 =  vcat([get_theo_t(seq.ADC[i])  .+ T0[i] for i=1:N]...)
	D =   vcat([get_theo_A(d;off_val) for d = seq.ADC]...)
	#Shapes
	shapes = []
	if show_seq_blocks
		aux = [line(
			xref="x", yref="paper",
			x0=T0[i]*1e3, y0=0,
			x1=T0[i]*1e3, y1=1,
			line=attr(color=sep_color, width=1),
			) for i = 1:N+1]
		append!(shapes, aux)
	end
	# The commented code here is to visually check the simulation blocks
	# if show_sim_blocks
	# 	aux = [line(
	# 		xref="x", yref="paper",
	# 		x0=t_sim_parts[i]*1e3, y0=0,
	# 		x1=t_sim_parts[i]*1e3, y1=1,
	# 		line=attr(color="Red", width=1),
	# 		) for i = 1:length(t_sim_parts)]
	# 		append!(shapes, aux)
	# end
	l = PlotlyJS.Layout(; hovermode="closest",
			xaxis_title="",
			modebar=attr(orientation="h",yanchor="bottom",xanchor="right",y=1,x=0,bgcolor=bgcolor,color=text_color,activecolor=plot_bgcolor),
			legend=attr(orientation="h",yanchor="bottom",xanchor="left",y=1,x=0),
			plot_bgcolor=plot_bgcolor,
			paper_bgcolor=bgcolor,
			xaxis_gridcolor=grid_color,
			yaxis_gridcolor=grid_color,
			xaxis_zerolinecolor=grid_color,
			yaxis_zerolinecolor=grid_color,
			font_color=text_color,
			yaxis_fixedrange = false,
			xaxis=attr(
				ticksuffix=" ms",
				range=[0.,min(20,1e3*dur(seq))],
				rangeslider=attr(visible=slider),
				rangeselector=attr(
					buttons=[
						attr(count=1,
						label="1m",
						step=10,
						stepmode="backward"),
						attr(step="all")
						]
					),
				),
			shapes = shapes,
			margin=attr(t=0,l=0,r=0,b=0),
			height=height, width=width
			)
	plotter = PlotlyJS.scatter #using scattergl speeds up the plotting but does not show the sequence in the slider below
	p = [plotter() for j=1:(3+2O+1)]
	#GR
	p[1] = plotter(x=t1x*1e3, y=Gx*1e3,name=idx[1],hovertemplate="(%{x:.4f} ms, %{y:.2f} mT/m)")
	p[2] = plotter(x=t1y*1e3, y=Gy*1e3,name=idx[2],hovertemplate="(%{x:.4f} ms, %{y:.2f} mT/m)")
	p[3] = plotter(x=t1z*1e3, y=Gz*1e3,name=idx[3],hovertemplate="(%{x:.4f} ms, %{y:.2f} mT/m)")
	#RF
	for j=1:O
		phase = angle.(R[:,j])
		phase[R[:,j] .== Inf] .= Inf
		p[2j-1+3] = plotter(x=t2*1e3, y=abs.(R[:,j])*1e6,name="|RF_$j|",hovertemplate="(%{x:.4f} ms, %{y:.2f} μT)")
		p[2j+3] =   plotter(x=t2*1e3, y=phase, text=ones(size(t2)), name="<RF_$j",hovertemplate="(%{x:.4f} ms, ∠B1: %{y:.4f} rad)", visible="legendonly")
	end
	#ADC
	p[2O+3+1] = plotter(x=t3*1e3, y=D*1., name="ADC",hovertemplate="(%{x:.4f} ms, %{y:i})")
	config = PlotConfig(
		displaylogo=false,
		toImageButtonOptions=attr(
			format="svg", # one of png, svg, jpeg, webp
		).fields,
		modeBarButtonsToRemove=["zoom", "autoScale", "resetScale2d", "pan", "tableRotation", "resetCameraLastSave", "zoomIn", "zoomOut"]
	)
	PlotlyJS.plot(p, l; config) #, options=Dict(:displayModeBar => false))
end

function plot_image(image; height=750, width=nothing, zmin=minimum(abs.(image[:])), zmax=maximum(abs.(image[:])), darkmode=false, title="")
	#Layout
	bgcolor, text_color, plot_bgcolor, grid_color, sep_color = theme_chooser(darkmode)
	l = PlotlyJS.Layout(;title=title,yaxis_title="y",
    xaxis_title="x",height=height,wifth=width,
    yaxis=attr(scaleanchor="x"),
	font_color=text_color,
    modebar=attr(orientation="v",bgcolor=bgcolor,color=text_color,activecolor=plot_bgcolor),xaxis=attr(constrain="domain"),hovermode="closest",
	paper_bgcolor=bgcolor,
	plot_bgcolor=plot_bgcolor,
	xaxis_gridcolor=grid_color,
	yaxis_gridcolor=grid_color,
	xaxis_zerolinecolor=grid_color,
	yaxis_zerolinecolor=grid_color)
	#Plot
	p = PlotlyJS.heatmap(z=image,transpose=false,zmin=zmin,zmax=zmax,colorscale="Greys")
	config = PlotConfig(
		displaylogo=false,
		toImageButtonOptions=attr(
			format="svg", # one of png, svg, jpeg, webp
		).fields,
		modeBarButtonsToRemove=["zoom", "autoScale", "resetScale2d", "pan", "tableRotation", "resetCameraLastSave", "zoomIn", "zoomOut"]
	)
	PlotlyJS.plot(p,l;config)
end

function plot_kspace(seq; width=nothing, height=nothing, darkmode=false)
	bgcolor, text_color, plot_bgcolor, grid_color, sep_color = theme_chooser(darkmode)
	#Calculations of theoretical k-space
	kspace, kspace_adc = get_kspace(seq; Δt=1) #simParams["Δt"])
	t_adc = get_sample_times(seq)
	#Colormap
	c_map = [[t, "hsv($(floor(Int,(1-t)*255)), 100, 50)"] for t=range(0,1,10)]
	c = "gray"
	c2_idx = []
	counter = 0
	for s in seq
		if is_ADC_on(s)
			N = s.ADC.N[1]
			append!(c2_idx, counter:N+counter-1)
			counter += N
		end
	end
	c2 = interp_map(c_map, c2_idx ./ maximum(c2_idx))
	#Layout
	mink = minimum(kspace_adc,dims=1)
	maxk = maximum(kspace_adc,dims=1)
	dW = maximum(maxk .- mink, dims=2) * .3
	mink .-= dW
	maxk .+= dW
	#Layout
	l = PlotlyJS.Layout(;
		paper_bgcolor=bgcolor,
		scene=attr(xaxis=attr(title="kx [m⁻¹]",range=[mink[1],maxk[1]],backgroundcolor=plot_bgcolor,gridcolor=grid_color,zerolinecolor=grid_color),
				   yaxis=attr(title="ky [m⁻¹]",range=[mink[2],maxk[2]],backgroundcolor=plot_bgcolor,gridcolor=grid_color,zerolinecolor=grid_color),
				   zaxis=attr(title="kz [m⁻¹]",range=[mink[3],maxk[3]],backgroundcolor=plot_bgcolor,gridcolor=grid_color,zerolinecolor=grid_color)),
		modebar=attr(orientation="h",yanchor="bottom",xanchor="right",y=1,x=0,bgcolor=bgcolor,color=text_color,activecolor=plot_bgcolor),
		legend=attr(orientation="h",yanchor="bottom",xanchor="left",y=1,x=0),
		height=height, width=width,font_color=text_color,
		scene_camera_eye=attr(x=0, y=0, z=1.7),
		scene_camera_up=attr(x=0, y=1., z=0),
		scene_aspectmode="cube",
		margin=attr(t=0,l=0,r=0))
	#Plot
	p = [PlotlyJS.scatter() for j=1:3]
	p[1] = PlotlyJS.scatter3d(x=kspace[:,1],y=kspace[:,2],z=kspace[:,3],mode="lines",
			line=attr(color=c),name="Trajectory",hoverinfo="skip")
	p[2] = PlotlyJS.scatter3d(x=kspace_adc[:,1],y=kspace_adc[:,2],z=kspace_adc[:,3],text=round.(t_adc*1e3,digits=3),mode="markers",
			line=attr(color=c2),marker=attr(size=2),name="ADC",hovertemplate="kx: %{x:.1f} m⁻¹
ky: %{y:.1f} m⁻¹
kz: %{z:.1f} m⁻¹
t_acq: %{text} ms<extra></extra>")
	p[3] = PlotlyJS.scatter3d(x=[0],y=[0],z=[0],name="k=0",marker=attr(symbol="cross",size=10,color="red"))
	config = PlotConfig(
		displaylogo=false,
		toImageButtonOptions=attr(
			format="svg", # one of png, svg, jpeg, webp
		).fields,
		modeBarButtonsToRemove=["zoom", "pan", "tableRotation", "resetCameraLastSave3d", "orbitRotation", "resetCameraDefault3d"]
	)
	PlotlyJS.plot(p,l; config)
end

function plot_M0(seq; height=nothing, width=nothing, slider=true, darkmode=false)
	#Times
	dt = 1
	t, Δt = Koma.get_uniform_times(seq, dt)
	#kx,ky
	ts = t .+ Δt
	rf_idx, rf_type = Koma.get_RF_types(seq, t)
	k, _ =  Koma.get_kspace(seq; Δt=dt)

	#plots k(t)
	bgcolor, text_color, plot_bgcolor, grid_color, sep_color = theme_chooser(darkmode)
	l = PlotlyJS.Layout(;yaxis_title="M0", hovermode="closest",
			xaxis_title="",
			modebar=attr(orientation="h",yanchor="bottom",xanchor="right",y=1,x=0,bgcolor=bgcolor,color=text_color,activecolor=plot_bgcolor),
			legend=attr(orientation="h",yanchor="bottom",xanchor="left",y=1,x=0),
			plot_bgcolor=plot_bgcolor,
			paper_bgcolor=bgcolor,
			xaxis_gridcolor=grid_color,
			yaxis_gridcolor=grid_color,
			xaxis_zerolinecolor=grid_color,
			yaxis_zerolinecolor=grid_color,
			font_color=text_color,
			yaxis_fixedrange = false,
			xaxis=attr(
				ticksuffix=" ms",
				range=[0.,min(20,1e3*dur(seq))],
				rangeslider=attr(visible=slider),
				rangeselector=attr(
					buttons=[
						attr(count=1,
						label="1m",
						step=10,
						stepmode="backward"),
						attr(step="all")
						]
					),
				),
			margin=attr(t=0,l=0,r=0),
			height=height, width=width
			)
	plotter = PlotlyJS.scatter #using scattergl speeds up the plotting but does not show the sequence in the slider below
	p = [plotter() for j=1:4]
	p[1] = plotter(x=ts*1e3, y=k[:,1], hovertemplate="(%{x:.4f} ms, %{y:.2f} mT/m⋅ms)", name="x")
	p[2] = plotter(x=ts*1e3, y=k[:,2], hovertemplate="(%{x:.4f} ms, %{y:.2f} mT/m⋅ms)", name="y")
	p[3] = plotter(x=ts*1e3, y=k[:,3], hovertemplate="(%{x:.4f} ms, %{y:.2f} mT/m⋅ms)", name="z")
	p[4] = plotter(x=t[rf_idx]*1e3,y=rf_type,name="RFs",marker=attr(symbol="cross",size=8,color="orange"),mode="markers")
	config = PlotConfig(
		displaylogo=false,
		toImageButtonOptions=attr(
			format="svg", # one of png, svg, jpeg, webp
		).fields,
		modeBarButtonsToRemove=["zoom", "select", "select2d","lasso2d", "autoScale", "resetScale2d", "pan", "tableRotation", "resetCameraLastSave", "zoomIn", "zoomOut"]
	)
	PlotlyJS.plot(p, l; config)
end

"""Plot phantom map, the keys could be {:ρ, :T1, :T2, :T2s, :x, :y, :z}. Optional parameter: t0 [ms] to see displacement.

Colormaps from https://github.com/markgriswold/MRFColormaps
Towards Unified Colormaps for Quantitative MRF Data, Mark Griswold, et al. (2018)."""
function plot_phantom_map(ph::Phantom, key::Symbol; t0=0, height=700, width=nothing, darkmode=false)
	cmin_key = minimum(getproperty(ph,key))
	cmax_key = maximum(getproperty(ph,key))
	if key == :T1 || key == :T2 || key == :T2s
		cmin_key = 0
		factor = 1e3
		unit = " ms"
		if key  == :T1
			cmax_key = 2500/factor
			colors = MAT.matread("src/ui/assets/T1cm.mat")["T1colormap"]
			N, _ = size(colors)
			idx = range(0,1,N)
			colormap = [[idx[n], "rgb($(floor(Int,colors[n,1]*255)),$(floor(Int,colors[n,2]*255)),$(floor(Int,colors[n,3]*255)))"] for n=1:N]
		elseif key == :T2 || key == :T2s
			if key == :T2
				cmax_key = 250/factor
			end
    		colors = MAT.matread("src/ui/assets/T2cm.mat")["T2colormap"]
			N, _ = size(colors)
			idx = range(0,1,N)
			colormap = [[idx[n], "rgb($(floor(Int,colors[n,1]*255)),$(floor(Int,colors[n,2]*255)),$(floor(Int,colors[n,3]*255)))"] for n=1:N]
		end
	elseif key == :x || key == :y || key == :z
		factor = 1e2
		unit = " cm"
		colormap="Greys"
	elseif key == :Δw
		factor = 1/(2π)
		unit = " Hz"
		colormap="Greys"
	else
		factor = 1
		cmin_key = 0
		unit=""
		colormap="Greys"
	end
	cmin_key *= factor
	cmax_key *= factor
	x0 = minimum([ph.x;ph.y;ph.z])*1e2
    xf = maximum([ph.x;ph.y;ph.z])*1e2
	#Layout
	bgcolor, text_color, plot_bgcolor, grid_color, sep_color = theme_chooser(darkmode)
	l = PlotlyJS.Layout(;title=ph.name, height=height, width=width,
		paper_bgcolor=bgcolor,font_color=text_color,
		scene=attr(
			xaxis=attr(title="x",range=[x0,xf],ticksuffix=" cm",backgroundcolor=plot_bgcolor,gridcolor=grid_color,zerolinecolor=grid_color),
			yaxis=attr(title="y",range=[x0,xf],ticksuffix=" cm",backgroundcolor=plot_bgcolor,gridcolor=grid_color,zerolinecolor=grid_color),
			zaxis=attr(title="z",range=[x0,xf],ticksuffix=" cm",backgroundcolor=plot_bgcolor,gridcolor=grid_color,zerolinecolor=grid_color),
		),
		modebar=attr(orientation="h",bgcolor=bgcolor,color=text_color,activecolor=plot_bgcolor),xaxis=attr(constrain="domain"),
		hovermode="closest",aspectmode="cube", margin=attr(t=25,l=0,r=0))
	h = PlotlyJS.scatter3d( x=(ph.x .+ ph.ux(ph.x,ph.y,ph.z,t0*1e-3))*1e2,
							y=(ph.y .+ ph.uy(ph.x,ph.y,ph.z,t0*1e-3))*1e2,
							z=(ph.z .+ ph.uz(ph.x,ph.y,ph.z,t0*1e-3))*1e2,
							mode="markers",
							marker=attr(color=getproperty(ph,key)*factor,
										showscale=true,
										colorscale=colormap,
										colorbar=attr(ticksuffix=unit, title=string(key)),
										cmin=cmin_key,
										cmax=cmax_key,
										size=2
										),
							text=round.(getproperty(ph,key)*factor,digits=4),
							hovertemplate="x: %{x:.1f} cm
y: %{y:.1f} cm
z: %{z:.1f} cm
$(string(key)): %{text}$unit<extra></extra>")
	config = PlotConfig(
		displaylogo=false,
		toImageButtonOptions=attr(
			format="svg", # one of png, svg, jpeg, webp
		).fields,
		modeBarButtonsToRemove=["zoom", "pan", "tableRotation", "resetCameraLastSave3d", "orbitRotation", "resetCameraDefault3d"]
	)
	p = PlotlyJS.plot(h,l;config)
end

function plot_signal(raw::RawAcquisitionData; height=nothing, width=nothing, darkmode=false)
	not_Koma = raw.params["systemVendor"] != "Koma.jl"
	t = []
	signal = []
	current_t0 = 0
	for p in raw.profiles
		dt = p.head.sample_time_us != 0 ? p.head.sample_time_us * 1e-3 : 1
		t0 = p.head.acquisition_time_stamp * 1e-3 #This parameter is used in Koma to store the time offset
		N =  p.head.number_of_samples != 0 ? p.head.number_of_samples : 1
		if not_Koma
			t0 = current_t0 * dt
			current_t0 += N
		end
		append!(t, t0.+(0:dt:dt*(N-1)))
		append!(signal, p.data[:,1]) #Just one coil
		#To generate gap
		append!(t, t[end])
		append!(signal, [Inf + Inf*1im])
	end
	#Time window
	Tmin = minimum(t)
	Tmax = maximum(t)
	Tacq = Tmax - Tmin
	#PLOT
	bgcolor, text_color, plot_bgcolor, grid_color, sep_color = theme_chooser(darkmode)
	l = PlotlyJS.Layout(;title="",
		font_color=text_color,
		xaxis=attr(ticksuffix=" ms",range=[Tmin+Tacq/2-.025*Tacq,Tmin+Tacq/2+.025*Tacq],
					rangeslider=attr(visible=true),
					gridcolor=grid_color,zerolinecolor=grid_color),
		yaxis=attr(gridcolor=grid_color,zerolinecolor=grid_color),
		paper_bgcolor=bgcolor,
		plot_bgcolor=plot_bgcolor, height=height, width=width,
		yaxis_fixedrange = false,
		legend=attr(orientation="h",yanchor="bottom",xanchor="left",y=1.02,x=0),
		modebar=attr(orientation="h",yanchor="bottom",xanchor="right",y=1.02,x=1,bgcolor=bgcolor,color=text_color,activecolor=plot_bgcolor))
	absS = PlotlyJS.scatter(x=t,y=abs.(signal), name="|S(t)|")
	reS =  PlotlyJS.scatter(x=t,y=real.(signal),name="Re{S(t)}")
	imS =  PlotlyJS.scatter(x=t,y=imag.(signal),name="Im{S(t)}")
	config = PlotConfig(
		displaylogo=false,
		toImageButtonOptions=attr(
			format="svg", # one of png, svg, jpeg, webp
		).fields,
		modeBarButtonsToRemove=["zoom", "select2d", "lasso2d", "autoScale", "resetScale2d", "pan", "tableRotation", "resetCameraLastSave", "zoomIn", "zoomOut"]
	)
	p = PlotlyJS.plot([absS,reS,imS],l;config)
end

function plot_dict(dict::Dict)
	html = """

			#
			Name
			Value

	"""
	i = 1
	for (key,val) = dict
		html *= """

				$i
				$(string(key))
				$(string(val))

		"""
		i += 1
	end
	html *= "

"
end

# """Plots gradient moments M0, M1 and M2 for specified axis. Does NOT consider RF pulses."""
# plot_grads_moments(seq::Sequence, idx::Int=1; title="", mode="quick") = begin
# 	if mode == "quick"
# 		plotter = PlotlyJS.scattergl
# 	else
# 		plotter = PlotlyJS.scatter
# 	end
# 	names = ["Gx", "Gy", "Gz"]
# 	M, N = size(seq.GR)
# 	G = [seq.GR[j,floor(Int,i/2)+1].A for i=0:2*N-1, j=1:M]
# 	T = [seq.GR[1,i].T for i=1:N]
# 	t = [sum(T[1:i]) for i=1:N]
# 	t = [t[floor(Int,i/2)+1] for i=0:2*N-1]
# 	t = [0; t[1:end-1]]

# 	M0, M1, M2 = get_M0_M1_M2(seq)
# 	M0t = [M0(t)[idx] for t=t][:]; M0max = maximum(abs.(M0t))*1.1e6
# 	M1t = [M1(t)[idx] for t=t][:]; M1max = maximum(abs.(M1t))*1.1e9
# 	M2t = [M2(t)[idx] for t=t][:]; M2max = maximum(abs.(M2t))*1.1e12
# 	Gmax = maximum(abs.(G))*1.1e3

# 	p = [plotter() for k = 1:4]

# 	l = PlotlyJS.Layout(;
# 	title=title,
# 	xaxis=attr(domain = [0, .75]),
# 	yaxis=attr(title=attr(text="G [mT/m]", standoff=0),
# 		side="left", tickfont=attr(color="#636efa"), titlefont=attr(color="#636efa"),range=[-Gmax, Gmax]),
# 	yaxis2=attr(title=attr(text="M0 [mT/m⋅ms]", standoff=1, position=1), showgrid=false,
# 		anchor="free", overlaying="y",position=0.75, side="right",
# 		tickfont=attr(color="#ef553b"), titlefont=attr(color="#ef553b"), range=[-M0max, M0max]),
# 	yaxis3=attr(title=attr(text="M1 [mT/m⋅ms²]", standoff=1), showgrid=false,
# 		anchor="free",overlaying="y",position=0.75+.25/3,side="right",
# 		tickfont=attr(color="#45d9b2"), titlefont=attr(color="#45d9b2"), range=[-M1max, M1max]),
# 	yaxis4=attr(title=attr(text="M2 [mT/m⋅ms³]", standoff=1), showgrid=false,
# 		anchor="free",overlaying="y",position=0.75+.5/3,side="right",
# 		tickfont=attr(color="#b373fa"), titlefont=attr(color="#b373fa"), range=[-M2max, M2max]),
# 	xaxis_title="t [ms]", height=300)
# 	p[1] = plotter(x=t*1e3, y=G[:,idx]*1e3, name=names[idx] , line_shape="hv")
# 	p[2] = plotter(x=t*1e3, y=M0t*1e6, name="M0", yaxis="y2",line=attr(dash="dash"))
# 	p[3] = plotter(x=t*1e3, y=M1t*1e9, name="M1", yaxis="y3",line=attr(dash="dash"))
# 	p[4] = plotter(x=t*1e3, y=M2t*1e12, name="M2", yaxis="y4",line=attr(dash="dash"))
# 	PlotlyJS.plot(p, l)
# end
# plot_Phantom(obj::Phantom,filename::String) = begin
# 	# Phantom
# 	p1 = heatmap(obj.x*1e2,obj.y*1e2,obj.ρ,aspect_ratio=:equal)#,clims=(0,maximum(T2[:])))
# 	xaxis!("x [cm]"); yaxis!("y [cm]")
# 	title!("Proton density")
# 	p2 = heatmap(obj.x*1e2,obj.y*1e2,obj.T2*1e3,aspect_ratio=:equal)#,clims=(0,maximum(T2[:])))
# 	xaxis!("x [cm]"); yaxis!("y [cm]")
# 	title!("T2 [ms]")
# 	p3 = heatmap(obj.x*1e2,obj.y*1e2,obj.Δw/2π,aspect_ratio=:equal)#,clims=(0,maximum(T2[:])))
# 	xaxis!("x [cm]"); yaxis!("y [cm]")
# 	title!("df [Hz]")
# 	P(i,j) = rotz(Dθ[i,j])[1:2,1:2]; D(i,j) = [obj.Dλ1[i,j] 0;0 obj.Dλ2[i,j]]
# 	nx = [1;0]; ny = [0;1]
# 	Dx = [nx'*P(i,j)'*D(i,j)*P(i,j)*nx for i=1:size(obj.Dλ1,1),j=1:size(obj.Dλ1,2)]
# 	Dy = [ny'*P(i,j)'*D(i,j)*P(i,j)*ny for i=1:size(obj.Dλ1,1),j=1:size(obj.Dλ1,2)]
# 	p4 = heatmap(obj.x*1e2,obj.y*1e2,Dx*1e12,aspect_ratio=:equal)#,clims=(0,maximum(T2[:])))
# 	xaxis!("x [cm]"); yaxis!("y [cm]")
# 	title!("Dx [um2/s]")
# 	p5 = heatmap(obj.x*1e2,obj.y*1e2,Dy*1e12,aspect_ratio=:equal)#,clims=(0,maximum(T2[:])))
# 	xaxis!("x [cm]"); yaxis!("y [cm]")
# 	title!("Dy [um2/s]")
# 	p = plot(p1,p2,p3,p4,p5,size=(1300,500),layout=@layout [a b c; d e])
# 	savefig(p,filename)
# end
# plot_sim_res(obj::Phantom,SEQ::Array{Sequence},S::Array{ComplexF64},
# 	t::Array{Float64,2},filename::String,t_k0::Float64,Ga::Float64) = begin
# 	T2m = sum(obj.T2.*obj.ρ)/sum(obj.ρ)
# 	b, n = get_bvalue(SEQ[1])
# 	println("Sequence with b-value: "*string(round(b*1e-6,digits=2))*" s/mm2") # s/mm2
# 	P(i) = rotz(obj.Dθ[i])[1:2,1:2]; D(i) = [obj.Dλ1[i] 0;0 obj.Dλ2[i]]
# 	Deq = [n'*P(i)'*D(i)*P(i)*n for i=1:prod(size(obj.Dλ1))]
# 	Dm = sum(Deq.*obj.ρ)/sum(obj.ρ) #MODIFYY

# 	p = plot_grads(SEQ,t,t_k0,Ga)
# 	p3 = plot([t_k0; t_k0]*1e3,[0; 1.1],linewidth=0.25,color=:black,label="k=0")
# 	plot!(t[:]*1e3,abs.(S),label="|S|")
# 	plot!(t[:]*1e3,exp.(-t[:]/T2m),linestyle=:dash,label="exp(-t/T2)")
# 	# exp(-b*D) <-> exp(-4*π^2*(Δ-δ/3)*q'*D*q)
# 	plot!(t[:]*1e3,exp.(-t[:]/T2m.-b*Dm).*(t[:].>=(Δ+δ)),linestyle=:dash,label="exp(-t/T2-bD)")

# 	xlabel!("Time [ms]")
# 	ylabel!("Signal [a.u]")
# 	xlims!((minimum(t),minimum(t).+dur(sum(SEQ))).*1e3)
# 	p = plot(p[1],p[2],p3,size=(800,600),layout=@layout [a ; b; c])
# 	savefig(p,filename)
# end
# plot_ksapce_trajectory(ACQ::Sequence,t::Array{Float64,2},filename::String) = begin
# 	k = get_designed_kspace(ACQ)
# 	p = plot(k[:,1],k[:,2],legend=:none,aspect_ratio=:equal)
# 	k = get_actual_kspace(ACQ,t)
# 	scatter!(k[:,1],k[:,2],legend=:none,markersize=1)
# 	xlabel!("kx [1/m]"); ylabel!("ky [1/m]")
# 	savefig(p,filename)
# end
# plot_recon(kdata::Array{ComplexF64},rec::Array{ComplexF64},
# 	Δx_pix::Float64,Δy_pix::Float64,filename::String,title::String) = begin
# 	Nx, Ny = size(rec)
# 	xr = -Δx_pix*(Nx-1)/2:Δx_pix:Δx_pix*(Nx-1)/2
# 	yr = -Δy_pix*(Ny-1)/2:Δy_pix:Δy_pix*(Ny-1)/2
# 	Wx, Wy = 1/Δx_pix, 1/Δy_pix
# 	kx = range(-Wx/2,stop=Wx/2,length=Nx)
# 	ky = range(-Wy/2,stop=Wy/2,length=Ny)
# 	p1 = heatmap(kx,ky,log.(abs.(kdata).+1),aspect_ratio=:equal,
# 				legend=:none,size=(400,400))
# 	xaxis!("kx [1/m]"); yaxis!("ky [1/m]")
# 	title!("k-space")
# 	p2 = heatmap(xr*1e2,yr*1e2,abs.(rec),aspect_ratio=:equal,
# 				legend=:none,size=(400,400))
# 	xaxis!("x [cm]"); yaxis!("y [cm]")
# 	title!(title)
# 	#savefig(p1,filename*"_ksp.pdf")
# 	savefig(p2,filename*".pdf")
# end
# plot_Eq(vx,vy,S,S0,SEQ) = begin
# 	Nq, Nθ = size(SEQ)
# 	rs = range(0,stop=1,length=Nq)
# 	θs = range(0,stop=π,length=Nθ); θs = [θs[1:end-1]; π.+θs[1:end-1]];
#     Eq(vx,vy) = begin
#         Eq = [S[i,j][vx,vy] for i=1:Nq,j=1:Nθ]
#         Eq = [Eq[:,1:end-1] Eq[:,1:end-1]]
#     end
#     c = RGB{Float64}(1,1,1)
#     α = real.(abs.(S0[vx,vy]))/maximum(real.(abs.(S0)))
#     pyplot()
#     hm = heatmap(θs,rs,Eq(vx,vy)*(α<.2 ? 0 : 1),proj=:polar,aspect_ratio=:equal,
#     	legend=:none,grid=false,xticks=:none,yticks=:none,
#     	background_color_subplot=α*c)
# end
# plot_Pr(vx,vy,S,S0,SEQ) = begin
# 	Nq, Nθ = size(SEQ)
# 	rs = range(0,stop=1,length=10)*25e-6
# 	θs = range(0,stop=2π,length=10)
#     Pr(vx,vy) = begin
#         Eq = [S[i,j][vx,vy] for i=1:Nq,j=1:Nθ]
# 		b = [get_bvalue(SEQ[i,j])[1] for i=1:Nq,j=1:Nθ]
# 		n = [get_bvalue(SEQ[i,j])[2] for i=1:Nq,j=1:Nθ]
# 		B = [-b[i,j]*[n[i,j][1]^2
# 					2*n[i,j][1]*n[i,j][2]
# 					  n[i,j][2]^2] for i=1:Nq,j=1:Nθ]
# 		B = [B[i][j] for i=1:Nq*Nθ,j=1:3] #reshape
# 		y = log.(abs.(Eq))[:]
# 		# exp(-b*n'D n) <-> exp(-4*π^2*(Δ-δ/3)*q'*D*q)
# 		# min D || -b*n'D*n - log(S) ||_2 + λ ||D||_2
# 		# min D || B D - y ||_2 + λ ||D||_2 => D = (B'B+λI)^-1 B' y
# 		λ = 0
# 		Id = Matrix{Float64}(I,3,3); D = [0;0;0]
# 		for n = 0:2 #Tikhonov regularized and iterativly weighted
# 			W = n==0 ? diagm(0=>Eq[:].^2) : diagm(0=>exp.(2*B*D))
# 			D = (B'*W*B + λ*Id)^-1*B'*W*y
# 		end
# 		# Diffusion propagator
# 		D_inv = [D[1] D[2];
# 				 D[2] D[3]]^-1
# 		pr = exp.(-[([r*cos(θ) r*sin(θ)]*D_inv*[r*cos(θ);r*sin(θ)])[1] for r=rs, θ=θs])
#     end
#     c = RGB{Float64}(1,1,1)
#     α = real.(abs.(S0[vx,vy]))/maximum(real.(abs.(S0)))
#     pyplot()
#     hm = heatmap(θs,rs,Pr(vx,vy)*(α<.2 ? 0 : 1),proj=:polar,aspect_ratio=:equal,
#     	legend=:none,grid=false,xticks=:none,yticks=:none,
#     	background_color_subplot=α*c)
# end


"""
    PowerNetwork <: AbstractPowerSystem

This structure contains constant parameters that define the topology and
physics of the power network.

The object `PowerNetwork` uses its own contiguous indexing for the buses.
The indexing is independent from those specified in the Matpower or the
PSSE input file. However, a correspondence between the two indexing
(Input indexing *to* `PowerNetwork` indexing) is stored inside the
attribute `bus_to_indexes`.

## Note
The object `PowerNetwork` is created in the main memory.
Use a `AbstractFormulation` object to move data to the target device.

"""
struct PowerNetwork <: AbstractPowerSystem
    vbus::Vector{Complex{Float64}}
    Ybus::SparseArrays.SparseMatrixCSC{Complex{Float64},Int64}
    data::Dict{String,Array}

    nbus::Int64
    ngen::Int64

    bustype::Vector{Int64}
    bus_to_indexes::Dict{Int, Int}
    ref::Vector{Int64}
    pv::Vector{Int64}
    pq::Vector{Int64}

    sbus::Vector{Complex{Float64}}
    sload::Vector{Complex{Float64}}

    function PowerNetwork(datafile::String, data_format::Int64)

        if data_format == 0
            println("Reading PSSE format")
            data_raw = ParsePSSE.parse_raw(datafile)
            data, bus_id_to_indexes = ParsePSSE.raw_to_exapf(data_raw)
        elseif data_format == 1
            data_mat = ParseMAT.parse_mat(datafile)
            data, bus_id_to_indexes = ParseMAT.mat_to_exapf(data_mat)
        end
        # Parsed data indexes
        BUS_I, BUS_TYPE, PD, QD, GS, BS, BUS_AREA, VM, VA, BASE_KV, ZONE, VMAX, VMIN,
        LAM_P, LAM_Q, MU_VMAX, MU_VMIN = IndexSet.idx_bus()

        # retrive required data
        bus = data["bus"]
        gen = data["gen"]
        SBASE = data["baseMVA"][1]

        # size of the system
        nbus = size(bus, 1)
        ngen = size(gen, 1)

        # obtain V0 from raw data
        vbus = zeros(Complex{Float64}, nbus)
        for i in 1:nbus
            vbus[i] = bus[i, VM]*exp(1im * pi/180 * bus[i, VA])
        end

        # form Y matrix
        Ybus = makeYbus(data, bus_id_to_indexes)

        # bus type indexing
        ref, pv, pq, bustype = bustypeindex(bus, gen, bus_id_to_indexes)

        sbus, sload = assembleSbus(gen, bus, SBASE, bus_id_to_indexes)

        new(vbus, Ybus, data, nbus, ngen, bustype, bus_id_to_indexes, ref, pv, pq, sbus, sload)
    end
end
function PowerNetwork(datafile::String)
    if endswith(datafile, ".raw")
        data_format = 0
    elseif endswith(datafile, ".m")
        data_format = 1
    else
        error("Unsupported format in file $(datafile): supported extensions are " *
              "Matpower (.m) or PSSE (.raw)")
    end
    return PowerNetwork(datafile, data_format)
end

# Getters
## Network attributes
get(pf::PowerNetwork, ::NumberOfBuses) = pf.nbus
get(pf::PowerNetwork, ::NumberOfLines) = size(pf.data["branch"], 2)
get(pf::PowerNetwork, ::NumberOfGenerators) = pf.ngen
get(pf::PowerNetwork, ::NumberOfPVBuses) = length(pf.pv)
get(pf::PowerNetwork, ::NumberOfPQBuses) = length(pf.pq)
get(pf::PowerNetwork, ::NumberOfSlackBuses) = length(pf.ref)

## Indexing
function get(pf::PowerNetwork, ::GeneratorIndexes)
    GEN_BUS = IndexSet.idx_gen()[1]
    gens = pf.data["gen"]
    ngens = size(gens, 1)
    # Create array on host memory
    indexing = zeros(Int, ngens)
    # Here, we keep the same ordering as specified in Matpower.
    for i in 1:ngens
        indexing[i] = pf.bus_to_indexes[gens[i, GEN_BUS]]
    end
    return indexing
end
get(pf::PowerNetwork, ::PVIndexes) = pf.pv
get(pf::PowerNetwork, ::PQIndexes) = pf.pq
get(pf::PowerNetwork, ::SlackIndexes) = pf.ref

# Pretty printing
function Base.show(io::IO, pf::PowerNetwork)
    println("Power Network characteristics:")
    @printf("\tBuses: %d. Slack: %d. PV: %d. PQ: %d\n", pf.nbus, length(pf.ref),
            length(pf.pv), length(pf.pq))
    println("\tGenerators: ", pf.ngen, ".")
    # Print system status
    @printf("\t==============================================\n")
    @printf("\tBUS \t TYPE \t VMAG \t VANG \t P \t Q\n")
    @printf("\t==============================================\n")

    for i=1:pf.nbus
        type = pf.bustype[i]
        vmag = abs(pf.vbus[i])
        vang = angle(pf.vbus[i])*(180.0/pi)
        pinj = real(pf.sbus[i])
        qinj = imag(pf.sbus[i])
        @printf("\t%d \t  %d \t %1.3f\t%3.2f\t%3.3f\t%3.3f\n", i,
                type, vmag, vang, pinj, qinj)
    end
end

# Some utils function

function bounds(pf::PowerNetwork, ::Buses, ::VoltageMagnitude)
    BUS_I, BUS_TYPE, PD, QD, GS, BS, BUS_AREA, VM, VA, BASE_KV, ZONE, VMAX, VMIN,
    LAM_P, LAM_Q, MU_VMAX, MU_VMIN = IndexSet.idx_bus()

    bus = pf.data["bus"]
    v_min = convert.(Float64, bus[:, VMIN])
    v_max = convert.(Float64, bus[:, VMAX])
    return v_min, v_max
end

function bounds(pf::PowerNetwork, ::Generator, ::ActivePower)
    GEN_BUS, PG, QG, QMAX, QMIN, VG, MBASE, GEN_STATUS, PMAX, PMIN, PC1, PC2, QC1MIN,
    QC2MIN, QC2MAX, RAMP_AGC, RAMP_10, RAMP_30, RAMP_Q, APF, MU_PMAG, MU_PMIN, MU_QMAX,
    MU_QMIN = IndexSet.idx_gen()

    gens = pf.data["gen"]
    baseMVA = pf.data["baseMVA"][1]

    p_min = convert.(Float64, gens[:, PMIN] / baseMVA)
    p_max = convert.(Float64, gens[:, PMAX] / baseMVA)
    return p_min, p_max
end

function bounds(pf::PowerNetwork, ::Generator, ::ReactivePower)
    GEN_BUS, PG, QG, QMAX, QMIN, VG, MBASE, GEN_STATUS, PMAX, PMIN, PC1, PC2, QC1MIN,
    QC2MIN, QC2MAX, RAMP_AGC, RAMP_10, RAMP_30, RAMP_Q, APF, MU_PMAG, MU_PMIN, MU_QMAX,
    MU_QMIN = IndexSet.idx_gen()

    gens = pf.data["gen"]
    baseMVA = pf.data["baseMVA"][1]

    q_min = convert.(Float64, gens[:, QMIN] / baseMVA)
    q_max = convert.(Float64, gens[:, QMAX] / baseMVA)
    return q_min, q_max
end

"""
    get_costs_coefficients(pf::PowerNetwork)

Return coefficients for costs function.

"""
function get_costs_coefficients(pf::PowerNetwork)
    # indexes
    BUS_I, BUS_TYPE, PD, QD, GS, BS, BUS_AREA, VM, VA, BASE_KV, ZONE, VMAX, VMIN,
    LAM_P, LAM_Q, MU_VMAX, MU_VMIN = IndexSet.idx_bus()
    GEN_BUS, PG, QG, QMAX, QMIN, VG, MBASE, GEN_STATUS, PMAX, PMIN, PC1, PC2, QC1MIN,
    QC2MIN, QC2MAX, RAMP_AGC, RAMP_10, RAMP_30, RAMP_Q, APF, MU_PMAG, MU_PMIN, MU_QMAX,
    MU_QMIN = IndexSet.idx_gen()
    MODEL, STARTUP, SHUTDOWN, NCOST, COST = IndexSet.idx_cost()

    ref = pf.ref
    pv = pf.pv
    pq = pf.pq
    b2i = pf.bus_to_indexes

    # Matpower assumes gens are ordered. Generator in row i has its cost on row i
    # of the cost table.
    gens = pf.data["gen"]
    baseMVA = pf.data["baseMVA"][1]
    bus = pf.data["bus"]
    ngens = size(gens)[1]
    nbus = size(bus)[1]

    # initialize cost
    # store coefficients in a Float64 array, with 4 columns:
    # - 1st column: bus type
    # - 2nd column: constant coefficient c0
    # - 3rd column: coefficient c1
    # - 4th column: coefficient c2
    coefficients = zeros(ngens, 4)
    # If cost is not specified, we return the array coefficients as is
    if !haskey(pf.data, "cost")
        @warn("PowerSystem: cost is not specified in PowerNetwork dataset")
        return coefficients
    end

    cost_data = pf.data["cost"]
    # iterate generators and check if pv or ref.
    for i = 1:ngens
        # only 2nd degree polynomial implemented for now.
        @assert cost_data[i, MODEL] == 2
        @assert cost_data[i, NCOST] == 3
        genbus = b2i[gens[i, GEN_BUS]]
        bustype = bus[genbus, BUS_TYPE]

        # polynomial coefficients
        # TODO: currently scale by baseMVA. Is it a good idea?
        c0 = cost_data[i, COST][3]
        c1 = cost_data[i, COST][2] * baseMVA
        c2 = cost_data[i, COST][1] * baseMVA^2
        coefficients[i, :] .= (bustype, c0, c1, c2)
    end
    return coefficients
end


# Tags:
Cassette.@context HessianSparsityContext

const TaggedOf{T} = Tagged{A, T} where A

const HTagType = Union{Input, TermCombination}
Cassette.metadatatype(::Type{<:HessianSparsityContext}, ::DataType) = HTagType

istainted(ctx::HessianSparsityContext, x) = ismetatype(x, ctx, TermCombination)

@inline Cassette.overdub(ctx::HessianSparsityContext, f::typeof(istainted), x) =
    istainted(ctx, x)
@inline Cassette.overdub(ctx::HessianSparsityContext, f::typeof(this_here_predicate!)) =
    this_here_predicate!(ctx.metadata)

# getindex on the input
@inline function Cassette.overdub(ctx::HessianSparsityContext,
                                  f::typeof(getindex),
                                  X::Tagged,
                                  idx::Tagged...)
    if any(i->ismetatype(i, ctx, TermCombination) && !isone(metadata(i, ctx)), idx)
        error("getindex call depends on input. Cannot determine Hessian sparsity")
    end
    Cassette.overdub(ctx, f, X, map(i->untag(i, ctx), idx)...)
end

# plugs an ambiguity
@inline function Cassette.overdub(ctx::HessianSparsityContext,
                                  f::typeof(getindex),
                                  X::Tagged)
    Cassette.recurse(ctx, f, X)
end

@inline function Cassette.overdub(ctx::HessianSparsityContext,
                                  f::typeof(getindex),
                                  X::Tagged,
                                  idx::Integer...)
    if ismetatype(X, ctx, Input)
        val = Cassette.fallback(ctx, f, X, idx...)
        i = LinearIndices(untag(X, ctx))[idx...]
        tag(val, ctx, TermCombination(Set([Dict(i=>1)])))
    else
        Cassette.recurse(ctx, f, X, idx...)
    end
end

@inline function Cassette.overdub(ctx::HessianSparsityContext,
                                  f::typeof(Base.unsafe_copyto!),
                                  X::Tagged,
                                  xstart,
                                  Y::Tagged,
                                  ystart,
                                  len)
    if ismetatype(Y, ctx, Input)
        val = Cassette.fallback(ctx, f, X, xstart, Y, ystart, len)
        nometa = Cassette.NoMetaMeta()
        X.meta.meta[xstart:xstart+len-1] .= (i->Cassette.Meta(TermCombination(Set([Dict(i=>1)])), nometa)).(ystart:ystart+len-1)
        val
    else
        Cassette.recurse(ctx, f, X, xstart, Y, ystart, len)
    end
end

@inline function Cassette.overdub(ctx::HessianSparsityContext,
                                  f::typeof(copy),
                                  X::Tagged)
    if ismetatype(X, ctx, Input)
        val = Cassette.fallback(ctx, f, X)
        tag(val, ctx, Input())
    else
        Cassette.recurse(ctx, f, X)
    end
end

combine_terms(::Nothing, terms...) = one(TermCombination)

# 1-arg functions
combine_terms(::Val{true}, term) = term
combine_terms(::Val{false}, term) = term * term

# 2-arg functions
function combine_terms(::Val{linearity}, term1, term2) where linearity

    linear11, linear22, linear12 = linearity
    term = zero(TermCombination)
    if linear11
        if !linear12
            term += term1
        end
    else
        term += term1 * term1
    end

    if linear22
        if !linear12
            term += term2
        end
    else
        term += term2 * term2
    end

    if linear12
        term += term1 + term2
    else
        term += term1 * term2
    end
    term
end

# Hessian overdub
#
function getterms(ctx, x)
    ismetatype(x, ctx, TermCombination) ? metadata(x, ctx) : one(TermCombination)
end

function hessian_overdub(ctx::HessianSparsityContext, f, linearity, args...)
    t = combine_terms(linearity, map(x->getterms(ctx, x), args)...)
    val = Cassette.fallback(ctx, f, args...)
    tag(val, ctx, t)
end

@inline function Cassette.overdub(ctx::HessianSparsityContext,
                                  f::typeof(getproperty),
                                  x::Tagged, prop)
    if ismetatype(x, ctx, TermCombination) && !isone(metadata(x, ctx))
        error("property of a non-constant term accessed")
    else
        Cassette.recurse(ctx, f, x, prop)
    end
end

haslinearity(ctx::HessianSparsityContext, f, nargs) = haslinearity(untag(f, ctx), nargs)
linearity(ctx::HessianSparsityContext, f, nargs) = linearity(untag(f, ctx), nargs)

@inline function Cassette.overdub(ctx::HessianSparsityContext,
                                  f,
                                  args...)
    tainted = any(x->ismetatype(x, ctx, TermCombination), args)
    val = if tainted && haslinearity(ctx, f, Val{nfields(args)}())
        l = linearity(ctx, f, Val{nfields(args)}())
        hessian_overdub(ctx, f, l, args...)
    else
        val = Cassette.recurse(ctx, f, args...)
       #if tainted && !ismetatype(val, ctx, TermCombination)
       #    @warn("Don't know the linearity of function $f")
       #end
        val
    end
    val
end


using BlackBoxOptim: general_stop_reason

@testset "general_stop_reason" begin
    problem = BlackBoxOptim.example_problems["Sphere"]

    @testset "Within fitness tolerance" begin
        # Set a low MinDeltaFitnessTolerance so it does not interfere.
        res = bboptimize(problem; NumDimensions = 1, TraceMode = :silent,
            FitnessTolerance = 1e-1, MinDeltaFitnessTolerance = 1e-100)
        @test general_stop_reason(res) == "Within fitness tolerance of optimum"
    end

    @testset "Delta fitness below tolerance" begin
        # Set a low FitnessTolerance so it does not interfere.
        res = bboptimize(problem; NumDimensions = 1, TraceMode = :silent,
            MinDeltaFitnessTolerance = 1e-1, FitnessTolerance = 1e-100)
        @test general_stop_reason(res) == "Delta fitness below tolerance"
    end
end


export NeighborInfo, Neighborhood, PackedNeighborhoods
export NearNeighborDataStructure, MetricNNDS, NullNNDS, MemoizedNNDS
export NeighborhoodCache, StreamingNC, MemoizedNC, OfflineNC
export includes_controls, neighbor_info_type, neighbors

# Utilities
const F_or_B = Union{Val{:F},Val{:B}}
const BoolVal = Union{Val{true},Val{false}}
@inline always_false(::SampleIndex) = false    # used as a filtering predicate
@inline keep_controls_if(nt::NamedTuple, cond::BoolVal) = cond === Val(true) ? nt : (nt..., controls=nothing)
# @inline lattice_neighbors(v::SVector{D}) where {D} = (setindex(v, v[i] + j, i) for i in 1:D for j in (-1, 1))
@inline lattice_neighbors(v::SVector{D}) where {D} = (v + SVector(Tuple(I)) - SVector(ntuple(i -> 2, D)) for
                                                      I in CartesianIndices(ntuple(i -> 3, D)) if
                                                      I != CartesianIndex(ntuple(i -> 2, D)))        # JuliaLang/julia#29440
@inline linear_index(i::Int) = i
@inline linear_index(i::TiledIndex) = i.off_lattice
@inline collect_if_not_Array(x::Array) = x
@inline collect_if_not_Array(x) = collect(x)
@inline includes_controls(x) = includes_controls(typeof(x))
@inline neighbor_info_type(x) = neighbor_info_type(typeof(x))

# Neighborhoods
const NeighborInfo{X<:SampleIndex,D,U} = NamedTuple{(:index,:cost,:controls),Tuple{X,D,U}}
const Neighborhood{X<:SampleIndex,D,U} = AbstractVector{NeighborInfo{X,D,U}}
includes_controls( ::Type{  NeighborInfo{X,D,U}}) where {X,D,U} = U !== Nothing
includes_controls( ::Type{<:Neighborhood{X,D,U}}) where {X,D,U} = U !== Nothing
neighbor_info_type(::Type{<:Neighborhood{X,D,U}}) where {X,D,U} = NeighborInfo{X,D,U}
# indextype(   ::Type{NeighborInfo{X,D,U}}) where {X,D,U} = X
# costtype(    ::Type{NeighborInfo{X,D,U}}) where {X,D,U} = D
# controlstype(::Type{NeighborInfo{X,D,U}}) where {X,D,U} = U
function _neighbor_info_type(::Type{X}, ::Type{NamedTuple{M,Tuple{D,U}}}, include_controls::BoolVal) where {X,M,D,U}
    include_controls === Val(true) ? NeighborInfo{X,D,U} : NeighborInfo{X,D,Nothing}
end
function neighbor_info_type(nodes::SampleSet, bvp::SteeringBVP; include_controls::BoolVal=Val(false))
    X = indextype(nodes)
    BVP = typeof(bvp(nodes.init, nodes.init))
    _neighbor_info_type(X, BVP, include_controls)
end
function shift_lattice_component(n::NeighborInfo{<:TiledIndex}, shift)
    (index=shift_lattice_component(n.index, shift), cost=n.cost, controls=n.controls)
end
function reverse_neighbor_dict(i, neighborhood, ::Type{NeighborInfo{X,D,U}}) where {X,D,U}
    Dict{X,NeighborInfo{X,D,U}}(n.index => (index=i, cost=n.cost, controls=n.controls) for n in neighborhood)
end

## CSC-style Neighborhoods
struct PackedNeighborhoods{X,D,U,N<:Neighborhood{X,D,U}} <: AbstractVector{N}
    ptr::Vector{Int}
    val::Vector{NeighborInfo{X,D,U}}
    neighborhoods::Vector{N}
end
function PackedNeighborhoods(neighborhoods::AbstractVector{<:Neighborhood})
    ptr = pushfirst!(cumsum([length(n) for n in neighborhoods]) + 1, 1)
    val = reduce(vcat, neighborhoods)
    PackedNeighborhoods(ptr, val, [view(val, ptr[i]:ptr[i+1]-1) for i in 1:length(neighborhoods)])
end
Base.length(PN::PackedNeighborhoods) = length(PN.neighborhoods)
Base.getindex(PN::PackedNeighborhoods, i::Int) = PN.neighborhoods[i]
includes_controls( ::Type{<:PackedNeighborhoods{X,D,U}}) where {X,D,U} = U !== Nothing
neighbor_info_type(::Type{<:PackedNeighborhoods{X,D,U}}) where {X,D,U} = NeighborInfo{X,D,U}

# Near Neighborhood Caches
abstract type NeighborhoodCache end
includes_controls(::Type{<:NeighborhoodCache}) = false
## StreamingNC
struct StreamingNC <: NeighborhoodCache end
function neighbors(cache::StreamingNC, nnds, nodes, bvp, v, r;
                   dir::F_or_B=Val(:F), include_controls::BoolVal=Val(false))
    neighbors(nnds, nodes, bvp, v, r, dir=dir, include_controls=include_controls)
end

## MemoizedNC
struct MemoizedNC{N<:Neighborhood} <: NeighborhoodCache
    neighborhoods_F::Vector{N}
    neighborhoods_B::Vector{N}
end
function MemoizedNC(::Type{T}, n::Int=0, issymmetric::Bool=false) where {T}
    if issymmetric
        neighborhoods = Vector{Vector{T}}(undef, n)
        MemoizedNC(neighborhoods, neighborhoods)
    else
        MemoizedNC(Vector{Vector{T}}(undef, n), Vector{Vector{T}}(undef, n))
    end
end
function MemoizedNC(nodes::SampleSet, bvp::SteeringBVP; include_controls::BoolVal=Val(false))
    MemoizedNC(neighbor_info_type(nodes, bvp, include_controls=include_controls), length(nodes), issymmetric(bvp))
end
includes_controls( ::Type{MemoizedNC{N}}) where {N} = includes_controls(N)
neighbor_info_type(::Type{MemoizedNC{N}}) where {N} = neighbor_info_type(N)
function neighbors(cache::MemoizedNC{N}, nnds, nodes, bvp, v, r;
                   dir::F_or_B=Val(:F), include_controls::BoolVal=Val(false)) where {N}
    neighborhoods = dir === Val(:F) ? cache.neighborhoods_F : cache.neighborhoods_B
    i = linear_index(v)
    if i <= 0 || !isassigned(neighborhoods, i)
        neighborhood = neighbors(nnds, nodes, bvp, v, r, dir=dir, include_controls=include_controls)
        i > 0 && (neighborhoods[i] = collect_if_not_Array(neighborhood))
        neighborhood
    else
        neighborhoods[i]
    end
end

## OfflineNC
struct OfflineNC{PN<:PackedNeighborhoods} <: NeighborhoodCache
    neighborhoods_F::PN
    neighborhoods_B::PN
end
function OfflineNC(cache::MemoizedNC)
    if cache.neighborhoods_F === cache.neighborhoods_B
        packed_neighborhoods = PackedNeighborhoods(cache.neighborhoods_F)
        OfflineNC(packed_neighborhoods, packed_neighborhoods)
    else
        OfflineNC(PackedNeighborhoods(cache.neighborhoods_F), PackedNeighborhoods(cache.neighborhoods_B))
    end
end
includes_controls( ::Type{OfflineNC{PN}}) where {PN} = includes_controls(PN)
neighbor_info_type(::Type{OfflineNC{PN}}) where {PN} = neighbor_info_type(PN)
function neighbors(cache::OfflineNC{PN}, nnds, nodes, bvp, v, r;
                   dir::F_or_B=Val(:F), include_controls::BoolVal=Val(false)) where {PN}
    dir === Val(:F) ? cache.neighborhoods_F[linear_index(v)] : cache.neighborhoods_B[linear_index(v)]
end

# Near Neighbor Data Structures
abstract type NearNeighborDataStructure end
abstract type MetricNNDS <: NearNeighborDataStructure end    # `using NearestNeighbors` or `using FLANN`
## NullNNDS (Streaming)
struct NullNNDS <: NearNeighborDataStructure end
function neighbors(nnds::NullNNDS, nodes, bvp, v, r;
                   dir::F_or_B=Val(:F), include_controls::BoolVal=Val(false))
    neighbors(nodes, bvp, v, r, dir=dir, include_controls=include_controls)
end

## MemoizedNNDS
struct MemoizedNNDS{S,D,U,F} <: NearNeighborDataStructure
    bvp_results::Dict{Tuple{S,S},NamedTuple{(:cost,:controls),Tuple{D,U}}}
    memoized_bvp::F
end
function MemoizedNNDS(nodes::ExplicitSampleSet, bvp::SteeringBVP; include_controls::BoolVal=Val(false))
    S = typeof(nodes.init)
    D = typeof(bvp(nodes.init, nodes.init).cost)
    U = include_controls === Val(true) ? typeof(bvp(nodes.init, nodes.init).controls) : Nothing
    bvp_results = Dict{Tuple{S,S},NamedTuple{(:cost,:controls),Tuple{D,U}}}()
    memoized_bvp = (x0, xf, cost_bound) -> begin    # TODO: check if closure performance is disastrous
        get!(bvp_results, (x0, xf)) do              #       also, make subtype of SteeringBVP?
            keep_controls_if(bvp(x0, xf, cost_bound), include_controls)
        end
    end
    MemoizedNNDS(bvp_results, memoized_bvp)
end
includes_controls(::MemoizedNNDS{S,D,U}) where {S,D,U} = U !== Nothing
function neighbors(nnds::MemoizedNNDS, nodes, bvp, v, r;
                   dir::F_or_B=Val(:F), include_controls::BoolVal=Val(false))
    neighbors(nodes, nnds.memoized_bvp, v, r, dir=dir, include_controls=include_controls)
end

# General Methods
function neighbors(nodes::ExplicitSampleSet, bvp, v::Int, r;
                   dir::F_or_B=Val(:F), include_controls::BoolVal=Val(false))
    neighbors(nodes, bvp, nodes[v], r, dir=dir, include_controls=include_controls, omit=isequal(v))
end
function neighbors(nodes::ExplicitSampleSet, bvp, x::State, r;
                   dir::F_or_B=Val(:F), include_controls::BoolVal=Val(false), omit::Function=always_false)
    gen = dir === Val(:F) ? ((index=i, bvp(x, nodes[i], r)...) for i in eachindex(nodes) if !omit(i)) :
                            ((index=i, bvp(nodes[i], x, r)...) for i in eachindex(nodes) if !omit(i))
    (keep_controls_if(n, include_controls) for n in gen if n.cost <= r)
end

function neighbors(nodes::TiledSampleSet, bvp, v::TiledIndex, r;
                   dir::F_or_B=Val(:F), include_controls::BoolVal=Val(false))
    if v.off_lattice <= 0
        neighbors(nodes, bvp, nodes[v], r, dir=dir, include_controls=include_controls)
    else
        if dir === Val(:F)
            x0_xf = i -> (nodes[zero_lattice_component(v)], nodes[shift_lattice_component(i, -v.on_lattice)])
        else
            x0_xf = i -> (nodes[zero_lattice_component(i)], nodes[shift_lattice_component(v, -i.on_lattice)])
        end
        neighbors(nodes, bvp, nodes[v], r, dir=dir, include_controls=include_controls, omit=isequal(v), x0_xf=x0_xf)
    end
end
function neighbors(nodes::TiledSampleSet, bvp, x::State, r;
                   dir::F_or_B=Val(:F), include_controls::BoolVal=Val(false), omit::Function=always_false,
                   x0_xf=(dir === Val(:F) ? i -> (x, nodes[i]) : i -> (nodes[i], x)))
    neighborhood = neighbor_info_type(nodes, bvp, include_controls=include_controls)[]
    close_idx = findclose(nodes, bvp, x, r, dir=dir, omit=omit, x0_xf=x0_xf)
    queue = [close_idx.on_lattice]
    visited = Set(queue)
    while !isempty(queue)
        on_lattice_idx = popfirst!(queue)
        found_a_neighbor = false
        for i in eachindex(nodes.Voff)
            idx = TiledIndex(on_lattice_idx, i)
            n = (index=idx, bvp(x0_xf(idx)..., r)...)
            if n.cost <= r
                found_a_neighbor = true
                !omit(idx) && push!(neighborhood, keep_controls_if(n, include_controls))
            end
        end
        found_a_neighbor && for oli in lattice_neighbors(on_lattice_idx)
            if !(oli in visited)
                push!(queue, oli)
                push!(visited, oli)
            end
        end
    end
    neighborhood
end
function findclose(nodes::TiledSampleSet, bvp, x::State, r; dir::F_or_B=Val(:F), omit::Function=always_false,
                   x0_xf=(dir === Val(:F) ? i -> (x, nodes[i]) : i -> (nodes[i], x)), cutoff::Int=30)
    start_lattice_idx = round.(Int, SVector{size(nodes.A,2)}(Matrix(nodes.A)\(x - nodes.b)))    # TODO: \ for non-square A
    sentinel = fill(typemin(Int), typeof(start_lattice_idx))
    queue = [start_lattice_idx, sentinel]
    visited = Set(queue)
    ct = 0
    while !isempty(queue)
        on_lattice_idx = popfirst!(queue)
        if on_lattice_idx == sentinel
            push!(queue, sentinel)
            ct += 1
            ct > cutoff && error("Either r is too small, or searching near the starting "*
                                 "state isn't working to find a close tiled sample!")
            continue
        end
        for i in eachindex(nodes.Voff)
            idx = TiledIndex(on_lattice_idx, i)
            omit(idx) && continue
            bvp(x0_xf(idx)..., r).cost <= r && return idx
        end
        for oli in lattice_neighbors(on_lattice_idx)
            if !(oli in visited)
                push!(queue, oli)
                push!(visited, oli)
            end
        end
    end
end


module CGP

# This defines what the import module of 'CGP is'

using Logging
using PaddedViews
using Distributions

Logging.configure(level=INFO)

# Various includes we wish to have in this module

include("config.jl")
include("chromosome.jl")
include("distance.jl")
include("mutation.jl")
include("crossover.jl")
include("chromosomes/cgp.jl")
include("chromosomes/pcgp.jl")
include("logging.jl")
include("EAs/oneplus.jl")
include("EAs/cgpneat.jl")
include("EAs/ga.jl")
include("EAs/cmaes.jl")

# The series of evolutionary strategies to be included and used
EAs = [oneplus, GA]

#  The types of chromosomes, normal CGP and positional CGP
chromosomes = [CGPChromo, PCGPChromo]

# The types of mutation strategies that are included in the package
mutations = [:gene_mutate, :mixed_node_mutate, :mixed_subtree_mutate, :adaptive_node_mutate,
             :adaptive_subtree_mutate]

# The types of crossover strategies included in the package
crossovers = [:single_point_crossover, :random_node_crossover, :aligned_node_crossover,
              :proportional_crossover, :output_graph_crossover, :subgraph_crossover]

# The distance measures included in the package
distances = [:genetic_distance, :positional_distance, :constant_functional_distance,
             :random_functional_distance, :active_distance]

end


using Turing, Random
using Turing: ParticleContainer, weights, resample!,
    effectiveSampleSize, Trace, current_trace, VarName,
    Sampler, consume, produce, copy, fork
using Turing.Core: logZ
using Test

dir = splitdir(splitdir(pathof(Turing))[1])[1]
include(dir*"/test/test_utils/AllUtils.jl")

@testset "container.jl" begin
    @turing_testset "copy particle container" begin
        pc = ParticleContainer{Trace}(x -> x * x)
        newpc = copy(pc)

        @test newpc.logE        == pc.logE
        @test newpc.logWs       == pc.logWs
        @test newpc.conditional == pc.conditional
        @test newpc.n_consume   == pc.n_consume
    end
    @turing_testset "particle container" begin
        n = Ref(0)

        alg = PG(5)
        spl = Turing.Sampler(alg, empty_model())
        dist = Normal(0, 1)

        function fpc()
            t = TArray(Float64, 1);
            t[1] = 0;
            while true
                ct = current_trace()
                vn = VarName(gensym(), :x, "[$n]", 1)
                Turing.assume(spl, dist, vn, ct.vi); n[] += 1;
                produce(0)
                vn = VarName(gensym(), :x, "[$n]", 1)
                Turing.assume(spl, dist, vn, ct.vi); n[] += 1;
                t[1] = 1 + t[1]
            end
        end

        pc = ParticleContainer{Trace}(fpc)
        model = Turing.Model{(:x,),()}(fpc, NamedTuple(), NamedTuple())
        tr = Trace(pc.model, model, spl, Turing.VarInfo())
        push!(pc, tr)
        tr = Trace(pc.model, model, spl, Turing.VarInfo())
        push!(pc, tr)
        tr = Trace(pc.model, model, spl, Turing.VarInfo())
        push!(pc, tr)

        @test weights(pc) == [1/3, 1/3, 1/3]
        @test logZ(pc) ≈ log(3)
        @test pc.logE ≈ log(1)

        @test consume(pc) == log(1)

        resample!(pc)
        @test pc.num_particles == length(pc)
        @test weights(pc) == [1/3, 1/3, 1/3]
        @test logZ(pc) ≈ log(3)
        @test pc.logE ≈ log(1)
        @test effectiveSampleSize(pc) == 3

        @test consume(pc) ≈ log(1)
        resample!(pc)
        @test consume(pc) ≈ log(1)
    end
    @turing_testset "trace" begin
        n = Ref(0)

        alg = PG(5)
        spl = Turing.Sampler(alg, empty_model())
        dist = Normal(0, 1)
        function f2()
            t = TArray(Int, 1);
            t[1] = 0;
            while true
                ct = current_trace()
                vn = VarName(gensym(), :x, "[$n]", 1)
                Turing.assume(spl, dist, vn, ct.vi); n[] += 1;
                produce(t[1]);
                vn = VarName(gensym(), :x, "[$n]", 1)
                Turing.assume(spl, dist, vn, ct.vi); n[] += 1;
                t[1] = 1 + t[1]
            end
        end

        # Test task copy version of trace
        model = Turing.Model{(:x,),()}(f2, NamedTuple(), NamedTuple())
        tr = Trace(f2, model, spl, Turing.VarInfo())

        consume(tr); consume(tr)
        a = fork(tr);
        consume(a); consume(a)

        @test consume(tr) == 2
        @test consume(a) == 4
    end
end


using DrWatson
@quickactivate
using ArgParse
using CSV, DataFrames
using HierarchicalAD
using Flux, CUDA

arg_table = ArgParseSettings()
@add_arg_table arg_table begin
    "latent_dim"
        arg_type = Int
        default = 2
        help = "dimensionality of latent spaces"
    "hdim"
        arg_type = Int
        default = 8
        help = "width of the hidden layers in the discriminator"
    "channels"
        arg_type = Int
        nargs = '*'
        default = [16, 32, 64]
        help = "channel sizes"
    "--savepath"
        arg_type = String
        default = "test"
        help = "subdir of data/models"
    "--kernelsizes"
        arg_type = Int
        nargs = '*'
        default = [5, 3, 3]
        help = "kernelsizes"
    "--stride"
        help = "stride length, *to be implemented*"
        arg_type = Int
        default = 1
    "--layer_depth"
        arg_type = Int
        default = 1
        help = "depth of individual Conv blocks between latent spaces, *to be implemented*"
    "--seed"
        default = nothing
        help = "data split seed"
    "--nepochs"
        default = 20
        arg_type = Int
        help = "number of epochs"
    "--gamma"
        default = 0.1f0
        arg_type = Float32
        help = "scaling constant of total correlation in loss"
    "--lambda"
        default = 0f0
        arg_type = Float32
        help = "L2 regularization constant"
    "--test"
        action = :store_true
        help = "test run with limited data"
    "--batchsize"
        default = 128
        arg_type = Int
        help = "batchsize"
    "--activation"
        default = "relu"
        arg_type = String
        help = "activation function"
    "--xdist"
        default = :gaussian
        arg_type = Symbol
        help = "decoder distribution, one of gaussian/bernoulli"
    "--last_conv"
        action = :store_true
        help = "should the last layer of decoder be a dense or a conv layer"
    "--lr"
        help = "learning rate"
        arg_type = Float32
        default = 0.001f0
    "--gpu_id"
        help = "GPU device switch"
        arg_type = Int
        default = 0
    "--epochsize"
        help = "number of samples used in each epoch"
        default = nothing
    "--digit"
        help = "which digits to include"
        arg_type = Int
        nargs = '*'
        default = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    "--area"
        help = "define the included area which is a number in [0,350] with mean of ~100, e.g. `--area >= 150`"
        arg_type = String
        nargs = 2
        default = [">=", "0"]
   "--length"
        help = "define the included digit length which is a number in [0,100] with mean of ~40"
        arg_type = String
        nargs = 2
        default = [">=", "0"]
   "--thickness"
        help = "define the included thickness which is a number in [0,10] with mean of ~3"
        arg_type = String
        nargs = 2
        default = [">=", "0"]
   "--slant"
        help = "define the included slant which is a number in [-1,1] with mean of ~0.1"
        arg_type = String
        nargs = 2
        default = [">=", "-1"]
   "--width"
        help = "define the included width which is a number in [0,25] with mean of ~14"
        arg_type = String
        nargs = 2
        default = [">=", "0"]
   "--height"
        help = "define the included height which is a number in [0,21] with mean of ~19.5"
        arg_type = String
        nargs = 2
        default = [">=", "0"]
end
args = parse_args(arg_table)
@unpack latent_dim, hdim, channels, kernelsizes, stride, layer_depth, last_conv, seed, test,
    lambda, batchsize, nepochs, gpu_id, epochsize, savepath, lr, activation, gamma, xdist = args
latent_count = length(channels)
(latent_count <= length(kernelsizes)) ? nothing : error("number of kernels and channels does not match.")
out_var = last_conv ? :conv : :dense
if seed != nothing
    seed = eval(Meta.parse(seed))
end
CUDA.device!(gpu_id)

# get filters
experiment_argnames = ["latent_dim", "hdim", "channels", "kernelsizes", "stride", "layer_depth", "last_conv",
    "seed", "lambda", "batchsize", "nepochs", "gpu_id", "epochsize", "savepath", "lr",
    "activation", "gamma", "xdist", "test"]
filter_dict, non_default_filters = HierarchicalAD.get_filter_info(experiment_argnames, args, arg_table)

# get the data
dataset = "morpho_mnist"
ratios = (0.8,0.199,0.001)
if test
    data, labels = HierarchicalAD.load_mnist("train")
    tr_x = HierarchicalAD.sample_tensor(data, 1000)
    val_x = HierarchicalAD.sample_tensor(data, 1000)
    tst_x = HierarchicalAD.sample_tensor(data, 100)
    tr_y, val_y, tst_y = nothing, nothing, nothing
    a_x, a_y = HierarchicalAD.sample_tensor(HierarchicalAD.load_mnist("test"), 1000), nothing
else
    (tr_x, tr_y), (val_x, val_y), (tst_x, tst_y), (a_x, a_y) =
        HierarchicalAD.load_train_val_test_data(dataset, filter_dict; ratios=ratios, seed=seed,
            categorical_key="digit")
end
if epochsize == nothing
    epochsize = size(tr_x, 4)
end

# now train the model
ncs = channels
ks = [(k,k) for k in kernelsizes][1:latent_count]
model, training_history, reconstructions, latent_representations =
    HierarchicalAD.train_fvlae(latent_dim, hdim, batchsize, ks, ncs, stride, nepochs, tr_x,
        val_x; λ=lambda, γ=gamma, epochsize=epochsize, layer_depth=layer_depth, lr=lr,
        var=out_var, activation=activation, xdist=xdist)

Flux.Zygote.ignore() do
    # compute scores
    tr_scores, val_scores, tst_scores, a_scores =
        map(x->HierarchicalAD.reconstruction_probability(gpu(model), x, 5, 16), (tr_x, val_x, tst_x, a_x))

    # compute encodings
    tr_encodings, val_encodings, tst_encodings, a_encodings =
        map(x->HierarchicalAD.encode_all(model,x,batchsize),(tr_x, val_x, tst_x, a_x))

    # now save everything
    model_id = HierarchicalAD.timetag()
    experiment_args = (model_id=model_id, data=dataset, latent_count=latent_count,
        latent_dim=latent_dim, channels=ncs, kernelsizes=ks, stride=stride, layer_depth=layer_depth,
        last_conv=last_conv, lr=lr, activation=activation, seed=seed, lambda=lambda,
        batchsize=batchsize, nepochs=nepochs, gpu_id=gpu_id, epochsize=epochsize, gamma=gamma,
        hdim=hdim, xdist=xdist, test=test)
    save_args = (model_id=model_id, data=dataset, model="fvlae", latent_dim=latent_dim,
        channels=channels,gamma=gamma, lambda=lambda, activation=activation, xdist=xdist)
    svn = HierarchicalAD.safe_savename(save_args, "bson", digits=5)
    svn = joinpath(datadir("models/$savepath"), svn)
    tagsave(svn, Dict(
            :model => cpu(model),
            :experiment_args => experiment_args,
            :filter_dict => filter_dict,
            :training_history => training_history,
            :reconstructions => reconstructions,
            :latent_representations => latent_representations,
            :ratios => ratios,
            :tr_scores => tr_scores,
            :val_scores => val_scores,
            :tst_scores => tst_scores,
            :a_scores => a_scores,
            :tr_encodings => tr_encodings,
            :val_encodings => val_encodings,
            :tst_encodings => tst_encodings,
            :a_encodings => a_encodings,
            :tr_labels => tr_y,
            :val_labels => val_y,
            :tst_labels => tst_y,
            :a_labels => a_y,
            :savepath => svn,
            :non_default_filters => non_default_filters
        ))
    @info "Results saved to $svn"
end


module AnalogDiscoveryNeuromore

using AnalogDiscovery

include("RingBuffer.jl")

export go

function float32msg!(buffer::Vector{UInt8},ch::Int64,val::Vector{Float32})
  buffer[1:5]=(ch==1) ? "/in/0".data : "/in/1".data
  @inbounds for i in 6:8
    buffer[i]=UInt8(0)
  end
  buffer[9:10]=",f".data
  @inbounds for i in 11:12
    buffer[i]=UInt8(0)
  end
  @inbounds for i in 1:4
    buffer[12+i]=unsafe_load(Ptr{UInt8}(pointer(val)),4-i+1)
  end
  #@show buffer
  return buffer
end

function go(fs,runForMin,twoChannels,lowVolt)
  cDev = enumDevices()
  @show ("Available",cDev)
  myDev = Int32(0)
  if cDev > 0
    if enumDeviceIsOpened(myDev)
      error("Error opening device, already in use.")
    end
    hdwf = deviceOpen(myDev)
  else
    error("No devices available.")
  end

  bmin,bmax = analogInBufferSizeInfo(hdwf)
  adbufsize = Int32(min(fs,bmax))
  adbufsize64 = Int64(adbufsize)
  analogInFrequencySet(hdwf,fs)
  analogInBufferSizeSet(hdwf,Int32(bmax)) # must be set at bmax for acqmodeRecord!
  analogInAcquisitionModeSet(hdwf,ACQMODERECORD)
  analogInRecordLengthSet(hdwf,0.0) # run indefinitely

  chVoltMin,chVoltMax,chSteps = analogInChannelRangeInfo(hdwf)
  chCount = analogInChannelCount(hdwf)
  for ch in Int32(0):(twoChannels?Int32(chCount-1):Int32(0))
    analogInChannelEnableSet(hdwf,ch,true)
    analogInChannelFilterSet(hdwf,ch,FILTERAVERAGE)
    analogInChannelRangeSet(hdwf,ch,lowVolt?chVoltMin:chVoltMax)
  end

  sock = UDPSocket(); const localh = ip"127.0.0.1"
  const fs_Int32 = round(Int32,fs)
  # Collect s-sec worth of data
  const s = 60*runForMin
  tBuf = Vector{Float64}(ceil(Int32,s*fs))

  # RingBuffers
  ch1buf = RingBuffer(Vector{Float64}(adbufsize),1,1)
  ch2buf = RingBuffer(Vector{Float64}(adbufsize),1,1)
  # Single-element buffer to reinterpret Float32 value later using unsafe_load
  valBuf::Vector{Float32} = Vector{Float32}(1)
  # Buffer for OSC packets sent over UDP
  udpBuf = Array{UInt8,1}(16)

  sleep(1.0) # gets rid of offset estimation transient
  analogInConfigure(hdwf,false,true) # Start acquistion

  # Prefill
  status = analogInStatus(hdwf,true)
  while status != DWFSTATETRIGGEREDRUNNING
    status = analogInStatus(hdwf,true)
  end
  cSamples::Int32 = 0
  chDataBuf = Vector{Float64}(adbufsize)
  sleep(0.25) # Deliberately introduce latency to improve throughput
  cAvailable,cLost,cCorrupted = analogInStatusRecord(hdwf)
  cSamples += cLost
  @inbounds for ch in Int32(0):(twoChannels?Int32(chCount-1):Int32(0))
    analogInStatusData!(chDataBuf,hdwf,ch,Int32(cAvailable))
    for n in 1:Int64(cAvailable)
      if ch==Int32(0)
        enque!(ch1buf,chDataBuf[n])
      else
        enque!(ch2buf,chDataBuf[n])
      end
    end
  end
  cSamples += cAvailable

  voltVal::Float64 = 0.0
  nSamples::Int32 = ceil(Int32,s*fs)
  startT = Libc.time(); t = -1
  display("Streaming data...")
  while t < nSamples-1
    if used(ch1buf) == 0
      status = analogInStatus(hdwf,true)
      if (cSamples == 0) && ((status == DWFSTATECONFIG) || (status == DWFSTATEPREFILL) || (status == DWFSTATEARMED))
        # Acquisition not yet started
        continue
      end
      cAvailable,cLost,cCorrupted = analogInStatusRecord(hdwf)
      # @show (cSamples,cAvailable,cLost,cCorrupted)
      cSamples += cLost
      if cAvailable==0
        continue
      end
      if cSamples+cAvailable > nSamples
        cAvailable = nSamples-cSamples
      end
      # get samples
      @inbounds for ch in Int32(0):(twoChannels?Int32(chCount-1):Int32(0))
        analogInStatusData!(chDataBuf,hdwf,ch,Int32(cAvailable))
        for n in 1:cAvailable
          if ch==0
            enque!(ch1buf,chDataBuf[n])
          else
            enque!(ch2buf,chDataBuf[n])
          end
        end
      end
      cSamples += cAvailable
    end

    # Now send UDP samples
    currT = Libc.time()
    theorT = t/fs+startT
    if (currT>=theorT)
      if used(ch1buf) > 0
        # @show used(ch1buf)
        t += 1
        tBuf[t+1] = theorT-currT
        for ch in Int32(0):(twoChannels?Int32(chCount-1):Int32(0))
          if ch==0
            voltVal = deque!(ch1buf)
          else
            voltVal = deque!(ch2buf)
          end
          valBuf[1] = Float32(voltVal)
          float32msg!(udpBuf,ch+1,valBuf)
          Base._send(sock,localh,UInt16(4545),udpBuf) # internal _send is faster and we don't care about failures
        end
      end
    end
  end

  deviceClose(hdwf)
  return tBuf
end

end # module


# Abstract interface for diffusion penalties

"""
  Abstract type for the different diffusion penalties that can be expressed
  in the framework of Yan et.al. "Interior Penalties for Summation-by-Parts
  Discretizations of Linear Second-Order Differential Equations".

  Each type must implement a method of:

  ```
   applyPenalty(penalty::AbstractDiffusionPenalty, sbp, sbpface,
                diffusion::AbstractDiffusion, iface::Interface,
                delta_w::AbstractMatrix, theta::AbstractMatrix,
                wL::AbstractMatrix, wR::AbstractMatrix,
                nrm_face::AbstractMatrix,
                alphas::AbstractVector,
                jacL::AbstractVector, jacR::AbstractVector,
                res1L::AbstractMatrix, res1R::AbstractMatrix,
                res2L::AbstractMatrix, res2R::AbstractMatrix
               ) where {Tsol, Tres}

  function applyDirichletPenalty(penalty::AbstractDiffusionPenalty,
                      sbp, sbpface,diffusion::AbstractDiffusion,
                      bndry::Boundary, delta_w::AbstractMatrix
                      wL::AbstractMatrix, nrm_face::AbstractMatrix,
                      alpha::Number, jacL::AbstractVector,
                      res1L::AbstractMatrix)

  ```
  specializing the `penalty` argument.  See the docstrings for these
  functions for details.

"""
abstract type AbstractDiffusionPenalty end

"""
  Applies the interface penalty for 2nd derivative terms.
  More specifically, it applies the penalty matrix
  for both sides of the face at the same time:

  [res1L   = [T1 T3  [delta_w   and   [res1R   = [T1 T3  [-delta_w
   res2L]     T2 T4]  theta]           res2R]     T2 T4]  theta]

  and similarly for the right element.  Note that delta_w is negated for
  the right element, and this function performs the required negation.

  **Inputs**

   * penalty: the [`AbstractDiffusionPenalty`](@ref) object
   * sbp
   * params
   * sbpface
   * diffusion: the [`AbstractDiffusion`](@ref) object
   * iface: `Interface` object
   * delta_w: difference of entropy variables at the face, as compuated by
              [`getFaceVariables`](@ref)
   * theta: Dgk * wL + Dgn * wR, as compuated by `getFaceVariables`
   * qL: conservative variables for left element, `numDofPerNode` x
         `numNodesPerElement`
   * qR: conservative variables for right element`, same size as `qL`
   * wL: entropy variables for the left element, `numDofPerNode` x
         `numNodesPerElement`
   * wR: entropy variables for the right element, same size as `wR`
   * coordsL: xyz coordinates of volume nodes of left element, `dim` x
              `numNodesPerElement`
   * coordsR: xyz coordinates of volumke nodes of right element, same size
              as `coordsL`
   * nrm_face: (scaled) normal vectors at the face, `dim` x `numNodesPerFace`
   * alphas: vector of length 2 containing alpha_gk and alpha_gn
   * dxidxL: (scaled) mapping jacobian for left element, `dim` x `dim` x
             `numNodesPerElement`
   * dxidxR: (scaled) mapping jacobian for right element, same size as `dxidxL`
   * jacL: mapping jacobian determinant for the left element,
           `numNodesPerElement`
   * jacR: mapping jacobian determinant for the right element, same size as
           `jacR`

  **Inputs/Outputs**

   * res1L: `numDofPerNode` x `numNodesPerFace`
   * res1R: same size as above
   * res2L: same size as above
   * res2R: same size as above

  All output arrays are overwritten
"""
function applyPenalty(penalty::AbstractDiffusionPenalty, sbp,
                      params::AbstractParamType, sbpface,
                      diffusion::AbstractDiffusion, iface::Interface,
                      delta_w::AbstractMatrix, theta::AbstractMatrix,
                      qL_el::AbstractMatrix, qR_el::AbstractMatrix,
                      wL::AbstractMatrix, wR::AbstractMatrix,
                      coordsL::AbstractMatrix, coordsR::AbstractMatrix,
                      nrm_face::AbstractMatrix,
                      alphas::AbstractVector,
                      dxidxL::Abstract3DArray, dxidxR::Abstract3DArray,
                      jacL::AbstractVector, jacR::AbstractVector,
                      res1L::AbstractMatrix, res1R::AbstractMatrix,
                      res2L::AbstractMatrix, res2R::AbstractMatrix
                     )

  error("generic fallback for applyPenalty() reached: did you forget to extend it with a new method for your AbstractDiffusionPenalty?")
end

"""
  **Inputs**

   * penalty: the [`AbstractDiffusionPenalty`](@ref) object
   * sbp
   * sbpface
   * diffusion: the [`AbstractDiffusion`](@ref) object
   * iface: `Interface` object
   * delta_w: difference of entropy variables at the face, as compuated by
              [`getFaceVariables`](@ref)
   * delta_w_dotL: array containing derivative of `delta_w` wrt `qL`,
                   `numDofPerNode` x `numDofPerNode` x `numNodesPerFace`
                   x `numNodesPerElement`
   * delta_w_dotR: derivative of `delta_w` wrt `qR, same shape as `delta_w_dotR`
   * theta: Dgk * wL + Dgn * wR, as compuated by `getFaceVariables`
   * theta_dotL derivative of `theta` wrt `qL`, same shape as `delta_w_dotL`
   * theta_dotR derivative of `theta` wrt `qR`, same shape as `delta_w_dotL`
   * qL: conservative variables for the left element, `numDofPerNode` x
         `numNodesPerElement`
   * qR: conservative variables for the right element, same size as `qL`
   * wL: entropy variables for the left element, `numDofPerNode` x
         `numNodesPerElement`
   * wR: entropy variables for the right element, same size as `wL`
   * coordsL: xyz coordinates of volume nodes of left element, `dim` x
              `numNodesPerElement`
   * coordsR: xyz coordinates of the volume nodes of the right element,
              same size as `coordsL`
   * nrm_face: (scaled) normal vectors at the face, `dim` x `numNodesPerFace`
   * alphas: vector of length 2 containing alpha_gk and alpha_gn
   * dxidxL: (scaled) mapping jacobian for left element, `dim` x `dim`,
             x `numNodesPerElement
   * dxidxR: (scaled) mapping jacobian for right element, same size as
             `dxidxL`
   * jacL: mapping jacobian determinant for the left element,
           `numNodesPerElement`
   * jacR: mapping jacobian determinant for the right element, same size as
           `jacR`

  **Inputs/Outputs**

   * res1L_dotL: derivative of `res1L` wrt `qL`, same size as `delta_w_dotL`
   * res1L_dotR: derivative of `res1L` wrt `qR`, same size as above`
   * res1R_dotL: derivative of `res1R` wrt `qL`, same size as `delta_w_dotL`
   * res1R_dotR: derivative of `res1R` wrt `qR`, same size as above`
   * res2L: see `applyPenalty`
   * res2R: see `applyPenalty
   * res2L_dotL: derivative of `res2L` wrt `qL`, same size as `delta_w_dotL`
   * res2L_dotR: derivative of `res2L` wrt `qR`, same size as above`
   * res2R_dotL: derivative of `res2R` wrt `qL`, same size as `delta_w_dotL`
   * res2R_dotR: derivative of `res2R` wrt `qR`, same size as above`

   All of these array are overwritten

  **Outputs**

   * T4_nonzero: true if the T4 penalty is zero, false otherwise. This is
                 used to determine the sparsity pattern of the Jacobian.
"""
function applyPenalty_diff(penalty::AbstractDiffusion, sbp,
                      params::AbstractParamType, sbpface,
                      diffusion::AbstractDiffusion, iface::Interface,
                      delta_w::AbstractMatrix,
                      delta_w_dotL::Abstract4DArray, delta_w_dotR::Abstract4DArray,
                      theta::AbstractMatrix,
                      theta_dotL::Abstract4DArray, theta_dotR::Abstract4DArray,
                      qL::AbstractMatrix, qR::AbstractMatrix,
                      wL::AbstractMatrix, wR::AbstractMatrix,
                      coordsL::AbstractMatrix, coordsR::AbstractMatrix,
                      nrm_face::AbstractMatrix,
                      alphas::AbstractVector,
                      dxidxL::Abstract3DArray, dxidxR::Abstract3DArray,
                      jacL::AbstractVector, jacR::AbstractVector,
                      res1L_dotL::Abstract4DArray, res1L_dotR::Abstract4DArray,
                      res1R_dotL::Abstract4DArray, res1R_dotR::Abstract4DArray,
                      res2L::AbstractMatrix, res2R::AbstractMatrix,
                      res2L_dotL::Abstract4DArray, res2L_dotR::Abstract4DArray,
                      res2R_dotL::Abstract4DArray, res2R_dotR::Abstract4DArray)

  error("generic fallback for applyPenalty_diff() reached: did you forget to extend it with a new method for your AbstractDiffusionPenalty?")
end

"""
  Applies the Dirichlet penalty matrix T_gammaD.

  **Inputs**

   * penalty: [`AbstractDiffusionPenalty`](@ref)
   * sbp
   * params
   * sbpface
   * diffusion: [`AbstractDiffusion`](@ref)
   * bndry: `Boundary` object
   * delta_w: values to multiply the penalty against, `numDofPerNode` x
              `numNodesPerFace`
   * qL: entropy variables for the element, `numDofPerNode` x `numNodesPerElement`
   * wL: entropy variables for the element, `numDofPerNode` x `numNodesPerElement`
   * coordsL: xyz coordinates of volume nodes of the element, `dim` x
              `numNodesPerElement`
   * nrm_face: (scaled) normal vector at each face node, `dim` x
               `numNodesPerFace`
   * alpha: the alpha_gk parameter (Number)
   * dxidxL: (scaled) mapping jacobian for the element, `dim` x `dim` x
             `numNodesPerElement`
   * jacL: mapping jacobian determinant, `numNodesPerElement`

  **Inputs/Outputs**

   * res1L: `numDofPerNode` x `numNodesPerFace` array to overwrite with
            the result

"""
function applyDirichletPenalty(penalty::AbstractDiffusionPenalty, sbp,
                      params::AbstractParamType, sbpface,
                      diffusion::AbstractDiffusion, bndry::Boundary,
                      delta_w::AbstractMatrix,
                      qL::AbstractMatrix,
                      wL::AbstractMatrix,
                      coordsL::AbstractMatrix,
                      nrm_face::AbstractMatrix,
                      alpha::Number,
                      dxidxL::Abstract3DArray,
                      jacL::AbstractVector,
                      res1L::AbstractMatrix)

  error("generic fallback for applDirichletPenalty() reached: did you forget to extend it with a new method for your AbstractDiffusionPenalty?")

end

#------------------------------------------------------------------------------
# Diffusion penalties

include("br2_penalty.jl")
include("inviscid_penalty.jl")

global const DiffusionPenaltyDict = Dict{String, Type{T} where T <: AbstractDiffusionPenalty}(
"BR2" => BR2Penalty,
"Inviscid" => InviscidPenalty,
)

"""
  Returns a fully constructed [`AbstractDiffusionPenalty`](@ref).

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * name: name of penalty to get, defaults to opts["DiffusionPenalty"]

  **Outputs**

   * obj: `AbstractDiffusionPenalty` object
"""
function getDiffusionPenalty(mesh, sbp, eqn::EulerData{Tsol, Tres}, opts,
                             name::String=opts["DiffusionPenalty"]
                            ) where {Tsol, Tres}

  obj = DiffusionPenaltyDict[name]{Tsol, Tres}(mesh, sbp, opts)
  assertArraysUnique(obj)

  return obj
end


# patches to metrics.jl of Distances.jl

const metrics = (SqEuclidean, Euclidean, Cityblock, Minkowski, Hamming, TotalVariation)
const UnionMetrics = Distances.UnionMetrics

# Before evaluation, unwrap the AbstractGray colorant and promote storage type
#
# (Gray{N0f8}, Gray{N0f8}) -> (Float32, Float32)
# (Gray{N0f8}, N0f8      ) -> (Float32, Float32)
# (N0f8      , Gray{N0f8}) -> (Float32, Float32)
# (N0f8      , N0f8      ) -> (Float32, Float32)
#
# For Color3, only promote storage type; basic operation such as `abs2`
# returns a `Number` for `Color3`(e.g., `RGB`)
#
# (RGB{N0f8}, RGB{N0f8}) -> (RGB{Float32}, RGB{Float32} )
#
# we don't extend `evaluate` here because it makes the dispatching rules too complicated
# we only need to extend `eval_op`, since other operations
# (e.g., `eval_start`, `eval_reduce`, and `eval_end`) will automatically get promoted
for M in metrics
    for (Ta, Tb) in ((AbstractGray, AbstractGray),
                     (AbstractGray, Number),
                     (Number, AbstractGray),
                     (PromoteType, PromoteType))
        @eval function Distances.eval_op(d::$M, a::$Ta, b::$Tb)
            T1 = eltype(floattype(typeof(a)))
            T2 = eltype(floattype(typeof(a)))
            Distances.eval_op(d, T1(a), T2(b))
        end
    end
    @eval function Distances.eval_op(d::$M,
                                     a::Color3{<:PromoteType},
                                     b::Color3{<:PromoteType})
        CT1 = floattype(typeof(a))
        CT2 = floattype(typeof(b))
        Distances.eval_op(d, CT1(a), CT2(b))
    end
end

# ambiguities
for (ATa, ATb) in ((AbstractGray, AbstractGray),
                   (AbstractGray, Number),
                   (Number, AbstractGray),
                   (PromoteType, PromoteType),
                   (Color3, Color3))
    @eval function result_type(dist::UnionMetrics, ::Type{Ta}, ::Type{Tb}) where {Ta <: $ATa,Tb <: $ATb}
        T1 = eltype(floattype(Ta))
        T2 = eltype(floattype(Tb))
        result_type(dist, T1, T2)
    end
end


# Script for the benchmark between Julia, Python, IDL, and Matlab.
#
# Hongyang Zhou, hyzhou@umich.edu 03/28/2020

using PyPlot

tJulia = [163.36*1e-6, 180.8*1e-3, 1.32]
tPython= [4390.95*1e-6, 179.5*1e-3, 1.34]
tIDL   = [1970.29*1e-6, 453.5*1e-3, 6.08]
tMatlab= [19273.25*1e-6, 698.4*1e-3, 10.60]

tTable = ones(4,3)
tTable[2,:] = tPython ./ tJulia
tTable[3,:] = tIDL ./ tJulia
tTable[4,:] = tMatlab ./ tJulia

testset = 1:3

figure()
plt.rc("font", family = "serif", size = 14)
scatter(testset, tTable[1,:], s=200, marker="*", label="Julia")
scatter(testset, tTable[2,:], s=100, alpha=0.5, label="Python+Numpy")
scatter(testset, tTable[3,:], s=100, label="IDL")
scatter(testset, tTable[4,:], s=100, label="Matlab")
legend()
my_xticks = ["65KB", "317MB", "2.4GB"]
plt.xticks(testset, my_xticks)
yscale("log")
grid(b=true, which="major", color="k", linestyle="-", alpha=0.5)
grid(b=true, which="minor", color="k", linestyle="--", linewidth=0.2, alpha=0.5)

xlabel("File Size")
ylabel("Normalized Timing to Julia")

tight_layout()

## Timing for reading files in Julia
#=
using Batsrus, BenchmarkTools
#filename = "3d_var_region0_0_t00001527_n00005528.out"
#filename = "z=0_raw_1_t25.60000_n00000258.out"
filename = "3d_var_region0_0_t00001205_n00037679.out"
dir = "../BATSRUS/VisAnaJulia"
#dir = "test/"
@btime data = readdata(filename, dir=dir);
=#


using Nettle

function Base.trunc(s::AbstractString, n::Integer)
    n > 0 || throw(DomainError())
    l = length(s)
    l > n || return s
    n > 3 || return s[1:n]
    return s[1:n-3] * "..."
end

tests = [""    => "d41d8cd98f00b204e9800998ecf8427e",
         "a"   => "0cc175b9c0f1b6a831c399e269772661",
         "abc" => "900150983cd24fb0d6963f7d28e17f72",
         "message digest" => "f96b697d7cb7938d525a2f31aaf161d0",
         "abcdefghijklmnopqrstuvwxyz" => "c3fcd3d76192e4007dfb496cca67e13b",
         "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" =>
         "d174ab98d277d9f5a5611c2c9f419d9f",
         "12345678901234567890123456789012345678901234567890123456789012345678901234567890" =>
         "57edf4a22be3c955ac49da2e2107b67a",
         "foobad" => "3858f62230ac3c915f300c664312c63f"]

println("Testing Julia's MD5 hash against RFC 1321.")
for (k, h) in sort(tests, by = length ∘ first)
    md5sum = hexdigest("md5", k)
    @printf("%20s → %s ", trunc(k, 15), md5sum)
    if md5sum == h
        println("MD5 OK")
    else
        println("MD5 Bad")
        println("* The sum should be  ", h)
    end
end


"""
Thermometer state machine:

@startuml
state ThermometerStateMachine {
    state OffState
    state OnState {
        state CelsiusState
        state FahrenheitState
        state KelvinState

        [*] --> CelsiusState
        CelsiusState : entry / set converter to celsius.
        CelsiusState -> FahrenheitState : UnitsEvent
        FahrenheitState: entry / set converter to fahrenheit.
        FahrenheitState --> KelvinState : UnitsEvent
        KelvinState : entry / set converter to kelvin.
        KelvinState --> CelsiusState : UnitsEvent
    }

    [*] --> OffState

    OffState : entry / set display to "off".
    OffState --> OnState[H*] : PowerEvent
    OnState : entry / set display to "--".
    OnState : TemperatureEvent / set display to converted temperature value.
    OnState --> OffState : PowerEvent
}
@enduml
"""

using Logging, Test, HSM

#####
# Temperature converter functions.
#####

celsius(t::Float64) = string(round(t, digits=1)) * "°C"
fahrenheit(t::Float64) = string(round((t * 1.8) + 32.0, digits=1)) * "°F"
kelvin(t::Float64) = string(round(t + 273.15, digits=1)) * "K"

#####
# Thermometer model.
#####

mutable struct Thermometer
    display::String
    temperature::Union{Float64, Nothing}
    convert::Function

    function Thermometer(display="", temperature=nothing, convert=celsius)
        new(display, temperature, convert)
    end
end

thermometer = Thermometer()

#####
# Events.
#####

struct PowerEvent <: HSM.AbstractHsmEvent
end

struct UnitsEvent <: HSM.AbstractHsmEvent
end

struct TemperatureEvent <:HSM.AbstractHsmEvent
    temperature::Float64
end

#####
# State machine.
#####

macro thermometer_state(name)
    return :(
        struct $name <: HSM.AbstractHsmState
            state_info::HSM.HsmStateInfo
            thermometer::Thermometer

            $name(parent, thermometer) = new(HSM.HsmStateInfo(parent), thermometer)
        end
    )
end

@thermometer_state(ThermometerStateMachine)
@thermometer_state(OffState)
@thermometer_state(OnState)
@thermometer_state(CelsiusState)
@thermometer_state(FahrenheitState)
@thermometer_state(KelvinState)

thermometer_state_machine = ThermometerStateMachine(nothing, thermometer)
off_state = OffState(thermometer_state_machine, thermometer)
on_state = OnState(thermometer_state_machine, thermometer)
celsius_state = CelsiusState(on_state, thermometer)
fahrenheit_state = FahrenheitState(on_state, thermometer)
kelvin_state = KelvinState(on_state, thermometer)

# Root state machine state event handlers.
# The root state machine state must handle all possible user-defined events --
# defining `on_event!()` functions for each -- which return `true` to avoid
# UnhandledExceptionErrors being thrown.

function HSM.on_initialize!(state::ThermometerStateMachine)
    @debug "on_initialize!(ThermometerStateMachine)"
    HSM.transition_to_state!(thermometer_state_machine, off_state)
end

function HSM.on_event!(state::ThermometerStateMachine, event::PowerEvent)
    @debug "on_event!(ThermometerStateMachine, PowerEvent)"
    return true
end

function HSM.on_event!(state::ThermometerStateMachine, event::UnitsEvent)
    @debug "on_event!(ThermometerStateMachine, UnitsEvent)"
    return true
end

function HSM.on_event!(state::ThermometerStateMachine, event::TemperatureEvent)
    @debug "on_event!(ThermometerStateMachine, TemperatureEvent)"
    return true
end

# Off state event handlers.

function HSM.on_entry!(state::OffState)
    @debug "on_entry!(OffState)"
    state.thermometer.display = "off"
    state.thermometer.temperature = nothing
end

function HSM.on_event!(state::OffState, event::PowerEvent)
    @debug "on_event!(OffState, PowerEvent)"
    HSM.transition_to_deep_history!(thermometer_state_machine, on_state)
    return true
end

# On state event handlers.

function HSM.on_initialize!(state::OnState)
    @debug "on_initialize!(OnState)"
    HSM.transition_to_state!(thermometer_state_machine, celsius_state)
end

function HSM.on_entry!(state::OnState)
    @debug "on_entry!(OnState)"
    state.thermometer.display = "--"
end

function HSM.on_event!(state::OnState, event::PowerEvent)
    @debug "on_event!(OnState, PowerEvent)"
    HSM.transition_to_state!(thermometer_state_machine, off_state)
    return true
end

function HSM.on_event!(state::OnState, event::TemperatureEvent)
    @debug "on_event!(OnState, TemperatureEvent)"
    state.thermometer.display = state.thermometer.convert(event.temperature)
    state.thermometer.temperature = event.temperature
    return true
end

# Celsius state event handlers.

function HSM.on_entry!(state::CelsiusState)
    @debug "on_entry!(CelsiusState)"
    state.thermometer.convert = celsius
    if !isnothing(state.thermometer.temperature)
        state.thermometer.display = state.thermometer.convert(state.thermometer.temperature)
    end
end

function HSM.on_event!(state::CelsiusState, event::UnitsEvent)
    @debug "on_event!(CelsiusState, UnitsEvent)"
    HSM.transition_to_state!(thermometer_state_machine, fahrenheit_state)
    return true
end

# Fahrenheit state event handlers.

function HSM.on_entry!(state::FahrenheitState)
    @debug "on_entry!(FahrenheitState)"
    state.thermometer.convert = fahrenheit
    if !isnothing(state.thermometer.temperature)
        state.thermometer.display = state.thermometer.convert(state.thermometer.temperature)
    end
end

function HSM.on_event!(state::FahrenheitState, event::UnitsEvent)
    @debug "on_event!(FahrenheitState, UnitsEvent)"
    HSM.transition_to_state!(thermometer_state_machine, kelvin_state)
    return true
end

# Kelvin state event handlers.

function HSM.on_entry!(state::KelvinState)
    @debug "on_entry!(KelvinState)"
    state.thermometer.convert = kelvin
    if !isnothing(state.thermometer.temperature)
        state.thermometer.display = state.thermometer.convert(state.thermometer.temperature)
    end
end

function HSM.on_event!(state::KelvinState, event::UnitsEvent)
    @debug "on_event!(KelvinState, UnitsEvent)"
    HSM.transition_to_state!(thermometer_state_machine, celsius_state)
    return true
end

#####
# Exercise / Test state machine.
#####

@testset "State machine initialization" begin
    # Initialize the state machine.
    HSM.transition_to_state!(thermometer_state_machine, thermometer_state_machine)

    @test HSM.active_state(thermometer_state_machine) == off_state
    @test thermometer.display == "off"
end

@testset "Off State tests" begin
    # Test ignored events.
    HSM.handle_event!(thermometer_state_machine, UnitsEvent())

    @test HSM.active_state(thermometer_state_machine) == off_state
    @test thermometer.display == "off"

    HSM.handle_event!(thermometer_state_machine, TemperatureEvent(10.0))

    @test HSM.active_state(thermometer_state_machine) == off_state
    @test thermometer.display == "off"

    # Test power event.
    HSM.handle_event!(thermometer_state_machine, PowerEvent())

    @test HSM.active_state(thermometer_state_machine) == celsius_state
    @test thermometer.display == "--"
end

@testset "On State tests" begin
    # Test units event.
    HSM.handle_event!(thermometer_state_machine, UnitsEvent())

    @test HSM.active_state(thermometer_state_machine) == fahrenheit_state
    @test thermometer.display == "--"

    HSM.handle_event!(thermometer_state_machine, UnitsEvent())

    @test HSM.active_state(thermometer_state_machine) == kelvin_state
    @test thermometer.display == "--"

    HSM.handle_event!(thermometer_state_machine, TemperatureEvent(0.0))

    @test thermometer.display == "273.2K"

    HSM.handle_event!(thermometer_state_machine, UnitsEvent())

    @test HSM.active_state(thermometer_state_machine) == celsius_state
    @test thermometer.display == "0.0°C"

    HSM.handle_event!(thermometer_state_machine, UnitsEvent())

    @test HSM.active_state(thermometer_state_machine) == fahrenheit_state
    @test thermometer.display == "32.0°F"

    HSM.handle_event!(thermometer_state_machine, PowerEvent())

    @test HSM.active_state(thermometer_state_machine) == off_state
    @test thermometer.display == "off"
end

@testset "On - Off - On retains memory" begin
    HSM.handle_event!(thermometer_state_machine, PowerEvent())

    @test HSM.active_state(thermometer_state_machine) == fahrenheit_state
    @test thermometer.display == "--"
end