export logψ!

# Definitions for batched evaluation of networks
# When the networrks are not cached and therefore allocate
# the result structure
@inline logψ!(out::AbstractMatrix, net::NeuralNetwork, σ::ADoubleStateBatch) where N =
    out .= net(σ...)
@inline logψ!(out::AbstractMatrix, net::NeuralNetwork, σ::AStateBatch) =
    out .= net(σ)

"""
    NNCache{N}

The base abstract type holding a cache for the network type `N<:NeuralNetwork`.
"""
abstract type NNBatchedCache{N} <: NNCache{N}
    #valid::Bool
end

cached(net::NeuralNetwork, batch_sz::Int) =
    CachedNet(net, cache(net, batch_sz))
cached(net::CachedNet, batch_sz::Int) =
    CachedNet(net.net, cache(net.net, batch_sz))

batch_size(net::CachedNet) = batch_size(net.cache)
batch_size(cache::NNBatchedCache) = throw("Not Implemented")

out_similar(net::CachedNet{N,C}) where{N,C<:NNBatchedCache} =
    similar(trainable_first(net), out_type(net), 1, batch_size(net))
out_similar(net::CachedNet{N,C}, dims::Vararg{T,I}) where{T<:Integer, I, N,C<:NNBatchedCache} =
    similar(trainable_first(net), out_type(net), 1, dims...)

# Definition for inplace evaluation of batched cached networks
@inline logψ!(out::AbstractMatrix, net::CachedNet, σ::Vararg{T,N}) where {N,T<:AbstractVecOrMat} =
    logψ!(out, net.net, net.cache, σ...)
@inline logψ!(out::AbstractMatrix, net::CachedNet, σ::ADoubleStateBatch) =
    logψ!(out, net.net, net.cache, row(σ), col(σ))
@inline logψ!(out::AbstractMatrix, cnet::CachedNet, σ::AStateBatch) =
    logψ!(out, cnet.net, cnet.cache, σ)

# Definition for allocating evaluation of batched cached networks
# Shadowing things at ~80 of base_cached_networks.jl
@inline logψ(net::CachedNet{NN,NC}, σr::AStateBatch, σc::AStateBatch) where {N,T,NN,NC<:NNBatchedCache} = begin
    logψ(net, (σr, σc))
end
@inline logψ(net::CachedNet{NN,NC}, σ::ADoubleStateBatch) where {N,T,NN,NC<:NNBatchedCache} = begin
    b_sz = last(size(first(σ)))
    out = similar(trainable_first(net), out_type(net), 1, b_sz)
    logψ!(out, net.net, net.cache, σ...)
end
@inline logψ(cnet::CachedNet{NN,NC}, σ::AStateBatch) where {NN,NC<:NNBatchedCache} = begin
    b_sz = last(size(σ))
    out = similar(trainable_first(cnet), out_type(cnet), 1, b_sz)
    logψ!(out, cnet.net, cnet.cache, σ)
end

# Declare the two functions, even if config(blabla)=blabla, because of a shitty
# Julia's performance bug #32761
# see https://github.com/JuliaLang/julia/issues/32761
@inline logψ_and_∇logψ!(der, n::CachedNet{NN,NC}, σ::AbstractArray) where {NN,NC<:NNBatchedCache} = begin
    b_sz = last(size(σ))
    out = similar(trainable_first(n), out_type(n), 1, b_sz)
    logψ_and_∇logψ!(der, out, n.net, n.cache, σ)
    return out, der
end
@inline logψ_and_∇logψ!(der, n::CachedNet{NN,NC}, σ::NTuple{N,<:AbstractArray}) where {N,NN,NC<:NNBatchedCache} = begin
    b_sz = last(size(first(σ)))
    out = similar(trainable_first(n), out_type(n), 1, b_sz)
    logψ_and_∇logψ!(der, out, n.net, n.cache, σ...)
    return out, der
end
@inline logψ_and_∇logψ!(der, n::CachedNet{NN,NC}, σ::Vararg{<:AbstractArray,N}) where {N,T,NN,NC<:NNBatchedCache} = begin
    b_sz = last(size(first(σ)))
    out = similar(trainable_first(n), out_type(n), 1, b_sz)
    logψ_and_∇logψ!(der, out, n.net, n.cache, σ...)
    return out, der
end

@inline function logψ_and_∇logψ!(der::AbstractDerivative, out, n::CachedNet, σ::NTuple{N,<:AbstractArray}) where N
    logψ_and_∇logψ!(der, out, n.net, n.cache, σ...)
    return (out, der)
end
@inline function logψ_and_∇logψ!(der::AbstractDerivative, out, n::CachedNet, σ::Vararg{<:AbstractArray,N}) where N
    logψ_and_∇logψ!(der, out, n.net, n.cache, σ...)
    return (out, der)
end
@inline function logψ_and_∇logψ!(der::AbstractDerivative, out, n::CachedNet, σ::AbstractArray) where N
    logψ_and_∇logψ!(der, out, n.net, n.cache, σ);
    return (out, der)
end

#
grad_cache(net::NeuralNetwork, batch_sz) =
    grad_cache(out_type(net), net, batch_sz)
function grad_cache(T::Type{<:Number}, net::NeuralNetwork, batch_sz)
    is_analytic(net) && return RealDerivative(T, net, batch_sz)
    return WirtingerDerivative(T, net, batch_sz)
end

#
grad_cache(T::Type{<:Number}, net::CachedNet{N,C}) where {N,C<:NNBatchedCache} =
    grad_cache(T, net, batch_size(net))

function RealDerivative(T::Type{<:Number}, net::NeuralNetwork, batch_sz::Int)
    pars = trainable(net)

    vec       = similar(trainable_first(pars), T, _tlen(pars), batch_sz)
    i, fields = batched_weight_tuple(net, vec)
    return RealDerivative(fields, [vec])
end

#
"""
    grad_cache_vec(net, batch_sz, vec_len)

Builds a vector of `grad_cache` batching `batch_sz` evaluations.
All share the same underlying vector.
"""
grad_cache(net::NeuralNetwork, batch_sz::Integer, vec_len::Integer) =
    grad_cache(out_type(net), net, batch_sz, vec_len)

function grad_cache(T::Type{<:Number}, net::NeuralNetwork, batch_sz, vec_len)
    pars  = trainable(net)
    npars = _tlen(pars)

    if is_analytic(net)
        vec       = similar(trainable_first(pars), T, npars, batch_sz, vec_len)
        # compute type
        j, fields = batched_weight_tuple(net, view(vec, :, :, 1))
        der       = RealDerivative(fields, (view(vec, :, :, 1),))
        ders      = Vector{typeof(der)}()
        for i=1:vec_len
            vec_rsp   = view(vec, :, :, i)
            j, fields = batched_weight_tuple(net, vec_rsp)
            der       = RealDerivative(fields, (vec_rsp,))
            push!(ders, der)
        end
        vecs = (vec,)
    else
        throw("To implement")
    end
    return ders, vecs
end

@inline store_state!(cache::Array,
             v::AbstractVector,
             i::Integer) = begin
    #@uviews cache v begin
        uview(cache, :, i) .= v
    #end
    return cache
end

@inline store_state!(cache::NTuple{2,<:Matrix},
             (vl, vr)::NTuple{2,<:AbstractVector},
             i::Integer) = begin
    cache_l, cache_r = cache
    #@uviews cache_l cache_r vl vr begin
        uview(cache_l, :,i) .= vl
        uview(cache_r, :,i) .= vr
    #end
    return cache
end

@inline @inbounds unsafe_get_batch(σ::AStateBatch, i::Integer) = uview(σ, :, i)
@inline @inbounds unsafe_get_batch(σ::ADoubleStateBatch, i::Integer) = (uview(row(σ), :, i), uview(col(σ), :, i))
@inline num_batches(σ::AStateBatch) = size(σ, 2)
@inline num_batches(σ::ADoubleStateBatch) = size(row(σ), 2)

# Automatic evaluation of batches
logψ!(out::AbstractArray{T,3}, cnet::CachedNet, σ::NTuple{2,AbstractArray{T2,3}}) where{T,T2} =
    logψ!(out, cnet, first(σ), last(σ))
function logψ!(out::AbstractArray{T,3}, cnet::Union{NeuralNetwork, CachedNet}, σr::AT, σc::AT) where {T, T2, AT<:AbstractArray{T2, 3}}
    n_batches = size(σr, 3)
    for i=1:n_batches
        σr_i  = state_uview(σr,  i)
        σc_i  = state_uview(σc,  i)
        out_i = state_uview(out, i)
        logψ!(out_i, cnet, σr_i, σc_i)
    end
    return out
end

function logψ!(out::AbstractArray{T,3}, cnet::NeuralNetwork, σ::AT) where {T, T2, AT<:AbstractArray{T2, 3}}
    n_batches = size(σ, 3)
    for i=1:n_batches
        σ_i   = state_uview(σ,  i)
        out_i = state_uview(out, i)
        logψ!(out_i, cnet, σ_i)
    end
    return out
end

# Compute batches of those things
function logψ_and_∇logψ!(der::Vector{<:AbstractDerivative}, out::AbstractArray{T,3},
                         net::Union{NeuralNetwork, CachedNet},
                         σ::Union{AStateBatchVec, ADoubleStateBatchVec}) where {T}
    n_batches  = length(der)
    for i=1:n_batches
        σ_i   = state_uview(σ,  i)
        out_i = state_uview(out, i)
        logψ_and_∇logψ!(der[i], out_i, net, σ_i)
    end
    return out
end

function logψ!(out::AbstractArray{T,3}, net::Union{NeuralNetwork, CachedNet},
               σ::Union{AStateBatchVec, ADoubleStateBatchVec}) where {T}
    n_batches  = length(der)
    for i=1:n_batches
        σ_i   = state_uview(σ,  i)
        out_i = state_uview(out, i)
        logψ!(out_i, net, σ_i)
    end
    return out
end

function logψ_and_∇logψ(net::Union{NeuralNetwork, CachedNet},
                        σ::Union{AStateBatchVec, ADoubleStateBatchVec}) where {T}
    ∇vals, ∇vec    = grad_cache(bnet, 1, batch_size(σ), chain_length(σ))
    out = similar(out_similar(net), batch_size(σ), chain_length(σ))
    logψ_and_∇logψ!(∇vals, out, net, σ)
    return ∇vals, out
end

function logψ(net::Union{NeuralNetwork, CachedNet},
               σ::Union{AStateBatch, ADoubleStateBatch}) where {T}
    out = out_similar(net, batch_size(σ))
    logψ!(out, net, σ)
end

function logψ(net::Union{NeuralNetwork, CachedNet},
               σ::Union{AStateBatchVec, ADoubleStateBatchVec}) where {T}
    out = similar(out_similar(net), 1, batch_size(σ), chain_length(σ))
    logψ!(out, net, σ)
end

# diagonal
logψ!(out::AbstractArray, net::MatrixNeuralNetwork, cache::NNCache, σ::AStateBatch) =
    logψ!(out, net, cache, σ, σ)

log_prob_ψ!(prob::AbstractArray, net::NeuralNetwork, σ) =
    log_prob_ψ!(prob, prob, net, σ)

function log_prob_ψ!(prob::AbstractArray, out_tmp::AbstractArray, net::MatrixNet, σ::AStateOrBatchOrVec)
    logψ!(out_tmp, net, σ, σ)
    prob .= abs.(out_tmp)
    return prob
end

function log_prob_ψ!(prob::AbstractArray, out_tmp::AbstractArray, net::NeuralNetwork, σ)
    logψ!(out_tmp, net, σ)
    prob .= 2 .* real.(out_tmp)
    return prob
end


#using Images
# Converts img to greyscale matrix representation of the image
#
# returns: mxn Float64 array
#function img_to_array(img_name)
#	img = load(img_name)
#	img = Gray.(img)
#	return convert(Array{Float64,2}, img)
#end

function load_array(file_name)
	value_arrays = []
	open(file_name) do file
		for ln in eachline(file)
			indices = [ c != "," for c in ln]
			values = split(ln, ",")
			push!(value_arrays, values)
		end
	end

	size = (parse(Int,value_arrays[1][1]),parse(Int,value_arrays[1][2]))
	vals = parse.(Int,value_arrays[2][1:end-1])

	v = reshape(vals,size[1],size[2])

end

