"""
    DofHandler(grid::Grid)

Construct a `DofHandler` based on the grid `grid`.
"""
struct DofHandler{dim,T}
    variables::Vector{TrialFunction}
    cell_dofs::Vector{Int}
    cell_dofs_offset::Vector{Int}
    mesh::AbstractPolytopalMesh{dim,T}
end

function DofHandler(mesh::AbstractPolytopalMesh{dim,T}, trialfuncs::Vector{TrialFunction{dim,T}}) where {dim,T}
    dofhandler = DofHandler{dim,T}(trialfuncs, Int[], Int[], mesh)
    _distribute_dofs(dofhandler)
end

function Base.show(io::IO, dh::DofHandler)
    println(io, "DofHandler")
    println(io, "  Number of Fields: ",nvariables(dh))
    println(io, "  Dofs per cell: ", ndofs_per_cell(dh))
    print(io, "  Total dofs: ", ndofs(dh))
end

ndofs(dh::DofHandler) = maximum(dh.cell_dofs)
ndofs_per_cell(dh::DofHandler, cell::Int=1) = dh.cell_dofs_offset[cell+1] - dh.cell_dofs_offset[cell]
@inline nvariables(dh::DofHandler) = length(dh.variables)

function find_field(dh::DofHandler, field::TrialFunction)
    j = findfirst(i->i == field, dh.variables)
    j == 0 && error("did not find field $field")
    return j
end

# Calculate the offset to the first local dof of a field
function field_offset(dh::DofHandler, field::TrialFunction, ci::Int)
    offset = 0
    for i in 1:find_field(dh, field)-1
        offset += getnlocaldofs(getfunctionspace(dh.variables[i]), getcells(dh.cells)[ci])::Int
    end
    return offset
end

"""
    dof_range(dh::DofHandler, u::Symbol)

Return the local dof range for `u`.
"""
function dof_range(dh::DofHandler, field::TrialFunction, ci::Int)
    f = find_field(dh, field)
    offset = field_offset(dh, field, ci)
    n_field_dofs = getnlocaldofs(getfunctionspace(dh.variables[f]),getcells(dh.mesh)[ci])
    return (offset+1):(offset+n_field_dofs)
end

# close the DofHandler and distribute all the dofs
#TODO: This only work for D < 3
function _distribute_dofs(dh::DofHandler{dim,T}) where {dim,T}
    # `vertexdict` keeps track of the visited vertices. We store the global vertex
    # number and the first dof we added to that vertex.
    vertexdicts = [Dict{entityeltype(dh.mesh,0),Int}() for _ in 1:nvariables(dh)]

    # `edgedict` keeps track of the visited edges, this will only be used for a 3D problem
    # We also need to store the direction
    # of the first edge we encounter and add dofs too. When we encounter the same edge
    # the next time we check if the direction is the same, otherwise we reuse the dofs
    # in the reverse order
    edgedicts = [Dict{entityeltype(dh.mesh,1),Int}() for _ in 1:nvariables(dh)]

    # `facedict` keeps track of the visited faces. We only need to store the first dof we
    # added to the face; if we encounter the same face again we *always* reverse the order
    # In 2D a face (i.e. a line) is uniquely determined by 2 vertices, and in 3D a
    # face (i.e. a surface) is uniquely determined by 3 vertices.
    facedicts = dim < 3 ? [] : [Dict{entityeltype(dh.mesh,2),Int}() for _ in 1:nvariables(dh)]

    topologyDicts = dim == 2 ? Dict(0=>vertexdicts, 1=>edgedicts) : Dict(0=>vertexdicts, 1=>edgedicts, 2 => facedicts)

    nextdof = 1 # next free dof to distribute
    push!(dh.cell_dofs_offset, 1) # dofs for the first cell start at 1

    # loop over all the cells, and distribute dofs for all the fields
    for (ci, cell) in enumerate(getcells(dh.mesh))
        # Get topologies
        n_max_topology_elements = maximum(keys(gettopology(dh.mesh,cell)))
        geometric_cell_topology = gettopology(dh.mesh,cell)
        for fi in 1:nvariables(dh)
            element = getelement(getfunctionspace(dh.variables[fi]))
            cell_topology = gettopology(dh.mesh, cell, element)
            for n_element in 0:n_max_topology_elements-1
                n_el_dofs_cell = cell_topology[n_element]
                n_el_cell = geometric_cell_topology[n_element]
                @assert mod(n_el_dofs_cell, n_el_cell) == 0
                nelementdofs = Int(n_el_dofs_cell/n_el_cell)
                if nelementdofs > 0
                    for element in getcellsubentities(dh.mesh,ci,n_element)
                        token = ht_keyindex2!(topologyDicts[n_element][fi], element)
                        if token > 0 # reuse dofs
                            reuse_dof = topologyDicts[n_element][fi].vals[token]
                            for d in 1:getncomponents(getfunctionspace(dh.variables[fi]))
                                push!(dh.cell_dofs, reuse_dof + (d-1))
                            end
                        else # token <= 0, distribute new dofs
                            for elementdof in 1:nelementdofs
                                Base._setindex!(topologyDicts[n_element][fi], nextdof, element, -token)
                                for d in 1:getncomponents(getfunctionspace(dh.variables[fi]))
                                    push!(dh.cell_dofs, nextdof)
                                    nextdof += 1
                                end
                            end
                        end
                    end # element loop
                end
            end #elements loop
            ncelldofs = cell_topology[dim]
            if ncelldofs > 0 # always distribute new dofs for cell
                for celldof in 1:ncelldofs
                    for d in 1:getncomponents(getfunctionspace(dh.variables[fi]))
                        push!(dh.cell_dofs, nextdof)
                        nextdof += 1
                    end
                end # cell loop
            end
        end # field loop
        # push! the first index of the next cell to the offset vector
        push!(dh.cell_dofs_offset, length(dh.cell_dofs)+1)
    end # cell loop
    return dh
end

function celldofs!(global_dofs::Vector{Int}, dh::DofHandler, i::Int)
    @assert length(global_dofs) == ndofs_per_cell(dh, i)
    unsafe_copyto!(global_dofs, 1, dh.cell_dofs, dh.cell_dofs_offset[i], length(global_dofs))
    return global_dofs
end

function vertexdofs(dh::DofHandler, u::TrialFunction)
    element = getelement(getfunctionspace(u))
    dofs = zeros(Int,getnvertices(dh.mesh))
    for (ci, cell) in enumerate(getcells(dh.mesh))
        nv = gettopology(dh.mesh,cell, element)[0]*getncomponents(getfunctionspace(u))
        for i in dh.cell_dofs_offset[ci]:(dh.cell_dofs_offset[ci]+nv-1)
            dof = dh.cell_dofs[i]
            if !(dof ∈ dofs)
                dofs[getverticesindices(dh.mesh,cell)[i-dh.cell_dofs_offset[ci]+1]] = dof
            end
        end
    end
    return dofs
end

# Creates a sparsity pattern from the dofs in a DofHandler.
# Returns a sparse matrix with the correct storage pattern.
"""
    create_sparsity_pattern(dh::DofHandler)

Create the sparsity pattern corresponding to the degree of freedom
numbering in the [`DofHandler`](@ref). Return a `SparseMatrixCSC`
with stored values in the correct places.

See the [Sparsity Pattern](@ref) section of the manual.
"""
@inline create_sparsity_pattern(dh::DofHandler) = _create_sparsity_pattern(dh, false)

"""
    create_symmetric_sparsity_pattern(dh::DofHandler)

Create the symmetric sparsity pattern corresponding to the degree of freedom
numbering in the [`DofHandler`](@ref) by only considering the upper
triangle of the matrix. Return a `Symmetric{SparseMatrixCSC}`.

See the [Sparsity Pattern](@ref) section of the manual.
"""
@inline create_symmetric_sparsity_pattern(dh::DofHandler) = Symmetric(_create_sparsity_pattern(dh, true), :U)

function _create_sparsity_pattern(dh::DofHandler, sym::Bool)
    ncells = getncells(dh.mesh)
    n = ndofs_per_cell(dh)
    N = sym ? div(n*(n+1), 2) * ncells : n^2 * ncells
    N += ndofs(dh) # always add the diagonal elements
    I = Int[]; resize!(I, N)
    J = Int[]; resize!(J, N)
    global_dofs = zeros(Int, n)
    cnt = 0
    for element_id in 1:ncells
        celldofs!(global_dofs, dh, element_id)
        @inbounds for j in 1:n, i in 1:n
            dofi = global_dofs[i]
            dofj = global_dofs[j]
            sym && (dofi > dofj && continue)
            cnt += 1
            if cnt > length(J)
                resize!(I, trunc(Int, length(I) * 1.5))
                resize!(J, trunc(Int, length(J) * 1.5))
            end
            I[cnt] = dofi
            J[cnt] = dofj

        end
    end
    @inbounds for d in 1:ndofs(dh)
        cnt += 1
        if cnt > length(J)
            resize!(I, trunc(Int, length(I) + ndofs(dh)))
            resize!(J, trunc(Int, length(J) + ndofs(dh)))
        end
        I[cnt] = d
        J[cnt] = d
    end
    resize!(I, cnt)
    resize!(J, cnt)
    V = zeros(length(I))
    K = sparse(I, J, V)
    return K
end


module SimulationSummary

using Statistics
using Serialization
using Printf
using Parameters
using TOML
using CSV

using ..SimulationParams: SimulationParameters
using ..Models: AbstractBond, SSHBond, Bond
using ..Models: AbstractModel, SSHModel, HolsteinModel, write_phonons!
using ..Models: write_phonons!, write_M_matrix!
using ..Measurements: measure_κ
using ..MuFinder: MuTuner, update_μ!, estimate_μ
using ..Utilities: reshaped

export write_simulation_summary!

"""
Write results of simulation to simulation summary file.
"""
function write_simulation_summary!(datafolder::String,Nbins::Int=10)

    # unpack last checkpoint from simulation
    chkpnt = deserialize(joinpath(datafolder,"checkpoint.jls"))
    @unpack model, μ_tuner, sim_stats = chkpnt

    # data foldername
    foldername = splitdir(datafolder)[end]

    # read in config file
    files    = readdir(datafolder)
    tomls    = findall(f -> endswith(f, r"\.toml|\.TOML"), files)
    tomlfile = joinpath(datafolder, files[tomls[1]])
    input    = TOML.parsefile(tomlfile)

    # write final phonon config to file
    write_phonons!(model,joinpath(datafolder,"$(foldername)_config.out"))

    # write M matrix to file
    if haskey(input["simulation"],"write_M_matrix")
        if input["simulation"]["write_M_matrix"]
            write_M_matrix!(model, joinpath(datafolder,"$(foldername)_matrix.out"))
        end
    end

    # construct filename for summary file
    filename = joinpath(datafolder, "$(foldername)_summary.out")

    # open summary file
    open(filename,"w") do fout

        # copy contents of input file to output file
        write(fout,"#########################\n")
        write(fout,"## INPUT FILE CONTENTS ##\n")
        write(fout,"#########################\n\n")

        TOML.print(fout, input)
        write(fout,"\n")

        # write bond defintions to summary file
        write(fout,"######################\n")
        write(fout,"## BOND DEFINITIONS ##\n")
        write(fout,"######################\n\n")

        write_bond_definitions!(fout,model)

        # write phonon definitions
        write(fout,"########################\n")
        write(fout,"## PHONON DEFINITIONS ##\n")
        write(fout,"########################\n\n")

        write_phonon_definitions!(fout,model)

        # write chemical potential defintions
        write(fout,"#########################\n")
        write(fout,"## CHEMICAL POTENTIALS ##\n")
        write(fout,"#########################\n\n")

        write_chemical_potential_definitions!(fout,model)

        # write simulation statistics to file
        write(fout, "#####################","\n")
        write(fout, "## SIMULATION INFO ##","\n")
        write(fout, "#####################","\n","\n")

        total_time = sim_stats["simulation_time"] + sim_stats["measurement_time"] + sim_stats["write_time"]
        @printf fout "Total Time (min)        = %.8f\n" total_time
        @printf fout "Simulation Time (min)   = %.8f\n" sim_stats["simulation_time"]
        @printf fout "Measurement Time (min)  = %.8f\n" sim_stats["measurement_time"]
        @printf fout "Write Time (min)        = %.8f\n" sim_stats["write_time"]
        @printf fout "Iterative Solver Steps  = %.8f\n" sim_stats["iters"]
        @printf fout "Acceptance Rate         = %.8f\n" sim_stats["acceptance_rate"]
        @printf fout "Reflect Acceptance Rate = %.8f\n" sim_stats["reflect_acceptance_rate"]
        @printf fout "Swap Acceptance Rate    = %.8f\n" sim_stats["swap_acceptance_rate"]

        # write global measurements to file
        write(fout,"\n")
        write(fout, "#########################","\n")
        write(fout, "## GLOBAL MEASUREMENTS ##","\n")
        write(fout, "#########################","\n","\n")

        write_global_measurements!(fout,model,μ_tuner,datafolder,Nbins)

        # write global measurements to file
        write(fout, "\n")
        write(fout, "##########################","\n")
        write(fout, "## ON-SITE MEASUREMENTS ##","\n")
        write(fout, "##########################","\n","\n")

        write_onsite_measurements!(fout,model,μ_tuner,datafolder,Nbins)

        # write global measurements to file
        write(fout,"\n")
        write(fout, "#############################","\n")
        write(fout, "## INTER-SITE MEASUREMENTS ##","\n")
        write(fout, "#############################","\n","\n")

        write_intersite_measurements!(fout,model,datafolder,Nbins)

        # write susceptibilities
        write(fout,"\n")
        write(fout, "######################","\n")
        write(fout, "## SUSCEPTIBILITIES ##","\n")
        write(fout, "######################","\n","\n")

        write_susceptibilities!(fout,model,datafolder,Nbins)

        # write on-site correlations
        write(fout,"\n")
        write(fout, "##################","\n")
        write(fout, "## CORRELATIONS ##","\n")
        write(fout, "##################","\n","\n")

        write_correlations!(fout,model,datafolder,Nbins)
    end

    return nothing
