"""
    IsStochasticInteger{T=Int}() <: StochasticStyle{T}

FCIQMC algorithm with integer walkers as in
[Booth *et al.* (2009)](https://doi.org/10.1063/1.3193710). During the vector matrix product
each individual diagonal and spawning step is rounded stochastically to a nearby integer
value.

See also [`StochasticStyle`](@ref).
"""
struct IsStochasticInteger{T<:Integer} <: StochasticStyle{T} end
IsStochasticInteger() =  IsStochasticInteger{Int}()

function step_stats(::IsStochasticInteger{T}) where {T}
    z = zero(T)
    return (
        (:spawn_attempts, :spawns, :deaths, :clones, :zombies, :annihilations),
        MultiScalar(0, z, z, z, z, z),
    )
end
function fciqmc_col!(
    ::IsStochasticInteger, w, ham, add, num::Real, shift, dτ
)
    clones, deaths, zombies, ann_diag = diagonal_step!(w, ham, add, num, dτ, shift, 0, true)
    attempts, spawns, ann_offdiag = spawn!(WithReplacement(), w, ham, add, num, dτ)
    return (attempts, spawns, deaths, clones, zombies, ann_diag + ann_offdiag)
end

"""
    IsStochastic2Pop{T=Complex{Int}}() <: StochasticStyle{T}

Stochastic propagation with complex walker numbers representing two populations of integer
walkers.

When using this style, make sure to set a complex target number walkers in the
[`ShiftStrategy`](@ref)!

This style is experimental.

See also [`StochasticStyle`](@ref).
"""
struct IsStochastic2Pop{T<:Complex{<:Integer}} <: StochasticStyle{T} end
IsStochastic2Pop() = IsStochastic2Pop{Complex{Int}}()

function step_stats(::IsStochastic2Pop{T}) where {T}
    z = zero(T)
    return (
        (:spawns, :deaths, :clones, :zombies, :annihilations),
        MultiScalar(z, z, z, z, z)
    )
end
function fciqmc_col!(::IsStochastic2Pop, w, ham, add, val, shift, dτ)
    offdiags = offdiagonals(ham, add)
    spawns = deaths = clones = zombies = ann_o = ann_d = 0 + 0im
    # off-diagonal real.
    s, a = spawn!(WithReplacement(), w, offdiags, add, real(val), dτ)
    spawns += s; ann_o += a
    # off-diagonal complex: complex dτ ensures spawning to the correct population.
    s, a = spawn!(WithReplacement(), w, offdiags, add, imag(val), dτ * im)
    spawns += s; ann_o += a

    clones, deaths, zombies, ann_d = diagonal_step!(w, ham, add, val, dτ, shift, 0, true)

    return (spawns, deaths, clones, zombies, ann_o + ann_d)
end

"""
    IsDeterministic{T=Float64}(compression=NoCompression()) <: StochasticStyle{T}

Propagate with deterministic vector matrix multiplications. Stochastic compression of the
resultant vector (after annihilations) can be triggered by setting the optional
`compression` argument to a relevant [`CompressionStrategy`](@ref).

See also [`StochasticStyle`](@ref).
"""
struct IsDeterministic{T<:AbstractFloat,C<:CompressionStrategy} <: StochasticStyle{T}
    compression::C
end
function IsDeterministic{T}(compression::C=NoCompression()) where {T,C}
    return IsDeterministic{T,C}(compression)
end
IsDeterministic(args...) = IsDeterministic{Float64}(args...)

CompressionStrategy(s::IsDeterministic) = s.compression

function Base.show(io::IO, s::IsDeterministic{T}) where {T}
    if s.compression isa NoCompression
        print(io, "IsDeterministic{$T}()")
    else
        print(io, "IsDeterministic{$T}($(s.compression))")
    end
end

function step_stats(::IsDeterministic)
    return (:exact_steps,), MultiScalar(0,)
end
function fciqmc_col!(::IsDeterministic, w, ham::AbstractMatrix, add, num, shift, dτ)
    for i in axes(ham, 1) # iterate through off-diagonal rows of `ham`
        i == add && continue
        deposit!(w, i, -dτ * ham[i, add] * num, add => num)
    end
    deposit!(w, add, (1 + dτ * (shift - ham[add, add])) * num, add => num) # diagonal
    return (1,)
end
function fciqmc_col!(::IsDeterministic, w, ham, add, val, shift, dτ)
    diagonal_step!(w, ham, add, val, dτ, shift)
    spawn!(Exact(), w, ham, add, val, dτ)
    return (1,)
end

"""
    IsStochasticWithThreshold{T=Float64}(threshold=1.0) <: StochasticStyle{T}

Stochastic propagation with floating point walker numbers. During the vector matrix product
each individual diagonal and spawning result is rounded stochastically if smaller than
`threshold` (before annihilations).

See also [`StochasticStyle`](@ref).
"""
struct IsStochasticWithThreshold{T<:AbstractFloat} <: StochasticStyle{T}
    threshold::T
end
IsStochasticWithThreshold(args...) = IsStochasticWithThreshold{Float64}(args...)
IsStochasticWithThreshold{T}(t=1.0) where {T} = IsStochasticWithThreshold{T}(T(t))

function step_stats(::IsStochasticWithThreshold{T}) where {T}
    z = zero(T)
    return (
        (:spawn_attempts, :spawns),
        MultiScalar(0, z)
    )
end
function fciqmc_col!(s::IsStochasticWithThreshold, w, ham, add, val, shift, dτ)
    diagonal_step!(w, ham, add, val, dτ, shift)
    attempts, spawns, _ = spawn!(WithReplacement(s.threshold), w, ham, add, val, dτ)
    return (attempts, spawns)
end

"""
    IsDynamicSemistochastic{T=Float64}(; kwargs...) <: StochasticStyle{T}

QMC propagation with floating-point walker numbers and reduced noise. All possible spawns
(offdiagonal elements in vector-matrix multiplication)
are performed deterministically when number of walkers in a configuration is high, as
controlled by the `rel_threshold` and `abs_threshold` keywords. Stochastic selection of
spawns is controlled  by the `spawning` keyword.

By default, a stochastic vector compression is applied after annihilations are completed.
This behaviour can be changed to on-the-fly projection (as in [`IsStochasticInteger`](@ref)
or [`IsStochasticWithThreshold`](@ref)) by setting a threshold to `spawning` and
setting `compression` to [`NoCompression`](@ref).

## Parameters:

* `rel_threshold = 1.0`: If the walker number on a configuration times this threshold
  is greater than the number of offdiagonals, spawning is done deterministically. Should be
  set to 1 or more for best performance.

* `abs_threshold = Inf`: If the walker number on a configuration is greater than this value,
  spawning is done deterministically. Can be set to e.g.
  `abs_threshold = 0.1 * target_walkers`.

* `proj_threshold = 1.0`: Values below this number are stochastically projected to this
  value or zero. See also [`ThresholdCompression`](@ref).

* `spawning = WithReplacement()`: [`SpawningStrategy`](@ref) to use for the non-exact
  spawns. See [`WithReplacement`](@ref) for behaviour and possible arguments.

* `compression = ThresholdCompression(proj_threshold)`: [`CompressionStartegy`](@ref) used
  to compress the vector after a step. Overrides `proj_threshold`.

See also [`StochasticStyle`](@ref).
"""
struct IsDynamicSemistochastic{
    T<:AbstractFloat,C<:CompressionStrategy,S<:DynamicSemistochastic
} <: StochasticStyle{T}
    compression::C
    spawning::S
end
function IsDynamicSemistochastic{T}(
    ;
    proj_threshold=1.0, strength=1.0, rel_threshold=1.0, abs_threshold=Inf,
    compression=ThresholdCompression(proj_threshold),
    spawning=WithReplacement(),
) where {T}
    s = DynamicSemistochastic(spawning, rel_threshold, abs_threshold)
    return IsDynamicSemistochastic{T,typeof(compression),typeof(s)}(compression, s)
end
IsDynamicSemistochastic(; kwargs...) = IsDynamicSemistochastic{Float64}(; kwargs...)

function Base.show(io::IO, s::IsDynamicSemistochastic{T,C,S}) where {T,C,S}
    print(io, "IsDynamicSemistochastic{$T,", nameof(C), ",", nameof(S), "}()")
end

CompressionStrategy(s::IsDynamicSemistochastic) = s.compression

function step_stats(::IsDynamicSemistochastic{T}) where {T}
    z = zero(T)
    return (
        (:exact_steps, :inexact_steps, :spawn_attempts, :spawns),
        MultiScalar(0, 0, 0, z),
    )
end
function fciqmc_col!(s::IsDynamicSemistochastic, w, ham, add, val, shift, dτ)
    diagonal_step!(w, ham, add, val, dτ, shift)
    exact, inexact, attempts, spawns, _ = spawn!(s.spawning, w, ham, add, val, dτ)
    return (exact, inexact, attempts, spawns)
end

"""
    IsExplosive{T=Float64}(;
        splatter_factor,
        explosion_threshold,
        compression
    ) <: StochasticStyle{T}

QMC propagation with explosive walkers. Walkers with small walker numbers do not perform the
standard death/spawning steps. Instead, a walker will either die completely and spawn with a
greater magnitude (refered to as explosion below), or stay unchanged and not spawn. The
probabilty of exploding is controlled by the shift and `dτ`.

Walkers with high walker numbers spawn as if [`IsDynamicSemistochastic`](@ref) was used.

Like [`IsDynamicSemistochastic`](@ref), the vector is compressed after all spawning is
performed.

This style is experimental.

## Parameters

* `splatter_factor = 1.0`: The spawning strength to use with exploded walkers.

* `explosion_threshold = 1.0`: Entries smaller or equal than this value will attempt to
  explode.

* `proj_threshold = 1.0`: Threshold to use in vector compression. See
  [`ThresholdCompression`](@ref).

* `compression = ThresholdCompression(proj_threshold)`: [`CompressionStrategy`](@ref) to use
  to compress the vector. Overrides `proj_threshold`.

See also [`StochasticStyle`](@ref).
"""
struct IsExplosive{T<:AbstractFloat,C<:CompressionStrategy} <: StochasticStyle{T}
    splatter_factor::T
    explosion_threshold::T
    compression::C
end
function IsExplosive{T}(
    ;
    splatter_factor=one(T),
    explosion_threshold=one(T),
    proj_threshold=one(T),
    compression=ThresholdCompression(proj_threshold),
) where {T}
    return IsExplosive(T(splatter_factor), T(explosion_threshold), compression)
end
IsExplosive(; kwargs...) = IsExplosive{Float64}(; kwargs...)

CompressionStrategy(s::IsExplosive) = s.compression

function step_stats(::IsExplosive{T}) where {T}
    z = zero(T)
    return (
        (:explosions, :ticks, :normal_steps, :spawn_attempts,
         :explosive_spawns, :normal_spawns,
         ),
        MultiScalar(0, 0, 0, 0, z, z),
    )
end
function fciqmc_col!(s::IsExplosive{T}, w, ham, add, val, shift, dτ) where {T}
    explosions = normal_steps = ticks = 0
    explosive_spawns = normal_spawns = zero(T)
    clones = deaths = zombies = zero(T)
    attempts = 0

    pd = dτ * (diagonal_element(ham, add) - shift)
    if abs(val) ≤ s.explosion_threshold && 0 ≤ pd < 1
        if cRand() < pd
            _, _, attempts, explosive_spawns, _ = spawn!(
                DynamicSemistochastic(strat=WithReplacement(zero(T), s.splatter_factor)),
                w, ham, add, val / pd, dτ,
            )
            explosions = 1
        else
            deposit!(w, add, val, add => val)
            ticks = 1
        end
    else
        _ = diagonal_step!(w, ham, add, val, dτ, shift)
        _, _, attempts, normal_spawns, _ = spawn!(
            DynamicSemistochastic(strat=WithReplacement()), w, ham, add, val, dτ,
        )
        normal_steps = 1
    end

    return (
        explosions, ticks, normal_steps, attempts, explosive_spawns, normal_spawns,
    )
end

default_style(::Type{T}) where {T<:Integer} = IsStochasticInteger{T}()
default_style(::Type{T}) where {T<:AbstractFloat} = IsDeterministic{T}()
default_style(::Type{T}) where {T<:Complex{<:Integer}} = IsStochastic2Pop{T}()