# Used to keep pixel intensity sums lower
switcheroo(x) = x = 1-x

man_norm(x) = sum(x.^2)^(1/2)

# Partitions image into num_cols * num_rows rectangles and returns the
# the partitions as t_images and the sum of the region's pixel intensity
# as sums
function get_regions(img,s,num_cols=2,num_rows=3)
	# initialize returned values
	sums = Float64[]
	t_images = []

	for col in 1:num_cols
		for row in 1:num_rows
			# Reference the (col * row)th partition of img
			region = @view img[max(1,Int(round((row-1)*s[1]/num_rows))):Int(round(row*s[1]/num_rows)),
												max(1,Int(round((col-1)*s[2]/num_cols))):Int(round(col*s[2]/num_cols))]

			push!(t_images, region)

			push!(sums, man_norm(region))
		end
	end
	return sums, t_images
end

# Take the sums matrix and output a matrix that represents the
# braille representation
function get_output_vector(sums,tol_div=4, zero_tol=5)
	largest = maximum(sums)
	tolerance = largest/tol_div
	vect = [ abs(largest - sum) < tolerance for sum in sums]
	if largest < zero_tol
		vect = zeros(6)
	end
	return Int.(reshape(vect,3,2))
end

# Given an image of a single char we return the regions of the
# image, the sums of those regions and the output vector
#
function process_char(img)
	s = size(img)
	sums, t_images = get_regions(img,s)
	output_vector = get_output_vector(sums)
	return output_vector
end

# Preliminary version of processing the image to feed the
# single characters to the process_char function
#
function process_chars(img,s,num_chars)
	chars = [ @view img[:,max(1,Int(round((col-1)*s[2]/num_chars))):Int(round(col*s[2]/num_chars))]
						for col in 1:num_chars ]

	outputs = [ process_char(char) for char in chars ]
	return outputs
end

function row_to_num_chars(l_row)
	last_col = 0
	distance = 0
	distances = Int64[]
	dot_size = 0
	dot_sizes = Int64[]

	for col in l_row
		distance += 1
		dot_size += 1
		if col > 0 && last_col == 0
			push!(distances, distance)
			dot_size = 0

		elseif col == 0 && last_col > 0
			push!(dot_sizes, dot_size)
			distance = 0
		end

		last_col = col
	end

	return distances, maximum(dot_sizes)
end

function optimize_x(t_sum, dot_size, inner_dist, outer_dist)
	x = 2
	y = 1
	z = 0

	while dot_size*x + y*inner_dist + z*outer_dist < t_sum
		x += 2
		y = x/2
		z = x/2 - 1
	end

	return (x, y, z)
end

function get_distances(img,s)
	l_row = img[1,:]
	l_row_sum = 0
	for row_num in 1:s[1]
		row = @view img[row_num, :]
		row_sum = sum(row)
		if row_sum >= l_row_sum
			l_row = row
			l_row_sum = row_sum
		end
	end
	distances, dot_size = row_to_num_chars(l_row)
	return distances, dot_size
end

function determine_chars(distances, dot_size)

	t_sum = sum([distance + dot_size for distance in distances[2:end]])

	inner_dist = minimum(distances)
	outer_dist = maximum(distances)
	opt_x = optimize_x(t_sum, dot_size, inner_dist, outer_dist)

	return t_sum, opt_x
end

function get_min_space(distances, tolerance=3)
	s_dists = sort(distances, rev=true)
	min_space = s_dists[1]
	for col in 2:size(s_dists)[1]
		gap = abs(min_space - s_dists[col])
		if gap > min_space/tolerance
			return min_space
		end
		min_space = s_dists[col]
	end
	return Inf
end

function get_words(img, distances, dot_size, indices)
	cur_region_strt = 1
	cur_region_end = cur_region_strt

	dot_ticker = 1

	words = []
	for col in 2:size(distances)[1]
		if dot_ticker > 0
			cur_region_end += dot_size
		end
		dot_ticker = mod(dot_ticker + 1,3)
		cur_region_end += distances[col]
		if indices[col]
			word = @view img[:,cur_region_strt:cur_region_end]
			push!(words,word)
			cur_region_strt = cur_region_end
		end
	end
	word = @view img[:,cur_region_strt:end]
	push!(words,word)
	return words
end

function vert_trim(word)
	while true
		if sum(word[1,:]) > 0
			break
		end
		word = @view word[2:end,:]
	end
	while true
		if sum(word[end,:]) > 0
			break
		end
		word = @view word[1:end-1,:]
	end
	return word
end

function horiz_trim(word)
	while true
		if sum(word[:,1]) > 0
			break
		end
		word = @view word[:,2:end]
	end
	while true
		if sum(word[:,end]) > 0
			break
		end
		word = @view word[:,1:end-1]
	end

	return word
end

function add_padding(word)
	s = size(word)
	word = [ zeros(5,s[2]); word; zeros(5,s[2])]

	word = [ zeros(s[1]+10,10) word zeros(s[1]+10,10) ]
	return word
end

function trim_word(word)
	word = vert_trim(word)
	word = horiz_trim(word)
	word = add_padding(word)
	return word
end

function image_to_braille(file_name)
	img = load_array(file_name)
	img = map(switcheroo,img)

	s = size(img)
	distances, dot_size = get_distances(img,s)
	braille_chars = []

	min_space = get_min_space(distances)
	indices = [ dist >= min_space for dist in distances]

	words = get_words(img, distances, dot_size, indices)
	for word in words
		word = trim_word(word)
		distances, dot_size = get_distances(word,size(word))
		t_sum, opt_x = determine_chars(distances,dot_size)
		num_chars = opt_x[1]/2
		outputs = process_chars(word, size(word), num_chars)
		if sum(sum(outputs)) > 0
			push!(braille_chars, outputs, [[0 0; 0 0; 0 0]])
		end
	end
	return braille_chars
end


"""
spin15.jl - T1 optimized pulses
"""

WDIR = get(ENV, "ROBUST_QOC_PATH", "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

using Altro
using HDF5
using LinearAlgebra
using RobotDynamics
using StaticArrays
using TrajectoryOptimization
const RD = RobotDynamics
const TO = TrajectoryOptimization

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin15"
const SAVE_PATH = abspath(joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME))

# problem
const CONTROL_COUNT = 1
const STATE_COUNT = 2
const ASTATE_SIZE_BASE = STATE_COUNT * HDIM_ISO + 3 * CONTROL_COUNT
# state indices
const STATE1_IDX = 1:HDIM_ISO
const STATE2_IDX = STATE1_IDX[end] + 1:STATE1_IDX[end] + HDIM_ISO
const INTCONTROLS_IDX = STATE2_IDX[end] + 1:STATE2_IDX[end] + CONTROL_COUNT
const CONTROLS_IDX = INTCONTROLS_IDX[end] + 1:INTCONTROLS_IDX[end] + CONTROL_COUNT
const DCONTROLS_IDX = CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + CONTROL_COUNT
const INTGAMMA_IDX = DCONTROLS_IDX[end] + 1:DCONTROLS_IDX[end] + 1
# control indices
const D2CONTROLS_IDX = 1:CONTROL_COUNT
const DT_IDX = D2CONTROLS_IDX[end] + 1:D2CONTROLS_IDX[end] + 1

# model
# DA = decay aware, TO = time optimal
struct Model{DA, TO} <: AbstractModel
    Model(DA::Bool=true, TO::Bool=true) = new{DA, TO}()
end
@inline RD.state_dim(::Model{false, TO}) where TO = ASTATE_SIZE_BASE
@inline RD.state_dim(::Model{true, TO}) where TO = ASTATE_SIZE_BASE + 1
@inline RD.control_dim(::Model{DA, false}) where DA = CONTROL_COUNT
@inline RD.control_dim(::Model{DA, true}) where DA = CONTROL_COUNT + 1

# dynamics
function RD.discrete_dynamics(::Type{RK3}, model::Model{DA, TO}, astate::StaticVector,
                              acontrols::StaticVector, time::Real, dt::Real) where {DA, TO}
    if TO
        dt = acontrols[DT_IDX[1]]^2
    end

    h_prop = exp(dt * (FQ_NEGI_H0_ISO + astate[CONTROLS_IDX[1]] * NEGI_H1_ISO))
    state1 = h_prop * astate[STATE1_IDX]
    state2 = h_prop * astate[STATE2_IDX]
    intcontrols = astate[INTCONTROLS_IDX[1]] + dt * astate[CONTROLS_IDX[1]]
    controls = astate[CONTROLS_IDX[1]] + dt * astate[DCONTROLS_IDX[1]]
    dcontrols = astate[DCONTROLS_IDX[1]] + dt * acontrols[D2CONTROLS_IDX[1]]
    astate_ = [
        state1; state2; intcontrols; controls; dcontrols;
    ]

    if DA
        intgamma = (astate[INTGAMMA_IDX[1]] +
                    dt * amp_t1_reduced_spline(astate[CONTROLS_IDX[1]])^(-1))
        push!(astate_, intgamma)
    end

    return astate_
end