end

######################
## PRIVATE FUNCTION ##
######################

"""
Write bond defintions to file.
"""
function write_bond_definitions!(fout,model::SSHModel{T1,T2,T3}) where {T1,T2,T3}

    bonds = model.bond_definitions

    for (id, bond) in enumerate(bonds)

        write(fout,"Bond ID       = ", string(id), "\n")
        write(fout,"name          = ",bond.name, "\n")
        write(fout,"t_avg         = ", string(bond.t),  "\n")
        write(fout,"t_std         = ", string(bond.σt), "\n")
        write(fout,"Initial Orbit = ", string(bond.o₁), "\n")
        write(fout,"Final Orbit   = ", string(bond.o₂), "\n")
        write(fout,"Displacement  = ", string(bond.v),  "\n", "\n")
    end

    return nothing
end

function write_bond_definitions!(fout,model::HolsteinModel{T1,T2,T3}) where {T1,T2,T3}

    bonds = model.bond_definitions

    for (id, bond) in enumerate(bonds)

        write(fout,"Bond ID       = ", string(id), "\n")
        write(fout,"t_avg         = ", string(bond.t),  "\n")
        write(fout,"t_std         = ", string(bond.σt), "\n")
        write(fout,"Initial Orbit = ", string(bond.o₁), "\n")
        write(fout,"Final Orbit   = ", string(bond.o₂), "\n")
        write(fout,"Displacement  = ", string(bond.v),  "\n", "\n")
    end

    return nothing
end

"""
Write Phonon definitions to file.
"""
function write_phonon_definitions!(fout,model::SSHModel{T1,T2,T3}) where {T1,T2,T3}

    bonds = model.bond_definitions

    id = 0
    for bond in bonds

        if bond.has_phonon
            id += 1
            write(fout,"SSH Phonon ID = ", string(id), "\n")
            write(fout,"name = ",bond.name, "\n")
            write(fout,"t_avg = ", string(bond.t),  "\n")
            write(fout,"t_std = ", string(bond.σt), "\n")
            write(fout,"alpha_avg = ", string(bond.α),  "\n")
            write(fout,"alpha_std = ", string(bond.σα), "\n")
            write(fout,"alpha2_avg = ", string(bond.α₂),  "\n")
            write(fout,"alpha2_std = ", string(bond.σα₂), "\n")
            write(fout,"omega_avg = ", string(bond.ω),  "\n")
            write(fout,"omega_std = ", string(bond.σω), "\n")
            write(fout,"omega4_avg = ", string(bond.ω₄),  "\n")
            write(fout,"omega4_std = ", string(bond.σω₄), "\n")
            write(fout,"Initial Orbit = ", string(bond.o₁), "\n")
            write(fout,"Final Orbit   = ", string(bond.o₂), "\n")
            write(fout,"Displacement  = ", string(bond.v),  "\n", "\n")
        end
    end

    return nothing
end

function write_phonon_definitions!(fout,model::HolsteinModel{T1,T2,T3}) where {T1,T2,T3}

    L₁ = model.lattice.L1
    L₂ = model.lattice.L2
    L₃ = model.lattice.L3
    n  = model.lattice.unit_cell.norbits

    ω  = reshaped(model.ω,(n,L₁,L₂,L₃))
    ω₄ = reshaped(model.ω₄,(n,L₁,L₂,L₃))
    λ  = reshaped(model.λ,(n,L₁,L₂,L₃))
    λ₂ = reshaped(model.λ₂,(n,L₁,L₂,L₃))

    for o in 1:n

        write(fout,"Orbit = ", string(o), "\n")

        vals  = @view ω[o,:,:,:]
        avg   = mean(vals)
        if length(vals)>1
            stdev = std(vals)
        else
            stdev = 0.0
        end
        write(fout,"Omega_avg   = ", string(avg), "\n")
        write(fout,"Omega_std   = ", string(stdev), "\n")

        vals  = @view ω₄[o,:,:,:]
        avg   = mean(vals)
        if length(vals)>1
            stdev = std(vals)
        else
            stdev = 0.0
        end
        write(fout,"Omega4_avg  = ", string(avg), "\n")
        write(fout,"Omega4_std  = ", string(stdev), "\n")

        vals  = @view λ[o,:,:,:]
        avg   = mean(vals)
        if length(vals)>1
            stdev = std(vals)
        else
            stdev = 0.0
        end
        write(fout,"Lambda_avg  = ", string(avg), "\n")
        write(fout,"Lambda_std  = ", string(stdev), "\n")

        vals  = @view λ₂[o,:,:,:]
        avg   = mean(vals)
        if length(vals)>1
            stdev = std(vals)
        else
            stdev = 0.0
        end
        write(fout,"Lambda2_avg = ", string(avg), "\n")
        write(fout,"Lambda2_std = ", string(stdev), "\n","\n")
    end

    return nothing
end

"""
Write chemical potential defintions to file.
"""
function write_chemical_potential_definitions!(fout,model::AbstractModel{T1,T2,T3}) where {T1,T2,T3}

    bonds = model.bond_definitions

    L₁ = model.lattice.L1
    L₂ = model.lattice.L2
    L₃ = model.lattice.L3
    n  = model.lattice.unit_cell.norbits

    μ  = reshaped(model.μ,(n,L₁,L₂,L₃))

    for o in 1:n

        vals  = @view μ[o,:,:,:]
        avg   = mean(vals)
        if length(vals)>1
            stdev = std(vals)
        else
            stdev = 0.0
        end
        write(fout,"Orbit  = ", string(o), "\n")
        write(fout,"Mu_avg = ", string(avg), "\n")
        write(fout,"Mu_std = ", string(stdev), "\n", "\n")
    end

    return nothing
end

"""
Write global measurements to file.
"""
function write_global_measurements!(fout, model::AbstractModel{T1,T2,T3}, μ_tuner::MuTuner{T1}, datafolder::String, Nbins::Int) where {T1,T2,T3}

    # file statistics will be written to
    statsfn = joinpath(datafolder, "global_measurements_stats.out")

    # data folder
    datafolder = joinpath(datafolder,"global_measurements_f")

    # get files in data folder
    files = readdir(datafolder,join=true)

    # number of data files
    Nmeas = length(files)

    # calculate number of measurements per bin
    Nbins = min(Nmeas,Nbins)
    N     = div(Nmeas,Nbins)
    @assert Nmeas%Nbins==0

    # get measurements
    measurements = String[]
    open(files[1],"r") do fin
        for line in readlines(fin)
            atoms = split(line)
            push!( measurements , atoms[1] )
        end
    end

    # container to hold binned data
    binned_data = Dict(string(k)=>zeros(T1,Nbins) for k in measurements)

    # container to contain statistics
    stats = Dict(string(k)=>zeros(T1,2) for k in measurements)

    # iterate over files in data folder
    for (i,file) in enumerate(files)
        # open file
        open(file,"r") do fin
            # iterate over each line in data file
            for line in eachline(fin)
                # split the line at white space
                atoms = split(line)
                # get the bin
                bin   = div(i-1,N)+1
                # record measurement
                binned_data[atoms[1]][bin] += parse(T1,atoms[2]) / N
            end
        end
    end

    # calculate average and error bar for each measurement
    for measurement in keys(binned_data)

        # deal with chemical potential seperately due to μ-tuning algorithm
        if measurement=="mu"
            estimate_μ(μ_tuner)
            stats[measurement][1] = μ_tuner.μ_avg
            stats[measurement][2] = μ_tuner.μ_err
        else
            avg, err = mean_and_error(binned_data[measurement])
            stats[measurement][1] = avg
            stats[measurement][2] = err
        end
    end

    # calculate the compressibility
    β   = model.β
    N   = model.Nsites
    n   = stats["density"][1]
    Δn  = stats["density"][2]
    N²  = stats["Nsqr"][1]
    ΔN² = stats["Nsqr"][2]
    κ, Δκ = measure_κ(β,N,N²,ΔN²,n,Δn)
    stats["compressibility"] = [κ,Δκ]

    # write averaged statistics to file
    open(statsfn,"w") do statsout

        # header
        header = "global_meas avg error\n"
        write(fout,    header)
        write(statsout,header)

        # iterate over measurements
        for measurement in keys(stats)

            # calculate mean and standard deviation of mean
            avg = stats[measurement][1]
            err = stats[measurement][2]

            # write statistics to file
            @printf fout     "%s %.8f %.8f\n" measurement avg err
            @printf statsout "%s %.8f %.8f\n" measurement avg err
        end
    end

    return nothing
end

"""
Write on-site measreuments to file.
"""
function write_onsite_measurements!(fout,model::AbstractModel{T1,T2,T3},μ_tuner::MuTuner{T1},datafolder::String,Nbins::Int) where {T1,T2,T3}

    # number of site per unit cell
    nₛ = model.lattice.unit_cell.norbits::Int

    # data folder
    folder = joinpath(datafolder, "onsite_measurements_f")

    # get files containing data
    files = readdir(folder,join=true)

    # number of data files
    Nmeas = length(files)

    # file statistics will be written to
    statsfn = joinpath(datafolder, "onsite_measurements_stats.out")

    # calculate number of measurements per bin
    Nbins = min(Nmeas,Nbins)
    N     = div(Nmeas,Nbins)
    @assert Nmeas%Nbins==0

    # get measurements
    measurements = String[]
    open(files[1],"r") do fin
        header = readline(fin)
        for line in readlines(fin)
            push!( measurements , split(line)[1] )
        end
    end
    measurements = unique(measurements)

    # container to hold binned data
    binned_data = Dict(string(k)=>zeros(T1,Nbins,nₛ) for k in measurements)

    # container to contain statistics
    stats = Dict(string(k)=>zeros(T1,2,nₛ) for k in measurements)

    # iterate over data files
    for (i,file) in enumerate(files)
        bin = div(i-1,N) + 1
        open(file,"r") do fin
            header = readline(fin)
            for line in eachline(fin)
                atoms = split(line)
                meas  = atoms[1]
                orbit = parse(Int,atoms[2])
                val   = parse(T1,atoms[3])
                binned_data[meas][bin,orbit] += val / N
            end
        end
    end

    # iterating over measurements
    for measurement in keys(stats)

        # deal with chemical potential seperately due to μ-tuning algorithm
        if measurement=="mu"
            μ  = reshaped(model.μ,(nₛ,div(model.Nsites,nₛ)))
            μ″ = mean(μ)
            μ′ = μ_tuner.μ_avg
            Δμ = μ_tuner.μ_err
            dμ = μ′-μ″
            # calculating average and error for on-site energy of each orbital
            for orbit in 1:nₛ
                μ₀ = @view μ[orbit,:]
                stats[measurement][1,orbit] = mean(μ₀) + dμ
                stats[measurement][2,orbit] = Δμ
            end
        else
            # iterate over orbits
            for orbit in 1:nₛ

                # calculating average and error of measurement
                data     = @view binned_data[measurement][:,orbit]
                avg, err = mean_and_error(data)
                stats[measurement][1,orbit] = avg
                stats[measurement][2,orbit] = err
            end
        end
    end

    # write averaged statistics to file
    open(statsfn,"w") do statsout

        # construct header
        header = "measurement orbit avg error\n"
        write(fout,    header)
        write(statsout,header)

        # iterate over measurements
        for measurement in keys(stats)

            # iterate over orbits
            for orbit in 1:nₛ

                # write out measurement
                avg  = stats[measurement][1,orbit]
                err  = stats[measurement][2,orbit]
                @printf fout     "%s %d %.6f %.6f\n" measurement orbit avg err
                @printf statsout "%s %d %.6f %.6f\n" measurement orbit avg err
            end
        end
    end

    return nothing