function edit{T <: Union(AbstractFixedVector, Real)}(style::Style{:Default}, numbertex::Texture{T, 1, 2}, customization::Dict{Symbol,Any})

  backgroundcolor = customization[:backgroundcolor]
  screen          = customization[:screen]
  camera          = screen.orthographiccam
  color           = customization[:color]
  maxdigits       = customization[:maxdigits]
  maxlength       = customization[:maxlength]
  model           = customization[:model]
  font            = customization[:font]

  numbers         = data(numbertex) # get data from texture/video memory
  text            = Array(GLGlyph{Uint16}, int(size(numbers,1)*maxdigits), size(numbers, 2))

  fill!(text, GLGlyph()) # Fill text array with blanks, as we don't need all of them

  # handle real values
  Base.stride(x::Real, i)       = 1
  # remove f0
  makestring(x::Integer)        = string(int(x))
  makestring(x::FloatingPoint)  = string(float64(x))
  makestring{T}(x::Vector1{T})  = string(float64(x[1]))
  makestring{T}(x::Vector1{T}, maxlen) = makestring(float64(x[1]), maxlen)

  makestring(x::Integer, maxlen) = begin
    tmp = string(int(x))
    len = length(tmp)
    if len > maxlen
        tmp = tmp[1:maxlen]
    elseif len < maxlen
      tmp = rpad(tmp, maxlen, " ")
    end
    tmp
  end
  makestring(x::FloatingPoint, maxlen) = begin
    tmp = string(float64(x))
    len = length(tmp)
    if len > maxlen
        tmp = tmp[1:maxlen]
    elseif len < maxlen
      tmp = rpad(tmp, maxlen, "0")
    end
    tmp
  end
  textgpu     = Texture(text, keepinram=true)
  customization[:style_group] = Texture([color])
  customization[:textlength]  = length(textgpu)
  obj         = visualize(style, textgpu, customization)

  startposition   = Uint16[0,0]
  positionrunner  = startposition

  maxlength = 5
  for i=1:size(numbers,1)
    for j=1:size(numbers, 2)
      number = numbers[i,j]
      i3 = ((i-1)*maxlength) + 1
      textgpu[i3:i3+maxlength-1, j:j] = GLGlyph{Uint16}[GLGlyph(c, positionrunner[1], positionrunner[2]+k-1, 0) for (k,c) in enumerate(makestring(number, maxlength))]
      positionrunner += [0, maxlength+1]
    end
    positionrunner = startposition + [i,0]
  end
  selectiondata = lift(first, SELECTION[:mouse_hover])
  # We allocated more space on the gpu then needed (length(numbers)*maxdigits)
  # So we need to update the render method, to render only length(numbers) * maxlength
  obj[:postrender, renderinstanced] = (obj.vertexarray, length(textgpu))
  # We allocated more space on the gpu then needed (length(numbers)*maxdigits)
  # So we need to update the render method, to render only length(numbers) * maxlength
  # ([numbers], zero(eltype(numbers)), -1, -1, -1, Vector2(0.0))
  number_lift = foldl(([numbers;], zero(eltype(numbers)), -1, -1, -1, Vector2(0.0)), screen.inputs[:mouseposition], screen.inputs[:mousebuttonspressed], selectiondata) do v0, mposition, mbuttons, selection
    numbers0, value0, inumbers0, igpu0, mbutton0, mposition0 = v0
    # if over a number           && nothing selected &&         only           left mousebutton clicked
    if selection[1] == obj.id && inumbers0 == -1 && length(mbuttons) == 1 && in(0, mbuttons)
      iorigin   = selection[2]
      inumbers  = div(iorigin, maxlength) + 1
      igpu      = int((iorigin - (iorigin%maxlength)) + 1) # get the first index of the number
      return (numbers0, numbers0[inumbers], inumbers, igpu, 0, mposition)
    end
    # if a number is selected && previous click was left && still only left button ist clicked
    if inumbers0 > 0 && mbutton0 == 0 && length(mbuttons) == 1 && in(0, mbuttons)
      xdiff                    = mposition[1] - mposition0[1]
      numbers0[inumbers0]      = value0 + (float32(xdiff)/ 50.0f0)
      numbertex[inumbers0]     = numbers0[inumbers0]
      a = mod(inumbers0-1, size(numbers0, 1))
      b = div(inumbers0-1, size(numbers0, 1))
      for (k,c) in enumerate(makestring(numbers0[inumbers0], maxlength))
        textgpu[igpu0+k-1] = GLGlyph(c, a, (b-1)+b*maxlength+k, 0)
      end
      return (numbers0, value0, inumbers0, igpu0, 0, mposition0)
    end
    return (numbers0, zero(eltype(numbers0)), -1, -1, -1, Vector2(0.0))
  end

  obj, numbertex
end

# remove f0
makestring(x::Integer)        = string(int(x))
makestring(x::FloatingPoint)  = string(float64(x))
makestring{T}(x::Vector1{T})  = string(float64(x[1]))
makestring{T}(x::Vector1{T}, maxlen) = makestring(float64(x[1]), maxlen)
makestring(x::Integer, maxlen) = begin
  tmp = string(int(x))
  len = length(tmp)
  if len > maxlen
      tmp = tmp[1:maxlen]
  elseif len < maxlen
    tmp = rpad(tmp, maxlen, " ")
  end
  tmp
end
makestring(x::FloatingPoint, maxlen) = begin
  tmp = string(float64(x))
  len = length(tmp)
  if len > maxlen
      tmp = tmp[1:maxlen]
  elseif len < maxlen
    tmp = rpad(tmp, maxlen, "0")
  end
  tmp
end

function edit{T <: Union(AbstractVector, Real)}(style::Style, numbers::T, customization::Dict{Symbol,Any})

  backgroundcolor = customization[:backgroundcolor]
  screen          = customization[:screen]
  camera          = screen.orthographiccam
  color           = customization[:color]
  maxdigits       = customization[:maxdigits]
  maxlength       = customization[:maxlength]
  model           = customization[:model]
  font            = customization[:font]

  if ndims(numbers) == 2
    xn, yn = size(numbers)
  elseif ndims(numbers) == 1 || ndims(numbers) == 0
    xn, yn = length(numbers), 1
  else
    error("Dimension Missmatch. Wanted: 0D,1D, 2D Vector. Got: ", typeof(numbers), " ndims: ", ndims(numbers))
  end

  text        = Array(GLGlyph{Uint16}, int(xn*maxdigits), yn)
  fill!(text, GLGlyph()) # Fill text array with blanks, as we don't need all of them

  textgpu     = Texture(text)
  customization[:style_group] = Texture([color])
  customization[:textlength]  = length(textgpu)
  obj         = visualize(style, textgpu, customization)

  startposition   = Uint16[0,0]
  positionrunner  = startposition

  maxlength = 5
  for i=1:xn
    for j=1:yn
      number = numbers[i,j]
      i3 = ((i-1)*maxlength) + 1
      textgpu[i3:i3+maxlength-1, j:j] = GLGlyph{Uint16}[GLGlyph(c, positionrunner[1], positionrunner[2]+k-1, 0) for (k,c) in enumerate(makestring(number, maxlength))]
      positionrunner += [0, maxlength+1]
    end
    positionrunner = startposition + [i,0]
  end
  selectiondata = lift(first, SELECTION[:mouse_hover])
  # We allocated more space on the gpu then needed (length(numbers)*maxdigits)
  # So we need to update the render method, to render only length(numbers) * maxlength
  obj[:postrender, renderinstanced] = (obj.vertexarray, length(textgpu))
  # We allocated more space on the gpu then needed (length(numbers)*maxdigits)
  # So we need to update the render method, to render only length(numbers) * maxlength
  # ([numbers], zero(eltype(numbers)), -1, -1, -1, Vector2(0.0))
  number_lift = foldl(([numbers;], zero(eltype(numbers)), -1, -1, -1, Vector2(0.0)), screen.inputs[:mouseposition], screen.inputs[:mousebuttonspressed], selectiondata) do v0, mposition, mbuttons, selection
    numbers0, value0, inumbers0::Int, igpu0, mbutton0, mposition0 = v0
    # if over a number           && nothing selected &&         only           left mousebutton clicked
    if selection[1] == obj.id && inumbers0::Int == -1 && length(mbuttons) == 1 && in(0, mbuttons)
      iorigin   = selection[2]
      inumbers  = div(iorigin, maxlength) + 1
      igpu      = int((iorigin - (iorigin%maxlength)) + 1) # get the first index of the number
      return (numbers0, numbers0[inumbers], inumbers, igpu, 0, mposition)
    end
    # if a number is selected && previous click was left && still only left button ist clicked
    if inumbers0::Int > 0 && mbutton0 == 0 && length(mbuttons) == 1 && in(0, mbuttons)
      xdiff                    = mposition[1] - mposition0[1]
      numbers0[inumbers0::Int]      = value0 + (float32(xdiff)/ 50.0f0)
      #numbertex[inumbers0]     = numbers0[inumbers0]
      a = mod(inumbers0::Int-1, size(numbers0::Int, 1))
      b = div(inumbers0::Int-1, size(numbers0::Int, 1))
      for (k,c) in enumerate(makestring(numbers0[inumbers0::Int], maxlength))
        textgpu[igpu0+k-1] = GLGlyph(c, a, (b-1)+b*maxlength+k, 0)
      end
      return (numbers0, value0, inumbers0::Int, igpu0, 0, mposition0)
    end
    return (numbers0, zero(eltype(numbers0)), -1, -1, -1, Vector2(0.0))
  end
  convert2{XX}(::Type{XX}, x::Vector{XX}) = first(x)
  convert2(t::Type, x) = convert(t, x)
  obj, lift(x->convert2(T,first(x)), number_lift)
end


module ChanVese
using DistanceTransforms
using Images
using ImageFiltering

include("initializers.jl")
include("classical_chan_vese.jl")

export
    checkerboard,
    disk
    classical_chan_vese

end


struct FRAPMesh <: Mesh
    Nodes::Array{Float64,1}
    NBases::Int
    Fil::IndexDict
end
# Convenience constructors
function FRAPMesh(
    model::SFFM.Model,
    Nodes::Array{<:Real,1},
    NBases::Int;
    Fil::IndexDict=IndexDict(),
    v::Bool = false,
)

    ## Construct the sets Fᵐ = ⋃ᵢ Fᵢᵐ, global index for sets of type m
    if isempty(Fil)
        Fil = MakeFil(model, Nodes)
    end

    mesh = FRAPMesh(
        Nodes,
        NBases,
        Fil,
    )
    v && println("UPDATE: DGMesh object created with fields ", fieldnames(SFFM.DGMesh))
    return mesh
end
function FRAPMesh()
    FRAPMesh(
        Array{Float64,1}(undef,0),
        0,
        Dict{String,BitArray{1}}(),
    )
end

"""

    NBases(mesh::FRAPMesh)

Number of bases in a cell
"""
NBases(mesh::FRAPMesh) = mesh.NBases