# main
function run_traj(;evolution_time=21.61, gate_type=zpiby2, time_optimal=false,
                  decay_aware=false, solver_type=altro, sqrtbp=false,
                  integrator_type=rk3, qs=[1e0, 1e0, 1e0, 1e-1, 1e2, 1e-1, 5e2],
                  smoke_test=false,
                  al_tol=1e-4, constraint_tol=1e-8, max_penalty=1e11,
                  pn_steps=2, save=true, verbose=true, dt_inv=Int64(1e1),
                  max_cost_value=1e8, max_iterations=Int64(2e5),
                  benchmark=false,)
    # model configuration
    dt = dt_inv^(-1)
    dt_max = (dt_inv / 2)^(-1)
    dt_min = (dt_inv * 2)^(-1)
    model = Model(decay_aware, time_optimal)
    n = state_dim(model)
    m = control_dim(model)
    t0 = 0.

    # initial state
    x0 = zeros(n)
    x0[STATE1_IDX] = IS1_ISO_
    x0[STATE2_IDX] = IS2_ISO_
    x0 = SVector{n}(x0)

    # final state
    if gate_type == xpiby2
        target_state1 = Array(XPIBY2_ISO_1)
        target_state2 = Array(XPIBY2_ISO_2)
    elseif gate_type == ypiby2
        target_state1 = Array(YPIBY2_ISO_1)
        target_state2 = Array(YPIBY2_ISO_2)
    elseif gate_type == zpiby2
        target_state1 = Array(ZPIBY2_ISO_1)
        target_state2 = Array(ZPIBY2_ISO_2)
    end
    xf = zeros(n)
    xf[STATE1_IDX] = target_state1
    xf[STATE2_IDX] = target_state2
    xf = SVector{n}(xf)

    # control amplitude constraint
    x_max = fill(Inf, n)
    x_max[CONTROLS_IDX] .= MAX_CONTROL_NORM_0
    x_max = SVector{n}(x_max)
    x_min = fill(-Inf, n)
    x_min[CONTROLS_IDX] .= -MAX_CONTROL_NORM_0
    x_min = SVector{n}(x_min)

    # control amplitude constraint at boundary
    x_max_boundary = fill(Inf, n)
    x_max_boundary[CONTROLS_IDX] .= 0
    x_max_boundary = SVector{n}(x_max_boundary)
    x_min_boundary = fill(-Inf, n)
    x_min_boundary[CONTROLS_IDX] .= 0
    x_min_boundary = SVector{n}(x_min_boundary)

    # dt bound
    u_max = SVector{m}([
        fill(Inf, CONTROL_COUNT);
        fill(sqrt(dt_max), eval(:($time_optimal ? 1 : 0))); #dt
    ])
    u_min = SVector{m}([
        fill(-Inf, CONTROL_COUNT);
        fill(sqrt(dt_min), eval(:($time_optimal ? 1 : 0))); #dt
    ])

    # initial trajectory
    N = Int(floor(evolution_time * dt_inv)) + 1
    U0 = [SVector{m}([
        fill(1e-4, CONTROL_COUNT);
        fill(sqrt(dt), eval(:($time_optimal ? 1 : 0)));
    ]) for k = 1:N - 1]
    X0 = [SVector{n}(
        fill(NaN, n)
    ) for k = 1:N]
    dt_ = time_optimal ? 1 : dt
    Z = Traj(X0, U0, dt_ * ones(N))

    # cost function
    Q = Diagonal(SVector{n}([
        fill(qs[1], STATE_COUNT * HDIM_ISO); # ψ1, ψ2
        fill(qs[2], CONTROL_COUNT); # ∫a
        fill(qs[3], CONTROL_COUNT); # a
        fill(qs[4], CONTROL_COUNT); # ∂a
        fill(qs[5], eval(:($decay_aware ? 1 : 0))); # ∫γ1
    ]))
    Qf = Q * N
    R = Diagonal(SVector{m}([
        fill(qs[6], CONTROL_COUNT); # ∂2a
        fill(qs[7], eval(:($time_optimal ? 1 : 0))); # Δt
    ]))
    objective = LQRObjective(Q, R, Qf, xf, N)

    # constraints
    # must satisfy control amplitude bound
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # must statisfy conrols start and end at 0
    control_bnd_boundary = BoundConstraint(n, m, x_max=x_max_boundary, x_min=x_min_boundary)
    # must satisfy dt bound
    dt_bnd = BoundConstraint(n, m, u_max=u_max, u_min=u_min)
    # must reach target state, must have zero net flux
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX; INTCONTROLS_IDX])
    # must obey unit norm
    norm_constraints = [NormConstraint(n, m, 1, TO.Equality(), idx)
                        for idx in [STATE1_IDX, STATE2_IDX]]

    constraints = ConstraintList(n, m, N)
    add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, control_bnd_boundary, N-1:N-1)
    add_constraint!(constraints, target_astate_constraint, N:N)
    for norm_constraint in norm_constraints
        add_constraint!(constraints, norm_constraint, 2:N-1)
    end
    if time_optimal
        add_constraint!(constraints, dt_bnd, 1:N-1)
    end

    # solve problem
    prob = Problem{IT_RDI[integrator_type]}(model, objective, constraints,
                                            x0, xf, Z, N, t0, evolution_time)
    solver = ALTROSolver(prob)
    verbose_pn = verbose ? true : false
    verbose_ = verbose ? 2 : 0
    projected_newton = solver_type == altro ? true : false
    constraint_tolerance = solver_type == altro ? constraint_tol : al_tol
    iterations_inner = smoke_test ? 1 : 300
    iterations_outer = smoke_test ? 1 : 30
    n_steps = smoke_test ? 1 : pn_steps
    set_options!(solver, square_root=sqrtbp, constraint_tolerance=constraint_tolerance,
                 projected_newton_tolerance=al_tol, n_steps=n_steps,
                 penalty_max=max_penalty, verbose_pn=verbose_pn, verbose=verbose_,
                 projected_newton=projected_newton, iterations_inner=iterations_inner,
                 iterations_outer=iterations_outer, iterations=max_iterations,
                 max_cost_value=max_cost_value)
    if benchmark
        benchmark_result = Altro.benchmark_solve!(solver)
    else
        benchmark_result = nothing
        Altro.solve!(solver)
    end

    # Post-process.
    acontrols_raw = TO.controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    Q_raw = Array(Q)
    Q_arr = [Q_raw[i, i] for i in 1:size(Q_raw)[1]]
    Qf_raw = Array(Qf)
    Qf_arr = [Qf_raw[i, i] for i in 1:size(Qf_raw)[1]]
    R_raw = Array(R)
    R_arr = [R_raw[i, i] for i in 1:size(R_raw)[1]]
    cidx_arr = Array(CONTROLS_IDX)
    d2cidx_arr = Array(D2CONTROLS_IDX)
    dtidx_arr = Array(DT_IDX)
    # Square the dts.
    if time_optimal
        acontrols_arr[:, DT_IDX] = acontrols_arr[:, DT_IDX] .^2
        # acontrols_arr[:, DT_IDX] = map(abs, acontrols_arr[:, DT_IDX])
    end
    cmax = TO.max_violation(solver)
    cmax_info = TO.findmax_violation(TO.get_constraints(solver))
    iterations_ = iterations(solver)

    result = Dict(
            "acontrols" => acontrols_arr,
            "controls_idx" => cidx_arr,
            "d2controls_dt2_idx" => d2cidx_arr,
            "dt_idx" => dtidx_arr,
            "evolution_time" => evolution_time,
            "astates" => astates_arr,
            "Q" => Q_arr,
            "Qf" => Qf_arr,
            "R" => R_arr,
            "cmax" => cmax,
            "cmax_info" => cmax_info,
            "dt" => dt,
            "solver_type" => Integer(solver_type),
            "sqrtbp" => Integer(sqrtbp),
            "max_penalty" => max_penalty,
            "constraint_tol" => constraint_tol,
            "al_tol" => al_tol,
            "integrator_type" => Integer(integrator_type),
            "gate_type" => Integer(gate_type),
            "save_type" => Integer(jl),
            "iterations" => iterations_,
            "pn_steps" => pn_steps,
            "max_iterations" => max_iterations,
            "max_cost_value" => max_cost_value,
    )

    if time_optimal
        # sample the important metrics
        (controls_sample, d2controls_dt2_sample, evolution_time_sample
         ) = sample_controls(save_file_path)
        result["controls_sample"] = controls_sample
        result["d2controls_dt2_sample"] = d2controls_dt2_sample
        result["evolution_time_sample"] = evolution_time_sample
        result["save_type"] = Integer(samplejl)
    end

    # save
    if save
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            for key in keys(result)
                write(save_file, key, result[key])
            end
        end
        result["save_file_path"] = save_file_path
    end

    result = benchmark ? benchmark_result : result

    return result
end


"""
    invgeotransform!(gt_in::Vector{Float64}, gt_out::Vector{Float64})

Invert Geotransform.

This function will invert a standard 3x2 set of GeoTransform coefficients. This
converts the equation from being pixel to geo to being geo to pixel.

### Parameters
* `gt_in`       Input geotransform (six doubles - unaltered).
* `gt_out`      Output geotransform (six doubles - updated).

### Returns
`gt_out`
"""
function invgeotransform!(
    gt_in::Vector{Float64},
    gt_out::Vector{Float64},
)::Vector{Float64}
    result = Bool(GDAL.gdalinvgeotransform(pointer(gt_in), pointer(gt_out)))
    result || error("Geotransform coefficients is uninvertable")
    return gt_out
end

invgeotransform(gt_in::Vector{Float64})::Vector{Float64} =
    invgeotransform!(gt_in, Array{Float64}(undef, 6))

"""
    applygeotransform(geotransform::Vector{Float64}, pixel::Float64,
        line::Float64)

Apply GeoTransform to x/y coordinate.

Applies the following computation, converting a (pixel,line) coordinate into a
georeferenced `(geo_x,geo_y)` location.
```
    geo_x = geotransform[1] + pixel*geotransform[2] + line*geotransform[3]
    geo_y = geotransform[4] + pixel*geotransform[5] + line*geotransform[6]
```
### Parameters
* `geotransform`  Six coefficient GeoTransform to apply.
* `pixel`           input pixel position.
* `line`            input line position.
"""
function applygeotransform(
    geotransform::Vector{Float64},
    pixel::Float64,
    line::Float64,
)::Vector{Float64}
    geo_xy = Vector{Float64}(undef, 2)
    geo_x = pointer(geo_xy)
    geo_y = geo_x + sizeof(Float64)
    GDAL.gdalapplygeotransform(pointer(geotransform), pixel, line, geo_x, geo_y)
    return geo_xy
end

"""
    composegeotransform!(gt1::Vector{Float64}, gt2::Vector{Float64},
        gtout::Vector{Float64})

Compose two geotransforms.

The resulting geotransform is the equivalent to `padfGT1` and then `padfGT2`
being applied to a point.

### Parameters
* `gt1`     the first geotransform, six values.
* `gt2`     the second geotransform, six values.
* `gtout`   the output geotransform, six values.
"""
function composegeotransform!(
    gt1::Vector{Float64},
    gt2::Vector{Float64},
    gtout::Vector{Float64},
)::Vector{Float64}
    GDAL.gdalcomposegeotransforms(pointer(gt1), pointer(gt2), pointer(gtout))
    return gtout
end

function composegeotransform(
    gt1::Vector{Float64},
    gt2::Vector{Float64},
)::Vector{Float64}
    return composegeotransform!(gt1, gt2, Vector{Float64}(undef, 6))
end


function LyapunovOrbit(arg0::CR3BPSystem, arg1::PVCoordinates, arg2::jdouble)
    return LyapunovOrbit((CR3BPSystem, PVCoordinates, jdouble), arg0, arg1, arg2)
end

function LyapunovOrbit(arg0::RichardsonExpansion, arg1::jdouble)
    return LyapunovOrbit((RichardsonExpansion, jdouble), arg0, arg1)
end


function step_through_env(env::NatureEnv, policy::T) where {T <: MultiPPOManager}
    RLBase.reset!(env)

    get_player_poses() = [p.pos for p in env.players if !p.dead]
    function get_food_poses()
        fps = Vector()
        for f in 1:env.food_types
            push!(fps, Vector())
            for pos in CartesianIndices(env.food_frames[f])
                if env.food_frames[f][pos] > 0f0
                    push!(fps[f], (Tuple(pos)..., env.food_frames[f][pos]))
                end
            end
        end
        fps
    end

    player_poses = []
    food_poses = []
    food_color_maps = (
        (:darkblue, :lightblue),
        (:yellow, :orange))
    nframes = env.episode_len
    frame = 0

    while !is_terminated(env) && frame < nframes
        env |> policy |> env
        push!(player_poses, get_player_poses())
        push!(food_poses, get_food_poses())
        frame += 1
    end

    fig = Figure()
    ax  = Axis(fig[1,1])
    limits!(ax, 0, env.world_size[1], 0, env.world_size[2])
    lsgrid = labelslidergrid!(fig, ["step"], [1:1:frame])
    fig[2, 1] = lsgrid.layout
    slider = lsgrid.sliders[1]

    for f in 1:env.food_types
        fps = @lift [(x, y) for (x, y, _) in food_poses[$(slider.value)][f]]
        fvs = @lift [z for (_, _, z) in food_poses[$(slider.value)][f]]
        GLMakie.scatter!(fps,color=fvs, colormap=food_color_maps[f], markersize=40, marker=:rect)
    end
    pps = @lift player_poses[$(slider.value)]
    GLMakie.scatter!(pps, color=:pink, markersize=60)

    display(fig)
    nothing
end


using WAV
using SpeechAugment: addEcho

T60 = 0.2
# read a wav file
wav,fs = wavread("xxx.wav")
# add reverberation effect on audio
echowav = addEcho(wav, fs, T60,
                  (3, 4, 2.5),      # room size
                  (1, 2, 1.2),      # sound src
                  (1.6, 2.6, 1.0)); # micphone


module Game2048

export Dirs, left, right, up, down, move!, move

@enum Dirs left right up down

include("uint_board/move-up.jl")
# include("gamesim/move-board.jl")
# include("gamesim/sim-game.jl")
# include("gamesim/play_game_with_policy.jl")

# include("identity-skip.jl")

include("rotate-mirror.jl")

include("bitboard/make_lookups.jl")
include("bitboard/bitboard.jl")
include("bitboard/simulate_bb.jl")
# include("bitboard/playahead.jl")
# include("bitboard/fingerprint.jl")

# include("monte-carlo-tree.jl")

# include("attempts2020/5_policy_gradient.jl")
end # module


using Distributed

@everywhere using Ipopt

include("code-2-lib/parsers.jl")
include("code-2-lib/lib.jl")