end

"""
Write on-site measreuments to file.
"""
function write_intersite_measurements!(fout,model::AbstractModel{T1,T2,T3},datafolder::String,Nbins::Int) where {T1,T2,T3}

    # number
    nᵥ = model.nbonds

    # data folder
    folder = joinpath(datafolder, "intersite_measurements_f")

    # get files containing data
    files = readdir(folder,join=true)

    # number of data files
    Nmeas = length(files)

    # file statistics will be written to
    statsfn = joinpath(datafolder, "intersite_measurements_stats.out")

    # calculate number of measurements per bin
    Nbins = min(Nmeas,Nbins)
    N     = div(Nmeas,Nbins)
    @assert Nmeas%Nbins==0

    # get measurements
    measurements = String[]
    open(files[1],"r") do fin
        header = readline(fin)
        for line in readlines(fin)
            push!( measurements , split(line)[1] )
        end
    end
    measurements = unique(measurements)

    # container to hold binned data
    binned_data = Dict(string(k)=>zeros(T1,Nbins,nᵥ) for k in measurements)

    # container to contain statistics
    stats = Dict(string(k)=>zeros(T1,2,nᵥ) for k in measurements)

    # iterate over data files
    for (i,file) in enumerate(files)
        bin = div(i-1,N) + 1
        open(file,"r") do fin
            header = readline(fin)
            for line in eachline(fin)
                atoms = split(line)
                meas  = atoms[1]
                bond  = parse(Int,atoms[2])
                val   = parse(T1,atoms[3])
                binned_data[meas][bin,bond] += val / N
            end
        end
    end

    # iterating over measurements
    for measurement in keys(stats)

        # iterate over orbits
        for vector in 1:nᵥ

            # calculating average and error of measurement
            data     = @view binned_data[measurement][:,vector]
            avg, err = mean_and_error(data)
            stats[measurement][1,vector] = avg
            stats[measurement][2,vector] = err
        end
    end

    # write averaged statistics to file
    open(statsfn,"w") do statsout

        # construct header
        header = "measurement bond avg error\n"
        write(fout,    header)
        write(statsout,header)

        # iterate over measurements
        for measurement in keys(stats)

            # iterate over orbits
            for b in 1:nᵥ

                # write out measurement
                avg  = stats[measurement][1,b]
                err  = stats[measurement][2,b]
                @printf fout     "%s %d %.6f %.6f\n" measurement b avg err
                @printf statsout "%s %d %.6f %.6f\n" measurement b avg err
            end
        end
    end

    return nothing
end

"""
Write correlations to file.
"""
function write_correlations!(fout,model::AbstractModel,datafolder::String,Nbins::Int)

    # get contents of data folder
    contents = readdir(datafolder)
    # iterate of names in data folder
    for name in contents
        # if name corresponds to folder
        if endswith(name,"_f")
            # if name of directory starts with capital letter
            if isuppercase(name[1])
                # split name by underscore
                atoms = split(name,"_")
                # if a correlation function
                if !endswith(atoms[1],"Susc")
                    # name of measurement
                    measurement = string(atoms[1])
                    # position or momentum space correlation function
                    space = string(atoms[2])
                    # process data
                    write_correlation!(fout,model,measurement,space,Nbins,datafolder)
                end
            end
        end
    end

    return nothing
end

"""
Write a correlation to file.
"""
function write_correlation!(fout,model::AbstractModel{T},measurement::String,space::String,Nbins::Int,datafolder::String) where {T<:AbstractFloat}

    # data folder
    folder = joinpath(datafolder,"$(measurement)_$(space)_f")
    files  = readdir(folder,join=true)

    # read info in from key file
    Lₜ      = 1
    pairs   = Vector{Tuple{Int,Int}}()
    header  = ""
    keyfile = joinpath(folder,"$(measurement)_$(space)_key.out")
    open(keyfile,"r") do fin
        # construct header for stats file, determining whether to label first two colmns
        # "orbit1 orbit2" or "bond1 bond2" depending on contents of key file
        keyheader = readline(fin)
        atoms     = split(keyheader)
        if space == "position"
            header = "$(atoms[2]) $(atoms[3]) r3 r2 r1 tau $(measurement) error\n"
        else
            header = "$(atoms[2]) $(atoms[3]) k3 k2 k1 tau $(measurement) error\n"
        end
        # get unique pairs of orbitals/bonds the correlation function was measured for
        # and determine if measurement was time dependent
        for line in eachline(fin)
            atoms = split(line)
            p     = ( parse(Int,atoms[2]) , parse(Int,atoms[3]) )
            if !(p in pairs)
                push!(pairs,p)
            end
            τ  = parse(Int,atoms[7])
            Lₜ = max(τ+1,Lₜ)
        end
    end

    # get number of pairs of orbitals/bonds
    nₚ = length(pairs)

    # get lattice size
    L₁ = model.lattice.L1::Int
    L₂ = model.lattice.L2::Int
    L₃ = model.lattice.L3::Int
    V  = L₁*L₂*L₃*nₚ*Lₜ

    # number of measurements
    Nmeas = length(files)-1

    # calculate number of measurements per bin
    Nbins = min(Nmeas,Nbins)
    N     = div(Nmeas,Nbins)
    @assert Nmeas % Nbins == 0

    # file statistics will be written to
    statsfn = joinpath(datafolder, "$(measurement)_$(space)_stats.out")

    # declare container for binned data
    container  = zeros(T,Nbins,Lₜ,L₁,L₂,L₃,nₚ)
    container′ = reshaped(container,Nbins,V)

    # iterate over number of files
    for m in 1:Nmeas
        # get bin number
        bin = div(m-1,N) + 1
        # iterate of rows in current data file
        for row in CSV.Rows(files[m], header=1, types=[Int,T], delim=" ", reusebuffer=true)
            # calculate binned averages
            container′[bin,row[1]] += row[2] / N
        end
    end

    # write stats file
    open(statsfn,"w") do statsout
        # write header
        write(fout,     header)
        write(statsout, header)
        # counter to indicate what chunk of data file to read
        n = 0
        # iterate of all position or momentum space variable
        for p in 1:nₚ
            n₁ = pairs[p][1]
            n₂ = pairs[p][2]
            for l₃ in 0:L₃-1
                for l₂ in 0:L₂-1
                    for l₁ in 0:L₁-1
                        # write averaged stats to file
                        for τ in 0:Lₜ-1
                            data     = @view container[:,τ+1,l₁+1,l₂+1,l₃+1,p]
                            avg, err = mean_and_error(data)
                            @printf fout     "%d %d %d %d %d %d %.8f %.8f\n" n₁ n₂ l₃ l₂ l₁ τ avg err
                            @printf statsout "%d %d %d %d %d %d %.8f %.8f\n" n₁ n₂ l₃ l₂ l₁ τ avg err
                        end
                        # increment to next chunk of file
                        n += 1
                    end
                end
            end
        end
    end

    write(fout,"\n")

    return nothing
end

"""
Write susceptibilities to file.
"""
function write_susceptibilities!(fout,model::AbstractModel, datafolder::String,Nbins::Int)

    # get contents of data folder
    contents = readdir(datafolder)
    # iterate of names in data folder
    for name in contents
        # if name corresponds to folder
        if endswith(name,"_f")
            # if name of directory starts with capital letter
            if isuppercase(name[1])
                # split name by underscore
                atoms = split(name,"_")
                # if a correlation function
                if endswith(atoms[1],"Susc")
                    # name of measurement
                    measurement = string(atoms[1])
                    # position or momentum space correlation function
                    space = string(atoms[2])
                    # process data
                    write_susceptibility!(fout,model,measurement,space,Nbins,datafolder)
                end
            end
        end
    end

    return nothing
end

"""
Write a susceptibility to file.
"""
function write_susceptibility!(fout,model::AbstractModel{T},measurement::String,space::String,Nbins::Int,datafolder::String) where {T<:AbstractFloat}

    # data folder
    folder = joinpath(datafolder,"$(measurement)_$(space)_f")
    files  = readdir(folder,join=true)

    # read info in from key file]
    pairs   = Vector{Tuple{Int,Int}}()
    header  = ""
    keyfile = joinpath(folder,"$(measurement)_$(space)_key.out")
    open(keyfile,"r") do fin
        # construct header for stats file, determining whether to label first two colmns
        # "orbit1 orbit2" or "bond1 bond2" depending on contents of key file
        keyheader = readline(fin)
        atoms     = split(keyheader)
        if space == "position"
            header = "$(atoms[2]) $(atoms[3]) r3 r2 r1 $(measurement) error\n"
        else
            header = "$(atoms[2]) $(atoms[3]) k3 k2 k1 $(measurement) error\n"
        end
        # get unique pairs of orbitals/bonds the correlation function was measured for
        for line in eachline(fin)
            atoms = split(line)
            p     = ( parse(Int,atoms[2]) , parse(Int,atoms[3]) )
            if !(p in pairs)
                push!(pairs,p)
            end
        end
    end

    # get number of pairs of orbitals/bonds
    nₚ = length(pairs)

    # get lattice size
    L₁ = model.lattice.L1::Int
    L₂ = model.lattice.L2::Int
    L₃ = model.lattice.L3::Int
    V  = L₁*L₂*L₃*nₚ

    # number of measurements
    Nmeas = length(files)-1

    # calculate number of measurements per bin
    Nbins = min(Nmeas,Nbins)
    N     = div(Nmeas,Nbins)
    @assert Nmeas % Nbins == 0

    # file statistics will be written to
    statsfn = joinpath(datafolder, "$(measurement)_$(space)_stats.out")

    # container to hold binned data
    container  = zeros(T,Nbins,L₁,L₂,L₃,nₚ)
    container′ = reshaped(container,Nbins,V)

    # iterate over files
    for m in 1:Nmeas
        # calculate bin
        bin = div(m-1,N) + 1
        # iterate over rows of csv file
        for row in CSV.Rows(files[m], types=[Int,T], reusebuffer=true, delim=' ')
            index = row[1]
            val   = row[2]
            container′[bin,index] += val / N
        end
    end

    # write avg and error of measurements to file
    open(statsfn,"w") do statsout
        write(statsout,header)
        write(fout,    header)
        for p in 1:nₚ
            n₁ = pairs[p][1]
            n₂ = pairs[p][2]
            for l₃ in 0:L₃-1
                for l₂ in 0:L₂-1
                    for l₁ in 0:L₁-1
                        vals     = @view container[:,l₁+1,l₂+1,l₃+1,p]
                        avg, err = mean_and_error(vals)
                        @printf fout     "%d %d %d %d %d %.8f %.8f\n" n₁ n₂ l₃ l₂ l₁ avg err
                        @printf statsout "%d %d %d %d %d %.8f %.8f\n" n₁ n₂ l₃ l₂ l₁ avg err
                    end
                end
            end
        end
    end

    write(fout,"\n")

    return nothing
end

"""
Calculate mean and standard deviation of the mean of a set of data.
"""
function mean_and_error(v::AbstractArray{T})::Tuple{T,T} where {T<:Number}

    N   = length(v)
    avg = mean(v)
    err = std(v,corrected=true,mean=avg)/sqrt(N)
    return avg, err
end

end


@testset "Cross validating ooHMM on CTB..." begin
        @time ctb = ChTreebank()
        #@time ctb = ZhuXian()

        nb = 10
        batches = KongYiji.foldbatch(ctb, nb)
        tbs = KongYiji.HmmScoreTable[]
        @showprogress 1 "Cross Validating ooHMM..." for (tr, te) in batches
                base = KongYiji.HMM(tr)
                hmm = KongYiji.HMM(;base=base)
                lm = LM(tr; rnnType=:lstm, numLayers=1, bidirectional=false)
                oohmm = hmm + lm
                x = KongYiji.rawsents(te)
                z = KongYiji.wordsents(te)
                @time y = oohmm(x; recover=false, withpos=false)
                push!(tbs, KongYiji.HmmScoreTable(z, y))
                println(tbs[end])
        end
        hmmscoretable = sum(tbs)
        println(hmmscoretable)
end


using Knet
using ArgParse
using JLD,MAT,JSON,HDF5

include("util.jl")
include("model.jl")