"""

    CellNodes(mesh::FRAPMesh)

The cell centre
"""
CellNodes(mesh::FRAPMesh) = Array(((mesh.Nodes[1:end-1] + mesh.Nodes[2:end]) / 2 )')

"""

    Basis(mesh::FRAPMesh)

Constant ""
"""
Basis(mesh::FRAPMesh) = ""

function MakeLazyGenerator(
    model::SFFM.Model,
    mesh::FRAPMesh,
    me::ME;
    v::Bool=false,
)
    blocks = (me.s*me.a, me.S, me.s*me.a)

    boundary_flux = (in = me.s[:], out = me.a[:])

    T = model.T
    C = model.C
    delta = Δ(mesh)
    D = me.D

    signChangeIndex = zeros(Bool,NPhases(model),NPhases(model))
    for i in 1:NPhases(model), j in 1:NPhases(model)
        if ((sign(model.C[i])!=0) && (sign(model.C[j])!=0))
            signChangeIndex[i,j] = (sign(model.C[i])!=sign(model.C[j]))
        elseif (sign(model.C[i])==0)
            signChangeIndex[i,j] = sign(model.C[j])>0
        elseif (sign(model.C[j])==0)
            signChangeIndex[i,j] = sign(model.C[i])>0
        end
    end

    out = SFFM.LazyGenerator(blocks,boundary_flux,T,C,delta,D,signChangeIndex,mesh.Fil)
    v && println("UPDATE: LazyGenerator object created with keys ", keys(out))
    return out
end
function MakeLazyGenerator(model::Model, mesh::FRAPMesh; v::Bool=false)
    me = SFFM.MakeME(SFFM.CMEParams[NBases(mesh)])
    return MakeLazyGenerator(model, mesh, me; v=v)
end
function MakeFullGenerator(model::Model, mesh::Mesh, me::ME; v::Bool=false)
    lazy = MakeLazyGenerator(model,mesh,me; v=v)
    return materialise(lazy)
end

# function MakeB(model::Model, mesh::FRAPMesh, me::ME)
#     N₊ = sum(model.C .>= 0)
#     N₋ = sum(model.C .<= 0)

#     F = Dict{String,SparseArrays.SparseMatrixCSC{Float64,Int64}}()
#     UpDiagBlock = me.s*me.a
#     LowDiagBlock = me.s*me.a
#     for i = ["+","-"]
#         F[i] = SparseArrays.spzeros(Float64, TotalNBases(mesh), TotalNBases(mesh))
#         for k = 1:NIntervals(mesh)
#             idx = (1:NBases(mesh)) .+ (k - 1) * NBases(mesh)
#             if k > 1
#                 idxup = (1:NBases(mesh)) .+ (k - 2) * NBases(mesh)
#                 if i=="+"
#                     F[i][idxup, idx] = UpDiagBlock
#                 elseif i=="-"
#                     F[i][idx, idxup] = LowDiagBlock
#                 end # end if C[i]
#             end # end if k>1
#         end # for k in ...
#     end # for i in ...

#     signChangeIndex = zeros(Bool,NPhases(model),NPhases(model))
#     for i in 1:NPhases(model), j in 1:NPhases(model)
#         if ((sign(model.C[i])!=0) && (sign(model.C[j])!=0))
#             signChangeIndex[i,j] = (sign(model.C[i])!=sign(model.C[j]))
#         elseif (sign(model.C[i])==0)
#             signChangeIndex[i,j] = sign(model.C[j])>0
#         elseif (sign(model.C[j])==0)
#             signChangeIndex[i,j] = sign(model.C[i])>0
#         end
#     end
#     B = SparseArrays.spzeros(
#         Float64,
#         NPhases(model) * TotalNBases(mesh) + N₋ + N₊,
#         NPhases(model) * TotalNBases(mesh) + N₋ + N₊,
#     )
#     B[N₋+1:end-N₊,N₋+1:end-N₊] = SparseArrays.kron(
#             model.T.*signChangeIndex,
#             SparseArrays.kron(SparseArrays.I(NIntervals(mesh)),me.D)
#         ) + SparseArrays.kron(
#             model.T.*(1 .- signChangeIndex),
#             SparseArrays.I(TotalNBases(mesh))
#         )

#     # Boundary conditions
#     T₋₋ = model.T[model.C.<=0,model.C.<=0]
#     T₊₋ = model.T[model.C.>=0,:].*((model.C.<0)')
#     T₋₊ = model.T[model.C.<=0,:].*((model.C.>0)')
#     T₊₊ = model.T[model.C.>=0,model.C.>=0]
#     # yuck
#     inLower = [
#         LinearAlgebra.kron(LinearAlgebra.diagm(abs.(model.C).*(model.C.<=0)),me.s)[:,model.C.<=0];
#         LinearAlgebra.zeros((NIntervals(mesh)-1)*NPhases(model)*NBases(mesh),N₋)
#     ]
#     outLower = [
#         LinearAlgebra.kron(T₋₊,me.a) LinearAlgebra.zeros(N₋,N₊+(NIntervals(mesh)-1)*NPhases(model)*NBases(mesh))
#     ]
#     inUpper = [
#         LinearAlgebra.zeros((NIntervals(mesh)-1)*NPhases(model)*NBases(mesh),N₊);
#         LinearAlgebra.kron(LinearAlgebra.diagm(abs.(model.C).*(model.C.>=0)),me.s)[:,model.C.>=0]
#     ]
#     outUpper = [
#         LinearAlgebra.zeros(N₊,N₋+(NIntervals(mesh)-1)*NPhases(model)*NBases(mesh)) LinearAlgebra.kron(T₊₋,me.a)
#     ]

#     QBDidx = MakeQBDidx(model,mesh)
#     B[1:N₋,QBDidx] = [T₋₋ outLower]
#     B[end-N₊+1:end,QBDidx] = [outUpper T₊₊]
#     B[QBDidx[N₋+1:end-N₊],1:N₋] = inLower
#     B[QBDidx[N₋+1:end-N₊],(end-N₊+1):end] = inUpper
#     for i = 1:NPhases(model)
#         idx = ((i-1)*TotalNBases(mesh)+1:i*TotalNBases(mesh)) .+ N₋
#         if model.C[i] > 0
#             B[idx, idx] += model.C[i] * (SparseArrays.kron(
#                     SparseArrays.I(NIntervals(mesh)), me.S
#                     ) + F["+"])
#         elseif model.C[i] < 0
#             B[idx, idx] += abs(model.C[i]) * (SparseArrays.kron(
#                 SparseArrays.I(NIntervals(mesh)), me.S
#                 ) + F["-"])
#         end
#     end

#     BDict = MakeDict(B, model, mesh)

#     return FullGenerator(BDict, B, mesh.Fil)
# end

# function MakeB(model::Model, mesh::FRAPMesh, order::Int)
#     me = SFFM.MakeME(SFFM.CMEParams[order], mean = 1)
#     return MakeB(model, mesh, me)
# end


using Test
#using BenchmarkTools
using LogicCircuits

function benchmark_pass_down()
    circuit = load_smooth_logical_circuit(zoo_psdd_file("plants.psdd"));
    data = train(dataset(twenty_datasets("plants")));
    upflow = UpFlowΔ(circuit, num_examples(data), Bool,  (max_factors = 0, compact⋀ = false, compact⋁ = false));
    pass_up(upflow, data);
    circuit(data);

    @time begin
        downflow_circuit = DownFlowΔ(upflow,  (max_factors = 0, compact⋀ = false, compact⋁ = false));
        pass_down(downflow_circuit);
    end
    # 547.595 ms (2416251 allocations: 474.71 MiB)

    @time Logical.pass_down2(circuit, data);
    # 379.885 ms (1687709 allocations: 419.05 MiB)

    for (i,n) in enumerate(downflow_circuit)
        if n isa HasDownFlow
            @test downflow(n) == n.origin.origin.data[1]
        else
            @assert n isa DownFlowLeaf "$(typeof(n))"
        end
    end
end


####################  Array Mask  ####################

struct GenericMask{N, M <:AbstractArray{Bool, N}} <: AbstractArrayMask
    mask::M
end

Base.ndims(::GenericMask{N}) where N = N

GenericMask(mask) = GenericMask(convert(AbstractArray{Bool}, mask))

adapt_structure(to, x::GenericMask) = GenericMask(adapt(to, x.mask))

Base.@propagate_inbounds Base.getindex(m::Indexer{<:GenericMask}, I::Integer...) = m.mask[I...]

AxesConstraint(m::GenericMask) = (NDimConstraint(ndims(m)), ntuple(i->DimConstraint(i, size(m.mask, i)), ndims(m))...)

struct SymLengthMask{N, L <: AbstractArray{Int32, N}, B<:StaticBool} <: AbstractArrayMask
    len::L
    one::B
end

Base.ndims(::SymLengthMask{N}) where N = N + 2

SymLengthMask(len) = SymLengthMask(convert(AbstractArray{Int32}, len), static(length(len) == 1))

adapt_structure(to, x::SymLengthMask) = SymLengthMask(adapt(to, x.len))

Base.@propagate_inbounds Base.getindex(m::Indexer{<:SymLengthMask}, i, j) = (l = m.len[]; i <= l && j <= l)
Base.@propagate_inbounds function Base.getindex(m::Indexer{<:SymLengthMask}, i, j, J::Integer...)
    l = m.len[J...]
    return i <= l && j <= l
end

AxesConstraint(m::SymLengthMask{N}) where N = as_bool(m.one) ? # only one mask
    (NDimConstraint(2, true), All1Constraint(3, ndims(m))) :
    (NDimConstraint(ndims(m)), ntuple(i->DimConstraint(i+2, size(m.len, i)), N)...)

struct BiLengthMask{N, L <: AbstractArray{Int32, N}, B<:StaticBool} <: AbstractArrayMask
    q_len::L
    k_len::L
    one::B
end

Base.ndims(::BiLengthMask{N}) where N = N + 2

function BiLengthMask(q_len, k_len)
    @assert size(q_len) == size(k_len)
    return BiLengthMask(convert(AbstractArray{Int32}, q_len), convert(AbstractArray{Int32}, k_len), static(length(q_len) == 1))
end

adapt_structure(to, x::BiLengthMask) = BiLengthMask(adapt(to, x.q_len), adapt(to, x.k_len))

Base.@propagate_inbounds Base.getindex(m::Indexer{<:BiLengthMask}, i, j) = i <= m.k_len[] && j <= m.q_len[]
Base.@propagate_inbounds function Base.getindex(m::Indexer{<:BiLengthMask}, i, j, J::Integer...)
    ql = m.q_len[J...]
    kl = m.k_len[J...]
    return i <= kl && j <= ql
end

AxesConstraint(m::BiLengthMask{N}) where N = as_bool(m.one) ? # only one mask
    (NDimConstraint(2, true), All1Constraint(3, ndims(m))) :
    (NDimConstraint(ndims(m)), ntuple(i->DimConstraint(i+2, size(m.q_len, i)), N)...)


"""
Tools for exploratory data analysis and debugging.
"""
module CelesteEDA

using DataFrames

import ..DeterministicVI:
    load_source_brightnesses, accumulate_source_pixel_brightness!, ElboArgs
import ..Model:
    linear_world_to_pix, load_bvn_mixtures, ids_names, ids, CatalogEntry,
    populate_fsm_vecs!
import ..SensitiveFloats.clear!
import ..DeterministicVIImagePSF:
    FSMSensitiveFloatMatrices
using ..DeterministicVIImagePSF:
    clear_brightness!, populate_star_fsm_image!, populate_gal_fsm_image!,
    accumulate_source_image_brightness!, load_gal_bvn_mixtures

import Base.print
function print(ce::CatalogEntry)
    for field in fieldnames(ce)
      println(field, ": ", getfield(ce, field))
    end
end

function print_vp(vp::Vector{Float64})
    df = DataFrame(ids=ids_names)
    s = 1
    df[Symbol(string("v", s))] = vp
    println(df)
end

function source_pixel_location(ea::ElboArgs, s::Int, n::Int)
    p = ea.patches[s, n]
    pix_loc = linear_world_to_pix(
        p.wcs_jacobian,
        p.center,
        p.pixel_center,
        ea.vp[s][ids.u])
    return pix_loc - p.bitmap_offset
end

function render_source(ea::ElboArgs, s::Int, n::Int;
                       include_epsilon=true, field=:E_G,
                       include_iota=true)
    local p = ea.patches[s, n]
    local image = fill(NaN, size(p.active_pixel_bitmap))
    local sbs = load_source_brightnesses(
        ea, calculate_gradient=false, calculate_hessian=false)
    star_mcs, gal_mcs = load_bvn_mixtures(ea.S, ea.patches,
                                ea.vp, ea.active_sources,
                                ea.psf_K, n,
                                calculate_gradient=false,
                                calculate_hessian=false)

    p = ea.patches[s,n]

    H2, W2 = size(p.active_pixel_bitmap)
    for w2 in 1:W2, h2 in 1:H2
        # (h2, w2) index the local patch, while (h, w) index the image
        h = p.bitmap_offset[1] + h2
        w = p.bitmap_offset[2] + w2

        if !p.active_pixel_bitmap[h2, w2]
            continue
        end

        clear!(ea.elbo_vars.E_G)
        clear!(ea.elbo_vars.var_G)

        populate_fsm_vecs!(ea.elbo_vars.bvn_derivs,
                           ea.elbo_vars.fs0m_vec,
                           ea.elbo_vars.fs1m_vec,
                           ea.patches,
                           ea.active_sources,
                           ea.num_allowed_sd,
                           n, h, w,
                           gal_mcs, star_mcs)

        accumulate_source_pixel_brightness!(
            ea.elbo_vars, ea, ea.elbo_vars.E_G, ea.elbo_vars.var_G,
            ea.elbo_vars.fs0m_vec[s], ea.elbo_vars.fs1m_vec[s],
            sbs[s], ea.images[n].b, s, false)

        if field == :E_G
            image[h2, w2] = ea.elbo_vars.E_G.v[]
        elseif field == :fs0m
            image[h2, w2] = ea.elbo_vars.fs0m_vec[s].v[]
        elseif field == :fs1m
            image[h2, w2] = ea.elbo_vars.fs1m_vec[s].v[]
        else
            error("Unknown field ", field)
        end
        if include_iota
            image[h2, w2] *= ea.images[n].iota_vec[h]
        end
        if include_epsilon
            image[h2, w2] += ea.images[n].epsilon_mat[h, w]
        end
    end

    return image
end

function render_source_fft(
    ea::ElboArgs,
    fsm_mat::Matrix{FSMSensitiveFloatMatrices},
    s::Int, n::Int;
    include_epsilon=true,
    field=:E_G, include_iota=true)

    local p = ea.patches[s, n]
    local image = fill(NaN, size(p.active_pixel_bitmap))
    local sbs = load_source_brightnesses(
        ea, calculate_gradient=false, calculate_hessian=false)
    local fsms = fsm_mat[s, n]

    local gal_mcs = load_gal_bvn_mixtures(
            ea.S, ea.patches, ea.vp, ea.active_sources, n,
            calculate_gradient=false,
            calculate_hessian=false);

    clear_brightness!(fsms)
    populate_star_fsm_image!(
        ea, s, n, fsms.psf, fsms.fs0m_conv,
        fsms.h_lower, fsms.w_lower, fsms.kernel_fun, fsms.kernel_width)
    populate_gal_fsm_image!(ea, s, n, gal_mcs, fsms)
    accumulate_source_image_brightness!(ea, s, n, fsms, sbs[s])

    H2, W2 = size(p.active_pixel_bitmap)
    for w2 in 1:W2, h2 in 1:H2
        # (h2, w2) index the local patch, while (h, w) index the image
        h = p.bitmap_offset[1] + h2
        w = p.bitmap_offset[2] + w2

        if !p.active_pixel_bitmap[h2, w2]
            continue
        end

        h_fsm = h - fsms.h_lower + 1
        w_fsm = w - fsms.w_lower + 1

        image[h2, w2] = getfield(fsms, field)[h_fsm, w_fsm].v[]
        if include_iota
            image[h2, w2] *= ea.images[n].iota_vec[h]
        end
        if include_epsilon
            image[h2, w2] += ea.images[n].epsilon_mat[h, w]
        end

    end

    return deepcopy(image)
end

function show_source_image(ea::ElboArgs, s::Int, n::Int)
    p = ea.patches[s, n]
    H2, W2 = size(p.active_pixel_bitmap)
    image = fill(NaN, H2, W2);
    for w2 in 1:W2, h2 in 1:H2
        # (h2, w2) index the local patch, while (h, w) index the image
        h = p.bitmap_offset[1] + h2
        w = p.bitmap_offset[2] + w2

        if !p.active_pixel_bitmap[h2, w2]
            continue
        end

        image[h2, w2] = ea.images[n].pixels[h, w]
    end
    return image
end

end


using Visualize: CairoWindow, AbstractCairoWindow, Canvas, Scene, ColorBuffer, DepthBuffer, clip2pixel_space, CairoContext
using Visualize.JLRasterization: JLRasterizer, edge_function, interpolate,standard_transparency, broadcastmin, broadcastmax
using FileIO, GeometryTypes, Colors
using Visualize: normalmesh, Light, Shading, MeshUniforms, Camera

x = CairoWindow()

function (r::JLRasterizer{Vert, Args, FragN}){Vert, Args, FragN}(
        window::AbstractCairoWindow, vertex_array::AbstractArray{Vert}, uniforms::Args
    )
    canvas = window[Canvas]
    framebuffers = canvas[ColorBuffer]; depthbuffer = canvas[DepthBuffer]
    resolution = Vec2f0(size(framebuffers[1]))
    # hoisting out functions... Seems to help inference a bit. Or not?
    vshader = r.vertexshader
    gshader = r.geometryshader
    fshader = r.fragmentshader
    FragNVal = Val{FragN}()
    fragments_drawn = 0
    clip_triangles = NTuple{3, Vec2f0}[]
    for face in vertex_array
        vertex_stage = map(reverse(face)) do f
            vshader(f, uniforms...)
        end
        geom_stage = if isa(r.geometryshader, Void)
            (vertex_stage,)
        else
            reset!(r.geometry_view)
            gshader(r.emit, vertex_stage, uniforms...)
            r.geometry_view.view
        end
        for geom_face in geom_stage
            fdepth = map(geom_face) do vert
                fv = first(vert)
                p = clip2pixel_space(fv, resolution)
                p[Vec(1, 2)], p[3]
            end
            f = map(first, fdepth)

            depths = map(last, fdepth)
            vertex_out = map(last, geom_face)
            # Bounding rectangle
            mini = floor(Int, max.(reduce(broadcastmin, f), 1f0))
            maxi = ceil(Int, min.(reduce(broadcastmax, f), resolution))
            area = edge_function(f[1], f[2], f[3])
            fragment_count = fragments_drawn
            for y = mini[2]:maxi[2]
                for x = mini[1]:maxi[1]
                    p = Vec(Float64(x), Float64(y))
                    w = Vec(
                        edge_function(f[2], f[3], p),
                        edge_function(f[3], f[1], p),
                        edge_function(f[1], f[2], p)
                    )
                    yi, xi = (y), (x)
                    if all(w-> w <= 0f0, w) && checkbounds(Bool, framebuffers[1], yi, xi)
                        bary = w / area
                        depth = bary[1] * depths[1] + bary[2] * depths[2] + bary[3] * depths[3]

                        if depth <= depthbuffer[yi, xi]
                            depthbuffer[yi, xi] = depth
                            fragment_in = interpolate(bary, vertex_out, FragNVal)
                            fragment_out = fshader(fragment_in, uniforms...)
                            for i = eachindex(fragment_out)
                                src_color = framebuffers[i][yi, xi]
                                dest_color = fragment_out[i]
                                fragments_drawn += 1
                                color = standard_transparency(
                                    RGBA{Float32}(src_color),
                                    RGBA{Float32}(dest_color[1], dest_color[2], dest_color[3], dest_color[4])
                                )
                                framebuffers[i][yi, xi] = mapc(x-> clamp(x, 0, 1), color)
                            end
                        end
                    end
                end
            end
            if fragments_drawn != fragment_count # only add triangle if it appeared on screen!
                push!(clip_triangles, f)
            end
        end
    end
    cr = canvas[CairoContext]
    Cairo.save(cr)
    Cairo.set_source_rgba(cr, 0.5, 0, 0, 1.0)
    for face in clip_triangles
        for (i, point) in enumerate(face)
            if i == 1
                Cairo.move_to(cr, point[1], point[2])
            else
                Cairo.line_to(cr, point[1], point[2])
            end
        end
        # Cairo.close_path(cr)
    end
    # Cairo.fill(cr)
    Cairo.clip(cr)
    # Cairo.reset_clip(cr)
    Cairo.set_source_surface(cr, Cairo.CairoImageSurface(rotl90(framebuffers[1])), 0, 0)
    Cairo.paint(cr)
    Cairo.restore(cr)
    println("fragments drawn: ", fragments_drawn)
    return
end

catmesh = normalmesh(load(Pkg.dir("GLVisualize", "assets", "cat.obj")))
catmesh = Base.view(
    reinterpret(Visualize.VertexN, catmesh.parent), catmesh.indexes[1]
)

args = MeshUniforms()

light = Light(
    Vec3f0(10),
    Vec3f0(0.0),
    Vec3f0(0.9),
    1.0f0,
    Vec3f0(0.2),
    1.0f0
)
shading = Shading(
    Vec3f0(0.0),
    Vec3f0(0.3),
    8.0f0
)
scene = Visualize.SceneUniforms(x[Camera])

uniforms = (scene, light, shading, args)
raster = JLRasterizer(
    catmesh,
    uniforms,
    Visualize.vert_mesh, Visualize.frag_mesh
)

sf = canvas[Visualize.CairoSurface]
canvas[Visualize.CairoContext] = Cairo.CairoContext(sf)
canvas[DepthBuffer] = zeros(Float32, 500, 500)

Cairo.save(cr);
Cairo.set_source_rgb(cr,1.0,1.0,1.0);    # light gray
Cairo.rectangle(cr,0.0,0.0, get(x, Visualize.Resolution)...); # background
Cairo.fill(cr);
Cairo.restore(cr);
Cairo.save(cr);
Cairo.reset_clip(cr)
raster(x, catmesh, uniforms)
Cairo.write_to_png(canvas[Visualize.CairoSurface], "test.png")
# save("test.png", canvas[Visualize.ColorBuffer][1])

#
# c = CairoPDFSurface("test.pdf", resolution...);
# cr = CairoContext(c);
# Cairo.save(cr);
# set_source_rgb(cr, 1.0, 1.0, 1.0);    # light gray
# rectangle(cr, 0.0, 0.0, resolution...); # background
# fill(cr);
# restore(cr);


cdir = pwd(); ii = findlast("MDP", cdir); MDPpath = cdir[1:ii[end]] * "/";
include(MDPpath * "Installation/setup.jl");

# Read xyz file to obtain configurations
filename = "/Users/ngoccuongnguyen/Dropbox (MIT)/MDP/data/LJCluster/curated_lj_cluster.xyz";
species = ["Ar"];
config = readEXTXYZ(filename, species);

# shift atom positions to put them in the box [0,10]x[0,10]x[0,10]
config.x = config.x .+ 5.0

# cumalative sum of numbers of atoms
natom = [0; cumsum(config.natom[:])];

ci = 1; # choose a configuration to start MD simulation
x = config.x[:,(natom[ci]+1):natom[ci+1]] # atom positions of configuration ci
t = config.t[:,(natom[ci]+1):natom[ci+1]] # atom types of configuration ci
q = [];                                   # atom charges of configuration ci
v = 0.0*x # zero velocity

# 3 vectors define the simulation box
a = [10.0; 0.0; 0.0]
b = [0.0; 10.0; 0.0]
c = [0.0; 0.0; 10.0]

# periodic boundary conditions
pbc = [0; 0; 0]

# Bessel parameters
mu = [9080.776024081284, 22740.90481568728, 22826.411376542143, 12272.607095393349, 3542.9517448941156, 437.16635703260755]
rcut = 3.294988099200801
alpha = 4.902241343306524

# define Bessel potential
potentials[1] = addpotential("pair", "nonbonded", [0], "scaledbessel", mu, rcut);
include(potentials[1].potentialfunction * ".jl");  # include the potentials file

# SNAP parameters
twojmax = 4
ntypes = 1
rcutfac = rcut;
rfac0 = 1.0;
rmin0 = 0;
bzeroflag = 0;
quadraticflag = 0;
switchflag = 1;
chemflag = 0;
bnormflag = 0;
wselfallflag = 0;

snaparam = [ntypes, twojmax, rcutfac, rfac0, rmin0, bzeroflag, switchflag, quadraticflag, chemflag, bnormflag, wselfallflag]
elemradius = 0.5*ones(Float64, ntypes)
elemweight = ones(Float64, ntypes)

snapcoeff = [0.0
9.51697531984316e-5
 -3.503209826706362e-6
 -0.00015477632687816527
  8.864246681893688e-5
 -6.228723743014278e-5
  0.0002572563742710612
  1.0535143289316965e-6
  9.529421498586589e-5
  9.123169496032785e-6
 -0.0001693491133185178
  0.00013696751121984198
  3.328433092023331e-5
 -6.500144488870558e-5
 -3.968313413722244e-5]

# define SNAP potential
potentials[2] = initsna(snaparam, elemradius, elemweight, snapcoeff)

app.unitstyle = "lj"; # unit system
app.ensemblemode = "nve";
app.time = 0.0;             # initial time
app.dt = 0.5e-3;            # time step size
app.ntimesteps = 1000;       # time steps
app.globalfreq = 10;        # time frequency to print out global outputs
app.neighskin = 0.5;        # add a skin distance for neighborlist
app.eta = [rcut alpha]
app.atommasses = [1.0 1.0];

# run MD simulation using velocity Verlet algorithm
VelocityVerlet(x, v, q, t, a, b, c, pbc, app, potentials)


importall MathProgBase.SolverInterface

export Presolver

type Presolver <: AbstractMathProgSolver
    realsolver :: AbstractMathProgSolver
    # Presolver() gives an error. says get_defualt_solver undefined, why?
    function Presolver()
        realsolver = get_defualt_solver()
        new(realsolver)
    end
    function Presolver(solver::AbstractMathProgSolver)
        realsolver = solver
        new(realsolver)
    end
end

type PresolveMathProgModel <: AbstractMathProgModel
    m::Int
    n::Int
    A::SparseMatrixCSC{Float64,Int}
    b::Vector{Float64}
    c::Vector{Float64}
    p::Presolve_Problem
    orig_sense::Symbol
    solve_stat::Symbol
    obj_val::Float64
    primal_sol::Array{Float64,1}
    dual_sol::Array{Float64,1}
    innermodel :: AbstractMathProgModel
    realsolver #gives an error if I specify the type. Can't converr from SCSSolver to AbstractMathProgSolver
    function PresolveMathProgModel()
        realsolver = get_default_solver()
        innermodel = LinearQuadraticModel(realsolver)
        p = Presolve_Problem()
        new(0,0,spzeros(0,0),Int[],Int[],:Min,p,:NotSolved,0.0,Int[],model,realsolver)
    end
    function PresolveMathProgModel(solver::AbstractMathProgSolver)
        realsolver=solver
        innermodel = LinearQuadraticModel(realsolver)
        p = Presolve_Problem()
        new(0,0,spzeros(0,0),Int[],Int[],:Min,p,:NotSolved,0.0,Int[],model,realsolver)
    end
end

function ConicModel(s::Presolver)
    return PresolveMathProgModel(s.realsolver)
end

LinearQuadraticModel(s::Presolver) = ConicToLPQPBridge(ConicModel(s))

status(model::PresolveMathProgModel) = model.solve_stat
getobjval(model::PresolveMathProgModel) = model.obj_val
getsolution(model::PresolveMathProgModel) = copy(model.primal_sol)
numvar(model::PresolveMathProgModel) = model.m  # my m doesn't change. in future it might.
numconstr(model::PresolveMathProgModel) = model.n
supportedcones(s::Presolver) = [:Free, :Zero, :NonNeg, :NonPos]

loadproblem!(model::PresolveMathProgModel,c,A,b,constr_cones,var_cones)
    = loadproblem!(model,c,sparse(A),b,constr_cones,var_cones)

function loadproblem!(model::PresolveMathProgModel, c, A::SparseMatrixCSC, b, constr_cones, var_cones)
    bad_cones = [:SOC,:SOCRotated,:SDP,:ExpPrimal,:ExpDual]
    for cone_vars in constr_cones
        cone_vars[1] in bad_cones && error("Cone type $(cone_vars[1]) is currently not supported")
        cone_vars[1] == :Free && error("There is a free constraint")
    end
    for cone_vars in var_cones
        cone_vars[1] in bad_cones && error("Cone type $(cone_vars[1]) is currently not supported")
    end

    # going to be used later. c_cones, and v_cones?
    cone_dict = Dict{Symbol,Int64}()
    cone_dict[:Zero] = 1
    cone_dict[:NonNeg] = 2
    cone_dict[:NonPos] = 3
    cone_dict[:Free] = 4

    c_cones = (Array{Int,1}(),Array{Int,1}(),Array{Int,1}(),Array{Int,1}())
    for (cone,id) in constr_cones
        if(isa(id,Number))
           push!(c_cones[cone_dict[c]],id)
        else
           append!(c_cones[cone_dict[c]],id)
        end
        push!(c_cones[cone_dict[cone]],id)
    end

    v_cones = (Array{Int,1}(),Array{Int,1}(),Array{Int,1}(),Array{Int,1}())
    for (cone,id) in var_cones
        if(isa(id,Number))
           push!(v_cones[cone_dict[c]],id)
        else
           append!(v_cones[cone_dict[c]],id)
        end
        push!(v_cones[cone_dict[cone]],id)
    end

    # Column or Variable bounds
    collb = fill(-Inf, length(c))
    colub = fill(Inf, length(c))
    for (cone,idxs) in var_cones
        l = (cone == :Free || cone == :NonPos) ? -Inf : 0.0
        u = (cone == :Free || cone == :NonNeg) ?  Inf : 0.0
        for idx in idxs
            collb[idx] = l
            colub[idx] = u
        end
    end

    # Row or Constraint bounds
    rowlb = Array(Float64,length(b))
    rowub = Array(Float64,length(b))
    for (cone,idxs) in constr_cones
        for idx in idxs
            rowlb[idx] = (cone == :Zero || cone == :NonPos) ? b[idx] : -Inf
            rowub[idx] = (cone == :Zero || cone == :NonNeg) ? b[idx] :  Inf
        end
    end

    model.m,model.n = size(A)
    model.p = Presolve_Problem(false,model.m,model.n)
    model.A = A
    model.b = b
    model.c = c

    #TODO.. change the variable orders in presolve routines.

    make_presolve!(false,model.p,A,collb,colub,c,rowlb,rowub)

    #TODO.. add code for Dual Variables.
    newA,newcollb,newcolub,newc,newrowlb,newrowub = presolver!(false,model.p,A,collb,colub,c,rowlb,rowub)

    loadproblem!(model.innermodel,A,collb,colub,c,rowlb,rowub,model.orig_sense)

    return model
end

function optimize!(m::PresolveMathProgModel)
    constr_primalsol = Array{Float64,1}()
    constr_dualsol = Array{Float64,1}()
    var_primalsol = Array{Float64,1}()
    var_dualsol = Array{Float64,1}()

    if(length(find(m.p.independentvar))!=0)
        optimize!(m.innermodel)
        if(m.innermodel.solve_stat!= :Optimal)
            error("Problem status is not optimal, solution might be inaccurate")
        end
        m.solve_stat = m.innermodel.solve_stat
        constr_primalsol = m.innermodel.getconstrsolution()
        constr_dualsol = m.innermodel.getconstrduals()
        var_primalsol = m.innermodel.getsolution()
        var_dualsol = m.innermodel.getreducedcosts()

        #TODO.. need to account for DUAL VARIABLES HERE.
    end

    sol = return_postsolved(constr_primalsol,constr_dualsol,var_primalsol,var_dualsol,m.p.independentvar,m.p.pstack)

    m.primal_sol = sol[1]
    # WHAT exactly is the dual sol in the conic problem format.
    # figure that out by reading the cblib or mosek manual. Might need to convert lp dual to conic dual.
    #TODO
    m.dual_sol = sol[2]
    m.obj_val = dot(m.c, m.primal_sol)

    # you  should get the dual solution and then convert it to the type required for Conic Models.
end

#getsolution        for x ,             col primal or var_primalsol
#getconstrsolution  for y (Ax),         row primal or constr_primalsol
#getreducedcosts    for gamma-delta,    col dual or var_dualsol
#getconstrduals     for alpha-beta,     row dual or constr_dualsol

# TODO: getdual and getvardaul, can be done once add the dualsolution support for all the presolve routines.
function getdual(m::PresolveMathProgModel)
    dual = m.dual_sol[m.row_map_ind]
    # flip sign for NonPos since it's treated as NonNeg by SCS
    for i in 1:length(m.row_map_type)
        if m.row_map_type[i] == :NonPos
            dual[i] = -dual[i]
        end
    end
    return dual
end

function getvardual(m::PresolveMathProgModel)
    dual = zeros(length(m.col_map_ind))
    for i in 1:length(m.col_map_type)
        if m.col_map_type[i] == :Free
            continue # dual is zero
        elseif m.col_map_type[i] == :NonPos
            # flip sign for NonPos since it's treated as NonNeg by SCS
            dual[i] = -m.dual_sol[m.col_map_ind[i]]
        else
            dual[i] = m.dual_sol[m.col_map_ind[i]]
        end
    end
    return dual
end


using ArgParse
using Random
using CUDAnative
using CuArrays
using Flux
using FileIO

include("DecisionStrain.jl")
using .DecisionStrain
using .DecisionStrain.Data
using .DecisionStrain.FluxFixes

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table s begin
        "--gpu"
            help = "Select GPU to use."
            arg_type = Int
            default = 0
        "--dat", "-i"
            help = "path to directory with data set"
            arg_type = String
            default = "dat/FB15K"
        "--binary_dat"
            help = "Data set is stored in binary files."
            action = :store_true
        "--model"
            help = "select model. must be either DecisionStrain or ComplEx"
            arg_type = String
            default = "DecisionStrain"
        "--head_prediction"
            help = "Only used for --model DecisionStrain: train head prediction rather than tail prediction"
            action = :store_true
        "--aux"
            help = "path to auxiliary model"
            arg_type = String
        "--epochs", "-E"
            help = "number of training epochs"
            arg_type = Int
            default = 100
        "--batchsize", "-B"
            help = "minibatch size"
            arg_type = Int
            default = 500
        "--eval_batchsize"
            help = "minibatch size for evaluation"
            arg_type = Int
            default = 500
        "--hidden_dim", "-k"
            help = "dimension of embedding space"
            arg_type = Int
            default = 20
        "--lmbda"
            help = "prefactor of regularizer"
            arg_type = Float64
            default = 0.01
        "--lr"
            help = "learning rate"
            arg_type = Float64
            default = 1.0
        "--neg_ratio"
            help = "Number of negative samples per positive sample"
            arg_type = Int
            default = 10
        "--rngseed"
            help = """
                seed for random number generator; if not provided, a seed will
                automatically be generated from system entropy and written to
                the output file."""
            arg_type = UInt32
        "--population", "-n"
            help = "number of simultaneously trained decision trees"
            arg_type = Int
            default = 1
        "--out", "-o"
            help = "output directory; must not exist"
            arg_type = String
        "--continue", "-c"
            help = "provide file name to continue training from."
            arg_type = String
        "--eval_every"
            help = "evaluate model every n epochs on validation set (set to zero to evaluate only at end)"
            arg_type = Int
            default = 10
        "--save_every"
            help = "save checkpoing every n epochs (set to zero to save only at end)"
            arg_type = Int
            default = 100
        "--copy_subtrees"
            help = "copy branching vectors of winning subtrees"
            action = :store_true
        "--precision"
            help = """
                floating point precision;  allowed values are "half", "single",
                and "double".  Default is "single"."""
            arg_type = Union{Type{Float16}, Type{Float32}, Type{Float64}}
            default = Float32
    end

    return parse_args(s, as_symbols = true)
end

function ArgParse.parse_item(
        ::Type{Union{Type{Float16}, Type{Float32}, Type{Float64}}},
        s::AbstractString)
    if s == "half"
        return Float16
    elseif s == "single"
        return Float32
    elseif s == "double"
        return Float64
    else
        throw("""Unknown float type "$s".""")
    end
end

function dicttostr(dict::Dict)
    stream = IOBuffer()
    show(stream, MIME("text/plain"), dict)
    lines = split(String(take!(stream)), "\n")
    "Dict(\n" * join("$line,\n" for line in lines[2:end]) * ")"
end

function main()
    args = parse_commandline()

    device!(args[:gpu])

    # initialize random number generator
    if args[:rngseed] === nothing
        args[:rngseed] = rand(RandomDevice(), UInt32)
    end
    rng = MersenneTwister(args[:rngseed]);

    if args[:out] === nothing
        args[:out] = joinpath("out", repr(hash(args))[end - 5 : end])
    end

    # create output directory and open log file
    print("Output directory: $(args[:out])\n")
    mkdir(args[:out])
    try
        open(joinpath(args[:out], "log"), write = true) do log
            # TODO: print date, time, hostname, git commit, and $CUDA_VISIBLE_DEVICES
            print(log, """
            args = $(dicttostr(args))

            # Running on $(Threads.nthreads()) threads.

            """)

            dat = Dataset(args[:dat], binary_files = args[:binary_dat], log = log,
                          reverse = args[:head_prediction])
            if args[:continue] === nothing
                if args[:model] == "DecisionStrain"
                    model = DecisionStrainModel(rng, dat, args, log = log)
                elseif args[:model] == "ComplEx"
                    aux_model = nothing
                    if args[:aux] !== nothing
                        aux_model = gpu(load(args[:aux], "model"), args, aux_model=true)
                    end
                    model = ComplExModel(rng, dat, args, aux_model)
                else
                    throw("Unknown model $(args[:model]).")
                end
                print(log, "# Initialized new model.\n\n")
            else
                model = gpu(load(args[:continue], "model"), args)
                print(log, "# Loaded model from file $(args[:continue])\n\n")
            end
            flush(log)

            # Run the training loop in a separate function to allow the compiler to generate
            # specialized code based on floattype(model).
            training_loop!(
                model, dat, rng, args[:epochs], args[:batchsize], args[:eval_every],
                args[:save_every], args[:out], log = log)
        end
    catch e
        open(joinpath(args[:out], "err"), write = true) do err
            trace = catch_backtrace()
            showerror(err, e, trace)
            print(err, "\n")
            rethrow()
        end
    end
end

function training_loop!(
        model::AbstractKnowledgGraphModel,
        dat::Dataset,
        rng::AbstractRNG,
        nepochs::Integer,
        batchsize::Integer,
        eval_every::Integer,
        save_every::Integer,
        out_path::AbstractString;
        log::IO = stdout)

    trainlength = size(dat.train_unsorted, 2)
    num_batches = trainlength ÷ batchsize

    # print(log, "initial_performance = ", evaluate(model, dat.val, dat), "\n\n")
    print(log, """
    col_headers = ["epoch", "avg_duration_continuous", "avg_duration_discrete",
        "mrr", "hits_at", "log_likelihood_per_data_point", "eval_fail_cnt"]

    data = [
    """)
    flush(log)

    sum_duration_cont = 0.0
    sum_duration_disc = 0.0
    duration_disc = 0.0
    last_print = 0

    for epoch in 1 : nepochs
        # print_stats(model)
        _, duration_cont = @timed begin
            shuffled = CuArray(shuffle_train(rng, dat))
            for batchstart = 1 : batchsize : batchsize * num_batches
                train_minibatch!(
                    rng, model, shuffled, batchstart, batchsize, epoch, log = log)
            end
        @sync_cuda("continuous updates done")
        end
        sum_duration_cont += duration_cont

        if needs_discrete_optimization(model)
            _, duration_disc = @timed begin
                train_epoch!(rng, model, dat, epoch, batchsize, log = log)
                @sync_cuda("discrete updates done")
            end
            sum_duration_disc += duration_disc
        else
            duration_disc = 0.0
        end

        GC.gc()

        print(log, "# Epoch $epoch took $duration_cont + $duration_disc seconds.\n")
        flush(log)

        if save_every != 0 && epoch % save_every == 0
            save_checkpoint(model, out_path, epoch)
            GC.gc()
        end

        # if eval_every != 0 && epoch % eval_every == 0
        #     print(log, (
        #         epoch,
        #         sum_duration_cont / (epoch - last_print),
        #         sum_duration_disc / (epoch - last_print),
        #         evaluate(model, dat.val, dat)...
        #         ), ",\n")
        #     flush(log)
        #     sum_duration_cont = sum_duration_disc = 0.0
        #     last_print = epoch
        #     GC.gc()
        # end
    end

    if save_every == 0 || nepochs % save_every != 0
        save_checkpoint(model, out_path, nepochs)
        GC.gc()
    end

    # if eval_every == 0 || nepochs % eval_every != 0
    #     print(
    #         log,
    #         (nepochs, nothing, nothing, evaluate(model, dat.val, dat)...),
    #         ",\n"
    #     )
    # end
    print(log, "]\n")
    flush(log)
end

@noinline function save_checkpoint(model, directory, epoch)
    save(joinpath(directory, "checkpoint-$epoch.jld2"), "model", cpu(model))
end

@noinline function evaluate(
            model::AbstractKnowledgGraphModel,
            heldout::AbstractArray{<:Integer, 2},
            fulldat::Dataset;
            batchsize = 100,
            windows=[1, 3, 10])
    hits_cnt_at_t = zeros(Int, length(windows), Threads.nthreads())
    hits_cnt_at_h = zeros(Int, length(windows), Threads.nthreads())
    sum_rr_t = zeros(Float64, Threads.nthreads())
    sum_rr_h = zeros(Float64, Threads.nthreads())
    cnt = zeros(Int, Threads.nthreads())
    total_ll = 0.0

    for batchstart = 1 : batchsize : size(heldout, 2)
        batchend = min(batchstart + batchsize - 1, size(heldout, 2))
        batch = batchstart : batchend
        hs = heldout[1, batch]
        ts = heldout[2, batch]
        rs = heldout[3, batch]
        all_ll, predicted_lls, raw_ranks, sum_ll = log_likelihoods_and_raw_ranks(model, hs, ts, rs)
        total_ll += sum_ll

        Threads.@threads for i = 1 : length(batch)
            predicted_ll = predicted_lls[i]
            raw_rank = raw_ranks[i]
            other_correct_ts = view(fulldat.all_sorted, fulldat.all_ranges[(hs[i], rs[i])])
            other_correct_lls = all_ll[other_correct_ts, i]
            filtered_rank_t = raw_rank - sum(other_correct_lls .> predicted_ll)

            threadid = Threads.threadid()

            for (ii, w) = enumerate(windows)
                if filtered_rank_t <= w
                    hits_cnt_at_t[ii, threadid] += 1
                end
            end

            sum_rr_t[threadid] += 1 / filtered_rank_t
            cnt[threadid] += 1
        end
    end

    success_cnt = sum(cnt)
    hits_cnt_at_t = sum(hits_cnt_at_t, dims = 2)[:, 1]
    hits_at_t = Dict(w => hits / success_cnt for (w, hits) = zip(windows, hits_cnt_at_t))
    mrr_t = sum(sum_rr_t) / success_cnt
    return mrr_t, hits_at_t, total_ll / size(heldout, 2), size(heldout, 2) - success_cnt
end

main()


module FresnelOptics

const VecI  = AbstractVector # Input  Vector
const VecO  = AbstractVector # Output Vector
const VecB  = AbstractVector # Buffer Vector
const VecIO = AbstractVector # In/Out Vector
const MatI  = AbstractMatrix # Input  Matrix
const MatO  = AbstractMatrix # Output Matrix
const MatB  = AbstractMatrix # Buffer Matrix
const MatIO = AbstractMatrix # In/Out Matrix

const 𝚷 = 2.0 * π

using LinearAlgebra.BLAS: gemv!

include("./utils/la.jl")
include("./fresnel.jl")
include("./lorentz.jl")
include("./dielectric.jl")
include("./layers.jl")

end # module


function beta(obj::GNSSFieldAttitudeContext, arg0::FieldAbsoluteDate)
    return jcall(obj, "beta", RealFieldElement, (FieldAbsoluteDate,), arg0)
end

function beta_ds(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "betaDS", FieldDerivativeStructure, ())
end

function compute_phi(obj::GNSSFieldAttitudeContext, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(obj, "computePhi", RealFieldElement, (RealFieldElement, RealFieldElement), arg0, arg1)
end

function get_date(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "getDate", FieldAbsoluteDate, ())
end

function get_delta_ds(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "getDeltaDS", FieldDerivativeStructure, ())
end

function get_mu_rate(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "getMuRate", RealFieldElement, ())
end

function get_orbit_angle_since_midnight(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "getOrbitAngleSinceMidnight", RealFieldElement, ())
end

function get_secured_beta(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "getSecuredBeta", RealFieldElement, ())
end

function get_sv_bcos(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "getSVBcos", RealFieldElement, ())
end

function get_turn_span(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "getTurnSpan", FieldTurnSpan, ())
end

function get_yaw_end(obj::GNSSFieldAttitudeContext, arg0::RealFieldElement)
    return jcall(obj, "getYawEnd", RealFieldElement, (RealFieldElement,), arg0)
end

function get_yaw_start(obj::GNSSFieldAttitudeContext, arg0::RealFieldElement)
    return jcall(obj, "getYawStart", RealFieldElement, (RealFieldElement,), arg0)
end

function in_orbit_plane_absolute_angle(obj::GNSSFieldAttitudeContext, arg0::RealFieldElement)
    return jcall(obj, "inOrbitPlaneAbsoluteAngle", RealFieldElement, (RealFieldElement,), arg0)
end

function in_sun_side(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "inSunSide", jboolean, ())
end

function in_turn_time_range(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "inTurnTimeRange", jboolean, ())
end

function linear_model_still_active(obj::GNSSFieldAttitudeContext, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(obj, "linearModelStillActive", jboolean, (RealFieldElement, RealFieldElement), arg0, arg1)
end

function nominal_yaw(obj::GNSSFieldAttitudeContext, arg0::FieldAbsoluteDate)
    return jcall(obj, "nominalYaw", TimeStampedFieldAngularCoordinates, (FieldAbsoluteDate,), arg0)
end

function orbit_normal_yaw(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "orbitNormalYaw", TimeStampedFieldAngularCoordinates, ())
end

function set_half_span(obj::GNSSFieldAttitudeContext, arg0::RealFieldElement, arg1::jdouble)
    return jcall(obj, "setHalfSpan", void, (RealFieldElement, jdouble), arg0, arg1)
end

function set_up_turn_region(obj::GNSSFieldAttitudeContext, arg0::jdouble, arg1::jdouble)
    return jcall(obj, "setUpTurnRegion", jboolean, (jdouble, jdouble), arg0, arg1)
end

function time_since_turn_start(obj::GNSSFieldAttitudeContext)
    return jcall(obj, "timeSinceTurnStart", RealFieldElement, ())
end

function turn_corrected_attitude(obj::GNSSFieldAttitudeContext, arg0::FieldDerivativeStructure)
    return jcall(obj, "turnCorrectedAttitude", TimeStampedFieldAngularCoordinates, (FieldDerivativeStructure,), arg0)
end

function turn_corrected_attitude(obj::GNSSFieldAttitudeContext, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(obj, "turnCorrectedAttitude", TimeStampedFieldAngularCoordinates, (RealFieldElement, RealFieldElement), arg0, arg1)
end

function yaw_rate(obj::GNSSFieldAttitudeContext, arg0::RealFieldElement)
    return jcall(obj, "yawRate", RealFieldElement, (RealFieldElement,), arg0)
end


importall MathProgBase.SolverInterface
importall SemidefiniteModels

export CSDPMathProgModel, CSDPSolver

immutable CSDPSolver <: AbstractMathProgSolver
    options::Dict{Symbol,Any}
end
CSDPSolver(;kwargs...) = CSDPSolver(Dict{Symbol,Any}(kwargs))

type CSDPMathProgModel <: AbstractSDModel
    C
    b
    As
    X
    y
    Z
    status::Cint
    pobj::Cdouble
    dobj::Cdouble
    options::Dict{Symbol,Any}
    function CSDPMathProgModel(; kwargs...)
        new(nothing, nothing, nothing, nothing, nothing, nothing,
            -1, 0.0, 0.0, Dict{Symbol, Any}(kwargs))
    end
end
SDModel(s::CSDPSolver) = CSDPMathProgModel(; s.options...)
ConicModel(s::CSDPSolver) = SDtoConicBridge(SDModel(s))

supportedcones(s::CSDPSolver) = [:Free,:Zero,:NonNeg,:NonPos,:SDP]
function setvartype!(m::CSDPMathProgModel, vtype, blk, i, j)
    if vtype != :Cont
        error("Unsupported variable type $vtype by CSDP")
    end
end

function loadproblem!(m::CSDPMathProgModel, filename::AbstractString)
    if endswith(filename,".dat-s")
       m.C, m.b, As = read_prob(filename)
       m.As = [ConstraintMatrix(As[i], i) for i in 1:length(As)]
    else
       error("unrecognized input format extension in $filename")
    end
end
#writeproblem(m, filename::String)
function loadproblem!(m::CSDPMathProgModel, blkdims::Vector{Int}, constr::Int)
    m.C = blockmatzeros(blkdims)
    m.b = zeros(Cdouble, constr)
    m.As = [constrmatzeros(i, blkdims) for i in 1:constr]
end

function setconstrB!(m::CSDPMathProgModel, val, constr::Integer)
    m.b[constr] = val
end
function setconstrentry!(m::CSDPMathProgModel, coef, constr::Integer, blk::Integer, i::Integer, j::Integer)
    m.As[constr][blk][i,j] = coef
end
function setobjentry!(m::CSDPMathProgModel, coef, blk::Integer, i::Integer, j::Integer)
    m.C[blk][i,j] = coef
end

function optimize!(m::CSDPMathProgModel)
    As = map(A->A.csdp, m.As)

    let wrt = string(get(m.options, :write_prob, ""))
        if length(wrt) > 0
            k = 1
            wrtf = "$wrt.$k"
            while isfile(wrtf)
                wrtf = "$wrt.$k"
                k += 1
            end
            info("Writing problem to $(pwd())/$(wrtf)")
            write_prob(wrtf, m.C, m.b, As)
        end
    end

    m.X, m.y, m.Z = initsoln(m.C, m.b, As)
    m.status, m.pobj, m.dobj = easy_sdp(m.C, m.b, As, m.X, m.y, m.Z)
end

function status(m::CSDPMathProgModel)
    status = m.status
    if status == 0
        return :Optimal
    elseif status == 1
        return :Infeasible
    elseif status == 2
        return :Unbounded
    elseif 3 <= status <= 7
        return :Unknown
    elseif 8 <= status <= 9
        return :Error
    elseif status == -1
        return :Uninitialized
    else
        error("Internal library error: status=$status")
    end
end

function getobjval(m::CSDPMathProgModel)
    m.pobj
end
function getsolution(m::CSDPMathProgModel)
    m.X
end
function getdual(m::CSDPMathProgModel)
    m.y
end
function getvardual(m::CSDPMathProgModel)
    m.Z
end


# Defines structus and functions relating to display.
# Defines `render` for all game objects.

struct ScreenPixelPos  # 0,0 == top-left
    x::Int
    y::Int
end
ScreenPixelPos(x::Number, y::Number) = ScreenPixelPos(convert.(Int, floor.((x,y)))...)
struct UIPixelPos  # 0,0 == top-left (Same as ScreenPixelPos but not scaled.)
    x::Int
    y::Int
end
UIPixelPos(x::Number, y::Number) = UIPixelPos(convert.(Int, floor.((x,y)))...)

+(a::UIPixelPos, b::UIPixelPos) = UIPixelPos(a.x+b.x, a.y+b.y)
# TODO: 1.0?
+(a::UIPixelPos, x::Number) = UIPixelPos(a.x+x, a.y+x)

mutable struct Camera
    pos::WorldPos
    w::Threads.Atomic{Int32}   # Note: These are Atomics, since they can be modified by the
    h::Threads.Atomic{Int32}   # windowEventWatcher callback, which can run in another thread!
end
Camera() = Camera(WorldPos(0,0),100,100)

screenCenter() = UIPixelPos(winWidth[]/2, winHeight[]/2)
screenCenterX() = winWidth[]/2
screenCenterY() = winHeight[]/2
screenOffsetFromCenter(x::Int,y::Int) = UIPixelPos(screenCenterX()+x,screenCenterY()+y)

worldScale(c::Camera) = cam.w[] / winWidth[];
function toScreenPos(p::WorldPos, c::Camera)
    scale = worldScale(c)
    ScreenPixelPos(
        floor(c.w[]/2.0 + scale*p.x), floor(c.h[]/2.0 - scale*p.y))
end
function toScreenPos(p::UIPixelPos, c::Camera)
    scale = worldScale(c)
    ScreenPixelPos(floor(scale*p.x), floor(scale*p.y))
end
function screenToUI(p::ScreenPixelPos, c::Camera)
    scale = worldScale(c)
    ScreenPixelPos(floor(p.x/scale), floor(p.y/scale))
end
function screenScaleDims(w,h,c::Camera)
    scale = worldScale(c)
    scale*w, scale*h
end

function render(o::Ball, cam::Camera, renderer)
    ballW = ballWidth; ballH = ballWidth;
    topLeft = WorldPos(o.pos.x - ballW/2.0, o.pos.y + ballH/2.0)
    screenPos = toScreenPos(topLeft, cam)
    rect = SDL2.Rect(screenPos.x, screenPos.y, screenScaleDims(ballW, ballH, cam)...)
    color = kBallColor
    SDL2.SetRenderDrawColor(renderer, Int64(color.r), Int64(color.g), Int64(color.b), Int64(color.a))
    SDL2.RenderFillRect(renderer, Ref(rect) )
end
function render(o::Paddle, cam::Camera, renderer)
    paddleW = o.length; paddleH = kPaddleRenderH;
    # Move up/down so the edge matches the "center" of paddle.
    edgeShift = if o.pos.y > 0; paddleH/2.0; else -paddleH/2.0; end
    topLeft = WorldPos(o.pos.x - paddleW/2.0, o.pos.y + paddleH/2.0 + edgeShift)
    screenPos = toScreenPos(topLeft, cam)
    rect = SDL2.Rect(screenPos.x, screenPos.y, screenScaleDims(paddleW, paddleH, cam)...)
    color = kPaddleColor
    SDL2.SetRenderDrawColor(renderer, Int64(color.r), Int64(color.g), Int64(color.b), Int64(color.a))
    SDL2.RenderFillRect(renderer, Ref(rect) )
end

abstract type AbstractButton end
mutable struct MenuButton <: AbstractButton
    enabled::Bool
    pos::UIPixelPos
    w::Int
    h::Int
    text::String
    callBack
end
mutable struct KeyButton <: AbstractButton
    enabled::Bool
    pos::UIPixelPos
    w::Int
    h::Int
    text::String
    callBack
end

mutable struct CheckboxButton
    toggled::Bool
    button::MenuButton
end

import Base.run
run(b::AbstractButton) = b.callBack()
function run(b::CheckboxButton)
    b.toggled = !b.toggled
    b.button.callBack(b.toggled)
end

# pointwise subtraction with bounds checking (floors to 0)
-(a::SDL2.Color, b::Int) = SDL2.Color(a.r-min(b,a.r), a.g-min(b,a.g), a.b-min(b,a.b), a.a-min(b,a.a))
SDL2.Color(1,5,1,1) - 2 == SDL2.Color(0,3,0,0)

function render(b::AbstractButton, cam::Camera, renderer, color, fontSize)
    if (!b.enabled)
         return
    end
    topLeft = UIPixelPos(b.pos.x - b.w/2.0, b.pos.y - b.h/2.0)
    screenPos = toScreenPos(topLeft, cam)
    rect = SDL2.Rect(screenPos.x, screenPos.y, screenScaleDims(b.w, b.h, cam)...)
    x,y = Int[0], Int[0]
    SDL2.GetMouseState(pointer(x), pointer(y))
    if clickedButton == b
        if mouseOnButton(UIPixelPos(x[],y[]),b,cam)
            color = color - 50
        else
            color = color - 30
        end
    else
        if mouseOnButton(UIPixelPos(x[],y[]),b,cam)
            color = color - 10
        end
    end
    SDL2.SetRenderDrawColor(renderer, Int64(color.r), Int64(color.g), Int64(color.b), 255)
    SDL2.RenderFillRect(renderer, Ref(rect) )
    renderText(renderer, cam, b.text, b.pos; fontSize = fontSize)
end

function render(b::MenuButton, cam::Camera, renderer)
    render(b, cam, renderer, kMenuButtonColor, kMenuButtonFontSize)
end
function render(b::KeyButton, cam::Camera, renderer)
    render(b, cam, renderer, kKeySettingButtonColor, kKeyButtonFontSize)
end

function render(b::CheckboxButton, cam::Camera, renderer)
    # Hack: move button text offcenter before rendering to accomodate checkbox
    offsetText = " "
    text_backup = b.button.text
    b.button.text = offsetText * b.button.text
    render(b.button, cam, renderer)
    b.button.text = text_backup

    # Render checkbox
    render_checkbox_square(b.button, 6, SDL2.Color(200,200,200, 255), cam, renderer)

    if b.toggled
        # Inside checkbox "fill"
        render_checkbox_square(b.button, 8, SDL2.Color(100,100,100, 255), cam, renderer)
    end
end

function render_checkbox_square(b::AbstractButton, border, color, cam, renderer)
    checkbox_radius = b.h/2.0 - border  # (checkbox is a square)
    topLeft = UIPixelPos(b.pos.x - b.w/2.0, b.pos.y - b.h/2.0)
    topLeft = topLeft + border
    screenPos = toScreenPos(topLeft, cam)
    rect = SDL2.Rect(screenPos.x, screenPos.y, screenScaleDims(checkbox_radius*2, checkbox_radius*2, cam)...)
    SDL2.SetRenderDrawColor(renderer, Int64(color.r), Int64(color.g), Int64(color.b), 255)
    SDL2.RenderFillRect(renderer, Ref(rect) )
end

# ---- Text Rendering ----

fonts_cache = Dict()
txt_cache = Dict()
function sizeText(cam, txt, fontName, fontSize)
    scale = worldScale(cam)
    sizeText(scale, txt, loadFont(scale, fontName, fontSize))
end
function sizeText(scale, txt, font::Ptr{SDL2.TTF_Font})
   fw,fh = Cint[1], Cint[1]
   SDL2.TTF_SizeText(font, txt, pointer(fw), pointer(fh))
   return fw[1]/scale,fh[1]/scale
end
function loadFont(scale, fontName, fontSize)
   fontSize = scale*fontSize
   fontKey = (fontName, fontSize)
   if haskey(fonts_cache, fontKey)
       font = fonts_cache[fontKey]
   else
       font = SDL2.TTF_OpenFont(fontKey...)
       font == C_NULL && throw(ErrorException("Failed to load font '$fontKey'"))
       fonts_cache[fontKey] = font
   end
   return font
end
function createText(renderer, cam, txt, fontName, fontSize)
   font = loadFont(worldScale(cam), fontName, fontSize)
   txtKey = (font, txt)
   if haskey(txt_cache, txtKey)
       tex = txt_cache[txtKey]
   else
       text = SDL2.TTF_RenderText_Blended(font, txt, SDL2.Color(20,20,20,255))
       tex = SDL2.CreateTextureFromSurface(renderer,text)
       #SDL2.FreeSurface(text)
       txt_cache[txtKey] = tex
   end

   fw,fh = Cint[1], Cint[1]
   SDL2.TTF_SizeText(font, txt, pointer(fw), pointer(fh))
   fw,fh = fw[1],fh[1]

   return tex, fw, fh
end
@enum TextAlign centered leftJustified rightJustified
function renderText(renderer, cam::Camera, txt::String, pos::UIPixelPos
                    ; fontName = defaultFontName,
                     fontSize=defaultFontSize, align::TextAlign = centered)
   tex, fw, fh = createText(renderer, cam, txt, fontName, fontSize)
   renderTextSurface(renderer, cam, pos, tex, fw, fh, align)
end

function renderTextSurface(renderer, cam::Camera, pos::UIPixelPos,
                           tex::Ptr{SDL2.Texture}, fw::Integer, fh::Integer, align::TextAlign)
   screenPos = toScreenPos(pos, cam)
   renderPos = SDL2.Rect(0,0,0,0)
   if align == centered
       renderPos = SDL2.Rect(Int(floor(screenPos.x-fw/2.0)), Int(floor(screenPos.y-fh/2.0)), fw,fh)
   elseif align == leftJustified
       renderPos = SDL2.Rect(Int(floor(screenPos.x)), Int(floor(screenPos.y-fh/2.0)), fw,fh)
   else # align == rightJustified
       renderPos = SDL2.Rect(Int(floor(screenPos.x-fw)), Int(floor(screenPos.y-fh/2.0)), fw,fh)
   end
   SDL2.RenderCopy(renderer, tex, C_NULL, pointer_from_objref(renderPos))
   #SDL2.DestroyTexture(tex)
end

function hcat_render_text(lines, renderer, cam, gap, pos::UIPixelPos;
         fixedWidth=nothing, fontName=defaultFontName, fontSize=defaultFontSize)
    numLines = size(lines)[1]
    if fixedWidth != nothing
        widths = fill(fixedWidth,size(lines))
    else
        widths = [sizeText(cam, line, fontName, fontSize)[1] for line in lines]
    end
    totalWidth = sum(widths) + gap*(numLines-1)
    runningWidth = 0
    leftMostPos = pos.x - totalWidth/2.0
    text_centers = []
    for i in 1:numLines
        linePos = leftMostPos + runningWidth
        leftPos = UIPixelPos(linePos, pos.y)
        renderText(renderer, cam, lines[i], leftPos; fontName=fontName, fontSize=fontSize, align=leftJustified)
        runningWidth += widths[i] + gap
        push!(text_centers, UIPixelPos(leftPos.x + widths[i]÷2, leftPos.y))
    end
    return text_centers
end

#  ------- Image rendering ---------

function render(t::Ptr{SDL2.Texture}, pos, cam::Camera, renderer; size = nothing)
    if (t == C_NULL) return end
    pos = toScreenPos(pos, cam)
    if size != nothing
        size = toScreenPos(size, cam)
        w = size.x
        h = size.y
    else
        w,h,access = Cint[1], Cint[1], Cint[1]
        format = Cuint[1]
        SDL2.QueryTexture( t, format, access, w, h );
        w,h = w[], h[]
    end
    rect = SDL2.Rect(pos.x - w÷2,pos.y - h÷2,w,h)
    SDL2.RenderCopy(renderer, t, C_NULL, pointer_from_objref(rect))
end


#! /usr/bin/env julia

println("Start test 'tests/train.jl'")

include("../src/DCellC.jl")
using DCellC

model = Multiscale3(GreyscaleImage, bn = true, seed = 42)
println("Test train.jl: Creating instance of FCRNA works")

lmgs  = [ synthesize(256, 256, (1500, 2000)) for i in 1:2 ]
println("Test train.jl: Creating synthesized images works")

train!(model, lmgs, patchsize = 128, epochs = 1, patchmode = 2, lr = 1e-3)
train!(model, lmgs, epochs = 1, lr = 1e-3)
println("Test train.jl: Training of model was sucessful")

println("Test train.jl completed")


using StaticArrays

@test begin
    x0 = SVector(0., 0.)
    x1 = SVector(1., 1.)

    curve = Curve([x0, x1])
    is_consistent(curve)
end

@test begin
    x0 = SVector(0., 0.)
    x1 = SVector(1., 1.)
    x2 = SVector(1., 1.)

    curve = Curve([x0, x1, x2])
    ! is_consistent(curve)
end

@test begin
    x0 = SVector(0., 0.)
    x1 = SVector(0., 0.)

    curve = Curve([x0, x1])
    !is_consistent(curve)
end

@test begin
    x0 = SVector(0., 0.)
    x1 = SVector(1., 1.)
    x2 = SVector(2., 2.)

    curve = Curve([x0, x1, x2], [(1, 2), (2, 3), (3, 1)])
    is_consistent(curve)
end

@test begin
    x0 = SVector(0., 0.)
    x1 = SVector(1., 1.)
    x2 = SVector(2., 2.)

    curve = Curve([x0, x1, x2], [(1, 1), (2, 3), (3, 1)])
    !is_consistent(curve)
end

@test begin
    x0 = SVector(0., 0.)
    x1 = SVector(1., 1.)
    x2 = SVector(2., 2.)

    curve = Curve([x0, x1, x2], [(1, 9), (2, 3), (3, 1)])
    !is_consistent(curve)
end

# IO
@test begin
    ext = ".vtu"
    all(length(Curve(vtu_path*ext).points) > 0 for vtu_path in ("mesh2d", "mesh2din3d", "tet_mesh"))
end

@test begin
    ext = ".xml"
    all(length(Curve(vtu_path*ext).points) > 0 for vtu_path in ("mesh2d", "mesh2din3d", "tet_mesh"))
end

@test begin
    ext = ".msh"
    all(length(Curve(vtu_path*ext).points) > 0 for vtu_path in ("mesh2d", "mesh2din3d", "tet_mesh"))
end


mutable struct SimulationState
    g::Global
    population::Population
    setup::Evaluation
    evaluate!::Function
    function SimulationState(g::Global)
        p = Population(g, g.cf.checkpoint_filename)
        setup = Evaluation(g.cf.dataset_id, g.cf.fitness_function_id, normalise=g.cf.normalise_input,
                           balance=g.cf.balance_dataset)
        fitness_function = (f::Function) -> f.(p.population)
        new(g, p, setup, fitness_function)
    end
    function SimulationState(g::Global, p::Population)
        setup = Evaluation(g.cf.dataset_id, g.cf.fitness_function_id, normalise=g.cf.normalise_input,
                           balance=g.cf.balance_dataset)
        fitness_function = (f::Function) -> f.(p.population)
        new(g, p, setup, fitness_function)
    end
end

function Base.show(io::IO, state::SimulationState)
    print(io, "State: \n")
    println(io, state.population)
    print(io, state.setup.dataset_setup)
end


# --------------------------------------------------------------------------- #
# Prepare layout for batches of images, weights, states and prmaps
# For suitable arraytypes, this will transfer the data into the graphics
# memory

function packimage(img :: Image; at = Array{Float32})
 pack = reshape(imgdata(img), imgsize(img)..., imgchannels(img), 1)
 return convert(at, pack)
end

function packimage(imgs :: Vector{<: Image}; at = Array{Float32})
  pack = cat(4, map(packimage, imgs)...)
  return convert(at, pack)
end

function packweights(w :: Vector{Any}; at = Array{Float32})
  return Any[ convert(at, x) for x in w ]
end

function packbnmoments(m; at = Array{Float32})
  mean = (m.mean != nothing) ? convert(at, m.mean) : nothing
  var  = (m.var  != nothing) ? convert(at, m.var)  : nothing
  return bnmoments(momentum = m.momentum, mean = mean, var = var)
end

function packstate(s :: Vector{Any}; at = Array{Float32})
  return Any[ packbnmoments(x, at = at) for x in s ]
end

function packproxymap(prmap; at = Array{Float32})
  if length(size(prmap)) < 3
    return convert(at, reshape(prmap, size(prmap)..., 1))
  else
    return convert(at, prmap)
  end
end

unpackweights(w) = packweights(w, at = Array{Float32})
unpackstate(s) = packstate(s, at = Array{Float32})

# --------------------------------------------------------------------------- #
# Check consistency between Array types

# TODO: This should be possible without relying on explicit
# knowledge of KnetArrays
arraytype(a :: Array{Float32}) = Array{Float32}
arraytype(a :: KnetArray{Float32}) = KnetArray{Float32}

# --------------------------------------------------------------------------- #
# Loss function

function loss(w, s, imgdata, prmap,
              mt :: Type{<: Model}; at = nothing, kwargs...)

  # Loss that can be applied to raw data without conversion
  # if `at` is not specified. This can be used to efficiently
  # calculate the loss for already gpu-packed data. If `at`
  # is specified, conversion might take place.

  if at != nothing
    w       = packweights(w, at = at)
    s       = packstate(s,  at = at)
    prmap   = packproxymap(prmap, at = at)
    imgdata = packimage(imgdata, at = at)
  end

  return sum(abs2, density(w, s, imgdata, mt) .- prmap) / length(prmap)
end

function loss(w, s, imgdata, lbls :: Vector{Label},
              mt :: Type{<: Model}; at = nothing, kwargs...)

  # This function will build proximity maps from the given labels.
  # If the argument `at` is not specified, these proximity maps
  # inherit their arraytype from `imgdata`

  _at = (at != nothing) ? at : arraytype(imgdata)

  imgdata = packimage(imgdata, at = _at)
  prmap   = proxymap(imgsize(imgdata)..., lbls; kwargs...)
  prmap   = packproxymap(prmap)
  return loss(w, s, imgdata, prmap, mt; at = at, kwargs...)
end

function loss{I <: Image}(w, s, img :: I, lbl :: Label,
                          mt :: Type{<: Model{I}}; at = nothing, kwargs...)

  # Loss calculated for an image
  # Creates proximity maps and then calculates the loss

  _at = (at != nothing) ? at : arraytype(imgdata(img))

  x = packimage(img, at = _at)
  return loss(w, s, x, [lbl], mt; at = at, kwargs...)
end

function loss(m :: Model, img :: GreyscaleImage, lbl :: Label)
  return loss(weights(m), state(m), img, lbl, typeof(m))
end

# --------------------------------------------------------------------------- #
# Loss gradient

_gradloss = gradloss(loss)  # Calculate both the gradient and loss value
_grad     = grad(loss)

lossgrad(args...; kwargs...) = _gradloss(args...; kwargs...)[2:-1:1]

function lossgrad{I <: Image}(m :: Model{I}, img :: I, lbl :: Label; kwargs...)
  return lossgrad(weights(m), state(m), img, lbl, typeof(m); kwargs...)
end

# --------------------------------------------------------------------------- #
# Divide dataset in batches

function makebatches(imgs, lbls, batchsize; shuffle = false)
  n = length(imgs)
  idx = shuffle ? Base.shuffle(collect(1:n)) : collect(1:n)
  sel = collect(1:batchsize:n)
  m = length(sel)

  idxs = [ idx[sel[i] : (i < m ? sel[i+1] - 1 : n)] for i in 1:m ]
  return [ (imgs[range], lbls[range]) for range in idxs ]
end

function packbatches(batches; at = Array{Float32}, kwargs...)
  return map(batches) do batch
    imgs   = packimage(batch[1], at = at)
    prmaps = proxymap(imgsize(imgs)..., batch[2]; kwargs...)
    prmaps = packproxymap(prmaps, at = at)
    return (imgs, prmaps)
  end
end

# --------------------------------------------------------------------------- #
# High level training function

"""
    train!(model, imgs, lbls; kwargs...)

    train!(model, lmgs; kwargs...)

Train a model `model` using the images `imgs` and labels `lbls`. Alternatively, an iterable `lmgs` of labeled images can be provided.

# Arguments
- `testset`: an iterable of labeled images to be used for testing the network.
- `epochs = 10`: the number of epochs (iterations of the training images)
  used for training.
- `batchsize = 1`: the size of a batch used for one update step of the model.
- `shuffle = true`: whether to shuffle the training images for each epoche
- `log = true`: whether to print information about the model's training
  performance.
- `logpath`: redirect logging to this file.
- `modelpath`: file to store the trained model in
- `record = true`: record the training performance in a `Dict` and return it.
- `patchsize`: providing an integer for this option causes that the
  effective training images are patches of this size, generated from the
  input image during each epoche.
- `patchmode`: if `patchmode < 1`, create patches by cutting ordered
  patches from the given input images. If `patchmode > 0`, create
  `patchmode` patches by random selection.
- `at = Array{Float32}`: the arraytyped used for training and evaluating.
  Currently, this is restricted to `Array{Float32}` for cpu computation or
  `KnetArray{Float32}` for gpu computation.
- `kernelsize = 7`: size of the kernel used to construct proximity maps.
- `peakheight = 100`: height of the kernels used to construct proximity
  maps.
- `opt = Adam`: Knet optimizer to be used. All other keyword arguments are
  passed to the constructor of the optimizer.
"""

function train!(model :: Model, imgs, lbls;
                epochs = 10,
                opt = Adam,
                log = true,
                logpath = nothing,
                record = true,
                batchsize = 1,
                patchsize = nothing,
                patchmode = 0,
                shuffle = false,
                modelpath = nothing,
                testset = nothing,
                kernelsize = 7,
                peakheight = 100.,
                imageop = Id(),
                at = nothing,
                kwargs...)

  # Sanity checks

  @assert (length(imgs) == length(lbls))
  if patchsize == nothing
    @assert all(x -> x == imgsize(imgs[1]), imgsize.(imgs))
  end

  # Pack the model weights and state.
  # In practice this means to transfer them to
  # gpu memory if a gpu based arratype is given.

  if at == nothing
    at = gpu() >= 0 ? KnetArray{Float32} : Array{Float32}
  end

  w = packweights(weights(model), at = at)
  s = packstate(state(model), at = at)

  # Initialize the optimizers and conduct some
  # general preparations

  optim = optimizers(w, opt; kwargs...)
  test  = (testset != nothing) && (log || record)

  if patchsize == nothing
    n = ceil(Int, length(imgs) / batchsize)
  elseif patchmode > 0
    n = ceil(Int, length(imgs) * patchmode / batchsize)
  else
    # WARNING: This will not work if margins or manual offsets are given to
    # `ordered_patches`
    m1 = floor(imgsize(imgs[1])[1] / patchsize)
    m2 = floor(imgsize(imgs[1])[2] / patchsize)
    n = ceil(Int, length(imgs) * m1 * m2 / batchsize)
  end

  if record
    rec = Dict{Symbol, Array{Float64}}(
               :loss_train    => zeros(n, epochs),
               :count_train   => zeros(n, epochs),
               :meanadj_train => zeros(n, epochs),
               :maxadj_train  => zeros(n, epochs))
    if test
      rec[:loss_test]    = zeros(epochs)
      rec[:count_test]   = zeros(epochs)
      rec[:meanadj_test] = zeros(epochs)
      rec[:maxadj_test]  = zeros(epochs)
    end
  end

  # Bring the testset in the correct format

  if test
    if typeof(testset) <: NTuple{2}
      timgs, tlbls = testset
    else
      timgs = [ t[1] for t in testset ]
      tlbls = [ t[2] for t in testset ]
    end
    tbatches = makebatches(timgs, tlbls, batchsize, shuffle = false)
    tpacks   = packbatches(tbatches, at = at,
                           peakheight = peakheight,
                           kernelsize = kernelsize)
  end

  # Print overview

  if log
    @printf "# CellCC training progress\n"
    @printf "#\n"
    @printf "# model      %s\n" typeof(model)
    @printf "# date       %s\n" now()
    @printf "# images     %d\n" length(imgs)
    @printf "# imagesize  %s\n" imgsize(imgs[1])
    @printf "# patchsize  %s\n" patchsize
    @printf "# patchmode  %s\n" patchmode
    @printf "# batches    %d\n" n
    @printf "# batchsize  %d\n" batchsize
    @printf "# epochs     %d\n" epochs
    @printf "# testset    %s\n" test ? length(timgs) : "-"
    @printf "# shuffle    %s\n" shuffle
    @printf "# record     %s\n" record
    @printf "# optimizer  %s\n" optim[1]
    @printf "# modelpath  %s\n" modelpath == nothing ? "-" : modelpath
    @printf "# logpath    %s\n" logpath == nothing ? "-" : logpath
  end

  # Start the training loop

  time = @elapsed for i in 1:epochs

    # Generate batches used for training. These batches
    # are also 'packed', meaning that they are brought in
    # a dense format consisting of image and proximity
    # map data arrays. Depending on `at`, this might
    # transfer the batch-data to gpu memory.

    # TODO: make the batch-acquicision process more dynamic!

    if patchsize != nothing && patchmode > 0
      patches = random_patches.(imgs, lbls, patchmode,
                                imageop = imageop,
                                size = (patchsize, patchsize))

    elseif patchsize != nothing && patchmode <= 0
      patches = ordered_patches.(imgs, lbls,
                                 imageop = imageop,
                                 size = (patchsize, patchsize))

    else
      patches = apply.(imageop, imgs, lbls)
    end

    timgs, tlbls = vcat(first.(patches)...), vcat(second.(patches)...)

    batches = makebatches(timgs, tlbls, batchsize, shuffle = shuffle)
    packs   = packbatches(batches, at = at,
                          peakheight = peakheight,
                          kernelsize = kernelsize)

    if log
      @printf "\n"
      @printf "# Epoche  Batch   Loss  Count  MeanAdj   MaxAdj \n"
      @printf "# ----------------------------------------------\n"
    end

    # Training epoche

    for j in 1:n

      # Make sure that everything is printed

      flush(STDOUT)
      flush(STDERR)

      # Extract image data and proximity maps

      data, prmaps = packs[j]
      lbl = batches[j][2]

      l, lg = lossgrad(w, s, data, prmaps, typeof(model))
      update!(w, lg, optim)

      if log || record

        # Let the model predict labels

        dlbl = label(w, s, data, typeof(model))

        # Evaluate the quality of the prediction

        c      = mean(length.(dlbl) ./ length.(lbl))
        adj    = adjacency.(dlbl, lbl)
        mm, mx = mean(mean, adj), mean(maximum, adj)

        if log
          @printf("  %6d  %5d  %5.1f  %5.3f  %7.2f  %7.2f\n",
                  i, j, l, c, mm, mx)
        end

        if record
          rec[:loss_train][j, i]    = l
          rec[:count_train][j, i]   = c
          rec[:meanadj_train][j, i] = mm
          rec[:maxadj_train][j, i]  = mx
        end
      end
    end

    # Evaluate performance on the test set

    if test

      l = c = mm = mx = 0.

      for j in 1:length(tpacks)
        data, prmaps = tpacks[j]
        timg, tlbl = tbatches[j]

        dlbl = label(w, s, data, typeof(model))
        adj  = adjacency.(dlbl, tlbl)

        l  += loss(w, s, data, prmaps, typeof(model)) * length(timg)
        c  += sum(length.(dlbl) ./ length.(tlbl))
        mm += sum(mean, adj)
        mx += sum(maximum, adj)
      end

      l  /= length(timgs)
      c  /= length(timgs)
      mm /= length(timgs)
      mx /= length(timgs)

      if log
          @printf "\n"
          @printf("  %6d      0  %5.1f  %5.3f  %7.2f  %7.2f\n",
                   i, l, c, mm, mx)
      end

      if record
        rec[:loss_test][i]    = l
        rec[:count_test][i]   = c
        rec[:meanadj_test][i] = mm
        rec[:maxadj_test][i]  = mx
      end
    end
  end

  # Update the model weights and its state

  weights(model)[:] = unpackweights(w)
  state(model)[:]   = unpackstate(s)

  if log
    @printf "\n"
    @printf "# End of training process after %.1f seconds\n" time
  end

  if logpath != nothing
    # TODO
    @printf "# Warning: Logpath not yet implemented\n"
    @printf "# Did not save a version of this output in file %s\n" logpath
  end

  if modelpath != nothing
    modelsave(modelpath, model)
  end

  return record ? rec : nothing
end

function train!(model :: Model, data; kwargs...)
  imgs = [ d[1] for d in data ]
  lbls = [ d[2] for d in data ]
  train!(model, imgs, lbls; kwargs...)
end

function train!(model :: Model, gen :: Function; samples = 10, kwargs...)
  data = [gen() for i in 1:samples]
  train!(model, data; kwargs...)
end

# --------------------------------------------------------------------------- #
# Function for training with a single image frame
# Intended for testing purposes

function train!(model :: Model, img :: Image, lbl :: Label;
                epochs = 10, record = true, opt = Adam,
                log = true, logpath = nothing, modelpath = nothing, kwargs...)

  optim = optimizers(weights(model), opt; kwargs...)

  if record
    rec = Dict{Symbol, Vector{Float64}}(
               :loss    => zeros(epochs),
               :count   => zeros(epochs),
               :meanadj => zeros(epochs),
               :maxadj  => zeros(epochs))
  end

  if log
    @printf "# CellCC training progress (single image mode)\n"
    @printf "#\n"
    @printf "# model      %s\n" typeof(model)
    @printf "# date       %s\n" now()
    @printf "# epochs     %d\n" epochs
    @printf "# record     %s\n" record
    @printf "# optimizer  %s\n" optim[1]
    @printf "# modelpath  %s\n" modelpath == nothing ? "-" : modelpath
    @printf "# logpath    %s\n" logpath == nothing ? "-" : logpath
    @printf "\n"
    @printf "# Epoche   Loss  Count  MeanAdj   MaxAdj\n"
    @printf "# --------------------------------------\n"
  end

  time = @elapsed for i in 1:epochs
    lg, l = lossgrad(model, img, lbl)
    update!(weights(model), lg, optim)

    if log || record

      dens = density(model, img)
      dlbl = label(dens)

      c      = length.(dlbl) ./ length.(lbl)
      adj    = adjacency(dlbl, lbl)
      mm, mx = mean(adj), maximum(adj)

      if log
        @printf("  %6d  %5.1f  %5.3f  %7.2f  %7.2f\n",
                 i, l, c, mm, mx)
      end

      if record
        rec[:loss][i]    = l
        rec[:count][i]   = c
        rec[:meanadj][i] = mm
        rec[:maxadj][i]  = mx
      end
    end
  end

  @printf "\n"
  @printf "# End of training process after %.1f seconds\n" time

  if logpath != nothing
    # TODO
    @printf "# Warning: Logpath not yet implemented\n"
    @printf "# Did not save a version of this output in file %s\n" logpath
  end

  if modelpath != nothing
    modelsave(model, modelpath)
    @printf "# Model saved at %s\n" modelpath
  end

  return record ? rec : nothing
end


export initNet

function initNet(layersNames,layersDims,netType=:classification)
  net=Layer[]
  if netType==:classification
    for (i,layerName) in enumerate(layersNames)
      if layerName==:d
        net=[net;dataLayer(layersDims[i,:]...)]
      elseif layerName==:h
        net=[net;hiddenLayer(layersDims[i,:]...)]
      elseif layerName==:p
        net=[net;predictLayer(layersDims[i,:]...)]
      elseif layerName==:l
        net=[net;lossLayer(layersDims[i,1])]
        net[i].Cost=crossEntropy;net[i].∂Cost=∂crossEntropy
      end
    end
  elseif netType==:regression
    for (i,layerName) in enumerate(layersNames)
      if layerName==:d
        net=[net;dataLayer(layersDims[i,:]...)]
      elseif layerName==:h
        net=[net;hiddenLayer(layersDims[i,:]...)]
      elseif layerName==:p
        net=[net;predictLayer(layersDims[i,:]...)]
      elseif layerName==:l
        net=[net;lossLayer(layersDims[i,1])]
      end
    end
  end
  return net
end