function compute_solution2(con_file::String, inl_file::String, raw_file::String, rop_file::String, time_limit::Int, scoring_method::Int, network_model::String; output_dir::String="", scenario_id::String="none")
    time_data_start = time()
    PowerModels.silence()
    goc_data = parse_goc_files(con_file, inl_file, raw_file, rop_file, scenario_id=scenario_id)
    network = build_pm_model(goc_data)

    sol = read_solution1(network, output_dir=output_dir)
    PowerModels.update_data!(network, sol)

    check_network_solution(network)

    network_tmp = deepcopy(network)
    balance = compute_power_balance_deltas!(network_tmp)

    if balance.p_delta_abs_max > 0.01 || balance.q_delta_abs_max > 0.01
        error(LOGGER, "solution1 power balance requirements not satified (all power balance values should be below 0.01). $(balance)")
    end
    load_time = time() - time_data_start

    ###### Prepare Solution 2 ######

    time_contingencies_start = time()

    gen_cont_total = length(network["gen_contingencies"])
    branch_cont_total = length(network["branch_contingencies"])
    cont_total = gen_cont_total + branch_cont_total

    cont_order = contingency_order(network)

    workers = Distributed.workers()

    process_data = []

    cont_per_proc = cont_total/length(workers)

    for p in 1:length(workers)
        cont_start = trunc(Int, ceil(1+(p-1)*cont_per_proc))
        cont_end = min(cont_total, trunc(Int,ceil(p*cont_per_proc)))
        pd = (
            pid = p,
            processes = length(workers),
            con_file = con_file,
            inl_file = inl_file,
            raw_file = raw_file,
            rop_file = rop_file,
            scenario_id = scenario_id,
            output_dir = output_dir,
            cont_range = cont_start:cont_end,
        )
        push!(process_data, pd)
    end

    for (i,pd) in enumerate(process_data)
        info(LOGGER, "worker task $(pd.pid): $(length(pd.cont_range)) / $(pd.cont_range)")
    end

    solution2_files = pmap(solution2_solver, process_data, retry_delays = zeros(3))
    sort!(solution2_files)
    #println("pmap result: $(solution2_files)")

    time_contingencies = time() - time_contingencies_start
    info(LOGGER, "contingency eval time: $(time_contingencies)")
    info(LOGGER, "time per contingency: $(time_contingencies/cont_total)")

    info(LOGGER, "combine $(length(solution2_files)) solution2 files")
    combine_files(solution2_files, "solution2.txt"; output_dir=output_dir)
    remove_files(solution2_files)

    println("")

    data = [
        "----",
        "scenario id",
        "bus",
        "branch",
        "gen_cont",
        "branch_cont",
        "runtime (sec.)",
    ]
    println(join(data, ", "))

    data = [
        "DATA_SSS",
        goc_data.scenario,
        length(network["bus"]),
        length(network["branch"]),
        length(network["gen_contingencies"]),
        length(network["branch_contingencies"]),
        time_contingencies,
    ]
    println(join(data, ", "))
end

@everywhere function solution2_solver(process_data)
    #println(process_data)
    time_data_start = time()
    PowerModels.silence()
    goc_data = parse_goc_files(
        process_data.con_file, process_data.inl_file, process_data.raw_file,
        process_data.rop_file, scenario_id=process_data.scenario_id)
    network = build_pm_model(goc_data)

    sol = read_solution1(network, output_dir=process_data.output_dir)
    PowerModels.update_data!(network, sol)
    correct_voltage_angles!(network)

    time_data = time() - time_data_start

    for (i,bus) in network["bus"]
        if haskey(bus, "evhi")
            bus["vmax"] = bus["evhi"]
        end
        if haskey(bus, "evlo")
            bus["vmin"] = bus["evlo"]
        end
    end

    for (i,branch) in network["branch"]
        if haskey(branch, "rate_c")
            branch["rate_a"] = branch["rate_c"]
        end
    end

    contingencies = contingency_order(network)[process_data.cont_range]

    for (i,branch) in network["branch"]
        g, b = PowerModels.calc_branch_y(branch)
        tr, ti = PowerModels.calc_branch_t(branch)
        branch["g"] = g
        branch["b"] = b
        branch["tr"] = tr
        branch["ti"] = ti
    end

    bus_gens = gens_by_bus(network)

    network["delta"] = 0
    for (i,bus) in network["bus"]
        bus["vm_base"] = bus["vm"]
        bus["vm_start"] = bus["vm"]
        bus["va_start"] = bus["va"]
        bus["vm_fixed"] = length(bus_gens[i]) != 0

    end

    for (i,gen) in network["gen"]
        gen["pg_base"] = gen["pg"]
        gen["pg_start"] = gen["pg"]
        gen["qg_start"] = gen["qg"]
        gen["pg_fixed"] = false
        gen["qg_fixed"] = false
    end

    #nlp_solver = JuMP.with_optimizer(Ipopt.Optimizer, tol=1e-6, mu_init=1e-6, hessian_approximation="limited-memory", print_level=0)
    nlp_solver = JuMP.with_optimizer(Ipopt.Optimizer, tol=1e-6, print_level=0)
    #nlp_solver = JuMP.with_optimizer(Ipopt.Optimizer, tol=1e-6)
    #nlp_solver = JuMP.with_optimizer(Ipopt.Optimizer, tol=1e-6, hessian_approximation="limited-memory")

    pad_size = trunc(Int, ceil(log(10,process_data.processes)))
    padded_pid = lpad(string(process_data.pid), pad_size, "0")
    solution_filename = "solution2-$(padded_pid).txt"

    if length(process_data.output_dir) > 0
        solution_path = joinpath(process_data.output_dir, solution_filename)
    else
        solution_path = solution_filename
    end
    if isfile(solution_path)
        warn(LOGGER, "removing existing solution2 file $(solution_path)")
        rm(solution_path)
    end
    open(solution_path, "w") do sol_file
        # creates an empty file in the case of workers without contingencies
    end

    #network_tmp = deepcopy(network)
    for cont in contingencies
        if cont.type == "gen"
            info(LOGGER, "working on: $(cont.label)")
            time_start = time()
            network_tmp = deepcopy(network)
            debug(LOGGER, "contingency copy time: $(time() - time_start)")
            network_tmp["cont_label"] = cont.label

            cont_gen = network_tmp["gen"]["$(cont.idx)"]
            cont_gen["contingency"] = true
            cont_gen["gen_status"] = 0
            pg_lost = cont_gen["pg"]

            time_start = time()
            result = run_fixpoint_pf_v2_2!(network_tmp, pg_lost, ACRPowerModel, nlp_solver, iteration_limit=5)
            debug(LOGGER, "second-stage contingency solve time: $(time() - time_start)")

            result["solution"]["label"] = cont.label
            result["solution"]["feasible"] = (result["termination_status"] == LOCALLY_SOLVED)
            result["solution"]["cont_type"] = "gen"
            result["solution"]["cont_comp_id"] = cont.idx

            result["solution"]["gen"]["$(cont.idx)"]["pg"] = 0.0
            result["solution"]["gen"]["$(cont.idx)"]["qg"] = 0.0

            correct_contingency_solution!(network, result["solution"])
            open(solution_path, "a") do sol_file
                sol2 = write_solution2_contingency(sol_file, network, result["solution"])
            end

            network_tmp["gen"]["$(cont.idx)"]["gen_status"] = 1
        elseif cont.type == "branch"
            info(LOGGER, "working on: $(cont.label)")
            time_start = time()
            network_tmp = deepcopy(network)
            debug(LOGGER, "contingency copy time: $(time() - time_start)")
            network_tmp["cont_label"] = cont.label
            network_tmp["branch"]["$(cont.idx)"]["br_status"] = 0

            time_start = time()
            result = run_fixpoint_pf_v2_2!(network_tmp, 0.0, ACRPowerModel, nlp_solver, iteration_limit=5)
            debug(LOGGER, "second-stage contingency solve time: $(time() - time_start)")

            result["solution"]["label"] = cont.label
            result["solution"]["feasible"] = (result["termination_status"] == LOCALLY_SOLVED)
            result["solution"]["cont_type"] = "branch"
            result["solution"]["cont_comp_id"] = cont.idx

            correct_contingency_solution!(network, result["solution"])
            open(solution_path, "a") do sol_file
                sol2 = write_solution2_contingency(sol_file, network, result["solution"])
            end

            network_tmp["branch"]["$(cont.idx)"]["br_status"] = 1
        else
            @assert("contingency type $(cont.type) not known")
        end
    end

    return solution_path
end


abstract type AbstractLeft end
abstract type AbstractRight end
abstract type AbstractUnit <: AbstractRight; end

struct Px <:   AbstractUnit;   val::Real; end
struct Em <:   AbstractUnit;   val::Real; end
struct Cm <:   AbstractUnit;   val::Real; end
struct In <:   AbstractUnit;   val::Real; end
struct Deg <:  AbstractUnit;   val::Real; end
struct Mm <:   AbstractUnit;   val::Real; end
struct Pt <:   AbstractUnit;   val::Real; end
struct Pc <:   AbstractUnit;   val::Real; end
struct Ex <:   AbstractUnit;   val::Real; end
struct Ch <:   AbstractUnit;   val::Real; end
struct Rem <:  AbstractUnit;   val::Real; end
struct Vw <:   AbstractUnit;   val::Real; end
struct Vh <:   AbstractUnit;   val::Real; end
struct Vmin <: AbstractUnit;   val::Real; end
struct Vmax <: AbstractUnit;   val::Real; end
struct Pcnt <: AbstractUnit;   val::Real; end

struct Hex <:   AbstractRight; val::Array{Int64}; end
struct Int_  <: AbstractRight; val::Int64;       end # These do seem redundant but actually help.
struct Float <: AbstractRight; val::Float64;     end
# ==============================================================================
struct Class <: AbstractLeft;     val::String; end
struct Tag <: AbstractLeft;       val::String; end
struct Ident <: AbstractLeft;       val::String; end
struct Attr <: AbstractLeft;      val::String; end

#"digit", "float", "value", "reference", "hex", "string", "unit"
struct Digit <: AbstractRight;     val::String; end
struct Value <: AbstractRight;     val::String; end
struct Reference <: AbstractRight; val::String; end
struct Str <: AbstractRight;       val::String; end
struct Unit <: AbstractRight;       val::String; end

struct Alpha <: AbstractRight;     val::String; end

struct Open;      val::String; end
struct Close;     val::String; end
struct End;       val::String; end

struct Hash;      val::String; end
struct Dot;       val::String; end
struct Dotdot;    val::String; end
struct Dollar;    val::String; end
struct At;        val::String; end
struct Space;     val::String; end

# ==============================================================================
units = [ "px", "em", "cm", "in", "deg", "mm", "pt", "pc",
               "ex", "ch", "rem", "vw", "vh", "vmin", "vmax"] # % & #


# Binary functions for Nullables ...

# imports
import Base: convert
import Base: +, -, *, /, \, ^, %, &, |, $, >>>, >>, <<, ==, !=, <, <=, >, >=, //,  div, fld, rem, divrem, mod, copysign, flipsign, hypot, ldexp, log, beta, lbeta, airy, airyx, besselj, besseljx, bessely, besselyx, hankelh1, hankelh1x, hankelh2, hankelh2x, besseli, besselix, besselk, besselkx

# Convert to Nullable
function as_nullable{T}(x::T)
  T <: Nullable ? x : Nullable{T}(x)
end

# Convert Any to Nullable
function convert{T}(x::Type{Nullable{T}}, y)
  Nullable{T}(convert(T, y))
end

function convert(x::Type{Nullable}, y)
  X = typeof(y)
  Nullable{X}(y)
end

# Convert Nullable to Nullable
function convert{T <: Nullable, U <: Nullable}(x::Type{T}, y::U)
  X = eltype(T)
  Nullable{X}(convert(X, eltype(y)))
end

# Convert Nullable to Any
function convert{T <: Nullable}(U::Type, y::T)
  convert(U, get(y))
end