function main(args)
    s = ArgParseSettings()
    s.description = "CNN-HMM hybrid ASR system in Knet by Ilker Kesen."

    @add_arg_table s begin
        ("--features"; required=true; help="MFSC features file")
        ("--jsonfile"; required=true)
        ("--model"; default=nothing)
        ("--savefile"; default=nothing)
        ("--atype"; default=(gpu()>=0?"KnetArray{Float32}":"Array{Float32}"))
        ("--window"; default=10; arg_type=Int64)
        ("--period"; default=20; arg_type=Int64)
    end

    isa(args, AbstractString) && (args=split(args))
    o = parse_args(args, s; as_symbols=true); println(o); flush(STDOUT)
    o[:atype] = eval(parse(o[:atype]))
    atype = o[:atype]
    # sr = o[:seed] > 0 ? srand(o[:seed]) : srand()

    # load data
    jsondata = JSON.parsefile(o[:jsonfile])
    p2i, i2p = build_vocab(jsondata)

    # make data
    trnids, tstids = [], []
    # trn = make_data(o[:features], jsondata; extra=trnids)
    tst = make_data(o[:features], jsondata, "test"; extra=tstids)

    # load model
    w = load(o[:model], "w")

    # convert atype
    w = map(wi->convert(atype,wi), w)

    solutions = Dict()

    data = tst
    getid(i) = tstids[i]
    for (i,sample) in enumerate(data)
        x = make_input(sample)
        # ygold = make_output(sample, p2i)
        ygold = map(si->si[2], sample)
        ygold = make_ygold(ygold)
        ypred = predict(w,convert(atype,x))
        ysoft = softmax(ypred)

        generation = generate(ysoft, o[:window])
        generation = map(t->i2p[t], generation)
        solutions[getid(i)] = Dict("ygold"=>ygold, "ypred"=>generation)

        # println("$i generation so far.")
        if i % o[:period] == 0
            println("$i generation so far ", now())
        end
    end

    ff = open(abspath(o[:savefile]), "w")
    write(ff, JSON.json(solutions))
    close(ff)
    println("done")
end

!isinteractive() && !isdefined(Core.Main, :load_only) && main(ARGS)


export WindingGeometry, Windings, windings, winding_geometry

"""
    WindingGeometry

2D representation of one layer of a winding.

**Fields**
- `width`   -- width of trace (m)
- `length`  -- length of trace (m)
- `turns`   -- number of turns

`WindingGeometry` objects are typicaly created with `winding_geometry`
"""
struct WindingGeometry
  width :: Float64
  length :: Float64
  turns :: Int
end
"""
    winding_geometry(pcb, core, turns)

Create a `WindingGeometry` object.
"""
function winding_geometry(pcb :: PCB_Specification,
                          core :: CoreGeometry,
                          turns :: Int)
  trace_width = (core.winding_aperture-2*pcb.trace_edge_gap-
    (turns-1)*pcb.trace_trace_gap)/turns
  trace_length = 0.0
  for i in 0:turns-1
    r = core.half_center_width+pcb.trace_edge_gap+
        0.5*trace_width+i*(trace_width+pcb.trace_trace_gap)
    trace_length += 2π*r
  end
  trace_length += 2*core.center_length
  if trace_width<=0.0
    warn("$(turns) on $(name(core)) will result in negative trace width.  Setting to NaN.")
    trace_width = NaN
  end
  WindingGeometry(trace_width, trace_length, turns)
end

"""
    Windings

**Fields**
- `pcb`                     -- `PCB_Specification`
- `core`                    -- `CoreGeometry`
- `primarywindinggeometry`  -- `WindingGeometry` for all layers of primary
- `secondarywindinggeometry`-- `WindingGeometry` for all layers of secondary
- `isprimary`               -- tuple of `Bool` for each conductor layer
- `sides`                   -- for internal use
- `isprimaryseries`         -- true if primary is connected in series
- `issecondaryseries`       -- true if secondary is connected in series

Create `Windings` objects with `windings`.
"""
struct Windings
  pcb :: PCB_Specification
  core :: CoreGeometry
  primarywindinggeometry :: WindingGeometry
  secondarywindinggeometry :: WindingGeometry
  isprimary :: Tuple
  sides :: Tuple
  isprimaryseries :: Bool
  issecondaryseries :: Bool
end
"""
    windings(pcb, core,
            primaryturnsperlayer, secondaryturnsperlayer,
            isprimary,
            isprimaryeries, issecondaryseries)

Create a `Windings` object.
"""
function windings(pcb::PCB_Specification, core::CoreGeometry,
                  primaryturnsperlayer::Int, secondaryturnsperlayer::Int,
                  isprimary,
                  isprimaryeries::Bool, issecondaryseries::Bool)
  if length(isprimary)!=div(length(pcb.stackup.material),2)+1
    throw(ArgumentError("length of isprimary must be the same as number of conductor layers"))
  end
  x = 0
  for ip in isprimary
    x += ip ? 1 : -1
    if abs(x)>1
      throw(ArgumentError("cannot analyze due to layer connections to primay and secondary"))
    end
  end
  Windings(pcb,core,
           winding_geometry(pcb,core,primaryturnsperlayer),
           winding_geometry(pcb,core,secondaryturnsperlayer),
           isprimary,
           ntuple(i->sides(isprimary,i),length(isprimary)),
           isprimaryeries, issecondaryseries
          )
end
"""
    sides(isprimary,index)

For a given layer (index), returns the number (1 or 2) of sides which are part
of the same winding.

This function is for internal use calculating effective_resistance and
 leakage_inductance.
"""
function sides(isprimary,i)::Float64
  i==1 && return 1.0
  i==length(isprimary) && return 1.0
  s = 0.0
  isprimary[i] != isprimary[i-1] && (s+=1.0)
  isprimary[i] != isprimary[i+1] && (s+=1.0)
  return s
end

effective_volume(w::Windings) = effective_volume(core(w))
effective_area(w::Windings) = effective_area(core(w))
effective_length(w::Windings) = effective_length(core(w))
winding_aperture_height(w::Windings) = winding_aperture_height(core(w))
winding_aperture(w::Windings) = winding_aperture(core(w))
core_name(w::Windings) = name(core(w))


#=== rec ===#
function wIB(
    fid::IOStream,
    rec::receiver,
    nT::Int64,
    OptCpnt::Array{String})

    CpntDict = Dict("P"=>1,"Vx"=>2,"Vz"=>3,
                "Txx"=>4,"Tzz"=>5,"Txz"=>6,"V"=>23,"Tii"=>45)

    write(fid,rec.nr) # Total Receiver Number
    write(fid,length(OptCpnt))
    for cpnt in OptCpnt
        write(fid,CpntDict[cpnt])
    end
    write(fid,nT) # Total Time Sampling Tn
    write(fid,rec.loc) #physical location
    write(fid,rec.nloc) #discrete location without PML
    write(fid,rec.BDnloc) #discrete location with PML
end

#=== elastic ===#

#=== wavefield ===#
function wIB(
    fid::IOStream,
    model::elmod2d,
    Slices::Union{UnitRange{Int64},StepRange{Int64,Int64},Array{Int64,1},Int64},
    OptCpnt::Array{String})

    CpntDict = Dict("P"=>1,"Vx"=>2,"Vz"=>3,
                "Txx"=>4,"Tzz"=>5,"Txz"=>6,
                "V"=>23,"Tii"=>45)

    write(fid,model.medium.iflag)
    write(fid,length(OptCpnt))
    for cpnt in OptCpnt
        write(fid,CpntDict[cpnt])
    end
    write(fid,model.nwf.BDntpp)
    write(fid,model.medium.ext)
    write(fid,model.medium.dx)
    write(fid,model.medium.dz)
    write(fid,length(Slices))
    write(fid,Slices.*model.medium.dt)

end

#=== acoustic ===#
#=== wavefield ===#
function wIB(
    fid::IOStream,
    model::acmod2d,
    Slices::Union{UnitRange{Int64},StepRange{Int64,Int64},Array{Int64,1},Int64},
    OptCpnt::Array{String})

    CpntDict = Dict("P"=>1,"Vx"=>2,"Vz"=>3,
                "Txx"=>4,"Tzz"=>5,
                "V"=>23,"Tii"=>45)

    write(fid,model.medium.iflag)
    write(fid,length(OptCpnt))
    for cpnt in OptCpnt
        write(fid,CpntDict[cpnt])
    end
    write(fid,model.nwf.BDntpp)
    write(fid,model.medium.ext)
    write(fid,model.medium.dx)
    write(fid,model.medium.dz)
    write(fid,length(Slices))
    write(fid,Slices.*model.medium.dt)

end


#
println("Ag) Test of the QED model corrections to the level structure of atoms and ions.")
#
wa = Atomic.Computation("QED estimates for carbon-like Xe",  Nuclear.Model(54.);
                        configs=[Configuration("1s^2 2s^2 2p^6")],
                        asfSettings=AsfSettings(true, false, "meanDFS", "hydrogenic", "xxx", [1],    40, 1.0e-6, JAC.Subshell[],
                                                true, false, NoneQed(), "yyy", LSjjSettings(true),
                                                false, [1,2,3,4], false, JAC.LevelSymmetry[] )  )

wb = perform(wa)


# Measures without Canvas Units
# -----------------------------

# Rather than introducing a new type to represent measures that
# have no canvas unit dimension, we just parameterize Measure
# over a special type that promotes appropriately.
immutable MeasureNil
end

const measure_nil = MeasureNil()

# All measures in Compose are specified as the sum
immutable Measure{S, T}
    abs::Float64 # absolute measurement in millimeters
    cx::S        # canvas x-units
    cy::T        # canvas y-units
    cw::Float64  # proportion of canvas width
    ch::Float64  # proportion of canvas height

    function Measure(abs::Number, cx::S, cy::T, cw::Number, ch::Number)
        new(float64(abs), cx, cy, float64(cw), float64(ch))
    end

    function Measure(; abs=0.0, cx::S=zero(S), cy::T=zero(T), cw=0.0, ch=0.0)
        new(float64(abs), cx, cy, float64(cw), float64(ch))
    end
end

typealias MeasureOrNumber Union(Measure, Number)

# MeasureNil semantics
zero(::Type{MeasureNil}) = measure_nil
zero(::MeasureNil) = measure_nil

iszero{T}(a::T) = a == zero(T)

# ambiguity warning
max(::Function, b::MeasureNil) = measure_nil
min(::Function, b::MeasureNil) = measure_nil

max(a::MeasureNil, b::MeasureNil) = measure_nil
max(a::MeasureNil, b)             = b
max(a, b::MeasureNil)             = a

min(a::MeasureNil, b::MeasureNil) = measure_nil
min(a::MeasureNil, b)             = measure_nil
min(a, b::MeasureNil)             = measure_nil

isless(a::MeasureNil, b::MeasureNil) = false

function -{S,T}(a::Measure{S, T})
    return Measure{S, T}(-a.abs, -a.cx, -a.cy, -a.cw, -a.ch)
end

function zero(::Type{Measure})
    return Measure()
end

function zero(::Measure)
    return Measure()
end

function Measure(abs, cx=measure_nil, cy=measure_nil, cw=0.0, ch=0.0)
    Measure{typeof(cx), typeof(cy)}(abs, cx, cy, cw, ch)
end

function Measure(; abs=0.0, cx=measure_nil, cy=measure_nil, cw=0.0, ch=0.0)
    Measure{typeof(cx), typeof(cy)}(abs, cx, cy, cw, ch)
end

const zero_measure = Measure()

# copy with substitutions
function Measure{S, T}(u::Measure{S, T};
                       abs=nothing, cx=nothing, cy=nothing,
                       cw=nothing, ch=nothing)
    return Measure(abs === nothing ? u.abs : abs,
                   cx  === nothing ? u.cx  : cx,
                   cy  === nothing ? u.cy  : cy,
                   cw  === nothing ? u.cw  : cw,
                   ch  === nothing ? u.ch  : ch)
end

function zero{S, T}(::Type{Measure{S, T}})
    Measure{S, T}()
end

function isless{T, S}(a::Measure{T, S}, b::Measure{T, S})
    return (a.abs < b.abs ||
            a.cx  < b.cx ||
            a.cy  < b.cy ||
            a.cw  < b.cw ||
            a.ch  < b.ch) &&
           a.abs <= b.abs &&
           a.cx  <= b.cx &&
           a.cy  <= b.cy &&
           a.cw  <= b.cw &&
           a.ch  <= b.ch
end

function isabsolute{S, T}(u::Measure{S, T})
    iszero(u.cx) && iszero(u.cy) && u.cw == 0.0 && u.ch == 0.0
end

function copy{S, T}(a::Measure{S, T})
    Measure{S, T}(abs=a.abs, cx=a.cx, cy=a.cy, cw=a.cw, ch=a.ch)
end

function show{S, T}(io::IO, a::Measure{S, T})
    first = true
    xz = zero(S)
    yz = zero(T)

    if a.abs != 0.0
        print(io, a.abs, "mm")
        first = false
    end

    if a.cx != xz
        if !first
            print(io, " + ")
        end
        print(io, a.cx, "cx")
        first = false
    end

    if a.cy != yz
        if !first
            print(io, " + ")
        end
        print(io, a.cy, "cy")
        first = false
    end

    if a.cw != 0.0
        if !first
            print(io, " + ")
        end
        print(io, a.cw, "w")
        first = false
    end

    if a.ch != 0.0
        if !first
            print(io, " + ")
        end
        print(io, a.ch, "h")
        first = false
    end

    if first
        print(io, "Measure()")
    end
end

# Measure Arithmatic
# ------------------

# Avoid ambiguity errors
add_measure_part(a::MeasureNil, b::MeasureNil) = measure_nil
add_measure_part(a::MeasureNil, b)             = b
add_measure_part(a, b::MeasureNil)             = a
add_measure_part(a, b)                         = a + b

sub_measure_part(a::MeasureNil, b::MeasureNil) = measure_nil
sub_measure_part(a::MeasureNil, b)             = b
sub_measure_part(a, b::MeasureNil)             = a
sub_measure_part(a, b)                         = a - b

*(a::MeasureNil, b::MeasureNil) = error("Two measure_nil objects multiplied")
*(a::MeasureNil, b)             = measure_nil
*(a, b::MeasureNil)             = measure_nil

/(a::MeasureNil, b::MeasureNil) = error("Division by a measure_nil")
/(a, b::MeasureNil)             = error("Division by a measure_nil")
/(a::MeasureNil, b)             = measure_nil

function +(a::Measure, b::Measure)
    Measure(a.abs + b.abs,
            add_measure_part(a.cx, b.cx),
            add_measure_part(a.cy, b.cy),
            a.cw  + b.cw,
            a.ch  + b.ch)
end

function -(a::Measure, b::Measure)
    Measure(a.abs - b.abs,
            sub_measure_part(a.cx, b.cx),
            sub_measure_part(a.cy, b.cy),
            a.cw  - b.cw,
            a.ch  - b.ch)
end

function *(a::Measure, b::Measure)
    error("A Measure can not be multiplied by another Measure.")
end

function /{T, S}(a::Measure{T}, b::Measure{S})
    factor = nothing

    for unit in (:abs, :cx, :cy, :cw, :ch)
        vala = getfield(a, unit)
        valb = getfield(b, unit)

        if valb == zero(typeof(valb))
            continue
        end

        if factor === nothing
            factor = vala / valb
        else
            if vala / valb != factor
                error("Measure {a} is not divisible by measure {b}")
            end
        end
    end

    if factor === nothing
        error("An empty measure can not be divided by another empty measure.")
    else
        factor
    end
end

# Return a measure that is at least as large as anything else
function max(measures::Measure...)
    # current maximums
    abs = 0.0
    cx = measure_nil
    cy = measure_nil
    cw = 0.0
    ch = 0.0

    for measure in measures
        abs = max(abs, measure.abs)
        cx  = max(cx, measure.cx)
        cy  = max(cy, measure.cy)
        cw  = max(cw, measure.cw)
        ch  = max(ch, measure.ch)
    end

    Measure(abs, cx, cy, cw, ch)
end

# Return a measure that at least as small as everything else
function min(measures::Measure...)
    # current maximums
    abs = Inf
    cx = Inf
    cy = Inf
    cw = Inf
    ch = Inf

    for measure in measures
        abs = min(abs, measure.abs)
        cx  = min(cx, measure.cx)
        cy  = min(cy, measure.cy)
        cw  = min(cw, measure.cw)
        ch  = min(ch, measure.ch)
    end

    Measure(abs, cx, cy, cw, ch)
end

# Versus plain numbers

function *(a::Measure, b::Number)
    Measure(a.abs*b, a.cx*b, a.cy*b, a.cw*b, a.ch*b)
end

function *(a::Number, b::Measure)
    b * a
end

function /{T}(a::Measure{T}, b::Number)
    Measure(a.abs/b, a.cx/b, a.cy/b, a.cw/b, a.ch/b)
end

# Measure Constants
# -----------------

# Measure instances are usually constructed by multiplying by one of these
# constants. E.g. "10mm"

const mm   = Measure(abs=1.0)
const cm   = Measure(abs=10.0)
const inch = Measure(abs=25.4)
const pt   = inch/72.0
const w    = Measure(cw=1.0)
const h    = Measure(ch=1.0)
const cx   = Measure(cx=1.0)
const cy   = Measure(cy=1.0)

# Pixels are not typically used in Compose in preference of absolute
# measurements or measurements relative to parent canvases. So for the
# 'px' constant, we just punt and give something do something vagually
# reasonable.

const assumed_ppmm = 3.78 # equivalent to 96 DPI
const px = Measure(abs=1.0/assumed_ppmm)

# Interpretation of bare numbers
# ------------------------------

function x_measure(a::Measure)
    a
end

function x_measure(a)
    Measure(0.0, a)
end

function y_measure(a::Measure)
    a
end

function y_measure(a)
    Measure(0.0, measure_nil, a)
end

function size_measure(a::Measure)
    a
end

function size_measure(a)
    Measure(a)
end

# Higher-order measures
# ---------------------

immutable Point
    x::Measure
    y::Measure

    function Point()
        new(Measure(), Measure())
    end

    function Point(x, y)
        new(x_measure(x), y_measure(y))
    end

    function Point(point::Point)
        new(point.x, point.y)
    end

    function Point(xy::NTuple{2})
        Point(x_measure(xy[1]), y_measure(xy[2]))
    end
end

function copy(point::Point)
    Point(point)
end

function isabsolute(point::Point)
    isabsolute(point.x) && isabsolute(point.y)
end

typealias XYTupleOrPoint Union(NTuple{2}, Point)

function convert(::Type{Point}, xy::NTuple{2})
    Point(xy[1], xy[2])
end

# Bounding Boxes
# --------------

# Measure based bounding box used to specify the extent and positioning
# of canvases.
immutable BoundingBox
    x0::Measure
    y0::Measure
    width::Measure
    height::Measure

    function BoundingBox(x0, y0, width, height)
        new(x_measure(x0), y_measure(y0),
            x_measure(width), y_measure(height))
    end

    function BoundingBox(box::BoundingBox)
        new(box.x0, box.y0, box.width, box.height)
    end

    function BoundingBox(box::BoundingBox;
                         x0=nothing,
                         y0=nothing,
                         width=nothing,
                         height=nothing)
        BoundingBox(x0     === nothing ? box.x0     : x0,
                    y0     === nothing ? box.y0     : y0,
                    width  === nothing ? box.width  : width,
                    height === nothing ? box.height : height)
    end

    function BoundingBox()
        new(0cx, 0cy, 1cx, 1cy)
    end
end

function copy(box::BoundingBox)
    BoundingBox(box)
end

# Compute the union of two bounding boxes.
#
# In other words, given two bounding boxes, return a new bounding box that
# contains both.
#
# Unfortunately this is in general uncomputable without knowing the absolute
# size of the parent canvas which may be passed in via the last two parameters.
# If not passed, this throws an error if they would have been required.
#
function union(a::BoundingBox, b::BoundingBox, units=nothing, parent_abs_width=nothing, parent_abs_height=nothing)
    (a.width == Measure() || a.height == Measure()) && return b
    (b.width == Measure() || b.height == Measure()) && return a
    x0 = min(a.x0, b.x0)
    y0 = min(a.y0, b.y0)
    x1 = max(a.x0 + a.width, b.x0 + b.width)
    y1 = max(a.y0 + a.height, b.y0 + b.height)
    # Check whether we had any problematic computations
    for m in (x0,y0,x1,y1)
        # Pure absolute or pure relative points are fine. When they are mixed,
        # there are problems
        if !isabsolute(m) && m.abs != 0.0
            if units == nothing || parent_abs_width == nothing || parent_abs_height == nothing
                error("""Bounding boxes are uncomputable without knowledge of the
                         absolute dimensions of the top canvase due to mixing of relative
                         and absolute coordinates. Either pass the dimension as a parameter
                         or restrict the context to one kind of coordinates.""")
            end
            parent_box = AbsoluteBoundingBox(0.0,0.0,parent_abs_width,parent_abs_height)
            abb = union(absolute_units(a,IdentityTransform(),units,parent_box),
                        absolute_units(b,IdentityTransform(),units,parent_box))
            return BoundingBox(Measure(;abs = abb.x0), Measure(;abs = abb.x0),
                               Measure(;abs = abb.width), Measure(;abs = abb.height))
        end
    end
    return BoundingBox(x0, y0, x1 - x0, y1 - y0)
end

# Absolute measure bounding box. This is used to keep track of
# coordinates when traversing canvas trees to perform drawing.
immutable AbsoluteBoundingBox
    x0::Float64
    y0::Float64
    width::Float64
    height::Float64

    function AbsoluteBoundingBox(x0::Number, y0::Number, width::Number, height::Number)
        return new(x0, y0, width, height)
    end

    function AbsoluteBoundingBox()
        return new(0.0, 0.0, 1.0, 1.0)
    end
end

function union(a::AbsoluteBoundingBox, b::AbsoluteBoundingBox)
    (a.width == 0.0 || a.height == 0.0) && return b
    (b.width == 0.0 || b.height == 0.0) && return a
    x0 = min(a.x0, b.x0)
    y0 = min(a.y0, b.y0)
    x1 = max(a.x0 + a.width, b.x0 + b.width)
    y1 = max(a.y0 + a.height, b.y0 + b.height)
    return AbsoluteBoundingBox(x0, y0, x1 - x0, y1 - y0)
end

# The same type-signature is used for a box used to assign
# a custom coordinate system to a canvas.
# TODO: There should not be one fixed type for every parameter
immutable UnitBox{S, T, U, V}
    x0::S
    y0::T
    width::U
    height::V

    leftpad::Measure
    rightpad::Measure
    toppad::Measure
    bottompad::Measure

    function UnitBox(x0::S, y0::T, width::U, height::V;
                     leftpad::Measure=Measure(),
                     rightpad::Measure=Measure(),
                     toppad::Measure=Measure(),
                     bottompad::Measure=Measure())
        return new(x0, y0, width, height, leftpad, rightpad, toppad, bottompad)
    end
end

function UnitBox{S,T}(width::S, height::T;
                      leftpad::Measure=Measure(),
                      rightpad::Measure=Measure(),
                      toppad::Measure=Measure(),
                      bottompad::Measure=Measure())
    x0 = zero(S)
    y0 = zero(T)

    return UnitBox{S, T, S, T}(x0, y0, width, height,
                               leftpad=leftpad, rightpad=rightpad,
                               toppad=toppad, bottompad=bottompad)
end

function UnitBox(x0, y0, width, height;
                 leftpad::Measure=Measure(),
                 rightpad::Measure=Measure(),
                 toppad::Measure=Measure(),
                 bottompad::Measure=Measure())
    x0, width  = promote(x0, width)
    y0, height = promote(y0, height)
    return UnitBox{typeof(x0), typeof(y0), typeof(width), typeof(height)}(
                   x0, y0, width, height,
                   leftpad=leftpad, rightpad=rightpad,
                   toppad=toppad, bottompad=bottompad)
end

function UnitBox()
    return UnitBox{Float64, Float64, Float64, Float64}(0.0, 0.0, 1.0, 1.0)
end

# copy with substitution
function UnitBox(units::UnitBox;
                 x0=nothing, y0=nothing, width=nothing, height=nothing,
                 leftpad=nothing, rightpad=nothing, toppad=nothing,
                 bottompad=nothing)
    return UnitBox(x0 === nothing ? units.x0 : x0,
                   y0 === nothing ? units.y0 : y0,
                   width === nothing ? units.width : width,
                   height === nothing ? units.height : height,
                   leftpad   = leftpad   === nothing ? units.leftpad : leftpad,
                   rightpad  = rightpad  === nothing ? units.rightpad : rightpad,
                   toppad    = toppad    === nothing ? units.toppad : toppad,
                   bottompad = bottompad === nothing ? units.bottompad : bottompad)
end

function NilUnitBox()
    return UnitBox{Nothing, Nothing, Nothing, Nothing}(
                nothing, nothing, nothing, nothing)
end

const nil_unit_box = NilUnitBox()

function isabsolute(units::UnitBox)
    return units.leftpad == zero_measure && units.rightpad == zero_measure &&
           units.toppad == zero_measure && units.bottompad == zero_measure
end

# Canvas Transforms
# -----------------

# Transform matrix in absolute coordinates

abstract Transform

immutable IdentityTransform <: Transform
end

immutable MatrixTransform <: Transform
    M::Matrix{Float64}

    function Transform()
        new([1.0 0.0 0.0
             0.0 1.0 0.0
             0.0 0.0 1.0])
    end

    function MatrixTransform(M::Matrix{Float64})
        new(M)
    end