# Define binary operators
arith = [:+, :-, :*, :/, :\, :^, :%, ://]
bit = [:&, :|, :$, :>>>, :>>, :<<]
update = [:+=, :-=, :/=, :\=, :%=, :^=, :|=, :>>>=, :>>=, :<<=]
comp = [:(==), :!=, :<, :<=, :>, :>=]
funs = [:div, :fld, :rem, :divrem, :mod, :copysign, :flipsign, :hypot, :ldexp, :log, :beta, :lbeta, :airy, :airyx, :besselj, :besseljx, :bessely, :besselyx, :hankelh1, :hankelh1x, :hankelh2, :hankelh2x, :besseli, :besselix, :besselk, :besselkx]

# Combine to one vector
syms = [arith; bit; update; comp; funs]

# Symbol expression union type
sym_expr = Union{Symbol, Expr}

# Nullable{T} & Nullable{T} family
function exnn{Q<:Union{Symbol, Expr}}(sym::Q)
  quote
    function $sym{T<:Nullable, U<:Nullable}(x::T, y::U)
      try
        ret = $sym(get(x), get(y))
        X = typeof(ret)
        Nullable{X}(ret)
      catch
        Nullable()
      end
    end
  end
end

# For the Nullable{T} & U family
function exnu{Q<:Union{Symbol, Expr}}(sym::Q)
  quote
    function $sym{T<:Nullable, U}(x::T, y::U)
       try
         ret = $sym(get(x), y)
         X = typeof(ret)
         Nullable{X}(ret)
       catch
          Nullable()
       end
    end
  end
end

# For the U & Nullable{T} family
function exun{Q<:Union{Symbol, Expr}}(sym::Q)
  quote
    function $sym{U, T<:Nullable}(x::U, y::T)
      try
        ret = $sym(x, get(y))
        X = typeof(ret)
        Nullable{X}(ret)
      catch
        Nullable()
      end
    end
  end
end

# Source the functions
for i in [exnn, exnu, exun]
  for j in syms
    eval(i(j))
  end
end


using Base.Test
using CGP
CGP.Config.init("cfg/test.yaml")

function test_mutate(c::Chromosome, child::Chromosome)
      @test child != c
      @test child.genes != c.genes
      @test child.nin == c.nin
      @test child.nout == c.nout
      @test distance(c, child) > 0
end

@testset "Mutation tests" begin
    for ct in CGP.chromosomes
        println(ct)
        nin = rand(1:100); nout = rand(1:100)
        c = ct(nin, nout)
        cgenes = deepcopy(c.genes)
        @testset "Clone $ct" begin
            copy = clone(c)
            @test copy.genes == c.genes
        end
        @testset "Constructor $ct" begin
            child = ct(c)
            test_mutate(c, child)
        end
        @testset "Mutate genes $ct" begin
            child = gene_mutate(c)
            test_mutate(c, child)
            @test length(child.genes) == length(c.genes)
        end
        @testset "Active genes $ct" begin
            child = active_gene_mutate(c)
            test_mutate(c, child)
            @test length(child.genes) == length(c.genes)
        end
        @testset "Add nodes $ct" begin
            child = add_nodes(c)
            test_mutate(c, child)
            @test length(child.genes) > length(c.genes)
            @test forward_connections(c) != forward_connections(child)
            @test issubset([n.p for n in c.nodes], [n.p for n in child.nodes])
            @test issubset([n.f for n in c.nodes], [n.f for n in child.nodes])
        end
        @testset "Delete nodes $ct" begin
            child = delete_nodes(c)
            test_mutate(c, child)
            @test length(child.genes) < length(c.genes)
            @test forward_connections(c) != forward_connections(child)
            @test issubset([n.p for n in child.nodes], [n.p for n in c.nodes])
            @test issubset([n.f for n in child.nodes], [n.f for n in c.nodes])
        end
        @testset "Mixed node mutate $ct" begin
            child = mixed_node_mutate(c)
            test_mutate(c, child)
        end
        @testset "Add subtree $ct" begin
            child = add_nodes(c)
            test_mutate(c, child)
            @test length(child.genes) > length(c.genes)
            @test forward_connections(c) != forward_connections(child)
            @test issubset([n.p for n in c.nodes], [n.p for n in child.nodes])
            @test issubset([n.f for n in c.nodes], [n.f for n in child.nodes])
        end
        @testset "Delete subtree $ct" begin
            child = delete_nodes(c)
            test_mutate(c, child)
            @test length(child.genes) < length(c.genes)
            @test forward_connections(c) != forward_connections(child)
            @test issubset([n.p for n in child.nodes], [n.p for n in c.nodes])
            @test issubset([n.f for n in child.nodes], [n.f for n in c.nodes])
        end
        @testset "Mixed subtree mutate $ct" begin
            child = mixed_subtree_mutate(c)
            test_mutate(c, child)
        end
    end
end

nothing


module OneDTMazeUtils

using SparseArrays
using Distributions

import ..TMazeCumulantSchedules
import ..OneDTmazeConst
import ..OneDTMaze
import ..Learner
import ..check_goal
import ..range_check
import ..get_action_probs
import ..GVFHordes
import ..update
import ..Curiosity
import ..GVFSRHordes
import ..SRCreationUtils
import ..FeatureCreator
using ...StatsBase

const TMCS = TMazeCumulantSchedules
const ODTMC = OneDTmazeConst
const SRCU = Curiosity.SRCreationUtils

#####
# GVF Parameter Functions
####

struct GoalTermination <: GVFHordes.GVFParamFuncs.AbstractDiscount
    γ::Float64
end

function Base.get(gt::GoalTermination; state_tp1, kwargs...)
    any([check_goal(OneDTMaze, i, state_tp1) for i in 1:4]) ? 0.0 : gt.γ
end

struct GoalPolicy <: GVFHordes.GVFParamFuncs.AbstractPolicy
    goal::Int
end
(π::GoalPolicy)(s) = sample(Weights([get(π, state_t=s, action_t=a) for a ∈ 1:4]))

function Base.get(π::GoalPolicy; state_t, action_t, kwargs...)
    cur_x = state_t[1]
    cur_y = state_t[2]
    if π.goal == 1
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 1.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 0.0 - ODTMC.EPSILON, 0.0 + ODTMC.EPSILON)
            ODTMC.UP == action_t
        else
            ODTMC.LEFT == action_t
        end
    elseif π.goal == 2
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 1.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.LEFT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 0.0 - ODTMC.EPSILON, 0.0 + ODTMC.EPSILON)
            ODTMC.DOWN == action_t
        else
            ODTMC.LEFT == action_t
        end
    elseif π.goal == 3
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 0.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 1.0 - ODTMC.EPSILON, 1.0 + ODTMC.EPSILON)
            ODTMC.UP == action_t
        else
            ODTMC.RIGHT == action_t
        end
    elseif π.goal == 4
        if cur_x == 0.5
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            else
                ODTMC.UP == action_t
            end
        elseif cur_x == 0.0
            if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON)
                ODTMC.RIGHT == action_t
            elseif cur_y <= 0.8 - ODTMC.EPSILON
                ODTMC.UP == action_t
            elseif cur_y >= 0.8 + ODTMC.EPSILON
                ODTMC.DOWN == action_t
            end
        elseif range_check(cur_x, 1.0 - ODTMC.EPSILON, 1.0 + ODTMC.EPSILON)
            ODTMC.DOWN == action_t
        else
            ODTMC.RIGHT == action_t
        end
    end
end

####
# Demon Creation
####

function create_demons(parsed, demon_projected_fc = nothing)
    action_space = 4
    demons = if parsed["demon_learner"] != "SR"
        GVFHordes.Horde(
            [GVFHordes.GVF(GVFHordes.GVFParamFuncs.FeatureCumulant(i+2),
                           GoalTermination(0.9),
                           GoalPolicy(i)) for i in 1:4])
    elseif parsed["demon_learner"] == "SR"
        @assert demon_projected_fc != nothing
        pred_horde =  GVFHordes.Horde(
                [GVFHordes.GVF(GVFHordes.GVFParamFuncs.FeatureCumulant(i+2),
                     GVFHordes.GVFParamFuncs.ConstantDiscount(0.0),
                     GoalPolicy(i)) for i in 1:4])

        SF_policies = [GoalPolicy(i) for i in 1:4]
        SF_discounts = [GoalTermination(0.9) for i in 1:4]
        num_SFs = length(SF_policies)
        # SF_horde = if parsed["demon_reward_feature_type"] == "state"
        #     SRCU.create_state_SF_horde(SF_policies, SF_discounts, demon_projected_fc,1:action_space)
        # else
        # SF_horde = SRCU.create_SF_horde(SF_policies, SF_discounts, demon_projected_fc,1:action_space)
        SF_horde = SRCU.create_SF_horde(SF_policies, SF_discounts, demon_projected_fc,1:action_space)

        # end
        GVFSRHordes.SRHorde(pred_horde, SF_horde, num_SFs, demon_projected_fc)
    else
        throw(ArgumentError("Cannot create demons"))
    end
    return demons
end

function make_behaviour_gvf(learner, γ, state_constructor, expl_strat) #discount, state_constructor_func, learner, exploration_strategy)
    function b_π(state_constructor, learner, exploration_strategy; kwargs...)
        s = state_constructor(kwargs[:state_t])
        preds = learner(s)
        return exploration_strategy(preds)[kwargs[:action_t]]
    end
    GVF_policy = GVFHordes.GVFParamFuncs.FunctionalPolicy((;kwargs...) -> b_π(state_constructor, learner, expl_strat; kwargs...))
    BehaviourGVF = GVFHordes.GVF(GVFHordes.GVFParamFuncs.RewardCumulant(),
                       GoalTermination(γ),
                       GVF_policy)
end

# ####
struct IdealStateActionDemonFeatures <: FeatureCreator
    num_actions::Int
end
function project_features(fc::IdealStateActionDemonFeatures, s_t, a_t, s_tp1)
    # new_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, s_tp1) for i in 1:4]))
    goal_ind = findfirst([check_goal(OneDTMaze, i, s_tp1) for i in 1:4])
    reward_feature = zeros(Int(fc.num_actions*4))
    if !(goal_ind isa Nothing)
        reward_feature[Int((goal_ind-1)*fc.num_actions + a_t)] = 1
    end
    return sparsevec(reward_feature)
end
(FP::IdealStateActionDemonFeatures)(s_t,a_t,s_tp1) = project_features(FP, s_t,a_t,s_tp1)

Base.size(FP::IdealStateActionDemonFeatures) = 16
####
# Ideal Feature Creator
####
struct IdealDemonFeatures <: FeatureCreator
end

function project_features(fc::IdealDemonFeatures, state,action,state_tp1)
    new_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, state_tp1) for i in 1:4]))
    return new_state
end

(FP::IdealDemonFeatures)(state, action, next_state) = project_features(FP, state, action, next_state)
Base.size(FP::IdealDemonFeatures) = 4

struct MarthaIdealDemonFeatures <: FeatureCreator
    num_actions::Int
end

function project_features(fc::MarthaIdealDemonFeatures, state)
    new_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, state, ODTMC.ACTION_STEP + ODTMC.EPSILON + 0.00001) for i in 1:4]))
    # alt_state = sparsevec(convert(Array{Int,1}, [check_goal(OneDTMaze, i, state_tp1) for i in 1:4]))
    return new_state
end

(FP::MarthaIdealDemonFeatures)(state) = project_features(FP, state)
Base.size(FP::MarthaIdealDemonFeatures) = 4

####
# GPI Feature Creation
####
struct TMazeEncoding <: FeatureCreator
    num_segments::Int
    num_partitions::Int
    num_actions::Int
    function TMazeEncoding()
        new(7,3,4)
    end
end
Base.size(FP::TMazeEncoding) = FP.num_segments * FP.num_partitions * FP.num_actions
function project_features(fc::FeatureCreator, obs, a_t, obs_tp1)
    segment = Inf
    partition = Inf
    x,y = obs

    #Segment 1: Top Right
    if (x == 0.0 && y >= 0.8)
        segment_length = (1.0 + ODTMC.ACTION_STEP) - (0.8)
        segment = 1
        partition = (y - 0.8) / segment_length
    #Segment 2: Bottom Right
    elseif (x == 0.0 && y < 0.8)
        segment_length = (0.8) - (0.6 - ODTMC.ACTION_STEP)
        segment = 2
        partition = (y - 0.8) / segment_length
    #Segment 3: Middle Left Branch
    elseif (y == 0.8 && range_check(x,0.0,0.5))
        segment_length = (0.5 - 0.0)
        segment = 3
        partition = (x) / segment_length
    # Segment 4: Middle Right Branch
    elseif (y == 0.8 && range_check(x,0.5,1.0))
        segment_length = (1.0-0.5)
        segment = 4
        partition = (x - 0.5)/segment_length
    #Segment 5: Bottom Branch
    elseif (x == 0.5)
        segment_length = (0.8 - 0.0)
        segment = 5
        partition = (y - 0.0)/segment_length
    # #Segment 6: Top Left Branch
    elseif (x == 1.0 && y >= 0.8)
        segment_length = (1.0 + ODTMC.ACTION_STEP) - (0.8)
        segment = 6
        partition = (y - 0.8) / segment_length
        #Segment 7: Bottom Left Branch
    elseif (x == 1.0 && y < 0.8)
        segment_length = (0.8) - (0.6 - ODTMC.ACTION_STEP)
        segment = 7
        partition = (y - 0.8) / segment_length
    else
        @warn "Not A Valid State (x,y)"
    end

    state_action = zeros(fc.num_segments * fc.num_partitions * fc.num_actions)
    # Need maximum since if agent is right on boundary, ceil (0.0) is 0.
    partition_offset = maximum([Int(ceil(partition * fc.num_partitions)),1])
    if partition_offset == 0
        @show x,y,segment
    end
    state_ind = (segment-1)*fc.num_partitions + partition_offset
    state_action_ind = state_ind * fc.num_actions + a_t

    state_action[state_action_ind] = 1
    return sparsevec(state_action)
end
(FP::TMazeEncoding)(state, action, state_tp1) = project_features(FP, state, action, state_tp1)

####
# Behaviour policies
####

Base.@kwdef struct RoundRobinPolicy <: Learner
    update = Nothing
end

Curiosity.update!(learner::RoundRobinPolicy, args...) = nothing

Base.get(π::RoundRobinPolicy; state_t, action_t, kwargs...) =
    get_action_probs(π, state_t, nothing)[action_t]

function Curiosity.get_action_probs(π::RoundRobinPolicy, features, state)
    cur_x = state[1]
    cur_y = state[2]
    ret = zeros(4)

    if cur_x == 0.5
        if range_check(cur_y, 0.8 - ODTMC.EPSILON, 0.8 + ODTMC.EPSILON) # Middle Junction
            ret[ODTMC.LEFT] = 0.5
            ret[ODTMC.RIGHT] = 0.5
        else # Middle Hallway
            ret[ODTMC.UP] = 1.0
        end
    elseif cur_y == 0.8 && range_check(cur_x, 0.0 - ODTMC.EPSILON, 0.0 + ODTMC.EPSILON) # Left Junction
        ret[ODTMC.UP] = 0.5
        ret[ODTMC.DOWN] = 0.5
    elseif cur_y == 0.8 && range_check(cur_x, 1.0 - ODTMC.EPSILON, 1.0 + ODTMC.EPSILON)
        ret[ODTMC.UP] = 0.5
        ret[ODTMC.DOWN] = 0.5
    elseif cur_x == 0.0
        if cur_y > 0.8
            ret[ODTMC.UP] = 1.0
        else
            ret[ODTMC.DOWN] = 1.0
        end
    elseif cur_x == 1.0
        if cur_y > 0.8
            ret[ODTMC.UP] = 1.0
        else
            ret[ODTMC.DOWN] = 1.0
        end
    elseif cur_x < 0.5
        ret[ODTMC.LEFT] = 1.0
    else
        ret[ODTMC.RIGHT] = 1.0
    end
    ret
end

####
# Cumulant Schedules
####
DrifterDistractor(parsed) = begin
    c_dist = Uniform(parsed["constant_target"][1],parsed["constant_target"][2])
    c1,c2 = rand(c_dist,2)
    if "drifter" ∈ keys(parsed)
        TMCS.DrifterDistractor(
            c1,
            c2,
            parsed["drifter"][1],
            parsed["drifter"][2],
            parsed["distractor"][1],
            parsed["distractor"][2])
    else
        TMCS.DrifterDistractor(
            c1,
            c2,
            parsed["drifter_init"],
            parsed["drifter_std"],
            parsed["distractor_mean"],
            parsed["distractor_std"])
    end
end

function get_cumulant_schedule(parsed)
    sched = parsed["cumulant_schedule"]
    if parsed["cumulant_schedule"] == "DrifterDistractor"
        DrifterDistractor(parsed)
    elseif parsed["cumulant_schedule"] == "Constant"
        if parsed["cumulant"] isa Number
            TMCS.Constant(parsed["cumulant"])
        else
            TMCS.Constant(parsed["cumulant"]...)
        end

    else
        throw("$(sched) Not Implemented")
    end
end

function get_true_values(env::Curiosity.OneDTMaze, eval_set)
    copy_eval_est = deepcopy(eval_set)
    num_gvfs = 4
    goal_cumulants = TMCS.get_cumulant_eval_values(env.cumulant_schedule)
    for i in 1:num_gvfs
        copy_eval_est[i, :] .*= goal_cumulants[i]
    end
    return copy_eval_est
end
function get_true_values(env::Curiosity.OneDTMaze, eval_set, gvf_idx)
    copy_eval_est = deepcopy(eval_set)
    goal_cumulants = TMCS.get_cumulant_eval_values(env.cumulant_schedule)
    copy_eval_est .*= goal_cumulants[gvf_idx]
    return copy_eval_est
end

end


"""
`Basics.tabulate()`  ... tabulates the data from various objects due to different criteria.

  + `("multiplet: energies", multiplet::Multiplet; stream::IO=stdout)`
    ... to tabulate the energies of all levels of the given multiplet into a neat format; nothing is returned.
  + `("multiplet: energy relative to immediately lower level", multiplet::Multiplet; stream::IO=stdout)`
    ... to tabulate the energy splitting between neighboured levels of all levels of the given multiplet into a neat
        format; nothing is returned.
  + `("multiplet: energy of each level relative to lowest level", multiplet::Multiplet; stream::IO=stdout)`
    ... to tabulate the energy splitting of all levels with regard to the lowest level of the given multiplet into
        a neat format; nothing is returned.
"""
function Basics.tabulate(sa::String, multiplet::Multiplet; stream::IO=stdout)
    if        sa == "multiplet: energies"
        println(stream, "\n  Eigenenergies:")
        sb = "  Level  J Parity          Hartrees       " * "             eV                   " *  JAC.TableStrings.inUnits("energy")
        println(stream, "\n", sb, "\n")
        for  i = 1:length(multiplet.levels)
            lev = multiplet.levels[i]
            en  = lev.energy;    en_eV = Defaults.convertUnits("energy: from atomic to eV", en);    en_requested = Defaults.convertUnits("energy: from atomic", en)
            sc  = " " * JAC.TableStrings.level(i) * "    " * string(LevelSymmetry(lev.J, lev.parity)) * "    "
            @printf(stream, "%s %.15e %s %.15e %s %.15e %s", sc, en, "  ", en_eV, "  ", en_requested, "\n")
        end

    elseif    sa == "multiplet: energy relative to immediately lower level"
        println(stream, "\n  Energy of each level relative to immediately lower level:")
        sb = "  Level  J Parity          Hartrees       " * "             eV                   " * JAC.TableStrings.inUnits("energy")
        println(stream, "\n", sb, "\n")
        for  i = 2:length(multiplet.levels)
            lev = multiplet.levels[i]
            en    = lev.energy - multiplet.levels[i-1].energy;
            en_eV = Defaults.convertUnits("energy: from atomic to eV", en);    en_requested = Defaults.convertUnits("energy: from atomic", en)
            sc  = " " * JAC.TableStrings.level(i) * "    " * string(LevelSymmetry(lev.J, lev.parity)) * "    "
            @printf(stream, "%s %.15e %s %.15e %s %.15e %s", sc, en, "  ", en_eV, "  ", en_requested, "\n")
        end

    elseif    sa == "multiplet: energy of each level relative to lowest level"
        println(stream, "\n  Energy of each level relative to lowest level:")
        sb = "  Level  J Parity          Hartrees       " * "             eV                   " * JAC.TableStrings.inUnits("energy")
        println(stream, "\n", sb, "\n")
        for  i = 2:length(multiplet.levels)
            lev = multiplet.levels[i]
            en    = lev.energy - multiplet.levels[1].energy;
            en_eV = Defaults.convertUnits("energy: from atomic to eV", en);    en_requested = Defaults.convertUnits("energy: from atomic", en)
            sc  = " " * JAC.TableStrings.level(i) * "    " * string(LevelSymmetry(lev.J, lev.parity)) * "    "
            @printf(stream, "%s %.15e %s %.15e %s %.15e %s", sc, en, "  ", en_eV, "  ", en_requested, "\n")
        end
    else
        error("Unsupported keystring.")
    end

    return( nothing )
end

"""
  + `("multiplet: energies", multiplet::JAC.Hfs.IJF_Multiplet; stream::IO=stdout)`
    ... to tabulate the energies of all hyperfine levels of the given multiplet into a neat format; nothing is returned.
  + `("multiplet: energy of each level relative to lowest level", multiplet::JAC.Hfs.IJF_Multiplet; stream::IO=stdout)`
    ... to tabulate the energy splitting of all levels with regard to the lowest level of the given multiplet into
        a neat format; nothing is returned.
"""
function Basics.tabulate(sa::String, multiplet::JAC.Hfs.IJF_Multiplet; stream::IO=stdout)
    if        sa == "multiplet: energies"
        println(stream, "\n  Eigenenergies for nuclear spin I = $(multiplet.levelFs[1].I):")
        sb = "  Level  F Parity          Hartrees       " * "             eV                   " *  JAC.TableStrings.inUnits("energy")
        println(stream, "\n", sb, "\n")
        for  i = 1:length(multiplet.levelFs)
            lev = multiplet.levelFs[i]
            en  = lev.energy;    en_eV = Defaults.convertUnits("energy: from atomic to eV", en);    en_requested = Defaults.convertUnits("energy: from atomic", en)
            sc  = " " * JAC.TableStrings.level(i) * "    " * string(LevelSymmetry(lev.F, lev.parity)) * "    "
            @printf(stream, "%s %.15e %s %.15e %s %.15e %s", sc, en, "  ", en_eV, "  ", en_requested, "\n")
        end

    elseif    sa == "multiplet: energy of each level relative to lowest level"
        println(stream, "\n  Energy of each level relative to lowest level for nuclear spin I = $(multiplet.levelFs[1].I):")
        sb = "  Level  F Parity          Hartrees       " * "             eV                   " * JAC.TableStrings.inUnits("energy")
        println(stream, "\n", sb, "\n")
        for  i = 2:length(multiplet.levelFs)
            lev = multiplet.levelFs[i]
            en    = lev.energy - multiplet.levelFs[1].energy;
            en_eV = Defaults.convertUnits("energy: from atomic to eV", en);    en_requested = Defaults.convertUnits("energy: from atomic", en)
            sc  = " " * JAC.TableStrings.level(i) * "    " * string(LevelSymmetry(lev.F, lev.parity))  * "    "
            @printf(stream, "%s %.15e %s %.15e %s %.15e %s", sc, en, "  ", en_eV, "  ", en_requested, "\n")
        end
    else
        error("Unsupported keystring.")
    end

    return( nothing )
end

"""
`Basics.tabulateKappaSymmetryEnergiesDirac(kappa::Int64, evalues::Array{Float64,1}, ns::Int64, nuclearModel::Nuclear.Model)`
    ... tabulates the eigenenergies for a given symmetry block kappa together with the corresponding Dirac energies for a
        point-like nucleus. The index ns tells the number of 'negative-continnum' energies in the given evalues. nothing is
        returned.
"""
function Basics.tabulateKappaSymmetryEnergiesDirac(kappa::Int64, evalues::Array{Float64,1}, ns::Int64, nuclearModel::Nuclear.Model)
    Z = nuclearModel.Z
    # Determine the allowed principal quantum numbers n
    l = Basics.subshell_l(Subshell(101,kappa))
    println("  ", JAC.TableStrings.hLine(77))
    sa = "  "
    sa = sa * JAC.TableStrings.center( 7, "Index";            na=2)
    sa = sa * JAC.TableStrings.center(10, "Subshell";         na=3)
    sa = sa * JAC.TableStrings.center(17, "Energies [a.u.]";  na=2)
    sa = sa * JAC.TableStrings.center(17, "Dirac-E  [a.u.]";  na=2)
    sa = sa * JAC.TableStrings.center(17, "Delta-E / |E|";    na=2)
    println(sa)
    println("  ", JAC.TableStrings.hLine(77))
    for  i = ns+1:ns+8
        sa = " " * JAC.TableStrings.center( 6, JAC.TableStrings.level(i-ns); na=2)
        sa = sa *  JAC.TableStrings.flushright(10, string(Subshell(i-ns+l, kappa)); na=6)
        en = Basics.computeDiracEnergy(Subshell(i-ns+l, kappa), Z)
        if  evalues[i] >= 0.      sa = sa * "+"   end;    sa = sa * @sprintf("%.8e", evalues[i])                       * "    "
        if  en         >= 0.      sa = sa * "+"   end;    sa = sa * @sprintf("%.8e", en)                               * "    "
        if  evalues[i]-en >= 0.   sa = sa * "+"   end;    sa = sa * @sprintf("%.8e", (evalues[i]-en)/abs(evalues[i])) * "    "
        println(sa)
    end
    println("      :       :    ")
    for  i = length(evalues)-5:length(evalues)
        sa = " " * JAC.TableStrings.center( 6, JAC.TableStrings.level(i-ns); na=2)
        sa = sa *  JAC.TableStrings.flushright(10, string(Subshell(i-ns+l, kappa)); na=6)
        en = Basics.computeDiracEnergy(Subshell(i-ns+l, kappa), Z)
        if  evalues[i] >= 0.      sa = sa * "+"   end;    sa = sa * @sprintf("%.8e", evalues[i])                       * "    "
        if  en         >= 0.      sa = sa * "+"   end;    sa = sa * @sprintf("%.8e", en)                               * "    "
        if  evalues[i]-en >= 0.   sa = sa * "+"   end;    sa = sa * @sprintf("%.8e", (evalues[i]-en)/abs(evalues[i])) * "    "
        println(sa)
    end
    println("  ", JAC.TableStrings.hLine(77))

    return( nothing )