end

function combine(a::IdentityTransform, b::IdentityTransform)
    return a
end

function combine(a::IdentityTransform, b::MatrixTransform)
    return b
end

function combine(a::MatrixTransform, b::IdentityTransform)
    return a
end

function combine(a::MatrixTransform, b::MatrixTransform)
    MatrixTransform(a.M * b.M)
end

# Rotation about a point.
type Rotation
    theta::Float64
    offset::Point

    function Rotation()
        new(0.0, Point(0.5w, 0.5h))
    end

    function Rotation(theta::Number)
        Rotation(theta, 0.5w, 0.5h)
    end

    function Rotation(theta::Number, offset::XYTupleOrPoint)
        new(convert(Float64, theta), convert(Point, offset))
    end

    function Rotation(theta::Number, offset_x, offset_y)
        new(convert(Float64, theta), Point(offset_x, offset_y))
    end

    # copy constructor
    function Rotation(rot::Rotation)
        new(copy(rot.theta),
            copy(rot.offset))
    end
end

copy(rot::Rotation) = Rotation(rot)

function convert(::Type{Transform}, rot::Rotation)
    if rot.theta == 0.0
        return IdentityTransform()
    else
        ct = cos(rot.theta)
        st = sin(rot.theta)
        x0 = rot.offset.x - (ct * rot.offset.x - st * rot.offset.y)
        y0 = rot.offset.y - (st * rot.offset.x + ct * rot.offset.y)
        return MatrixTransform([ct  -st  x0.abs
                                st   ct  y0.abs
                                0.0 0.0  1.0])
    end
end

# Conversion to absolute units
# ----------------------------

# Covert a Measure to a Float64
function absolute_units(u::Measure,
                        t::Transform,
                        unit_box::UnitBox,
                        parent_box::AbsoluteBoundingBox)
    add_measure_part(u.abs,
      add_measure_part(
        add_measure_part((u.cx / unit_box.width) * parent_box.width,
                         (u.cy / unit_box.height) * parent_box.height),
        u.cw * parent_box.width +
        u.ch * parent_box.height))
end

function absolute_units(units::UnitBox, t::Transform, parent_units::UnitBox,
                        box::AbsoluteBoundingBox)
    if isabsolute(units)
        return units
    else
        leftpad   = absolute_units(units.leftpad, t, parent_units, box)
        rightpad  = absolute_units(units.rightpad, t, parent_units, box)
        toppad    = absolute_units(units.toppad, t, parent_units, box)
        bottompad = absolute_units(units.bottompad, t, parent_units, box)

        # just give up trying to pad the units if it's impossible
        if leftpad + rightpad >= box.width ||
           toppad + bottompad >= box.height
            return UnitBox(units.x0, units.y0, units.width, units.height)
       end

        width = units.width * (box.width / (box.width - leftpad - rightpad))
        height = units.height * (box.height / (box.height - toppad - bottompad))
        x0 = units.x0 - width * (leftpad / box.width)
        y0 = units.y0 - height * (toppad / box.height)

        return UnitBox(x0, y0, width, height)
    end
end

# Convert a Rotation to a Transform
function absolute_units(rot::Rotation,
                        t::MatrixTransform,
                        unit_box::UnitBox,
                        parent_box::AbsoluteBoundingBox)

    absrot = Rotation(rot.theta,
                      absolute_units(rot.offset, t, unit_box, parent_box))

    rott = convert(Transform, absrot)
    t = combine(rott, t)
    theta = atan2(t.M[2,1], t.M[1,1])
    return Rotation(theta, absrot.offset)
end

function absolute_units(rot::Rotation,
                        t::IdentityTransform,
                        unit_box::UnitBox,
                        parent_box::AbsoluteBoundingBox)

    return Rotation(rot.theta, absolute_units(rot.offset, t, unit_box, parent_box))
end

function absolute_position_cx(::MeasureNil, unit_box::UnitBox,
                              parent_box::AbsoluteBoundingBox)
    0.0
end

function absolute_position_cx(cx, unit_box::UnitBox,
                              parent_box::AbsoluteBoundingBox)
    float64(((cx - unit_box.x0) / unit_box.width)) * parent_box.width
end

function absolute_position_cy(::MeasureNil, unit_box::UnitBox,
                              parent_box::AbsoluteBoundingBox)
    0.0
end

function absolute_position_cy(cy, unit_box::UnitBox,
                              parent_box::AbsoluteBoundingBox)
    float64(((cy - unit_box.y0) / unit_box.height)) * parent_box.height
end

function absolute_x_position(u::Measure,
                             t::Transform,
                             unit_box::UnitBox,
                             parent_box::AbsoluteBoundingBox)
    parent_box.x0 +
      u.abs +
      absolute_position_cx(u.cx, unit_box, parent_box) +
      absolute_position_cy(u.cy, unit_box, parent_box) +
      u.cw * parent_box.width +
      u.ch * parent_box.height
end

function absolute_y_position(u::Measure,
                             t::Transform,
                             unit_box::UnitBox,
                             parent_box::AbsoluteBoundingBox)
    parent_box.y0 +
      u.abs +
      absolute_position_cx(u.cx, unit_box, parent_box) +
      absolute_position_cy(u.cy, unit_box, parent_box) +
      u.cw * parent_box.width +
      u.ch * parent_box.height
end

# Convert a BoundingBox to a AbsoluteBoundingBox
function absolute_units(box::BoundingBox,
                        t::Transform,
                        unit_box::UnitBox,
                        parent_box::AbsoluteBoundingBox)
    AbsoluteBoundingBox(
        absolute_x_position(box.x0, t, unit_box, parent_box),
        absolute_y_position(box.y0, t, unit_box, parent_box),
        absolute_units(box.width, t, unit_box, parent_box),
        absolute_units(box.height, t, unit_box, parent_box))
end

# Convert a Point to a Point in absolute units
function absolute_units(point::Point,
                        t::MatrixTransform,
                        unit_box::UnitBox,
                        parent_box::AbsoluteBoundingBox)
    x = absolute_x_position(point.x, t, unit_box, parent_box)
    y = absolute_y_position(point.y, t, unit_box, parent_box)
    xyt = t.M * [x, y, 1.0]
    return Point(Measure(xyt[1]), Measure(xyt[2]))
end

function absolute_units(point::Point,
                        t::IdentityTransform,
                        unit_box::UnitBox,
                        parent_box::AbsoluteBoundingBox)
    x = absolute_x_position(point.x, t, unit_box, parent_box)
    y = absolute_y_position(point.y, t, unit_box, parent_box)
    return Point(Measure(x), Measure(y))
end


@recipe(Errorbars) do scene
    Theme(
        whiskerwidth = 10,
        color = :black,
        linewidth = 1,
        direction = :y,
        visible = theme(scene, :visible)
    )
end

function AbstractPlotting.plot!(plot::Errorbars{T}) where T <: Tuple{<:Any, <:Any, <:Any, <:Any}

    xs, ys, low, high = plot[1:4]
    xys = lift((x, y) -> Point2.(x, y), xs, ys)
    _plot_errorbars!(plot, xys, low, high)
end

function AbstractPlotting.plot!(plot::Errorbars{T}) where T <: Tuple{<:Any, <:Any, <:Any}
    xys, low, high = plot[1:3]
    _plot_errorbars!(plot, xys, low, high)
end

function AbstractPlotting.plot!(plot::Errorbars{T}) where T <: Tuple{<:Any, <:Any}
    xys, lowhigh = plot[1:2]
    _plot_errorbars!(plot, xys, lowhigh, lowhigh)
end

f_if(condition, f, arg) = condition ? f(arg) : arg

function _plot_errorbars!(plot, xys, low, high)

    @extract plot (whiskerwidth, color, linewidth, direction, visible)

    is_in_y_direction = lift(direction) do dir
        if dir == :y
            true
        elseif dir == :x
            false
        else
            error("Invalid direction $dir. Options are :x and :y.")
        end
    end

    linesegpairs = lift(xys, low, high, is_in_y_direction) do xys, low, high, is_in_y_direction

        [(xy .+ f_if(is_in_y_direction, reverse, Point2f0(-lo, 0)),
          xy .+ f_if(is_in_y_direction, reverse, Point2f0( hi, 0)),)
                for (xy, lo, hi) in zip(xys, low, high)]
    end

    scene = plot.parent

    whiskers = lift(linesegpairs, scene.camera.projectionview,
        scene.camera.pixel_space, whiskerwidth) do pairs, _, _, whiskerwidth

        endpoints = [p for pair in pairs for p in pair]

        screenendpoints = scene_to_screen(endpoints, scene)

        screenendpoints_shifted_pairs = map(screenendpoints) do sep
            (sep .+ f_if(is_in_y_direction[], reverse, Point(0, -whiskerwidth/2)),
             sep .+ f_if(is_in_y_direction[], reverse, Point(0,  whiskerwidth/2)))
        end

        screen_to_scene([p for pair in screenendpoints_shifted_pairs for p in pair], scene)
    end

    linesegments!(plot, linesegpairs, color = color, linewidth = linewidth, visible = visible)
    linesegments!(plot, whiskers, color = color, linewidth = linewidth, visible = visible)
    plot
end

function scene_to_screen(pts, scene)
    p4 = to_ndim.(Vec4f0, to_ndim.(Vec3f0, pts, 0.0), 1.0)
    p1m1 = Ref(scene.camera.projectionview[]) .* p4
    projected = Ref(inv(scene.camera.pixel_space[])) .* p1m1
    [Point2.(p[1:2]...) for p in projected]
end

function screen_to_scene(pts, scene)
    p4 = to_ndim.(Vec4f0, to_ndim.(Vec3f0, pts, 0.0), 1.0)
    p1m1 = Ref(scene.camera.pixel_space[]) .* p4
    projected = Ref(inv(scene.camera.projectionview[])) .* p1m1
    [Point2.(p[1:2]...) for p in projected]
end

function scene_to_screen(p::T, scene) where T <: Point
    p4 = to_ndim(Vec4f0, to_ndim(Vec3f0, p, 0.0), 1.0)
    p1m1 = scene.camera.projectionview[] * p4
    projected = inv(scene.camera.pixel_space[]) * p1m1
    T(projected[1:2]...)
end

function screen_to_scene(p::T, scene) where T <: Point
    p4 = to_ndim(Vec4f0, to_ndim(Vec3f0, p, 0.0), 1.0)
    p1m1 = scene.camera.pixel_space[] * p4
    projected = inv(scene.camera.projectionview[]) * p1m1
    T(projected[1:2]...)
end


# BitIndex
# --------
#
# Utils for indexing bits in a vector of unsigned integers (internal use only).
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/BioSequences.jl/blob/master/LICENSE.md
#
#         index(i)-1        index(i)        index(i)+1
# ....|................|..X.............|................|....
#                          |<-offset(i)-|
#                      |<--- 64 bits -->|

struct BitIndex{N, W}
    val::Int64
end

BitsPerSymbol(::BitIndex{N, W}) where {N,W} = BitsPerSymbol{N}()
bits_per_symbol(::BitIndex{N, W}) where {N,W} = N

@inline function bitindex(::BitsPerSymbol{N}, ::Type{W}, i) where {N,W}
    return BitIndex{N, W}((i - 1) << trailing_zeros(N))
end

index_shift(i::BitIndex{N, UInt64}) where N = 6
index_shift(i::BitIndex{N, UInt32}) where N = 5
index_shift(i::BitIndex{N, UInt16}) where N = 4
index_shift(i::BitIndex{N, UInt8}) where N = 3
offset_mask(i::BitIndex{N, W}) where {N, W} = UInt8(8 * sizeof(W)) - 0x01

index(i::BitIndex) = (i.val >> index_shift(i)) + 1
offset(i::BitIndex) = i.val & offset_mask(i)

Base.:+(i::BitIndex{N,W}, n::Integer) where {N,W} = BitIndex{N,W}(i.val + n)
Base.:-(i::BitIndex{N,W}, n::Integer) where {N,W} = BitIndex{N,W}(i.val - n)
Base.:-(i1::BitIndex, i2::BitIndex) = i1.val - i2.val
Base.:(==)(i1::BitIndex, i2::BitIndex) = i1.val == i2.val
Base.isless(i1::BitIndex, i2::BitIndex) = isless(i1.val, i2.val)
Base.cmp(i1::BitIndex, i2::BitIndex) = cmp(i1.val, i2.val)

@inline function nextposition(i::BitIndex{N,W}) where {N,W}
    return i + N
end

@inline function prevposition(i::BitIndex{N,W}) where {N,W}
    return i - N
end

function Base.iterate(i::BitIndex, s = 1)
    if s == 1
        return (index(i), 2)
    elseif s == 2
        return (offset(i), 3)
    else
        return nothing
    end
end