end


# Feature specifications:

# Individual dense features are represented with strings like "s1w".
typealias DFeature AbstractString

# DFvec is a String vector specifying a set of dense features.
typealias DFvec{T<:DFeature} Vector{T}

# Individual sparse features are specified using a vector of feature
# names indicating a conjunction, e.g. ["s0w", "s0p"]
typealias SFeature{T<:AbstractString} Vector{T}

# SFVec is an vector of SFeatures specifying a set of sparse features.
# e.g. [["sw"],["nw","np"]]
typealias SFvec{T<:SFeature} Vector{T}

# Fvec is a union of DFvec and SFvec, which specify dense and sparse
# features respectively.  The user controls what kind of features are
# generated by using the appropriate type of feature vector.
typealias Fvec Union{DFvec,SFvec}

"""

    features(p::Parser, s::Sentence, feats, [x, idx])

Return a feature vector given a [`Parser`](@ref) state, a
[`Sentence`](@ref), and a feature specification vector.  The last two
arguments (optional) can provide a preallocated output array and an
offset position in that array.

Primitive features are represented by strings of the following form:

    [sn]\d?([hlr]\d?)*[wvcpdLabAB]

The meaning of each character is described below. The letter i
represents an optional single digit integer. A default value is used
if i is not specified.

|char|meaning|
|:---|:------|
|si|i'th stack word, default i=0 means top|
|ni|i'th buffer word, default i=0 means first|
|hi|i'th degree head, default i=1 means direct head|
|li|i'th leftmost child, default i=1 means the leftmost child|
|ri|i'th rightmost child, default i=1 means the rightmost child|
|w|word|
|v|word vector (first half of wvec)|
|c|context vector (second half of wvec)|
|p|postag|
|d|distance to the right using encoding: 1,2,3,4,5-10,10+|
|L|dependency label (0 is ROOT or NONE)|
|a|number of left children.|
|b|number of right children.|
|A|set of left dependency labels|
|B|set of right dependency labels|

For example "sp" or "s0p" indicate the postag of the top stack word,
"n0lL" means the dependency label of the leftmost child of the first
buffer word, "s1d" is the s0s1 distance, "s0d" is the s0n0 distance.

The returned feature vector can be dense (`Vector{Float}`) or sparse
(`Vector{Int}` indicating non-zero positions in a sparse binary
vector).  The choice is marked by the type of the `feats` argument:
for dense output it should be a `Vector{String}`, for sparse output it
should be a `Vector{Vector{String}}`:

    ["s0v","n0lp","n1p"]  # dense feature spec
    [["sw"],["sw","sp"]]  # sparse feature spec

This allows sparse feature specifications to contain elements with
multiple primitive features indicating conjunctions. Some other
notable differences between dense and sparse features are:

* Dense vectors do not support "word", sparse vectors do not support "word vector" or "context vector".
* Sparse vectors represent number of children exactly, dense vectors use the encoding 0,1,...,8,9+.
* For dense features, `x` is a float matrix and `idx=1` is the column number to fill. For sparse features `x` is an int vector and `idx=0` is an offset.

"""
function features(p::Parser, s::Sentence, feats::DFvec,
                  x::AbstractMatrix=Array(wtype(s),flen(p,s,feats),1),
                  xcol::Integer=1)
    if xcol > size(x,2); error("xcol > size(x,2)"); end
    wrows = wdim(s)             #11
    xrows = size(x,1)
    xtype = eltype(x)
    x1 = one(xtype)
    x[:,xcol] = zero(xtype)     #173
    nx = 0                      # last entry in x
    nv = wrows # >> 1             # size of word/context vector, assumes word vec and context vec concatenated
    nd = p.ndeps
    np = length(s.vocab.postags)
    nw = p.nword
    ldep,rdep,lset,rset = getdeps(p) #3775
    for f in feats                   #48
        a = getanchor(f,p,ldep,rdep) #1402
        fn = f[end]                  #642
        if fn == 'v'                 #234
            if a>0                   #91
                copy!(x, (xcol-1)*xrows+nx+1, s.vocab.wvecs, (s.word[a]-1)*wrows+1, nv) #331
            end; nx += nv       #5
        elseif fn == 'c'        #200
            error("Context vectors not implemented.")
            # if a>0              #122
            #     copy!(x, (xcol-1)*xrows+nx+1, s.vocab.wvec, (a-1)*wrows+nv+1, nv) #492
            # end; nx += nv       #6
        elseif fn == 'p'        #114
            if a>0              #159
                if s.postag[a] > np; error("postag out of bound"); end #210
                x[nx+s.postag[a], xcol] = x1 #839
            end; nx += np                    #5
        elseif fn == 'd'
            if a>0
                d = getrdist(f,p,a)
                if d>0; x[nx+(d>10?6:d>5?5:d), xcol] = x1; end
            end; nx += 6
        elseif fn == 'L'
            if a>0
                r = p.deprel[a]
                if r > nd; error("deprel out of bound"); end
                x[nx+1+r, xcol] = x1 # first bit for deprel=0 (ROOT)
            end; nx += (nd+1)
        elseif fn == 'a'
            if a>0
                if isassigned(ldep,a)
                    lcnt=length(ldep[a])
                    if lcnt > 9; lcnt = 9; end
                else
                    lcnt = 0
                end
                x[nx+1+lcnt, xcol] = x1 # 0-9
            end; nx += 10
        elseif fn == 'b'
            if a>0
                if isassigned(ldep,a)
                    rcnt=length(ldep[a])
                    if rcnt > 9; rcnt = 9; end
                else
                    rcnt = 0
                end
                x[nx+1+rcnt, xcol] = x1 # 0-9
            end; nx += 10
        elseif fn == 'A'
            if a>0 && isassigned(lset,a)
                copy!(x, (xcol-1)*xrows+nx+1, Array{wtype(s)}(lset[a]), 1, nd) #
            end; nx += nd
        elseif fn == 'B'
            if a>0 && isassigned(rset,a)
                copy!(x, (xcol-1)*xrows+nx+1, Array{wtype(s)}(rset[a]), 1, nd) #
            end; nx += nd
        elseif fn == 'w'
            error("Dense features do not support 'w'")
        else
            error("Unknown feature $(fn)") # 3
        end
    end
    nx == xrows || error("Bad feature vector length $nx != $xrows")
    return x
end

# Sparse feature extractor:

# TODO: change default idx=1.

function features(p::Parser, s::Sentence, feats::SFvec,
                  x::AbstractMatrix=Array(SFtype, length(feats), 1), xcol=1)
    if xcol > size(x,2); error("xcol > size(x,2)"); end
    deps = getdeps(p)
    SFhash = s.vocab.fdict
    @inbounds for i = 1:length(feats)
        f = feats[i]
        v = Array(Any, length(f))                             # TODO: get rid of alloc here?
        for j=1:length(f)
            v[j] = features1(p,s,f[j],deps...) #1023
        end
        x[i,xcol] = get!(SFhash, (f,v), 1+length(SFhash)) #7782 TODO: need a better hash function here
    end
    return x
end

# Here is where the actual feature lookup happens.  Similar to the
# dense lookup.  But does not have to convert everything to a number,
# can return strings etc.  Returns 'nothing' if target word does not
# exist or the feature is not available.

function features1(p::Parser, s::Sentence, f::String, ldep, rdep, lset, rset)
    a = getanchor(f,p,ldep,rdep)
    if a == 0; return nothing; end
    fn = f[end]
    if fn == 'w'; s.word[a]
    elseif fn == 'p'; s.postag[a]
    elseif fn == 'L'; p.deprel[a]
    elseif fn == 'a'; if isassigned(ldep,a); length(ldep[a]); else; 0; end
    elseif fn == 'b'; if isassigned(rdep,a); length(rdep[a]); else; 0; end
    elseif fn == 'A'; if isassigned(lset,a); lset[a]; else; nothing; end
    elseif fn == 'B'; if isassigned(rset,a); rset[a]; else; nothing; end
    elseif fn == 'd'; (d=getrdist(f,p,a); if d>10; 6; elseif d>5; 5; else; d; end)
    elseif fn == 'v'; error("Sparse features do not support 'v'")
    elseif fn == 'c'; error("Sparse features do not support 'c'")
    else error("Unknown feature letter $fn")
    end
end

function flen(p::Parser, s::Sentence, feats::SFvec)
    length(feats) # one integer per feature
end

function flen(p::Parser, s::Sentence, feats::DFvec)
    nx = 0
    nw = wdim(s)  # >> 1
    nd = p.ndeps
    np = length(s.vocab.postags)
    for f in feats
        nx += flen1(f[end], nw, nd, np) # 1129
    end
    return nx
end

function flen1(c::Char, nw::Int, nd::Int, np::Int)
    if c == 'v'; nw
    elseif c == 'c'; nw
    elseif c == 'p'; np
    elseif c == 'd'; 6
    elseif c == 'L'; (nd+1) # ROOT is not included in nd, TODO: fix this
    elseif c == 'a'; 10
    elseif c == 'b'; 10
    elseif c == 'A'; nd
    elseif c == 'B'; nd
    else error("Unknown feature character $c")
    end
end

# Utility functions to calculate the size of the feature matrix

xsize(p::Parser, s::Sentence, f::Fvec)=(flen(p,s,f),nmoves(p,s))
xsize(p::Parser, c::Corpus, f::Fvec)=(flen(p,c[1],f),nmoves(p,c))
xsize{T<:Parser}(p::Vector{T}, c::Corpus, f::Fvec)=xsize(p[1],c,f)
ysize(p::Parser, s::Sentence)=(nmoves(p),nmoves(p,s))
ysize(p::Parser, c::Corpus)=(nmoves(p),nmoves(p,c))
ysize{T<:Parser}(p::Vector{T}, c::Corpus)=ysize(p[1],c)

function getdeps{T<:Parser}(p::T)
    nw = p.nword
    nd = p.ndeps
    dr = p.deprel
    ldep = Array(Any,nw)
    rdep = Array(Any,nw)
    lset = Array(Any,nw)
    rset = Array(Any,nw)
    @inbounds for d=1:nw
        h=Int(p.head[d])
        if h==0
            continue
        elseif d<h
            if !isassigned(ldep,h)
                ldep[h]=[d]
                lset[h]=zeros(UInt8,nd) # falses(nd) is slower, zeros(WVtype,nd) same speed more space
                lset[h][dr[d]]=1
            else
                push!(ldep[h],d)
                lset[h][dr[d]]=1
            end
        elseif d>h
            if !isassigned(rdep,h)
                rdep[h]=[d]
                rset[h]=zeros(UInt8,nd)
                rset[h][dr[d]]=1
            else
                unshift!(rdep[h],d)
                rset[h][dr[d]]=1
            end
        else
            error("h==d")
        end
    end
    return (ldep,rdep,lset,rset)
end

function getanchor(f::String, p::Parser, ldep, rdep)
    f1 = f[1]; f2 = f[2]; flen = length(f)
    if isdigit(f2)
        i = f2 - '0'            # index
        n = 3                   # next character
    else
        i = 0
        n = 2
    end
    a = 0                       # target word
    if f1 == 's'
        if (p.sptr - i >= 1) # 25
            a = Int(p.stack[p.sptr - i])             # 456
        end
    elseif f1 == 'n'
        if p.wptr + i <= p.nword
            a = p.wptr + i
        end
    else
        error("feature string should start with [sn]")
    end
    if a==0; return 0; end
    while n < flen
        f1 = f[n]; f2 = f[n+1]
        if isdigit(f2)
            i = f2 - '0'
            n = n+2
        else
            i = 1
            n = n+1
        end
        if i <= 0
            error("hlr indexing is one based") # 3 [lrh] is one based, [sn] was zero based
        end
        if f1 == 'l'                              # 2
            if a > p.wptr; error("buffer words other than n0 do not have ldeps"); end # 252
            if isassigned(ldep,a) && i <= length(ldep[a])
                a = Int(ldep[a][i])
            else
                return 0
            end
        elseif f1 == 'r'
            if a >= p.wptr; error("buffer words do not have rdeps"); end
            if isassigned(rdep,a) && i <= length(rdep[a])
                a = Int(rdep[a][i])
            else
                return 0
            end
        elseif f1 == 'h'
            for j=1:i       # 5
                a = Int(p.head[a]) # 147
                if a == 0
                    return 0
                end
            end
        else
            break
        end
    end
    if n != length(f); error("n!=length(f)"); end
    return a
end

function getrdist(f::AbstractString, p::Parser, a::Integer)
    if f[1]=='s'
        if isdigit(f[2])
            i=f[2]-'0'
        else
            i=0
        end
        if i>0
            return p.stack[p.sptr - i + 1] - a
        elseif p.wptr <= p.nword
            return p.wptr - a
        else
            return 0
        end
    else
        return 0 # dist only defined for stack words
    end
end


"""
```
struct VarName{sym}
    indexing  ::    String
end
```

A variable identifier. Every variable has a symbol `sym` and `indices `indexing`.
The Julia variable in the model corresponding to `sym` can refer to a single value or
to a hierarchical array structure of univariate, multivariate or matrix variables. `indexing` stores the indices that can access the random variable from the Julia
variable.

Examples:

- `x[1] ~ Normal()` will generate a `VarName` with `sym == :x` and `indexing == "[1]"`.
- `x[:,1] ~ MvNormal(zeros(2))` will generate a `VarName` with `sym == :x` and
 `indexing == "[Colon(),1]"`.
- `x[:,1][2] ~ Normal()` will generate a `VarName` with `sym == :x` and
 `indexing == "[Colon(),1][2]"`.
"""
struct VarName{sym}
    indexing::String
end

"""
    @varname(var)

A macro that returns an instance of `VarName` given the symbol or expression of a Julia variable, e.g. `@varname x[1,2][1+5][45][3]` returns `VarName{:x}("[1,2][6][45][3]")`.
"""
macro varname(expr::Union{Expr, Symbol})
    expr |> varname |> esc
end
function varname(expr)
    ex = deepcopy(expr)
    (ex isa Symbol) && return quote
        DynamicPPL.VarName{$(QuoteNode(ex))}("")
    end
    (ex.head == :ref) || throw("VarName: Mis-formed variable name $(expr)!")
    inds = :(())
    while ex.head == :ref
        if length(ex.args) >= 2
            strs = map(x -> :($x === (:) ? "Colon()" : string($x)), ex.args[2:end])
            pushfirst!(inds.args, :("[" * join($(Expr(:vect, strs...)), ",") * "]"))
        end
        ex = ex.args[1]
        isa(ex, Symbol) && return quote
            DynamicPPL.VarName{$(QuoteNode(ex))}(foldl(*, $inds, init = ""))
        end
    end
    throw("VarName: Mis-formed variable name $(expr)!")
end

macro vsym(expr::Union{Expr, Symbol})
    expr |> vsym
end

"""
    vsym(expr::Union{Expr, Symbol})

Returns the variable symbol given the input variable expression `expr`. For example, if the input `expr = :(x[1])`, the output is `:x`.
"""
function vsym(expr::Union{Expr, Symbol})
    ex = deepcopy(expr)
    (ex isa Symbol) && return QuoteNode(ex)
    (ex.head == :ref) || throw("VarName: Mis-formed variable name $(expr)!")
    while ex.head == :ref
        ex = ex.args[1]
        isa(ex, Symbol) && return QuoteNode(ex)
    end
    throw("VarName: Mis-formed variable name $(expr)!")
end

"""
    @vinds(expr)

Returns a tuple of tuples of the indices in `expr`. For example, `@vinds x[1,:][2]` returns
`((1, Colon()), (2,))`.
"""
macro vinds(expr::Union{Expr, Symbol})
    expr |> vinds |> esc
end
function vinds(expr::Union{Expr, Symbol})
    ex = deepcopy(expr)
    inds = Expr(:tuple)
    (ex isa Symbol) && return inds
    (ex.head == :ref) || throw("VarName: Mis-formed variable name $(expr)!")
    while ex.head == :ref
        pushfirst!(inds.args, Expr(:tuple, ex.args[2:end]...))
        ex = ex.args[1]
        isa(ex, Symbol) && return inds
    end
    throw("VarName: Mis-formed variable name $(expr)!")
end

"""
    split_var_str(var_str, inds_as = Vector)

This function splits a variable string, e.g. `"x[1:3,1:2][3,2]"` to the variable's symbol `"x"` and the indexing `"[1:3,1:2][3,2]"`. If `inds_as = String`, the indices are returned as a string, e.g. `"[1:3,1:2][3,2]"`. If `inds_as = Vector`, the indices are returned as a vector of vectors of strings, e.g. `[["1:3", "1:2"], ["3", "2"]]`.
"""
function split_var_str(var_str, inds_as = Vector)
    ind = findfirst(c -> c == '[', var_str)
    if inds_as === String
        if ind === nothing
            return var_str, ""
        else
            return var_str[1:ind-1], var_str[ind:end]
        end
    end
    @assert inds_as === Vector
    inds = Vector{String}[]
    if ind === nothing
        return var_str, inds
    end
    sym = var_str[1:ind-1]
    ind = length(sym)
    while ind < length(var_str)
        ind += 1
        @assert var_str[ind] == '['
        push!(inds, String[])
        while var_str[ind] != ']'
            ind += 1
            if var_str[ind] == '['
                ind2 = findnext(c -> c == ']', var_str, ind)
                push!(inds[end], strip(var_str[ind:ind2]))
                ind = ind2+1
            else
                ind2 = findnext(c -> c == ',' || c == ']', var_str, ind)
                push!(inds[end], strip(var_str[ind:ind2-1]))
                ind = ind2
            end
        end
    end
    return sym, inds
end

@generated function inargnames(::VarName{s}, ::Model{_F, argnames}) where {s, argnames, _F}
    return s in argnames
end

@generated function inmissings(::VarName{s}, ::Model{_F, _a, _T, missings}) where {s, missings, _F, _a, _T}
    return s in missings
end


# --------------------------------------------------------------------------- #
# General image type

abstract type Image end

# --------------------------------------------------------------------------- #
# Greyscale Images

struct GreyscaleImage <: Image
  data :: Array{Float32, 2}
  GreyscaleImage(data) = new(autoscale(convert(Array{Float32}, data)))
end

function GreyscaleImage(data :: Array{Float32, 3})
  @assert size(data, 3) == 1
  GreyscaleImage(reshape(data, size(data)[1:2]))
end

# --------------------------------------------------------------------------- #
# RGB Images

struct RGBImage <: Image
  data :: Array{Float32, 3}
  RGBImage(data) = new(autoscale(convert(Array{Float32}, data)))
end

function greyscale(img :: RGBImage)
  # Using rec 601 luma standard
  data = 0.299img.data[:,:,1] + 0.587img.data[:,:,2] + 0.114img.data[:,:,3]
  return GreyscaleImage(data)
end

# --------------------------------------------------------------------------- #
# Create datasets from high-res images

function intervals(L, l, m)
  grid = Int[]
  i = 1

  while i + l - 1 <= L
    push!(grid, i)
    i += l + m
  end

  return grid
end

function ordered_patches{I <: Image}(img :: I, lbl;
                                     size = (128, 128),
                                     offset = nothing,
                                     margin :: Integer = 0,
                                     shuffle :: Bool = true,
                                     imageop = Id(),
                                     multitude = 1)

  # Note: components of argument size should be divisible by 8,
  # otherwise some (multiresolution) networks will not like the input

  n, m = imgsize(img)

  @assert (n >= size[1] && m >= size[2])
  @assert (-margin < min(size...))

  ygrid = intervals(n, size[1], margin)
  xgrid = intervals(m, size[2], margin)

  if offset == nothing
    dy = div(n - (ygrid[end] + size[1] - 1), 2)
    dx = div(m - (xgrid[end] + size[2] - 1), 2)
    offset = (dy, dx)
  end

  ygrid = Int[ i + offset[1] for i in ygrid ]
  xgrid = Int[ j + offset[2] for j in xgrid ]

  imgs = I[]
  lbls = Label[]
  ld = convert(Matrix{Int}, lbl)

  for i in ygrid, j in xgrid
    ii = i + size[1] - 1
    jj = j + size[2] - 1
    push!(imgs, I(img.data[i:ii, j:jj, :]))
    if !isempty(lbl)
      sel = (i .<= ld[2,:] .<= ii) .& (j .<= ld[1,:] .<= jj)
      push!(lbls, Label(ld[:,sel] .- [j-1, i-1]))
    else
      push!(lbls, Label())
    end
  end

  data = vcat( (apply.(imageop, imgs, lbls) for i in 1:multitude)... )
  imgs = [ img for (img, _) in data ]
  lbls = [ lbl for (_, lbl) in data ]

  return imgs, lbls
end

function random_patches{I <: Image}(img :: I, lbl, number;
                                    size = (128, 128),
                                    imageop = Id(),
                                    multitude = 1)

  # Note: components of argument size should be divisible by 8,
  # otherwise some (multiresolution) networks will not like the input

  n, m = imgsize(img)
  @assert (n >= size[1] && m >= size[2])

  imgs = I[]
  lbls = Label[]
  ld = convert(Matrix{Int}, lbl)

  for k in 1:number
    i, j = rand(1:n-size[1]), rand(1:m-size[2])
    ii = i + size[1] - 1
    jj = j + size[2] - 1
    push!(imgs, I(img.data[i:ii, j:jj, :]))
    if !isempty(lbl)
      sel = (i .<= ld[2,:] .<= ii) .& (j .<= ld[1,:] .<= jj)
      push!(lbls, Label(ld[:,sel] .- [j-1, i-1]))
    else
      push!(lbls, Label())
    end
  end

  data = vcat( (apply.(imageop, imgs, lbls) for i in 1:multitude)... )
  imgs = [ img for (img, _) in data ]
  lbls = [ lbl for (_, lbl) in data ]

  return imgs, lbls
end

function ordered_patches(img; kwargs...)
  return ordered_patches(img, Label(); kwargs...)[1]
end

function random_patches(img, number; kwargs...)
  return random_patches(img, Label(), number; kwargs...)[1]
end

# --------------------------------------------------------------------------- #
# Convenience functions

imgdata(img :: Image) = img.data
imgsize(a) = size(a)[1:2]
imgtype{I<:Image}(a::I) = I

imgchannels(::GreyscaleImage) = 1
imgchannels(::RGBImage) = 3
imgchannels(::Type{<: GreyscaleImage}) = 1
imgchannels(::Type{<: RGBImage}) = 3

Base.size(img::Image, args...; kwargs...) = size(img.data, args..., kwargs...)

function imgconvert{T <: Images.Gray}(imgarray :: Array{T})
  data = permutedims(Images.channelview(imgarray), (2, 3, 1))
  return GreyscaleImage(reshape(data, size(data)[1:2]...))
end

function imgconvert{T <: Images.RGB}(imgarray :: Array{T})
  data = permutedims(Images.channelview(imgarray), (2, 3, 1))
  return RGBImage(data)
end

function crop(data :: Matrix{Float32}, x :: Int, y :: Int, w :: Int, h :: Int)
  return data[y:y+h-1,x:x+w-1]
end

function crop(img :: RGBImage, x :: Int, y :: Int, w :: Int, h :: Int)
  return RGBImage(imgdata(img)[y:y+h-1,x:x+w-1,:])
end

function crop(img :: GreyscaleImage, x :: Int, y :: Int, w :: Int, h :: Int)
  return GreyscaleImage(imgdata(img)[y:y+h-1,x:x+w-1])
end

Base.similar(img::GreyscaleImage) = GreyscaleImage(zeros(Float32, imsize(img)))