Base.show(io::IO, i::BitIndex) = print(io, '(', index(i), ", ", offset(i), ')')

"Extract the element stored in a packed bitarray referred to by bidx."
@inline function extract_encoded_element(bidx::BitIndex{N,W}, data::AbstractArray{W}) where {N,W}
    @inbounds chunk = data[index(bidx)]
    offchunk = chunk >> offset(bidx)
    return offchunk & bitmask(bidx)
end

# Create a bit mask that fills least significant `n` bits (`n` must be a
# non-negative integer).
"Create a bit mask covering the least significant `n` bits."
bitmask(::Type{T}, n::Integer) where {T} = (one(T) << n) - one(T)

# Create a bit mask filling least significant N bits.
# This is used in the extract_encoded_element function.
bitmask(bidx::BitIndex{N,W}) where {N, W} = bitmask(W, N)
bitmask(n::Integer) = bitmask(UInt64, n)
bitmask(::Type{T}, ::Val{N}) where {T, N} = (one(T) << N) - one(T)

# TODO: Work out places this is used and see if it is really nessecery given the
# bitmask methods above.
# TODO: Resolve this use of bits_per_symbol and A().
bitmask(::A) where {A<:Alphabet} = bitmask(bits_per_symbol(A()))


using ArgParse, Distributed

function parse_cli(args)
    settings = ArgParseSettings("ABC-MK CLI")

    @add_arg_table! settings begin
        # CMD 1
        "rates", "R"
            help = "Function to solve fixation and polymorphic rates analitically. The function will create N random models from prior values. Use the arguments to defined the input range for each parameter.\nIf rho and/or theta are set to nothing, the function will input random values given the range 0.0005:0.0005:0.01. Otherwise you can fix the values.\nIf gL is set to nothing, the function will not account the role of the weakly selected alleles in the estimation.\nThe function returns a HDF5 file containing models solved and rates. The rates will be used to compute summary statistics required at ABC.\nPlease check the documentation to get more info about models parameters or detailed arguments description https://jmurga.github.io/Analytical.jl/dev/cli/ to check model
            "
            action = :command
        # CMD 2
        "parse_data", "P"
            help = "Function to parse polymorphic and divergence data from Uricchio et. al (2019) and Murga-Moreno et al (2019). Please input a path to create a new analysis folder. You can filter the dataset using a file containing a list of Ensembl IDs. The function returns files containing raw polymorphic and divergence data, parsed SFS and parsed divegence required to estimate summary statistics. Please check the documentation to get more info https://jmurga.github.io/Analytical.jl/dev/cli/"
            action = :command
        # CMD 3
        "summaries", "S"
            help = "Function to parse polymorphic and divergence data from Uricchio et. al (2019) and Murga-Moreno et al (2019). Please input a path to create a new analysis folder. You can filter the dataset using a file containing a list of Ensembl IDs. The function returns files containing raw polymorphic and divergence data, parsed SFS and parsed divegence required to estimate summary statistics. Please check the documentation to get more info https://jmurga.github.io/Analytical.jl/dev/cli/"
            action = :command
        # CMD 4
        "inference", "I"
            help = "ABCreg inference. The function returns posterior distributions from ABC inference. Each posterior file contains information about alpha_w, alpha_s, alpha, gamNeg and shape parameter. The number of posterior distributions will depend on the number of bootstrap replicas. Check the documentation to get more info https://jmurga.github.io/Analytical.jl/dev/cli"
            action = :command

    end

    add_arg_table!(settings["rates"],
        ["--pop_size"],
        Dict(
            :help => "Population size",
            :arg_type => Int64,
            :default => 1000
        ),
        ["--sample_size"],
        Dict(
            :help => "Sample size",
            :arg_type => Int64,
            :required => true
        ),
        ["--dac"],
        Dict(
            :help => "Derived Allele Count",
            :arg_type => String,
            :required => true
        ),
        ["--gam_neg"],
        Dict(
            :help => "Selection coefficient for deleterious alleles",
            :arg_type => String,
            :required => true
        ),
        ["--positive_strong"],
        Dict(
            :help => "Selection coefficient for strongly beneficial alleles",
            :arg_type => String,
            :required => true
        ),
        ["--positive_weak"],
        Dict(
            :help => "Selection coefficient for weakly beneficial alleles",
            :arg_type => String,
            :required => true
        ),
        ["--shape"],
        Dict(
            :help => "Shape value modeling Gamma distribution for deleterious alleles",
            :arg_type => Float64,
            :default => 0.184
        ),
        ["--rho"],
        Dict(
            :help => "Recombination rate",
            :arg_type => Float64,
            :default => 0.001
        ),
        ["--theta"],
        Dict(
            :help => "Mutation rate on coding locus",
            :arg_type => Float64,
            :default => 0.001
        ),
        ["--solutions"],
        Dict(
            :help => "Mutation rate on coding locus",
            :arg_type => Int64,
            :default => 100_000
        ),
        ["--output"],
        Dict(
            :help => "Output file",
            :arg_type => String,
            :default => "rates.jld2"
        ),
        ["--scheduler"],
        Dict(
            :help => "Select scheduler manager",
            :arg_type => String,
            :default => "local"
        ),
        ["--nthreads"],
        Dict(
            :help => "Select number of threads to parallelize",
            :arg_type => Int64,
            :default => 1
        ),
    )

    add_arg_table!(settings["parse_data"],
        "folder",
        Dict(
            :help => "a positional argument",
            :required => true,
            :arg_type => String
        ),
        ["--dataset","-d"],
        Dict(
            :help => "a positional argument",
            :default => "tgp",
            :arg_type => String
        ),
        ["--gene_list","-g"],
        Dict(
            :help => "a positional argument",
            :arg_type => Union{Bool,String},
            :default => false
        ),
        ["--bins","-b"],
        Dict(
            :help => "a positional argument",
            :arg_type => Union{Bool,Int64},
            :default => false
        )
    )

    add_arg_table!(settings["summaries"],
        "folder",
        Dict(
            :help => "Folder path containing SFS and divergence files to run the analysis",
            :required => true,
            :arg_type => String
        ),
        ["--rates"],
        Dict(
            :help => "H5 file containing precomputed rates",
            :required => true,
            :arg_type => String
        ),
        ["--sample_size"],
        Dict(
            :help => "Sample size",
            :required => true,
            :arg_type => Int64
        ),
        ["--dac"],
        Dict(
            :help => "Derived allele count",
            :required => true,
            :arg_type => String
        ),
        ["--summstat_size"],
        Dict(
            :help => "Define number of summary estatistic to perform ABC",
            :required => true,
            :arg_type => Int64
        ),
        ["--bootstrap"],
        Dict(
            :help => "Allow bootstrap following polyDFE manual",
            :arg_type => Bool,
            :default => false
        ),
        ["--replicas"],
        Dict(
            :help => "Number of bootstrap replicas",
            :arg_type => Int64,
            :default => 1
        ),
        ["--nthreads"],
        Dict(
            :help => "Select scheduler manager",
            :arg_type => Int64,
            :default => 1
        ),
        ["--scheduler"],
        Dict(
            :help => "Select scheduler manager",
            :arg_type => String,
            :default => "local"
        )
    )
    add_arg_table!(settings["inference"],
        "folder",
        Dict(
            :help => "Folder path containing SFS and divergence files to run the analysis",
            :required => true,
            :arg_type => String
        ),
        ["--S"],
        Dict(
            :help => "Define number of summary estatistic to perform ABC",
            :required => true,
            :arg_type => Int64
        ),
        ["--tol"],
        Dict(
            :help => "Tolerance",
            :required => true,
            :arg_type => Float64
        ),
        ["--abcreg"],
        Dict(
            :help => "ABCreg static binary",
            :required => true,
            :arg_type => String
        ),
        ["--nthreads"],
        Dict(
            :help => "Select scheduler manager",
            :arg_type => Int64,
            :default => 1
        ),
        ["--scheduler"],
        Dict(
            :help => "Select scheduler manager",
            :arg_type => String,
            :default => "local"
        )
    )

    return parse_args(settings)
end

cli = parse_cli(ARGS)

for cmd in keys(cli)
    if(cmd == "parse_data")

        @eval using Analytical, DataFrames, CSV

        folder = cli[cmd]["folder"]

        run(`mkdir -p $folder`)

        dataset = lowercase(cli[cmd]["dataset"])
        data    = folder * "/" * dataset * ".txt"

        @eval download("https://raw.githubusercontent.com/jmurga/Analytical.jl/master/data/"* $dataset * ".txt",$data)

        # Check if bins or gene_list are defined
        gene_list = cli[cmd]["gene_list"]
        @eval if $gene_list != false
            @eval gList = CSV.read($gene_list,DataFrame,header=false) |> Array
        else
            @eval gList = nothing
        end

        bins = cli[cmd]["bins"]

        @eval if $bins != 0
            @eval bins_size = $bins
        else
            bins_size = nothing
        end

        # Parsing TGP data
        if dataset == "tgp"
            @eval α,sfs, divergence = Analytical.parse_sfs(sample_size=661,data=$data,gene_list=$gList,bins=$bins_size)
        elseif occursin("zi",dataset)
            @eval α,sfs, divergence = Analytical.parse_sfs(sample_size=154,data=$data,gene_list=$gList,bins=$bins_size,isolines=true)
        elseif occursin("ral",dataset)
            @eval α,sfs, divergence = Analytical.parse_sfs(sample_size=160,data=$data,gene_list=$gList,bins=$bins_size,isolines=true)
        end
        # Writting data to folder
        @eval sName = $folder * "/sfs.tsv"
        @eval dName = $folder * "/div.tsv"

        @eval CSV.write($sName,DataFrame($sfs,:auto),delim='\t',header=false)
        @eval CSV.write($dName,DataFrame($divergence',:auto),delim='\t',header=false)
    elseif (cmd == "rates")

        neg    = parse.(Int,split(cli[cmd]["gam_neg"],":"))
        strong = parse.(Int,split(cli[cmd]["positive_strong"],":"))

        if (cli[cmd]["positive_weak"] == false)
            weak = nothing
        else
            weak = parse.(Int,split(cli[cmd]["positive_weak"],":"))
        end

        dac = parse.(Int,split(cli[cmd]["dac"],","))

        if (cli[cmd]["rho"] == false)
            rho = nothing
        else
            rho = cli[cmd]["rho"]
        end

        if (cli[cmd]["theta"] == "nothing")
            theta = nothing
        else
            theta = cli[cmd]["theta"]
        end

        scheduler = cli[cmd]["scheduler"];nthreads = cli[cmd]["nthreads"]
        if scheduler == "slurm"
            @eval using ClusterManagers
            @eval addprocs_slurm($nthreads)
        elseif scheduler == "htcondor"
            @eval using ClusterManagers
            @eval addprocs_htc($nthreads)
        else
            @eval addprocs($nthreads)
        end

        ne = cli[cmd]["pop_size"];samples= cli[cmd]["sample_size"];shape = cli[cmd]["shape"]
        @eval @everywhere using Analytical, ParallelUtilities
        @eval adap = Analytical.parameters(N=$ne,n=$samples,dac=$dac,al=$shape)

        @eval cnv = Analytical.binomial_dict()
        @eval Analytical.binomOp!($adap,$cnv.bn);

        solutions = cli[cmd]["solutions"]; output = cli[cmd]["output"]
        @eval Analytical.rates(param = $adap,convoluted_samples=$cnv,gH=$strong[1]:$strong[2],gL=$weak[1]:$weak[2],gamNeg=$neg[1]:$neg[2],iterations = $solutions,rho=$rho,theta=$theta,shape=$adap.al,output=$output,scheduler=$scheduler);

        for i in workers()
            rmprocs(i)
        end
    elseif (cmd == "summaries")

        scheduler = cli[cmd]["scheduler"];nthreads = cli[cmd]["nthreads"];

        if scheduler == "slurm"
            @eval using ClusterManagers
            @eval addprocs_slurm($nthreads)
        elseif scheduler == "htcondor"
            @eval using ClusterManagers
            @eval addprocs_htc($nthreads)
        else
            @eval addprocs($nthreads)
        end

        samples       = cli[cmd]["sample_size"]
        dac           = parse.(Int,split(cli[cmd]["dac"],","))
        bootstrap     = cli[cmd]["bootstrap"]
        replicas      = cli[cmd]["replicas"]
        rates         = cli[cmd]["rates"]
        summstat_size = cli[cmd]["summstat_size"]
        folder        = cli[cmd]["folder"]

        @eval  using JLD2, DataFrames, CSV, ProgressMeter
        @eval @everywhere using  Analytical, ParallelUtilities
        @eval adap    = Analytical.parameters(n=$samples,dac = $dac)

        @eval if ($bootstrap == true)
            @eval summstat  = Analytical.summary_statistics(param=$adap,h5_file=$rates,analysis_folder=$folder,summstat_size=$summstat_size,replicas=$replicas,bootstrap=true)
        else
            @eval summstat  = Analytical.summary_statistics(param=$adap,h5_file=$rates,analysis_folder=$folder,summstat_size=$summstat_size,replicas=$replicas,bootstrap=false)
        end

        for i in workers()
            rmprocs(i)
        end
    elseif (cmd == "inference")

        scheduler = cli[cmd]["scheduler"]
        nthreads = cli[cmd]["nthreads"]
        folder   = cli[cmd]["folder"]
        S        = cli[cmd]["S"]
        tol      = cli[cmd]["tol"]
        abcreg   = cli[cmd]["abcreg"]

        if scheduler == "slurm"
            @eval using ClusterManagers
            @eval addprocs_slurm($nthreads)
        elseif scheduler == "htcondor"
            @eval using ClusterManagers
            @eval addprocs_htc($nthreads)
        else
            @eval addprocs($nthreads)
        end

        @eval @everywhere using Analytical,ParallelUtilities

        @eval Analytical.ABCreg(analysis_folder=$folder,S=$S,tol=$tol,abcreg=$abcreg)

        for i in workers()
            rmprocs(i)
        end
    end
end


"""
The constants representing the game variation being played.
They come from game engine and changing them has no effect.
They are strictly informational.
"""
struct Constants
    "The cost to build a single ship."
    SHIP_COST::Int
    "The cost to build a dropoff."
    DROPOFF_COST::Int
    "The maximum amount of halite a ship can carry."
    MAX_HALITE::Int
    """
    The maximum number of turns a game can last. This reflects the fact
    that smaller maps play for fewer turns.
    """
    MAX_TURNS::Int
    "1/EXTRACT_RATIO halite (truncated) is collected from a square per turn."
    EXTRACT_RATIO::Int
    "1/MOVE_COST_RATIO halite (truncated) is needed to move off a cell."
    MOVE_COST_RATIO::Int
    "Whether inspiration is enabled."
    INSPIRATION_ENABLED::Bool
    """
    A ship is inspired if at least INSPIRATION_SHIP_COUNT opponent
    ships are within this Manhattan distance.
    """
    INSPIRATION_RADIUS::Int
    """
    A ship is inspired if at least this many opponent ships are within
    INSPIRATION_RADIUS distance.
    """
    INSPIRATION_SHIP_COUNT::Int
    "An inspired ship mines 1/X halite from a cell per turn instead."
    INSPIRED_BONUS_MULTIPLIER::Float64
    "An inspired ship that removes Y halite from a cell collects X*Y additional halite."
    INSPIRED_EXTRACT_RATIO::Int
    "An inspired ship instead spends 1/X% halite to move."
    INSPIRED_MOVE_COST_RATIO::Int
end

"Load constants from JSON given by the game engine."
function load_constants(constants::Dict)
    SHIP_COST = constants["NEW_ENTITY_ENERGY_COST"]
    DROPOFF_COST = constants["DROPOFF_COST"]
    MAX_HALITE = constants["MAX_ENERGY"]
    MAX_TURNS = constants["MAX_TURNS"]
    EXTRACT_RATIO = constants["EXTRACT_RATIO"]
    MOVE_COST_RATIO = constants["MOVE_COST_RATIO"]
    INSPIRATION_ENABLED = constants["INSPIRATION_ENABLED"]
    INSPIRATION_RADIUS = constants["INSPIRATION_RADIUS"]
    INSPIRATION_SHIP_COUNT = constants["INSPIRATION_SHIP_COUNT"]
    INSPIRED_EXTRACT_RATIO = constants["INSPIRED_EXTRACT_RATIO"]
    INSPIRED_BONUS_MULTIPLIER = constants["INSPIRED_BONUS_MULTIPLIER"]
    INSPIRED_MOVE_COST_RATIO = constants["INSPIRED_MOVE_COST_RATIO"]

    Constants(SHIP_COST, DROPOFF_COST, MAX_HALITE, MAX_TURNS, EXTRACT_RATIO,
              MOVE_COST_RATIO, INSPIRATION_ENABLED, INSPIRATION_RADIUS,
              INSPIRATION_SHIP_COUNT, INSPIRED_BONUS_MULTIPLIER,
              INSPIRED_EXTRACT_RATIO, INSPIRED_MOVE_COST_RATIO)
end


module CVMAbstractModel

require("CommonTools.jl")

using DataTypes
using Corpora.Token
using CommonTools.CVMIntermediates
using CommonTools: mult_by_0!, sigmoid, loopdot, get_rand
using CommonTools: tanh, square!, l2sq, l2sq_distance

export AbstractModel
export AbstractParameters
export accumulate_gradients_batch!
export accumulate_closeness_gradients_batch!
export objective_gradients_batch
export regularize_l2!

abstract AbstractModel;
abstract AbstractParameters;

# Whenever we create a new model, these two methods need to be implemented
# (compose! and store_gradients!)

# Given the gradient of the error function w.r.t. the phrase representation,
# backpropagate this gradient to the parameters of the model
function store_gradients!(model::AbstractModel,
                    pre_grad::Vector{Float},
                    s::Vector{Token}, intermediate::CVMIntermediates;
                    intermediates_precomputed::Bool=false)
    error("This is an abstract method");
end

# Given a list of tokens, compose a phrase vector representation
function compose!(model::AbstractModel,
                                    s::Vector{Token},
                                    intermediate::CVMIntermediates)
    error("This is an abstract method");
end

# Adds l2 regularization to the given parameters
function regularize_l2!(model::AbstractModel, params::AbstractParameters;
                            sparse_updates::Bool=false)
    acc_reg = const_zero;
    if model.regularizer > const_zero
        if sparse_updates
            # Only add regularization to the gradients that are not 0
            for i=1:length(params.wordvecs)
                if params.gradients[i] != 0
                    params.gradients[i] += params.wordvecs[i] * model.regularizer
                    acc_reg += model.regularizer * params.wordvecs[i]^2
                end
            end
        else
            # Add l2 regularization
            Base.LinAlg.BLAS.axpy!(model.regularizer, params.wordvecs, params.gradients)
            # Compute l2 error
            acc_reg += l2sq(params.wordvecs) * model.regularizer
        end
    end
    return acc_reg*0.5f0;
end

# Accumulates the gradients of |s1-s2|^2
function accumulate_closeness_gradients_batch!(model::AbstractModel,
                s1::Vector{Token}, s2::Vector{Token},
                intermediates::Vector{CVMIntermediates}; scaling=1.0)
    err = 0.0
    # Compose vector representation sentence1 and sentence2
    o1 = compose!(model, s1, intermediates[1])
    o2 = compose!(model, s2, intermediates[2])

    # Cache distance of the pair of aligned sentences (valid pair)
    dist_o1_o2 = l2sq_distance(o1, o2)
    # Update error
    err += scaling * dist_o1_o2
    # Calculate the gradient of |o1-o2|^2 w.r.t. o1
    pre_grad_closeness = (o1-o2)
    # Propagate that gradient to the individual word representations
    # of sentence1
    Base.LinAlg.BLAS.scal!(model.wordvecdim,
                                    scaling,
                                    pre_grad_closeness, 1);
    store_gradients!(model, pre_grad_closeness, s1, intermediates[1])
    # The gradient of |o1-o2|^2 w.r.t. o2 is just the negative o1 gradient
    # => invert the sign
    Base.LinAlg.BLAS.scal!(model.wordvecdim,
                                    -1.0,
                                    pre_grad_closeness, 1);
    # Propagage the gradient down to the word representations of sentence 2
    store_gradients!(model, pre_grad_closeness, s2, intermediates[2])

    return err;
end

# Calculates the error of the hinge loss and the closeness
function objective_gradients_batch(model::AbstractModel,
                s1::Vector{Token}, s2::Vector{Token},
                sns::Vector{Vector{Token}},
                intermediates::Vector{CVMIntermediates};
                scaling::Float=const_one)

    # Calculate the representations of sentence1 and sentence2
    # In case of the inclusion criterion, sentence1 is the outer window
    # and sentence 2 is the inner window.
    o1 = compose!(model, s1, intermediates[1])
    o2 = compose!(model, s2, intermediates[2])

    # Cache distance
    dist_o1_o2 = l2sq_distance(o1, o2)
    # Update error
    err_closeness = dist_o1_o2 * 0.5

    # Calculate the contrastive error (hinge loss) for each noise sample
    err_contrastive = 0.0
    inv_len_sns = 1.0/size(sns,1)
    cn_scaling = scaling * model.gamma * inv_len_sns
    for sn in sns
        # Compose the vector representation of the noise sample
        on = compose!(model, sn, intermediates[3])

        # Compute the contastive part of the hinge loss
        contrastive = 0.5*(dist_o1_o2 - l2sq_distance(o1, on))

        # Calculate the hinge loss
        contrastive_err = max(0.0, model.margin + contrastive)

        # If hinge loss greater than 0: Update the accumulative error
        if  contrastive_err > 0
            err_contrastive += contrastive_err * cn_scaling;
        end
    end

    return err_contrastive + err_closeness;
end

# Accumulates the gradients of the hinge loss and the closeness function
function accumulate_gradients_batch!(model::AbstractModel,
                s1::Vector{Token}, s2::Vector{Token},
                sns::Vector{Vector{Token}},
                intermediates::Vector{CVMIntermediates};
                scaling::Float=const_one)

    # Calculate the representations of sentence1 and sentence2
    # In case of the inclusion criterion, sentence1 is the outer window
    # and sentence 2 is the inner window.
    o1 = compose!(model, s1, intermediates[1])
    o2 = compose!(model, s2, intermediates[2])

    # Cache distance of the pair of aligned sentences (valid pair)
    dist_o1_o2 = l2sq_distance(o1, o2)
    err_closeness = dist_o1_o2 * 0.5

    # Calculate the gradient of |o1-o2|^2 w.r.t. o1
    pre_grad_closeness = (o1-o2)
    Base.LinAlg.BLAS.scal!(model.wordvecdim,
                                    scaling,
                                    pre_grad_closeness, 1);
    # Propagage the gradient down to the word representations of sentence 2
    store_gradients!(model, pre_grad_closeness, s1, intermediates[1])
    # The gradient of |o1-o2|^2 w.r.t. o2 is just the negative o1 gradient
    # => invert the sign
    Base.LinAlg.BLAS.scal!(model.wordvecdim,
                                    -1.0,
                                    pre_grad_closeness, 1);
    # Propagage the gradient down to the word representations of sentence 2
    store_gradients!(model, pre_grad_closeness, s2, intermediates[2])

    err_contrastive = 0.0
    # Scale the number contribution of each individual noise sample
    # by 1/their number
    inv_len_sns = 1.0/size(sns,1)
    cn_scaling = scaling * model.gamma * inv_len_sns
    # Apply the scaling to the |o1-o2|^2 w.r.t. o2 gradient
    Base.LinAlg.BLAS.scal!(model.wordvecdim,
                                    model.gamma * inv_len_sns,
                                    pre_grad_closeness, 1);
    for sn in sns
        on = compose!(model, sn, intermediates[3])
        contrastive = 0.5*(dist_o1_o2 - l2sq_distance(o1, on))

        # Calculate the hinge loss
        contrastive_err = max(0.0, model.margin + contrastive)

        # If hinge loss greater than 0: Minimize the contrastive distances
        if  contrastive_err > 0

            # Compute gradients of |o1,o2|^2-|o1,on|^2 w.r.t. on
            pre_gradn = o1-on
            Base.LinAlg.BLAS.scal!(model.wordvecdim,
                                    cn_scaling,
                                    pre_gradn, 1);
            ### Dangerous! Overwriting "on" and reuising it is potentially
            # dangerous!!
            # Original, clean version
            #pre_grad1 = on-o2
            # hacky, high-speed version reusing "on"
            pre_grad1 = on
            Base.LinAlg.BLAS.axpy!(-1.0, o2, pre_grad1)
            Base.LinAlg.BLAS.scal!(model.wordvecdim,
                                    cn_scaling,
                                    pre_grad1, 1);

            # Add the gradients up in the gradient field of each token
            ## Hinge objective
            store_gradients!(model, pre_grad1, s1, intermediates[1],
                                    intermediates_precomputed=true)
            # Could optimize this by pulling it out of the loop
            store_gradients!(model, pre_grad_closeness, s2, intermediates[2],
                                    intermediates_precomputed=true)
            # Propagage the gradients down to the word representations
            # of the noise sentence
            store_gradients!(model, pre_gradn, sn, intermediates[3])

            # Update accumulative error
            err_contrastive += contrastive_err * model.gamma;
        end
    end

    return err_closeness, err_contrastive;
end

end