################################################################################

"""
    process(self::HMM, d::Array{T, 2}, splitSw::Bool;
    params::HMMParams) where T <: Number

# Description
Process hidden Markov model object.
Meant as an iterative mutating function, perform several steps:
  - reset model traceback.
  - feed frames into model data.
  - update model.
  - generate hypothesized new states.

See also: [`setup`](@ref), [`HMM`](@ref), [`HMMParams`](@ref)
"""
function process!(self::HMM, d::Array{T, 2}, splitSw::Bool; params::HMMParams) where T <: Number

  # reset
  reset!(self)

  # feed frame
  for ι ∈ axes(d, 1)
    feed!(self, ι, d, params = params)
  end

  # backtrace
  τ = backTrace(self)

  divider = fill(1, size(self.data, 1))
  orig = deepcopy(self.data)
  pp = StructArrays.StructArray{ScorePair}(undef, 0)

  # update model
  mdist = zeros(Float64, size(self.data))
  mcount = zeros(Float64, size(self.data))

  for ι ∈ axes(d, 1)
    self.data[τ[ι]] .+= d[ι, :]
    divider[τ[ι]] += 1
    pair = ScorePair(params.distance(orig[τ[ι]], d[ι, :]), ι)

    mdist[τ[ι]] += pair.score
    mcount[τ[ι]] += 1

    push!(pp, pair)
  end

  scores = sort(pp.score, rev = true)
  ixs = map(χ -> findall(χ .== pp.score), scores)
  pp = pp[vcat(ixs...)]

  for ι ∈ eachindex(self.data)
    self.data[ι] /= divider[ι]
  end

  if params.verbosity
    for ε ∈ eachindex(self.data)
      @info "Print state $(ε)"
      for ι ∈ eachindex(self.data[ε])
        println(round(self.data[ε][ι]; digits = 3))
      end
    end
  end

  sortHMM!(self)

  if !splitSw
    return τ, self.data
  end

  max = 0.
  toSplit = 1

  for ι ∈ eachindex(mdist)
    if mcount[ι] > params.minimumFrequency
      avdist = mdist[ι] / mcount[ι]
      if avdist > max
        max = avdist
        toSplit = ι
      end
    end
  end

  half = 1 + mcount[toSplit] / 4

  extra = fill(0, size(self.data[1], 1))

  count = 1
  for ι ∈ eachindex(pp)
    if τ[pp[ι].index] != toSplit
      continue
    end
    extra += d[pp[ι].index, :]
    count += 1
    if count >= half
      break
    end
  end

  extra ./= (count - 1)

  push!(self.data, extra)

  push!(self.tb, fill(0, size(self.tb[1], 1)))

  return τ, self.data

end

################################################################################


module _Derivatives

using LinearAlgebra
using SparseArrays

using Base.Meta
using ForwardDiff
import Calculus
using SpecialFunctions  # Needed for some of the derivatives in Calculus.
import MathOptInterface
const MOI = MathOptInterface
using ..JuMP

const TAG = :jump_tag

# Override basic math functions to return NaN instead of throwing errors.
# This is what NLP solvers expect, and
# sometimes the results aren't needed anyway,
# because the code may compute derivatives wrt constants.
import NaNMath:
    sin,
    cos,
    tan,
    asin,
    acos,
    acosh,
    atanh,
    log,
    log2,
    log10,
    lgamma,
    log1p,
    pow,
    sqrt

include("types.jl")
include("conversion.jl")
include("coloring.jl")
export Coloring
include("linearity.jl")
include("sparsity.jl")
include("forward.jl")
include("reverse.jl")
include("subexpressions.jl")

end # module


import Pkg
Pkg.activate(".")

using GBIF
using SimpleSDMLayers
using StatsBase
using Plots
using ArchGDAL
using Images
using StatsPlots.PlotMeasures
using ProgressMeter

include("bioclim.jl")
include("landcover.jl")

# Load raccoon emoji
img = load("images/raccoon_freepik_cropped.png")

## Get occurrences & climatic data

# Get the GBIF taxon
raccoon = GBIF.taxon("Procyon lotor", rank=:SPECIES, strict=true)

# Get the initial occurrences
raccoon_occ = GBIF.occurrences(raccoon, "limit" => 100)

# Get the rest of the occurrences
while length(raccoon_occ) <= 2000
    try
        GBIF.occurrences!(raccoon_occ)
    catch#
        break
    end
end

# Get the temperature and precipitation
temperature   = worldclim(1)[(bottom = -60.0)]
precipitation = worldclim(12)[(bottom = -60.0)]
# Get landcover data
tree  = landcover(1)[(bottom = -60.0)]
urban = landcover(2)[(bottom = -60.0)]
water = landcover(3)[(bottom = -60.0)]
# Combine variables
vars = [temperature, precipitation, tree, urban, water]

# Map climate variables
temp_map = heatmap(temperature, c = :inferno, xlab= "Longitude", ylab = "Latitude",
                   colorbar_title = "Temperature (° C)")
prec_map = heatmap(precipitation, c = :blues, xlab= "Longitude", ylab = "Latitude",
                   colorbar_title = "Precipitation (mm)")

# Map raccoon occurrences
# Option 1: Regular markers & raccoon emoji on the side
occ_map  = heatmap(temperature, c = :lightgrey, xlab= "Longitude", ylab = "Latitude",
                   colorbar = :none) |> x ->
    scatter!(x, longitudes(raccoon_occ), latitudes(raccoon_occ),
         lab = "Raccoons",
         legend = :outerbottomright,
         foreground_color_legend = nothing,
         size = (600, 400)
    ) |> x ->
    plot!(x, img,
        yflip = true,
        inset = bbox(0.4, -0.05, 150px, 75px, :center),
        subplot = 2,
        grid = false, axis = false,
        bg_inside = nothing
    )
# Option 2: Raccoon emojis as markers (max 200 for now)
function emojimap(layer, lons, lats, n)
    lons = lons[1:n]
    lats = lats[1:n]
    emoji_map = heatmap(layer, c = :lightgrey, # xlab= "Longitude", ylab = "Latitude",
        colorbar = :none, size = (360,150).*2,
        ticks = false, margin = -1.9mm, box = :none)
    @time @showprogress for (lon, lat, i) in zip(lons, lats, 1:n)
        plot!(emoji_map, img, yflip = true, grid = false, axis = false, bg_inside = nothing,
            inset = (1, bbox(((lon + 180)/360)w - 15px, ((lat+60)/150)h - 8px, 30px, 15px, :bottom, :left)),
            subplot = i+1
        )
    end
    return emoji_map
end
occ_map2 = emojimap(temperature, longitudes(raccoon_occ), latitudes(raccoon_occ), 200)

# Option 3: Unique sites
using DataFrames

df = DataFrame(raccoon_occ)
select!(df, :longitude, :latitude)
filter!(x -> !ismissing(x.latitude) || !ismissing(x.longitude), df)
df.gridlon = [SimpleSDMLayers._match_longitude(temperature, lon) for lon in df.longitude]
df.gridlat = [SimpleSDMLayers._match_latitude(temperature, lat) for lat in df.latitude]
uniquedf = unique(df, [:gridlon, :gridlat])

occ_map2 = emojimap(temperature, uniquedf.longitude, uniquedf.latitude, 125)

## Bioclim model

# Get prediction for each variable
vars_predictions = [bioclim(v, raccoon_occ) for v in vars]

# Get minimum prediction per site
sdm_raccoon = reduce(min, vars_predictions)

# Set value to NaN if prediction is zero
replace!(x -> iszero(x) ? NaN : x, sdm_raccoon.grid)

# Filter predictions with threshold
threshold = quantile(filter(!isnan, sdm_raccoon.grid), 0.05)
replace!(x -> x <= threshold ? NaN : x, sdm_raccoon.grid)

# Group predictions in categories
lim1 = 0.08
lim2 = 0.25
replace!(x -> x >= lim2 ? 1.0 : x, sdm_raccoon.grid)
replace!(x -> lim1 < x < lim2 ? 0.5 : x, sdm_raccoon.grid)
replace!(x -> x <= lim1 ? 0.0 : x, sdm_raccoon.grid)

# Custom colorpicking function
colorpick(cg::ColorGradient, n::Int) = RGB[cg[i] for i in LinRange(0, 1, n)]

# Map predictions
begin
    pred_map = heatmap(temperature, c = :lightgrey, colorbar = :none,
                       xlab= "Longitude", ylab = "Latitude",
                       framestyle = :box, size = (600, 300))
    heatmap!(pred_map, sdm_raccoon, c = :viridis, clim = (0,1), colorbar_title = "Suitability for raccoons")
    scatter!(pred_map, [NaN NaN NaN NaN],
                 c = [colorpick(cgrad(:viridis), 3)... :lightgrey],
                 labels = ["Low" "Medium" "High" "Not suitable"],
                 legend = :outerbottomright, legendtitle = "Suitability",
                 foreground_color_legend = nothing,
                 legendtitlefontsize = 10)
    plot!(pred_map, img,
            yflip = true,
            inset = bbox(0.36, -0.23, 150px, 75px, :center),
            subplot = 2,
            grid = false, axis = false,
            bg_inside = nothing
        )
end

## Export figures
# Occurrence figure
savefig(plot(occ_map2, dpi = 600), joinpath("fig", "occurrences_emojis.png"))
savefig(plot(occ_map2, dpi = 600), joinpath("fig", "occurrences_emojis.svg"))
# Predictions figure
savefig(plot(pred_map, dpi = 600), joinpath("fig", "predictions.png"))
savefig(plot(pred_map, dpi = 600), joinpath("fig", "predictions.svg"))
savefig(plot(pred_map, dpi = 1200), joinpath("fig", "predictions_1200.png"))
# Extra figures
savefig(plot(temp_map, dpi = 150), joinpath("fig", "temperature.png"))
savefig(plot(prec_map, dpi = 150), joinpath("fig", "precipitation.png"))
savefig(plot(occ_map,  dpi = 150), joinpath("fig", "occurrences.png"))


# -----------------------------------------------------------------------------

doc"Represents a `Constant Distribution`"
immutable Constant{T} <: DiscreteUnivariateDistribution
  c::T
end
StatsBase.sample(c::Constant) = c.c

immutable HiddenLayerSampler <: DiscreteMultivariateDistribution
  numlayers::AVecI
  numnodes::AVecI
end
function StatsBase.sample(hls::HiddenLayerSampler)
  numlayers = sample(hls.numlayers)
  Int[sample(hls.numnodes) for i in 1:numlayers]
end

immutable VectorSampler{T} <: DiscreteUnivariateDistribution
  vec::AVec{T}
  probs::VecF
end
VectorSampler(vec::AVec) = VectorSampler(vec, fill(1.0 / length(vec), length(vec)))
VectorSampler{T}(vec::AVec{Tuple{T,Float64}}) = VectorSampler(unzip(vec)...)
function StatsBase.sample(vs::VectorSampler)
  r = rand()
  totp = 0.0
  for (i,p) in enumerate(vs.probs)
    totp += p
    if r <= totp
      return vs.vec[i]
    end
  end
  error("probs don't sum to 1?? $r $totp $vs")
end

# -----------------------------------------------------------------------------

StatsBase.sample(d::Distribution) = rand(d)

doc"""
A `ParameterSampler` holds vectors of keyword parameter names and the distributions to sample from
when randomly choosing a parameter set.  Example:

```
  ps = ParameterSampler(Dict(:x => Constant(1), :y => Normal()))
```
"""
type ParameterSampler
  d::Dict
end
StatsBase.sample(ps::ParameterSampler) = [(s,sample(dist)) for (s,dist) in ps.d]

function Base.print(io::IO, ps::ParameterSampler)
  println("ParameterSampler{")
  for k in sort(keys(ps.d))
    println("  $k => $(typeof(ps.d[k]))")
  end
  print("} ")
end
Base.show(io::IO, ps::ParameterSampler) = print(io, ps)

# -----------------------------------------------------------------------------

doc"""
Creates a `VectorTransformer` by sampling (N times) both a transform type and the index that it maps to.

Note that we call `unique` so as not to duplicate data, so we may actually return less than N samples.
"""
function generateTransformer(transformSampler::VectorSampler{DataType}, indices::AVecI, N::Int)
  VectorTransformer(unique(Transformation[sample(transformSampler)(sample(indices)) for i in 1:N]))
end

doc"""
Create numModels models, all using the `buildFn`, which handles creating the model with the constant args/kwargs,
and a random sampling of the parameters from the parameter sampler.
Note we assume all parameters we care about setting are keyword params.
"""
function generateModels(buildFn::Function, ps::ParameterSampler, numModels::Integer)
  OnlineStat[buildFn(; sample(ps)...) for i in 1:numModels]
end

# -----------------------------------------------------------------------------

# TODO: ensemble tools, parameter optimization (grid search, etc)

doc"""
The Ensemble type represents a set of models which have been fit to the same data.
A call to fit! will update all the models with that data.  Then the outputs of
the models (yhatᵢⱼ) are stacked and get passed through to the `agg` aggregator. `agg`
is another OnlineStat which regresses yhats (model estimates) on y (targets).

Note: predict first generates estimates from each model, then uses those to predict y
with the `agg` model aggregator.
"""
type Ensemble{STAT} <: OnlineStat
  models::Vector{OnlineStat}
  aggs::Vector{STAT}  # one aggregation model for each target variable... each agg maps nummodels vars to numtargets vars
end

# the build function should take only keyword arguments to build, which can be populated with a ParameterSampler
function Ensemble{AGG<:OnlineStat}(::Type{AGG}, buildFn::Function, ps::ParameterSampler, numModels::Int, nout::Int, args...; kwargs...)
  Ensemble(generateModels(buildFn, ps, numModels), AGG[AGG(args...; kwargs...) for i in 1:nout])
end

# OnlineStats.statenames(o::Ensemble) = [:nobs]
# OnlineStats.state(o::Ensemble) = Any[nobs(o)]
nobs(o::Ensemble) = 0

StatsBase.predict(models::Vector{OnlineStat}, x::AVecF) = [predict(model, x) for model in models]

# get estimates from each model.  then for each i, grab the estimates of the iᵗʰ target variable
# and use those as the inputs to the iᵗʰ aggregation model
function StatsBase.predict(ensemble::Ensemble, x::AVecF)
  ests = predict(ensemble.models)
  nout = length(ensemble.aggs)
  yhat = zeros(nout)
  for i in 1:nout
    esti = Float64[est[i] for est in ests]
    yhat[i] = predict(ensemble.aggs[i], esti)
  end
end

# similar to predict, but fit! instead
function OnlineStats.fit!(ensemble::Ensemble, x::AVecF, y::AVecF)
  for model in ensemble.models
    fit!(model, x, y)
  end

  ests = predict(ensemble.models)

  yhat = similar(y)
  for i in 1:length(y)
    esti = Float64[est[i] for est in ests]
    fit!(ensemble.aggs[i], esti, y)
  end
end


export DiffCorrectAxisymmetric

#-------#
# TYPES #
#-------#

abstract type Abstract_AxisymmetricCorrector <: Abstract_DifferentialCorrector end

struct DiffCorrectAxisymmetric{D<:DiffEqBase.DEAlgorithm,F<:Abstract_ReferenceFrame} <: Abstract_AxisymmetricCorrector
    u0_free :: SVector{2,Integer}  # Indices of the free variables in the initial condition
    u1_free :: SVector{3,Integer}  # Indices of the free variables in the final condition
    alg :: D
    frame :: F

    # Search perturbations for axisymmetric case [Russell 2006]
    # NOTE: Only valid if the reference sol is already near-axisymmetric
    DiffCorrectAxisymmetric(u0_free, u1_free, alg, frame) = new{typeof(alg),typeof(frame)}(u0_free, u1_free, alg, frame)
    DiffCorrectAxisymmetric() = DiffCorrectAxisymmetric((5,6), (2,3,4), DEFAULT_ALG, SynodicFrame(true))
end

#-----------#
# FUNCTIONS #
#-----------#

abstract type AxisymmetricException <: Exception end
struct DidNotCross <: AxisymmetricException end
struct Crashed <: AxisymmetricException end
struct BadTspan <: AxisymmetricException end

function orbit_xcrossing(u, t, integrator)
    state = get(integrator.opts.userdata, :real_state, nothing)
    isnothing(state) && error("orbit_xcrossing expects a `real_state` to be provided in userdata.")
    new_state = remake(state; u0=u, tspan=(t, t))
    synodic_state = convert_to_frame(new_state, default_synodic_reference_frame(state.model))

    # XXX: This cannot return the Dual, or the downstream Duals (i.e. sol.u) get very broken when this triggers!
    return @inbounds ForwardDiff.value.(synodic_state.prob.u0[2])
end

function terminate_after_N_crossings!(integrator)
    t = integrator.t
    tspan = integrator.sol.prob.tspan
    if (t - tspan[begin]) / (tspan[end] - tspan[begin]) >= 0.01
        N = get!(integrator.opts.userdata, :crossings, 1)
        integrator.opts.userdata[:crossings] = N - 1
        if integrator.opts.userdata[:crossings] == 0
            terminate!(integrator, :NCrossings)
        end
    end
end

crossed(sol::Trajectory) = sol.retcode == :NCrossings

function corrector_callback(::Abstract_AxisymmetricCorrector, system::EphemerisNBP)
    # TODO: Play with interp_points, interp_points=0 halves runtime/memory, but might cause issues due to oscillation around y-axis
    # TODO: idxs=[2] makes things very slow, even though it seems it should speed it up.
    return VectorContinuousCallback(
        (integrator, idx) -> begin
            if idx == 3 # Crossing!
                terminate_after_N_crossings!(integrator)
            else
                terminate!(integrator, :Crashed)
            end
        end, 3;
        interp_points=10, save_positions=(false, false)) do du, u, t, integrator
            diam1 = maximum(bodvrd(String(primary_body(system)), "RADII"))  # km
            diam2 = maximum(bodvrd(String(secondary_body(system)), "RADII"))  # km
            du[1] = check_distance(u, t, system, primary_body(system), diam1)
            du[2] = check_distance(u, t, system, secondary_body(system), diam2)
            du[3] = orbit_xcrossing(u, t, integrator)

            # XXX: I think that some of the underlying SPICE functions are slightly discontinuous,
            # maybe there are differences in number of bits being used between here and in the
            # SPICE library. Basically, this causes issues in event callback detection, because it
            # gets stuck in a value close, but not quite equal to zero (e.g. on a value like -2.0e-14),
            # despite incrementing the timestep -- see https://github.com/SciML/DiffEqBase.jl/issues/646.
            # This fix is very hacky, but I have no idea what might be a better option for now.
            @. du[isapprox(du, 0; atol=1e-4)] = 0.
        end
end

function corrector_callback(::Abstract_AxisymmetricCorrector, ::Abstract_DynamicalModel)
    # TODO: Play with interp_points, interp_points=0 halves runtime/memory, but might cause issues due to oscillation around y-axis
    # TODO: idxs=[2] makes things very slow, even though it seems it should speed it up.
    return ContinuousCallback(orbit_xcrossing, terminate_after_N_crossings!; interp_points=10, save_positions=(false, false))
end

function corrector_solve(corrector::Abstract_AxisymmetricCorrector, state::State; strict=false, verbose=false, crossings=1, kwargs...)
    callback = strict ? corrector_callback(corrector, state.model) : nothing
    userdata = Dict{Symbol,Any}(:crossings => crossings)  # XXX: Dict type specified for easy merging

    sol = sensitivity_trace(AD, state, corrector.frame, corrector.alg; trace_time=true, abstol=1e-12, reltol=1e-12, callback, userdata, kwargs...)
    sol_tspan = ForwardDiff.value.((sol.t[begin], sol.t[end]))

    # Error-checking
    if strict
        try
            crashed(sol) && throw(Crashed())
            !crossed(sol) && throw(DidNotCross())
        catch ex
            if verbose && isa(ex, AxisymmetricException)
                display(plot(solve(remake(state, tspan=sol_tspan)), corrector.frame; title="$(ex)"))
            end
            rethrow(ex)
        end
    end

    if verbose
        display(plot(sol, corrector.frame))
    end

    return convert_to_frame(sol, corrector.frame), sol_tspan
end

function (corrector::Abstract_AxisymmetricCorrector)(state::State, F, J, x; kwargs...)
    tspan = (state.tspan[begin], state.tspan[end] + x[end])
    if (tspan[end] - tspan[begin]) < 0 || (tspan[end] - tspan[begin]) > 1e8
        # Don't like this tspan, it's either negative time direction, or too long
        throw(BadTspan())
    end

    # Recreate the state with the parameters provided in x
    correction_state = convert_to_frame(state, corrector.frame)
    u0 = deepcopy(correction_state.u0)
    u0[corrector.u0_free] .= x[1:length(corrector.u0_free)]
    correction_state = remake(correction_state; u0, tspan)

    # Compute sensitivities!
    sol, sol_tspan = corrector_solve(corrector, correction_state; strict=false, kwargs...)

    if !isnothing(F)
        # Calculate how far away we are from the desirable final state (i.e. the residual).
        end_u0 = ForwardDiff.value.(sol.u[end])
        F .= [end_u0[corrector.u1_free]...]
    end

    if !isnothing(J)
        # As per [Russell.2006, Eq.(10) to (13)], we need to add the time variation term,
        # which requires the values of [ẏ, ż, ẍ]. Here, these are provided in the STM
        # on the "end" column (since we're calling corrector_solve with trace_time=true).
        STM = extract_STMs([sol.u[end]])[1]
        J .= STM[[corrector.u1_free...], [corrector.u0_free..., end]]
    end
end


#
# Trailer path planning library
#
# author: Atsushi Sakai(@Atsushi_twi)
#

module trailerlib

using PyPlot
using NearestNeighbors

# Vehicle parameter
const WB = 3.7  #[m] wheel base: rear to front steer
const LT = 8.0 #[m] rear to trailer wheel
const W = 2.6 #[m] width of vehicle
const LF = 4.5 #[m] distance from rear to vehicle front end of vehicle
const LB = 1.0 #[m] distance from rear to vehicle back end of vehicle
const LTF = 1.0 #[m] distance from rear to vehicle front end of trailer
const LTB = 9.0 #[m] distance from rear to vehicle back end of trailer
const MAX_STEER = 0.6 #[rad] maximum steering angle
const TR = 0.5 # Tyre radius [m] for plot
const TW = 1.0 # Tyre width [m] for plot

# for collision check
const WBUBBLE_DIST = 3.5 #distance from rear and the center of whole bubble
const WBUBBLE_R = 10.0 # whole bubble radius
const B = 4.45 # distance from rear to vehicle back end
const C = 11.54 # distance from rear to vehicle front end
const I = 8.55 # width of vehicle
const VRX = [C, C, -B, -B, C ]
const VRY = [-I/2.0, I/2.0, I/2.0, -I/2.0, -I/2.0]

function check_collision(x::Array{Float64},
                         y::Array{Float64},
                         yaw::Array{Float64},
                         kdtree::KDTree,
                         ox::Array{Float64},
                         oy::Array{Float64},
                         wbd::Float64,
                         wbr::Float64,
                         vrx::Array{Float64},
                         vry::Array{Float64})::Bool

    for (ix, iy, iyaw) in zip(x, y, yaw)
        cx = ix + wbd*cos(iyaw)
        cy = iy + wbd*sin(iyaw)

        # Whole bubble check
        ids = inrange(kdtree, [cx, cy], wbr, true)
        # println(length(ids))
        if length(ids) == 0 continue end

        if !rect_check(ix, iy, iyaw, ox[ids], oy[ids], vrx, vry)
            # println("collision")
            return false #collision
        end

        # println(ids)
    end

    return true #OK

end

function rect_check(ix::Float64, iy::Float64, iyaw::Float64,
                    ox::Array{Float64}, oy::Array{Float64},
                    vrx::Array{Float64}, vry::Array{Float64}
                   )

    c = cos(-iyaw)
    s = sin(-iyaw)

    for (iox, ioy) in zip(ox, oy)
        tx = iox - ix
        ty = ioy - iy
        lx = (c*tx - s*ty)
        ly = (s*tx + c*ty)

        sumangle = 0.0
        for i in 1:length(vrx)-1
            x1 = vrx[i] - lx
            y1 = vry[i] - ly
            x2 = vrx[i+1] - lx
            y2 = vry[i+1] - ly
            d1 = hypot(x1,y1)
            d2 = hypot(x2,y2)
            theta1 = atan(y1,x1)
            tty = (-sin(theta1)*x2 + cos(theta1)*y2)
            tmp = (x1*x2+y1*y2)/(d1*d2)

            if tmp >= 1.0
                tmp = 1.0
            elseif tmp <= 0.0
                tmp = 0.0
            end

            if tty >= 0.0
                sumangle += acos(tmp)
            else
                sumangle -= acos(tmp)
            end
        end

        if sumangle >= pi
            return false #collision
        end
    end

    return true #OK
end

function calc_trailer_yaw_from_xyyaw(
                   x::Array{Float64},
                   y::Array{Float64},
                   yaw::Array{Float64},
                   init_tyaw::Float64,
                   steps::Array{Float64})::Array{Float64}
    """
    calc trailer yaw from x y yaw lists
    """

    tyaw = fill(0.0, length(x))
    tyaw[1] = init_tyaw

    ## for trailer truck
    # for i in 2:length(x)
    #     tyaw[i] += tyaw[i-1] + steps[i-1]/LT*sin(yaw[i-1] - tyaw[i-1])
    # end

    for i in 2:length(x)
        tyaw[i] = yaw[i]# tyaw[i-1] + steps[i-1]/LT*sin(yaw[i-1] - tyaw[i-1])
    end
    return tyaw
end

function trailer_motion_model(x, y, yaw0, yaw1, D, d, L, delta)
    """
    Motion model for trailer
    see:
    http://planning.cs.uiuc.edu/node661.html#77556
    """
    x += D*cos(yaw0)
    y += D*sin(yaw0)
    yaw0 += D/L*tan(delta)
    yaw1 += D/d*sin(yaw0 - yaw1)

    return x, y, yaw0, yaw1
end

function check_trailer_collision(
                    ox::Array{Float64},
                    oy::Array{Float64},
                    x::Array{Float64},
                    y::Array{Float64},
                    yaw0::Array{Float64},
                    yaw1::Array{Float64};
                    kdtree = nothing
                   )
    """
    collision check function for trailer

    """

    if kdtree == nothing
		kdtree = KDTree([ox'; oy'])
    end

    vrxt = [LTF, LTF, -LTB, -LTB, LTF]
    vryt = [-W/2.0, W/2.0, W/2.0, -W/2.0, -W/2.0]

    # bubble parameter
    DT = (LTF + LTB)/2.0 - LTB
    DTR = (LTF + LTB)/2.0 + 0.3

    # check trailer
    if !check_collision(x, y, yaw1, kdtree, ox, oy, DT, DTR, vrxt, vryt)
        return false
    end

    vrxf = [LF, LF, -LB, -LB, LF]
    vryf = [-W/2.0, W/2.0, W/2.0, -W/2.0, -W/2.0]

    # bubble parameter
    DF = (LF + LB)/2.0 - LB
    DFR = (LF + LB)/2.0 + 0.3

    # check front trailer
    if !check_collision(x, y, yaw0, kdtree, ox, oy, DF, DFR, vrxf, vryf)
        return false
    end

    return true #OK
end

function plot_trailer(x::Float64,
                      y::Float64,
                      yaw::Float64,
                      yaw1::Float64,
                      steer::Float64)

    truckcolor = "-k"

    LENGTH = LB+LF
    LENGTHt = LTB+LTF

    truckOutLine = [-LB (LENGTH - LB) (LENGTH - LB) (-LB) (-LB);
                    W/2 W/2 -W/2 -W/2 W / 2]
    trailerOutLine = [-LTB (LENGTHt - LTB) (LENGTHt - LTB) (-LTB) (-LTB);
                    W/2 W/2 -W/2 -W/2 W / 2]

    rr_wheel = [TR -TR -TR TR TR;
                -W/12.0+TW  -W/12.0+TW W/12.0+TW W/12.0+TW -W/12.0+TW]

    rl_wheel = [TR -TR -TR TR TR;
                -W/12.0-TW  -W/12.0-TW W/12.0-TW W/12.0-TW -W/12.0-TW]

    fr_wheel = [TR -TR -TR TR TR
                -W/12.0+TW  -W/12.0+TW W/12.0+TW W/12.0+TW -W/12.0+TW]

    fl_wheel = [TR -TR -TR TR TR;
                -W/12.0-TW  -W/12.0-TW W/12.0-TW W/12.0-TW -W/12.0-TW]
    tr_wheel = [TR -TR -TR TR TR
                -W/12.0+TW  -W/12.0+TW W/12.0+TW W/12.0+TW -W/12.0+TW]

    tl_wheel = [TR -TR -TR TR TR;
                -W/12.0-TW  -W/12.0-TW W/12.0-TW W/12.0-TW -W/12.0-TW]

    Rot1 = [cos(yaw) sin(yaw);
           -sin(yaw) cos(yaw)]
    Rot2 = [cos(steer) sin(steer);
           -sin(steer) cos(steer)]
    Rot3 = [cos(yaw1) sin(yaw1);
           -sin(yaw1) cos(yaw1)]

    fr_wheel = (fr_wheel' * Rot2)'
    fl_wheel = (fl_wheel' * Rot2)'
    fr_wheel[1,:] .+= WB
    fl_wheel[1,:] .+= WB
    fr_wheel = (fr_wheel' * Rot1)'
    fl_wheel = (fl_wheel' * Rot1)'

    tr_wheel[1,:] .-= LT
    tl_wheel[1,:] .-= LT
    tr_wheel = (tr_wheel' * Rot3)'
    tl_wheel = (tl_wheel' * Rot3)'

    truckOutLine = (truckOutLine' * Rot1)'
    trailerOutLine = (trailerOutLine' * Rot3)'
    rr_wheel = (rr_wheel' * Rot1)'
    rl_wheel = (rl_wheel' * Rot1)'

    truckOutLine[1, :] .+= x
    truckOutLine[2, :] .+= y
    trailerOutLine[1, :] .+= x
    trailerOutLine[2, :] .+= y
    fr_wheel[1, :] .+= x
    fr_wheel[2, :] .+= y
    rr_wheel[1, :] .+= x
    rr_wheel[2, :] .+= y
    fl_wheel[1, :] .+= x
    fl_wheel[2, :] .+= y
    rl_wheel[1, :] .+= x
    rl_wheel[2, :] .+= y

    tr_wheel[1, :] .+= x
    tr_wheel[2, :] .+= y
    tl_wheel[1, :] .+= x
    tl_wheel[2, :] .+= y

    plot(truckOutLine[1, :], truckOutLine[2, :], truckcolor)
    plot(trailerOutLine[1, :], trailerOutLine[2, :], truckcolor)
    plot(fr_wheel[1, :], fr_wheel[2, :], truckcolor)
    plot(rr_wheel[1, :], rr_wheel[2, :], truckcolor)
    plot(fl_wheel[1, :], fl_wheel[2, :], truckcolor)
    plot(rl_wheel[1, :], rl_wheel[2, :], truckcolor)

    plot(tr_wheel[1, :], tr_wheel[2, :], truckcolor)
    plot(tl_wheel[1, :], tl_wheel[2, :], truckcolor)
    plot(x, y, "*")
end

function main()
    x = 0.0
    y = 0.0
    yaw0 = deg2rad(-90.0)
    yaw1 = deg2rad(-90.0)

    plot_trailer(x, y, yaw0, yaw1, 0.0)

    DF = (LF + LB)/2.0 - LB
    DFR = (LF + LB)/2.0 + 0.3

    DT = (LTF + LTB)/2.0 - LTB
    DTR = (LTF + LTB)/2.0 + 0.3

    axis("equal")

    show()

end

# @time main()
if length(PROGRAM_FILE)!=0 &&
	occursin(PROGRAM_FILE, @__FILE__)
    @time main()
end

end


function ChiSquareTest()
    return ChiSquareTest(())
end

function chi_square(obj::ChiSquareTest, arg0::Vector{Vector{jlong}})
    return jcall(obj, "chiSquare", jdouble, (Vector{Vector{jlong}},), arg0)
end

function chi_square(obj::ChiSquareTest, arg0::Vector{jdouble}, arg1::Vector{jlong})
    return jcall(obj, "chiSquare", jdouble, (Vector{jdouble}, Vector{jlong}), arg0, arg1)
end

function chi_square_data_sets_comparison(obj::ChiSquareTest, arg0::Vector{jlong}, arg1::Vector{jlong})
    return jcall(obj, "chiSquareDataSetsComparison", jdouble, (Vector{jlong}, Vector{jlong}), arg0, arg1)
end

function chi_square_test(obj::ChiSquareTest, arg0::Vector{Vector{jlong}})
    return jcall(obj, "chiSquareTest", jdouble, (Vector{Vector{jlong}},), arg0)
end

function chi_square_test(obj::ChiSquareTest, arg0::Vector{Vector{jlong}}, arg1::jdouble)
    return jcall(obj, "chiSquareTest", jboolean, (Vector{Vector{jlong}}, jdouble), arg0, arg1)
end

function chi_square_test(obj::ChiSquareTest, arg0::Vector{jdouble}, arg1::Vector{jlong})
    return jcall(obj, "chiSquareTest", jdouble, (Vector{jdouble}, Vector{jlong}), arg0, arg1)
end

function chi_square_test(obj::ChiSquareTest, arg0::Vector{jdouble}, arg1::Vector{jlong}, arg2::jdouble)
    return jcall(obj, "chiSquareTest", jboolean, (Vector{jdouble}, Vector{jlong}, jdouble), arg0, arg1, arg2)
end

function chi_square_test_data_sets_comparison(obj::ChiSquareTest, arg0::Vector{jlong}, arg1::Vector{jlong})
    return jcall(obj, "chiSquareTestDataSetsComparison", jdouble, (Vector{jlong}, Vector{jlong}), arg0, arg1)
end

function chi_square_test_data_sets_comparison(obj::ChiSquareTest, arg0::Vector{jlong}, arg1::Vector{jlong}, arg2::jdouble)
    return jcall(obj, "chiSquareTestDataSetsComparison", jboolean, (Vector{jlong}, Vector{jlong}, jdouble), arg0, arg1, arg2)
end


#!/usr/bin/env julia

using Pkg
Pkg.activate(normpath(joinpath(@__DIR__, ".")))

using ColorSchemes
using Images
using Colors
using EchogramColorSchemes
using TestImages
using ImageQualityIndexes

zx = ColorScheme([colorant"#000000",
                  colorant"#0000ee",
                  colorant"#ee0000",
                  colorant"#ee00ee",
                  colorant"#00ee00",
                  colorant"#00eeee",
                  colorant"#eeee00",
                  colorant"#eeeeee",
                  colorant"#0000ff",
                  colorant"#ff0000",
                  colorant"#ff00ff",
                  colorant"#00ff00",
                  colorant"#00ffff",
                  colorant"#ffff00",
                  colorant"#ffffff" ],
                 "ZX Spectrum",
                 "The colors used on the Sinclair ZX Spectrum computer.")

spectra = ColorScheme(
    [RGB{N0f8}(r, g, b) for r in [0, 1/3 , 2/3 , 1] for g in [0, 1/3 , 2/3 , 1] for b in [0, 1/3 , 2/3 , 1]],
    "Spectra",
    "The colours used on the Sinclair ZX Spectrum SPECTRA interface.")

function show_colorscheme(cs)
    cs = Lab.(cs)
    cs = sort(cs, by=x->x.l)
    img = Array{RGB{N0f8}}(undef, 10, length(cs) *10)
    for i in 1:length(cs)
        for j in (i-1) * 10 + 1:(i-1) * 10 + 10
            for k in 1:10
                img[k,j] = cs[i]
            end
        end
    end
    return img
end

save("zx.png", show_colorscheme(zx))
save("spectra.png", show_colorscheme(spectra))

"""
    rgb_dist(c1, c2)

Return the squared euclidean distance between two colours, c1 and c2
in RGB space.
"""
function rgb_dist(c1, c2)
    rdiff = float(red(c1)) - float(red(c2))
    gdiff = float(green(c1)) - float(green(c2))
    bdiff = float(blue(c1)) - float(blue(c2))
    return rdiff * rdiff + gdiff * gdiff + bdiff * bdiff
end

"""
    nearest(c, cs, f=colordiff)

Return the nearest colour in colorscheme cs to colour c using the
given colour distance metric f (which defaults to CIEDE2000).
"""
function nearest(c, cs, f=colordiff)
    d = f.(c, cs.colors)
    return cs.colors[argmin(d)]
end

"""
    recolor(img, colorscheme::ColorScheme)

Recolour image using colours from colorscheme, returning the
recoloured image. Use the given colour distance metric f (which
defaults to CIEDE2000).
"""
function recolor(img, colorscheme::ColorScheme, f=colordiff)
    map(x -> nearest(x, colorscheme, f), img)
end

function test(name, img, cs=spectra)
    save("$(name).png", img)
    M1 =  colorfulness(img)

    recoloured =  recolor(img, cs, rgb_dist)
    M2 =  colorfulness(recoloured)
    save("$(name)1.png", recoloured)

    recoloured = recolor(img, cs)
    M3 =  colorfulness(recoloured)
    save("$(name)2.png", recoloured )

    println("Colourfulness $name $M1 $M2 $M3")
end

test("sblair", load("blair.png"))

test("smandril", testimage("mandril_color.tif"))

test("slighthouse", testimage("lighthouse.png"))

test("sbarbara", testimage("barbara_color.png"))

test("schelsea", testimage("chelsea.png"))

test("zxblair", load("blair.png"), zx)

test("zxmandril", testimage("mandril_color.tif"), zx)

test("zxlighthouse", testimage("lighthouse.png"), zx)

test("zxbarbara", testimage("barbara_color.png"), zx)

test("zxchelsea", testimage("chelsea.png"), zx)


mutable struct MapSimulateState{T,U}
    score::Float64
    noise::Float64
    subtraces::Vector{U}
    retval::Vector{T}
    num_nonempty::Int
end

function process!(gen_fn::Map{T,U}, args::Tuple,
                  key::Int, state::MapSimulateState{T,U}) where {T,U}
    local subtrace::U
    local retval::T
    kernel_args = get_args_for_key(args, key)
    subtrace = simulate(gen_fn.kernel, kernel_args)
    state.noise += project(subtrace, EmptySelection())
    state.num_nonempty += (isempty(get_choices(subtrace)) ? 0 : 1)
    state.score += get_score(subtrace)
    state.subtraces[key] = subtrace
    retval = get_retval(subtrace)
    state.retval[key] = retval
end

function simulate(gen_fn::Map{T,U}, args::Tuple) where {T,U}
    len = length(args[1])
    state = MapSimulateState{T,U}(0., 0., Vector{U}(undef,len), Vector{T}(undef,len), 0)
    for key=1:len
        process!(gen_fn, args, key, state)
    end
    VectorTrace{MapType,T,U}(gen_fn,
        PersistentVector{U}(state.subtraces), PersistentVector{T}(state.retval),
        args, state.score, state.noise, len, state.num_nonempty)
end


function add_additional_equations(obj::AbstractPropagatorBuilder, arg0::AdditionalEquations)
    return jcall(obj, "addAdditionalEquations", void, (AdditionalEquations,), arg0)
end

function build_propagator(obj::PropagatorBuilder, arg0::Vector{jdouble})
    return jcall(obj, "buildPropagator", Propagator, (Vector{jdouble},), arg0)
end

function equals(obj::Object, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function get_attitude_provider(obj::AbstractPropagatorBuilder)
    return jcall(obj, "getAttitudeProvider", AttitudeProvider, ())
end

function get_class(obj::Object)
    return jcall(obj, "getClass", Class, ())
end

function get_frame(obj::AbstractPropagatorBuilder)
    return jcall(obj, "getFrame", Frame, ())
end

function get_initial_orbit_date(obj::AbstractPropagatorBuilder)
    return jcall(obj, "getInitialOrbitDate", AbsoluteDate, ())
end

function get_mu(obj::AbstractPropagatorBuilder)
    return jcall(obj, "getMu", jdouble, ())
end

function get_orbit_type(obj::AbstractPropagatorBuilder)
    return jcall(obj, "getOrbitType", OrbitType, ())
end

function get_orbital_parameters_drivers(obj::AbstractPropagatorBuilder)
    return jcall(obj, "getOrbitalParametersDrivers", ParameterDriversList, ())
end

function get_position_angle(obj::AbstractPropagatorBuilder)
    return jcall(obj, "getPositionAngle", PositionAngle, ())
end

function get_position_scale(obj::AbstractPropagatorBuilder)
    return jcall(obj, "getPositionScale", jdouble, ())
end

function get_propagation_parameters_drivers(obj::AbstractPropagatorBuilder)
    return jcall(obj, "getPropagationParametersDrivers", ParameterDriversList, ())
end

function get_selected_normalized_parameters(obj::AbstractPropagatorBuilder)
    return jcall(obj, "getSelectedNormalizedParameters", Vector{jdouble}, ())
end

function hash_code(obj::Object)
    return jcall(obj, "hashCode", jint, ())
end

function notify(obj::Object)
    return jcall(obj, "notify", void, ())
end

function notify_all(obj::Object)
    return jcall(obj, "notifyAll", void, ())
end

function reset_orbit(obj::AbstractPropagatorBuilder, arg0::Orbit)
    return jcall(obj, "resetOrbit", void, (Orbit,), arg0)
end

function set_attitude_provider(obj::AbstractPropagatorBuilder, arg0::AttitudeProvider)
    return jcall(obj, "setAttitudeProvider", void, (AttitudeProvider,), arg0)
end

function to_string(obj::Object)
    return jcall(obj, "toString", JString, ())
end

function wait(obj::Object)
    return jcall(obj, "wait", void, ())
end

function wait(obj::Object, arg0::jlong)
    return jcall(obj, "wait", void, (jlong,), arg0)
end

function wait(obj::Object, arg0::jlong, arg1::jint)
    return jcall(obj, "wait", void, (jlong, jint), arg0, arg1)
end


struct Header
    paramptr::UInt8
    npoints::UInt16
    ampf::UInt16
    fframe::UInt16
    lframe::UInt16
    maxinterp::UInt16
    scale::Float32
    datastart::UInt16
    aspf::UInt16
    pointrate::Float32
    res1::Array{UInt8,1}
    labelrange::Union{UInt16,Nothing}
    res2::UInt16
    events::Dict{Symbol,Array{Float32,1}}
    res3::UInt16
end

function readheader(f::IOStream, FEND::Endian, FType::Type{T}) where T <: Union{Float32,VaxFloatF}
    seek(f,0)
    paramptr = read(f, UInt8)
    read(f, Int8)
    npoints = saferead(f, UInt16, FEND)
    ampf = saferead(f, UInt16, FEND)
    fframe = saferead(f, UInt16, FEND)
    lframe = saferead(f, UInt16, FEND)
    maxinterp = saferead(f, UInt16, FEND)
    scale = saferead(f, FType, FEND)
    datastart = saferead(f, UInt16, FEND)
    aspf = saferead(f, UInt16, FEND)
    pointrate = saferead(f, FType, FEND)
    res1 = read(f, 137*2)

    tmp = saferead(f, UInt16, FEND)
    tmp_lr = saferead(f, UInt16, FEND)
    labelrange = (tmp == 0x3039) ? tmp_lr : nothing

    char4 = (saferead(f, UInt16, FEND) == 0x3039)
    skip(f, 2)
    res2 = read(f, UInt16)

    eventtimes = saferead(f, FType, FEND, 18)
    eventflags = read!(f, Array{Bool}(undef, 18))
    validevents = findall(iszero, eventflags)

    res3 = read(f, UInt16)

    tdata = convert.(Char, read!(f, Array{UInt8}(undef, (4, 18))))
    eventlabels = [ Symbol(replace(strip(String(@view(tdata[((i - 1) * 4 + 1):(i * 4)]))),
                                   r"[^a-zA-Z0-9_]" => '_')) for i in 1:18]

    events = Dict{Symbol,Array{Float32,1}}(( ev => Float32[]
                for ev in unique(eventlabels[validevents])))
    for (event, time) in zip(eventlabels[validevents], eventtimes[validevents])
        push!(events[event], time)
    end

    return Header(paramptr, npoints, ampf, fframe, lframe, maxinterp, scale,
                     datastart, aspf, pointrate, res1, labelrange, res2, events, res3)
end


#using Soss
using HypercubeTransform
import Distributions
const Dists = Distributions
import MeasureTheory
const MT = MeasureTheory

using Test

@testset "HypercubeTransform.jl" begin
    # Write your tests here.
    tests = [
        "transform",
        "composite",
        #"soss"
    ]
    res = map(tests) do t
        @eval module $(Symbol("Test_", t))
            include($t*".jl")
        end
        return
    end

end


struct TrialFESpace{B} <: SingleFieldFESpace
  space::SingleFieldFESpace
  dirichlet_values::AbstractVector
  cell_basis::CellField
  constraint_style::Val{B}

  function TrialFESpace(dirichlet_values::AbstractVector,space::SingleFieldFESpace)
    cell_basis = trialize_cell_basis(get_cell_basis(space))
    cs = constraint_style(space)
    B = get_val_parameter(cs)
    new{B}(space,dirichlet_values,cell_basis,cs)
  end
end

"""
"""
function TrialFESpace(space::SingleFieldFESpace)
  dirichlet_values = get_dirichlet_values(space)
  TrialFESpace(dirichlet_values,space)
end

"""
"""
function TrialFESpace(space::SingleFieldFESpace,objects)
  dirichlet_values = compute_dirichlet_values_for_tags(space,objects)
  TrialFESpace(dirichlet_values,space)
end

"""
"""
function TrialFESpace!(dir_values::AbstractVector,space::SingleFieldFESpace,objects)
  dir_values = compute_dirichlet_values_for_tags!(dir_values,space,objects)
  TrialFESpace(dir_values,space)
end

"""
"""
function TrialFESpace!(space::TrialFESpace,objects)
  dir_values = get_dirichlet_values(space)
  dir_values = compute_dirichlet_values_for_tags!(dir_values,space,objects)
  space
end

function TrialFESpace(space::TrialFESpace)
  space
end

function HomogeneousTrialFESpace(U::FESpace)
  dirichlet_values = zero_dirichlet_values(U)
  TrialFESpace(dirichlet_values,U)
end

function HomogeneousTrialFESpace!(dirichlet_values::AbstractVector,U::FESpace)
  fill!(dirichlet_values,zero(eltype(dirichlet_values)))
  TrialFESpace(dirichlet_values,U)
end

# Genuine functions

get_dirichlet_values(f::TrialFESpace) = f.dirichlet_values

get_cell_basis(f::TrialFESpace) = f.cell_basis

# Delegated functions

get_cell_axes(t::TrialFESpace)= get_cell_axes(t.space)

get_cell_axes_with_constraints(t::TrialFESpace)= get_cell_axes_with_constraints(t.space)

CellData.CellField(t::TrialFESpace,cell_vals) = CellField(t.space,cell_vals)

constraint_style(::Type{<:TrialFESpace{B}}) where B = Val{B}()

get_cell_dof_basis(f::TrialFESpace) = get_cell_dof_basis(f.space)

num_free_dofs(f::TrialFESpace) = num_free_dofs(f.space)

zero_free_values(f::TrialFESpace) = zero_free_values(f.space)

get_cell_dofs(f::TrialFESpace) = get_cell_dofs(f.space)

num_dirichlet_dofs(f::TrialFESpace) = num_dirichlet_dofs(f.space)

zero_dirichlet_values(f::TrialFESpace) = zero_dirichlet_values(f.space)

num_dirichlet_tags(f::TrialFESpace) = num_dirichlet_tags(f.space)

get_dirichlet_dof_tag(f::TrialFESpace) = get_dirichlet_dof_tag(f.space)

scatter_free_and_dirichlet_values(f::TrialFESpace,fv,dv) = scatter_free_and_dirichlet_values(f.space,fv,dv)

gather_free_and_dirichlet_values(f::TrialFESpace,cv) = gather_free_and_dirichlet_values(f.space,cv)

gather_free_and_dirichlet_values!(fv,dv,f::TrialFESpace,cv) = gather_free_and_dirichlet_values!(fv,dv,f.space,cv)

gather_dirichlet_values(f::TrialFESpace,cv) = gather_dirichlet_values(f.space,cv)

gather_dirichlet_values!(dv,f::TrialFESpace,cv) = gather_dirichlet_values!(dv,f.space,cv)

gather_free_values(f::TrialFESpace,cv) = gather_free_values(f.space,cv)

gather_free_values!(fv,f::TrialFESpace,cv) = gather_free_values!(fv,f.space,cv)

get_cell_isconstrained(f::TrialFESpace) = get_cell_isconstrained(f.space)

get_cell_constraints(f::TrialFESpace) = get_cell_constraints(f.space)


module BioMakie

export  bondshape,
        resletterdict,
        viewmsa,
		atomcolors,
        atomcoords,
        atomradii,
		viewstruc,
        resatoms,
        resbonds,
        backbonebonds,
        defaultatom,
        defaultresidue,
        covrad,
        vdwrad

using BioStructures, MolecularGraph, MIToS
using Lazy, SplitApplyCombine, TensorCast
using DataStructures, DelimitedFiles, JLD2, JSON3, HTTP
using Distances, Distributions, GeometryBasics, Colors
using FileIO, FastaIO, OrderedCollections
using GraphMakie, JSServe, Meshes
using Makie, GLMakie, WGLMakie
GLMakie.activate!()

include("../src/utils.jl")
include("../src/chemdata.jl")
include("../src/bonds.jl")
include("../src/structure.jl")
include("../src/msa.jl")

end # BioMakie


using DeutschlandDaten
using Test

@testset "DeutschlandDaten.jl" begin

    @test [deaths(deathdat; alter=a, jahr=2020, monat=1, geschlecht="Ins")
           for a in DeutschlandDaten.Altersgruppen_Monate] ==
               [ 272, 329, 248, 374, 486, 856, 1943, 3159, 4414, 5820, 6653, 11018,  16556, 15577+12282+ 4993]

    @test [deaths(deathdat; alter=a, jahr=2018, monat=6, geschlecht="Ins")
           for a in DeutschlandDaten.Altersgruppen_Monate] ==
        [ 322, 345, 210, 317, 418, 873, 1845, 2748, 3863, 4842, 5555, 10125, 12784, 12196+ 9168+ 3717 ]

    @test [deaths(deathdat; alter=a, jahr=2019, monat=3, geschlecht="Männlich")
           for a in DeutschlandDaten.Altersgruppen_Monate] ==
               [ 181, 228, 149, 217, 338, 585, 1295, 2045, 2877, 3640, 3991, 6926, 8337, 6668 + 3916 + 1037 ]

    @test [ deaths(deathdat; alter=a, jahr=2020, kw=3, geschlecht="Ins")
           for a in DeutschlandDaten.Altersgruppen_KW ] ==
               [ 130, 62, 79, 102, 190, 449, 723, 1003, 1259, 1459, 2424, 3701, 3526 + 2746 + 1100 ]

    @test [ deaths(deathdat; alter=a, jahr=2018, kw=6, geschlecht="Männlich")
           for a in DeutschlandDaten.Altersgruppen_KW ] ==
               [ 98, 44, 43, 77, 122, 273, 463, 617, 793, 992, 1738, 1874, 1511 + 862 + 259 ]

    @test [ deaths(deathdat; alter=a, jahr=2017, bundesland="Schleswig-Holstein", kw=1, geschlecht="Männlich")
            for a in DeutschlandDaten.Altersgruppen_Bundesländer ] ==
                [  61, 75, 101, 80  ]

    @test [ deaths(deathdat; alter=a, jahr=2018, bundesland="Bayern", kw=6, geschlecht="Männlich")
           for a in DeutschlandDaten.Altersgruppen_Bundesländer ] ==
               [ 250, 256, 487, 406  ]

    @test [ deaths(deathdat; alter=a, jahr=2021, bundesland="Niedersachsen", kw=6, geschlecht="Weiblich")
           for a in DeutschlandDaten.Altersgruppen_Bundesländer ] ==
               [ 97, 124, 324, 503 ]

    @test [ deaths(deathdat; alter=a, jahr=2016, bundesland="Baden-Württemberg", kw=52, geschlecht="Weiblich")
           for a in DeutschlandDaten.Altersgruppen_Bundesländer ] ==
               [ 94, 124, 350, 595 ]

    # Write your tests here.
end


module TestMOIWrapper

using MathOptInterface
using NLopt
using Test

const MOI = MathOptInterface

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

function test_runtests()
    model = MOI.instantiate(NLopt.Optimizer; with_bridge_type = Float64)
    MOI.set(model, MOI.RawOptimizerAttribute("algorithm"), :LD_SLSQP)
    MOI.Test.runtests(
        model,
        MOI.Test.Config(
            optimal_status = MOI.LOCALLY_SOLVED,
            atol = 1e-2,
            rtol = 1e-2,
            exclude = Any[
                MOI.ConstraintBasisStatus,
                MOI.ConstraintDual,
                MOI.ConstraintName,
                MOI.DualObjectiveValue,
                MOI.ObjectiveBound,
                MOI.NLPBlockDual,
                MOI.SolverVersion,
                MOI.VariableBasisStatus,
                MOI.VariableName,
            ]
        ),
        exclude = String[
            # TODO(odow): investigate failures. A lot of these are probably just
            # suboptimal solutions.
            "test_conic_NormInfinityCone_3",
            "test_conic_NormInfinityCone_INFEASIBLE",
            "test_conic_NormOneCone",
            "test_conic_NormOneCone_INFEASIBLE",
            "test_conic_linear_INFEASIBLE",
            "test_conic_linear_INFEASIBLE_2",
            "test_constraint_qcp_duplicate_diagonal",  # Fails on 32-bit
            "test_infeasible_",
            "test_linear_DUAL_INFEASIBLE",
            "test_linear_INFEASIBLE",
            "test_linear_add_constraints",
            "test_linear_VectorAffineFunction_empty_row",
            "test_model_ScalarAffineFunction_ConstraintName",
            "test_model_duplicate_ScalarAffineFunction_ConstraintName",
            "test_nonlinear_invalid",
            "test_objective_ObjectiveFunction_VariableIndex",
            "test_quadratic_SecondOrderCone_basic",
            "test_quadratic_constraint_integration",
            "test_quadratic_nonconvex_constraint_basic",
            "test_quadratic_nonconvex_constraint_integration",
            "test_unbounded_",
            "test_solve_TerminationStatus_DUAL_INFEASIBLE",
            "test_solve_result_index",
        ]
    )
    return
end

end

TestMOIWrapper.runtests()


mutable struct DropoutDesc
    ptr::Cptr
    states

    function DropoutDesc(droprate::Float64)
        seed = 0
        ref = Ref{Cptr}()
        @cudnn :cudnnCreateDropoutDescriptor (Ptr{Cptr},) ref
        desc = new(ref[])

        h = gethandle()
        ref = Ref{Csize_t}()
        @cudnn :cudnnDropoutGetStatesSize (Cptr,Ptr{Csize_t}) h ref
        states = CuArray{UInt8}(Int(ref[]))
        @cudnn(:cudnnSetDropoutDescriptor,
            (Cptr,Cptr,Cfloat,Cptr,Csize_t,Culonglong),
            desc, h, droprate, states, length(states), seed)
        desc.states = states

        finalizer(desc) do x
            @cudnn :cudnnDestroyDropoutDescriptor (Cptr,) x.ptr
        end
        desc
    end
end

Base.cconvert(::Type{Cptr}, desc::DropoutDesc) = desc.ptr

const DICT_DropoutDesc = Dict()

function dropout(x::CuArray{T,N}, droprate::Float64) where {T,N}
    h = gethandle()
    dropdesc = get!(DICT_DropoutDesc, (h,droprate)) do
        DropoutDesc(droprate)
    end
    xdesc = TensorDesc(x, 4)

    ref = Ref{Csize_t}()
    @cudnn :cudnnDropoutGetReserveSpaceSize (Cptr,Ptr{Csize_t}) xdesc ref
    reserve_space = CuArray{UInt8}(Int(ref[]))

    y = similar(x)
    @cudnn(:cudnnDropoutForward,
        (Cptr,Cptr,Cptr,Cptr,Cptr,Cptr,Cptr,Csize_t),
        h, dropdesc, xdesc, x, xdesc, y, reserve_space, length(reserve_space))
    y, (dropdesc,xdesc,reserve_space)
end

function ∇dropout!(dy::CuArray, dx::CuArray, work)
    dropdesc, xdesc, reserve_space = work
    h = gethandle()
    @cudnn(:cudnnDropoutBackward,
        (Cptr,Cptr,Cptr,Cptr,Cptr,Cptr,Cptr,Csize_t),
        h, dropdesc, xdesc, dy, xdesc, dx, reserve_space, length(reserve_space))
end


###############################################################################
#
#   fq_rel_series.jl : Power series over flint finite fields
#
###############################################################################

export fq_rel_series, FqRelSeriesRing

###############################################################################
#
#   Data type and parent object methods
#
###############################################################################

function O(a::fq_rel_series)
   val = pol_length(a) + valuation(a) - 1
   val < 0 && throw(DomainError())
   z = fq_rel_series(base_ring(a), Array{fq}(0), 0, val, val)
   z.parent = parent(a)
   return z
end

elem_type(::Type{FqRelSeriesRing}) = fq_rel_series

parent_type(::Type{fq_rel_series}) = FqRelSeriesRing

base_ring(R::FqRelSeriesRing) = R.base_ring

var(a::FqRelSeriesRing) = a.S

###############################################################################
#
#   Basic manipulation
#
###############################################################################

max_precision(R::FqRelSeriesRing) = R.prec_max

function normalise(a::fq_rel_series, len::Int)
   ctx = base_ring(a)
   if len > 0
      c = base_ring(a)()
      ccall((:fq_poly_get_coeff, :libflint), Void,
         (Ref{fq}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
                                                         c, a, len - 1, ctx)
   end
   while len > 0 && iszero(c)
      len -= 1
      if len > 0
         ccall((:fq_poly_get_coeff, :libflint), Void,
            (Ref{fq}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
                                                         c, a, len - 1, ctx)
      end
   end
   return len
end

function pol_length(x::fq_rel_series)
   return ccall((:fq_poly_length, :libflint), Int,
                (Ref{fq_rel_series}, Ref{FqFiniteField}), x, base_ring(x))
end

precision(x::fq_rel_series) = x.prec

function polcoeff(x::fq_rel_series, n::Int)
   z = base_ring(x)()
   if n < 0
      return z
   end
   ccall((:fq_poly_get_coeff, :libflint), Void,
         (Ref{fq}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
                                                      z, x, n, base_ring(x))
   return z
end

zero(R::FqRelSeriesRing) = R(0)

one(R::FqRelSeriesRing) = R(1)

function gen(R::FqRelSeriesRing)
   z = fq_rel_series(base_ring(R), [base_ring(R)(1)], 1, max_precision(R) + 1, 1)
   z.parent = R
   return z
end

function deepcopy_internal(a::fq_rel_series, dict::ObjectIdDict)
   z = fq_rel_series(base_ring(a), a)
   z.prec = a.prec
   z.val = a.val
   z.parent = parent(a)
   return z
end

function renormalize!(z::fq_rel_series)
   i = 0
   zlen = pol_length(z)
   zval = valuation(z)
   zprec = precision(z)
   while i < zlen && iszero(polcoeff(z, i))
      i += 1
   end
   z.prec = zprec
   if i == zlen
      z.val = zprec
   else
      z.val = zval + i
      ccall((:fq_poly_shift_right, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
                                                      z, z, i, base_ring(z))
   end
   return nothing
end

###############################################################################
#
#   AbstractString I/O
#
###############################################################################

function show(io::IO, a::FqRelSeriesRing)
   print(io, "Univariate power series ring in ", var(a), " over ")
   show(io, base_ring(a))
end

show_minus_one(::Type{fq_rel_series}) = show_minus_one(fq)

###############################################################################
#
#   Unary operators
#
###############################################################################

function -(x::fq_rel_series)
   z = parent(x)()
   ccall((:fq_poly_neg, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series}, Ref{FqFiniteField}),
               z, x, base_ring(x))
   z.prec = x.prec
   z.val = x.val
   return z
end

###############################################################################
#
#   Binary operators
#
###############################################################################

function +(a::fq_rel_series, b::fq_rel_series)
   check_parent(a, b)
   lena = pol_length(a)
   lenb = pol_length(b)
   prec = min(a.prec, b.prec)
   val = min(a.val, b.val)
   lena = min(lena, prec - a.val)
   lenb = min(lenb, prec - b.val)
   z = parent(a)()
   ctx = base_ring(a)
   if a.val < b.val
      lenz = max(lena, lenb + b.val - a.val)
      ccall((:fq_poly_set_trunc, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            z, b, max(0, lenz - b.val + a.val), ctx)
      ccall((:fq_poly_shift_left, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            z, z, b.val - a.val, ctx)
      ccall((:fq_poly_add_series, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               z, z, a, lenz, ctx)
   elseif b.val < a.val
      lenz = max(lena + a.val - b.val, lenb)
      ccall((:fq_poly_set_trunc, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            z, a, max(0, lenz - a.val + b.val), ctx)
      ccall((:fq_poly_shift_left, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            z, z, a.val - b.val, ctx)
      ccall((:fq_poly_add_series, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               z, z, b, lenz, ctx)
   else
      lenz = max(lena, lenb)
      ccall((:fq_poly_add_series, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               z, a, b, lenz, ctx)
   end
   z.prec = prec
   z.val = val
   renormalize!(z)
   return z
end

function -(a::fq_rel_series, b::fq_rel_series)
   check_parent(a, b)
   lena = pol_length(a)
   lenb = pol_length(b)
   prec = min(a.prec, b.prec)
   val = min(a.val, b.val)
   lena = min(lena, prec - a.val)
   lenb = min(lenb, prec - b.val)
   lenz = max(lena, lenb)
   z = parent(a)()
   ctx = base_ring(a)
   if a.val < b.val
      lenz = max(lena, lenb + b.val - a.val)
      ccall((:fq_poly_set_trunc, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            z, b, max(0, lenz - b.val + a.val), ctx)
      ccall((:fq_poly_shift_left, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            z, z, b.val - a.val, ctx)
      ccall((:fq_poly_neg, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Ref{FqFiniteField}),
            z, z, ctx)
      ccall((:fq_poly_add_series, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               z, z, a, lenz, ctx)
   elseif b.val < a.val
      lenz = max(lena + a.val - b.val, lenb)
      ccall((:fq_poly_set_trunc, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            z, a, max(0, lenz - a.val + b.val), ctx)
      ccall((:fq_poly_shift_left, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            z, z, a.val - b.val, ctx)
      ccall((:fq_poly_sub_series, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               z, z, b, lenz, ctx)
   else
      lenz = max(lena, lenb)
      ccall((:fq_poly_sub_series, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               z, a, b, lenz, ctx)
   end
   z.prec = prec
   z.val = val
   renormalize!(z)
   return z
end

function *(a::fq_rel_series, b::fq_rel_series)
   check_parent(a, b)
   lena = pol_length(a)
   lenb = pol_length(b)
   aval = valuation(a)
   bval = valuation(b)
   prec = min(a.prec - aval, b.prec - bval)
   lena = min(lena, prec)
   lenb = min(lenb, prec)
   z = parent(a)()
   z.val = a.val + b.val
   z.prec = prec + z.val
   if lena == 0 || lenb == 0
      return z
   end
   lenz = min(lena + lenb - 1, prec)
   ccall((:fq_poly_mullow, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               z, a, b, lenz, base_ring(a))
   return z
end

###############################################################################
#
#   Ad hoc binary operators
#
###############################################################################

function *(x::fq, y::fq_rel_series)
   z = parent(y)()
   z.prec = y.prec
   z.val = y.val
   ccall((:fq_poly_scalar_mul_fq, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq}, Ref{FqFiniteField}),
               z, y, x, base_ring(y))
   return z
end

*(x::fq_rel_series, y::fq) = y * x

###############################################################################
#
#   Shifting
#
###############################################################################

function shift_left(x::fq_rel_series, len::Int)
   len < 0 && throw(DomainError())
   xlen = pol_length(x)
   z = fq_rel_series(base_ring(x), x)
   z.prec = x.prec + len
   z.val = x.val + len
   z.parent = parent(x)
   return z
end

function shift_right(x::fq_rel_series, len::Int)
   len < 0 && throw(DomainError())
   xlen = pol_length(x)
   xval = valuation(x)
   z = parent(x)()
   if len >= xlen + xval
      z.prec = max(0, x.prec - len)
      z.val = max(0, x.prec - len)
   else
      z.prec = max(0, x.prec - len)
      z.val = max(0, xval - len)
      zlen = min(xlen + xval - len, xlen)
      ccall((:fq_poly_shift_right, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Int, Ref{FqFiniteField}),
               z, x, xlen - zlen, base_ring(x))
      renormalize!(z)
   end
   return z
end

###############################################################################
#
#   Truncation
#
###############################################################################

function truncate(x::fq_rel_series, prec::Int)
   prec < 0 && throw(DomainError())
   xlen = pol_length(x)
   xprec = precision(x)
   xval = valuation(x)
   if xprec + xval <= prec
      return x
   end
   z = parent(x)()
   z.prec = prec
   if prec <= xval
      z = parent(x)()
      z.val = prec
      z.prec = prec
   else
      z.val = xval
      ccall((:fq_poly_set_trunc, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Int, Ref{FqFiniteField}),
               z, x, min(prec - xval, xlen), base_ring(x))
   end
   return z
end

###############################################################################
#
#   Powering
#
###############################################################################

function ^(a::fq_rel_series, b::Int)
   b < 0 && throw(DomainError())
   if isgen(a)
      z = parent(a)()
      z = setcoeff!(z, 0, base_ring(a)(1))
      z.prec = a.prec + b - 1
      z.val = b
   elseif pol_length(a) == 0
      z = parent(a)()
      z.prec = b*valuation(a)
      z.val = b*valuation(a)
   elseif pol_length(a) == 1
      return parent(a)([polcoeff(a, 0)^b], 1,
                           (b - 1)*valuation(a) + precision(a), b*valuation(a))
   elseif b == 0
      return one(parent(a))
   else
      bit = ~((~UInt(0)) >> 1)
      while (UInt(bit) & b) == 0
         bit >>= 1
      end
      z = a
      bit >>= 1
      while bit != 0
         z = z*z
         if (UInt(bit) & b) != 0
            z *= a
         end
         bit >>= 1
      end
   end
   return z
end

###############################################################################
#
#   Comparison
#
###############################################################################

function ==(x::fq_rel_series, y::fq_rel_series)
   check_parent(x, y)
   prec = min(x.prec, y.prec)
   if prec <= x.val && prec <= y.val
      return true
   end
   if x.val != y.val
      return false
   end
   xlen = normalise(x, min(pol_length(x), prec - x.val))
   ylen = normalise(y, min(pol_length(y), prec - y.val))
   if xlen != ylen
      return false
   end
   return Bool(ccall((:fq_poly_equal_trunc, :libflint), Cint,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Int, Ref{FqFiniteField}),
               x, y, xlen, base_ring(x)))
end

function isequal(x::fq_rel_series, y::fq_rel_series)
   if parent(x) != parent(y)
      return false
   end
   if x.prec != y.prec || x.val != y.val || pol_length(x) != pol_length(y)
      return false
   end
   return Bool(ccall((:fq_poly_equal, :libflint), Cint,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Int, Ref{FqFiniteField}),
               x, y, pol_length(x), base_ring(x)))
end

###############################################################################
#
#   Exact division
#
###############################################################################

function divexact(x::fq_rel_series, y::fq_rel_series)
   check_parent(x, y)
   iszero(y) && throw(DivideError())
   yval = valuation(y)
   xval = valuation(x)
   if yval != 0
      if xval >= yval
         x = shift_right(x, yval)
         y = shift_right(y, yval)
      end
   end
   !isunit(y) && error("Unable to invert power series")
   prec = min(x.prec - x.val, y.prec - y.val)
   z = parent(x)()
   z.val = xval - yval
   z.prec = prec + z.val
   if prec != 0
      ccall((:fq_poly_div_series, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               z, x, y, prec, base_ring(x))
   end
   return z
end

###############################################################################
#
#   Ad hoc exact division
#
###############################################################################

function divexact(x::fq_rel_series, y::fq)
   iszero(y) && throw(DivideError())
   z = parent(x)()
   z.prec = x.prec
   z.prec = x.prec
   z.val = x.val
   ccall((:fq_poly_scalar_div_fq, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq}, Ref{FqFiniteField}),
               z, x, y, base_ring(x))
   return z
end

###############################################################################
#
#   Inversion
#
###############################################################################

function inv(a::fq_rel_series)
   iszero(a) && throw(DivideError())
   !isunit(a) && error("Unable to invert power series")
   ainv = parent(a)()
   ainv.prec = a.prec
   ainv.val = 0
   ccall((:fq_poly_inv_series, :libflint), Void,
         (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               ainv, a, a.prec, base_ring(a))
   return ainv
end

###############################################################################
#
#   Unsafe functions
#
###############################################################################

function zero!(x::fq_rel_series)
  ccall((:fq_poly_zero, :libflint), Void,
                   (Ref{fq_rel_series}, Ref{FqFiniteField}), x, base_ring(x))
  x.prec = parent(x).prec_max
  return x
end

function fit!(z::fq_rel_series, n::Int)
   ccall((:fq_poly_fit_length, :libflint), Void,
         (Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
         z, n, base_ring(z))
   return nothing
end

function setcoeff!(z::fq_rel_series, n::Int, x::fmpz)
   ccall((:fq_poly_set_coeff_fmpz, :libflint), Void,
                (Ref{fq_rel_series}, Int, Ref{fmpz}, Ref{FqFiniteField}),
               z, n, x, base_ring(z))
   return z
end

function setcoeff!(z::fq_rel_series, n::Int, x::fq)
   ccall((:fq_poly_set_coeff, :libflint), Void,
                (Ref{fq_rel_series}, Int, Ref{fq}, Ref{FqFiniteField}),
               z, n, x, base_ring(z))
   return z
end

function mul!(z::fq_rel_series, a::fq_rel_series, b::fq_rel_series)
   lena = pol_length(a)
   lenb = pol_length(b)
   aval = valuation(a)
   bval = valuation(b)
   prec = min(a.prec - aval, b.prec - bval)
   lena = min(lena, prec)
   lenb = min(lenb, prec)
   z.val = a.val + b.val
   z.prec = prec + c.val
   lenz = min(lena + lenb - 1, prec)
   if lena <= 0 || lenb <= 0
      lenz = 0
   end
   ccall((:fq_poly_mullow, :libflint), Void,
         (Ref{fq_rel_series}, Ref{fq_rel_series},
          Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               z, a, b, lenz, base_ring(z))
   return z
end

function addeq!(a::fq_rel_series, b::fq_rel_series)
   lena = pol_length(a)
   lenb = pol_length(b)
   prec = min(a.prec, b.prec)
   val = min(a.val, b.val)
   lena = min(lena, prec - a.val)
   lenb = min(lenb, prec - b.val)
   ctx = base_ring(a)
   if a.val < b.val
      z = fq_rel_series(base_ring(a))
      lenz = max(lena, lenb + b.val - a.val)
      ccall((:fq_poly_set_trunc, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            z, b, max(0, lenz - b.val + a.val), ctx)
      ccall((:fq_poly_shift_left, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            z, z, b.val - a.val, ctx)
      ccall((:fq_poly_add_series, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               a, a, z, lenz, ctx)
   elseif b.val < a.val
      lenz = max(lena + a.val - b.val, lenb)
      ccall((:fq_poly_truncate, :libflint), Void,
            (Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            a, max(0, lenz - a.val + b.val), ctx)
      ccall((:fq_poly_shift_left, :libflint), Void,
            (Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
            a, a, a.val - b.val, ctx)
      ccall((:fq_poly_add_series, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               a, a, b, lenz, ctx)
   else
      lenz = max(lena, lenb)
      ccall((:fq_poly_add_series, :libflint), Void,
                (Ref{fq_rel_series}, Ref{fq_rel_series},
                 Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               a, a, b, lenz, ctx)
   end
   a.prec = prec
   a.val = val
   renormalize!(a)
   return a
end

function add!(c::fq_rel_series, a::fq_rel_series, b::fq_rel_series)
   ctx = base_ring(a)
   lena = length(a)
   lenb = length(b)

   prec = min(a.prec, b.prec)

   lena = min(lena, prec)
   lenb = min(lenb, prec)

   lenc = max(lena, lenb)
   c.prec = prec
   ccall((:fq_poly_add_series, :libflint), Void,
     (Ref{fq_rel_series}, Ref{fq_rel_series}, Ref{fq_rel_series}, Int, Ref{FqFiniteField}),
               c, a, b, lenc, ctx)
   return c
end

###############################################################################
#
#   Promotion rules
#
###############################################################################

promote_rule(::Type{fq_rel_series}, ::Type{T}) where {T <: Integer} = fq_rel_series

promote_rule(::Type{fq_rel_series}, ::Type{fmpz}) = fq_rel_series

promote_rule(::Type{fq_rel_series}, ::Type{fq}) = fq_rel_series

###############################################################################
#
#   Parent object call overload
#
###############################################################################

function (a::FqRelSeriesRing)()
   ctx = base_ring(a)
   z = fq_rel_series(ctx)
   z.prec = a.prec_max
   z.val = a.prec_max
   z.parent = a
   return z
end

function (a::FqRelSeriesRing)(b::Integer)
   ctx = base_ring(a)
   if b == 0
      z = fq_rel_series(ctx)
      z.prec = a.prec_max
   else
      z = fq_rel_series(ctx, [ctx(b)], 1, a.prec_max, 0)
   end
   z.parent = a
   return z
end

function (a::FqRelSeriesRing)(b::fmpz)
   ctx = base_ring(a)
   if iszero(b)
      z = fq_rel_series(ctx)
      z.prec = a.prec_max
   else
      z = fq_rel_series(ctx, [ctx(b)], 1, a.prec_max, 0)
   end
   z.parent = a
   return z
end

function (a::FqRelSeriesRing)(b::fq)
   ctx = base_ring(a)
   if iszero(b)
      z = fq_rel_series(ctx)
      z.prec = a.prec_max
   else
      z = fq_rel_series(ctx, [b], 1, a.prec_max, 0)
   end
   z.parent = a
   return z
end

function (a::FqRelSeriesRing)(b::fq_rel_series)
   parent(b) != a && error("Unable to coerce power series")
   return b
end

function (a::FqRelSeriesRing)(b::Array{fq, 1}, len::Int, prec::Int, val::Int)
   ctx = base_ring(a)
   z = fq_rel_series(ctx, b, len, prec, val)
   z.parent = a
   return z
end


##################################
## Setting up agents properties ##
##################################
"""
`pick_random_node` function is used to get starting and ending node of agents.
Nodes are randomly chosen from set of provided areas.

**Input parameters**
* `OSMmap` : OpenStreetMapX mapData object with road network data
* `rects` : vector of Rect types interpreted as a set of rectangle areas
"""
function pick_random_node(OSMmap::OpenStreetMapX.MapData,
                          rects::Vector{Rect})
    nodes_in_rects = Vector{Int}()
    for rect in rects
        p1 = ENU(LLA(rect.p1[1], rect.p1[2]), OSMmap.bounds)
        p2 = ENU(LLA(rect.p2[1], rect.p2[2]), OSMmap.bounds)
        exE = extrema([p1.east, p2.east])
        exN = extrema([p1.north, p2.north])
        for key in keys(OSMmap.v)
            if (exE[1] <= OSMmap.nodes[key].east <= exE[2] &&
                exN[1] <= OSMmap.nodes[key].north <= exN[2])
                push!(nodes_in_rects, key)
            end
        end
    end
    unique_nodes = unique!(nodes_in_rects)
    return rand(unique_nodes)
end

"""
`get_nodes_set` function is used to get set of nodes within passed area.

**Input parameters**
* `OSMmap` : OpenStreetMapX mapData object with road network data
* `rects` : vector of Rect types interpreted as a set of rectangle areas
"""
function get_nodes_set(OSMmap::OpenStreetMapX.MapData,
                          rects::Vector{Rect})
    nodes_in_rects = Vector{Int}()
    for rect in rects
        p1 = ENU(LLA(rect.p1[1], rect.p1[2]), OSMmap.bounds)
        p2 = ENU(LLA(rect.p2[1], rect.p2[2]), OSMmap.bounds)
        exE = extrema([p1.east, p2.east])
        exN = extrema([p1.north, p2.north])
        for key in keys(OSMmap.v)
            if (exE[1] <= OSMmap.nodes[key].east <= exE[2] &&
                exN[1] <= OSMmap.nodes[key].north <= exN[2])
                push!(nodes_in_rects, key)
            end
        end
    end
    unique_nodes = unique!(nodes_in_rects)
    return unique_nodes
end

"""
`generate_agents` function creating vector of agents and returning travel time
for initial routes travelled with maximal speed

**Input parameters**
* `OSMmap` : OpenStreetMapX MapData object with road network data
* `N` : number of agents to be generated
* `StartArea` : vector of areas from which agents randomly pick starting point
* `EndArea` : vector of areas from which agents randomly pick ending point
* `α` : percentage of smart agents
* `k` : number of fastest routes used in k-path algorithm
* `T` : control variable for k-path algorithm probability distribution
* `AvgStartTime` : average agents start time in minutes
* `k_routes_dict` : dictionary with multiple shortest paths (values) between vertices (keys)
* `seed` : set seed for random function
"""
function generate_agents(OSMmap::OpenStreetMapX.MapData,
                        N::Int,
                        StartArea::Vector{Rect},
                        EndArea::Vector{Rect},
                        α::Float64,
                        k::Int = 3,
                        T::Float64 = 1.0,
                        AvgStartTime::Float64 = 0.0,
                        k_routes_dict::Dict{Tuple{Int,Int},Array{Vector{Int}}}=
                        Dict{Tuple{Int,Int},Array{Vector{Int}}}())
    #Initialize empty working variables
    AgentsArr = Vector{Agent}()
    #Indicate smart agents
    N_int= Int(ceil(N*α))
    smart_ind = [trues(N_int); falses(N-N_int)]
    #Obtaining nodes set for given starting and ending areas
    start_set = get_nodes_set(OSMmap, StartArea)
    end_set = get_nodes_set(OSMmap, EndArea)
    #Generating N agents
    for i in 1:N
        dist = Inf
        start_node = end_node = counter =  0
        init_route = Array{Int64,1}()
        while dist == Inf
            start_node = rand(start_set)
            end_node = rand(end_set)
            init_route, dist = OpenStreetMapX.fastest_route(OSMmap, start_node, end_node)[1:2]
            counter +=1
            counter == 100 && error("Route from starting to ending point can't be calculated.")
        end
        #First edge in vertices notation
        firstEdge = (OSMmap.v[init_route[1]], OSMmap.v[init_route[2]])
        sTime = rand()*AvgStartTime*60*2
        NewAgent = Agent(smart_ind[i], start_node, end_node,
                        init_route, sTime, 0.0, firstEdge,
                        ifelse(sTime == 0.0,true,false))
        push!(AgentsArr, NewAgent)
    end
    #Get averages speeds from base scenario - simulating short term memory
    AverageSpeeds = simulation_run(:base, OSMmap,AgentsArr; track_avg_speeds=true).AvgSpeeds
    #Recalculate agents paths with k-shortest path algorithm and avg speeds
    max_speeds = OpenStreetMapX.get_velocities(OSMmap)
    for a in AgentsArr
        k_shortest_path_rerouting!(OSMmap, k_routes_dict, a, AverageSpeeds,max_speeds, k, T, 0)
    end
    return AgentsArr
end


type DArray{T,N,A} <: AbstractArray{T,N}
    dims::NTuple{N,Int}

    chunks::Array{RemoteRef,N}

    # pmap[i]==p ⇒ processor p has piece i
    pmap::Vector{Int}

    # indexes held by piece i
    indexes::Array{NTuple{N,UnitRange{Int}},N}
    # cuts[d][i] = first index of chunk i in dimension d
    cuts::Vector{Vector{Int}}

    function DArray(dims, chunks, pmap, indexes, cuts)
        # check invariants
        assert(size(chunks) == size(indexes))
        assert(length(chunks) == length(pmap))
        assert(dims == map(last,last(indexes)))
        new(dims, chunks, pmap, indexes, cuts)
    end
end

typealias SubDArray{T,N,D<:DArray} SubArray{T,N,D}
typealias SubOrDArray{T,N}         Union(DArray{T,N}, SubDArray{T,N})

## core constructors ##

# dist == size(chunks)
function DArray(init, dims, procs, dist)
    np = prod(dist)
    procs = procs[1:np]
    idxs, cuts = chunk_idxs([dims...], dist)
    chunks = Array(RemoteRef, dist...)
    for i = 1:np
        chunks[i] = remotecall(procs[i], init, idxs[i])
    end
    p = max(1, localpartindex(procs))
    A = remotecall_fetch(procs[p], r->typeof(fetch(r)), chunks[p])
    DArray{eltype(A),length(dims),A}(dims, chunks, procs, idxs, cuts)
end

function DArray(init, dims, procs)
    if isempty(procs)
        error("no processors")
    end
    DArray(init, dims, procs, defaultdist(dims,procs))
end
DArray(init, dims) = DArray(init, dims, workers()[1:min(nworkers(),maximum(dims))])

# new DArray similar to an existing one
DArray(init, d::DArray) = DArray(init, size(d), procs(d), [size(d.chunks)...])

similar(d::DArray, T, dims::Dims)= DArray(I->Array(T, map(length,I)), dims, procs(d))
similar(d::DArray, T)= similar(d, T, size(d))
similar{T}(d::DArray{T}, dims::Dims)= similar(d, T, dims)
similar{T}(d::DArray{T})= similar(d, T, size(d))

size(d::DArray) = d.dims
procs(d::DArray) = d.pmap

chunktype{T,N,A}(d::DArray{T,N,A}) = A

## chunk index utilities ##

# decide how to divide each dimension
# returns size of chunks array
function defaultdist(dims, procs)
    dims = [dims...]
    chunks = ones(Int, length(dims))
    np = length(procs)
    f = sort!(collect(keys(factor(np))), rev=true)
    k = 1
    while np > 1
        # repeatedly allocate largest factor to largest dim
        if np%f[k] != 0
            k += 1
            if k > length(f)
                break
            end
        end
        fac = f[k]
        (d, dno) = findmax(dims)
        # resolve ties to highest dim
        dno = last(find(dims .== d))
        if dims[dno] >= fac
            dims[dno] = div(dims[dno], fac)
            chunks[dno] *= fac
        end
        np = div(np,fac)
    end
    chunks
end

# get array of start indexes for dividing sz into nc chunks
function defaultdist(sz::Int, nc::Int)
    if sz >= nc
        iround(linspace(1, sz+1, nc+1))
    else
        [[1:(sz+1)], zeros(Int, nc-sz)]
    end
end

# compute indexes array for dividing dims into chunks
function chunk_idxs(dims, chunks)
    cuts = map(defaultdist, dims, chunks)
    n = length(dims)
    idxs = Array(NTuple{n,UnitRange{Int}},chunks...)
    cartesianmap(tuple(chunks...)) do cidx...
        idxs[cidx...] = ntuple(n, i->(cuts[i][cidx[i]]:cuts[i][cidx[i]+1]-1))
    end
    idxs, cuts
end

function localpartindex(pmap::Vector{Int})
    mi = myid()
    for i = 1:length(pmap)
        if pmap[i] == mi
            return i
        end
    end
    return 0
end

localpartindex(d::DArray) = localpartindex(d.pmap)

function localpart{T,N,A}(d::DArray{T,N,A})
    lpidx = localpartindex(d)
    if lpidx == 0
        convert(A, Array(T, ntuple(N,i->0)))::A
    else
        fetch(d.chunks[lpidx])::A
    end
end
function localindexes(d::DArray)
    lpidx = localpartindex(d)
    if lpidx == 0
        ntuple(ndims(d), i->1:0)
    else
        d.indexes[lpidx]
    end
end

# find which piece holds index (I...)
function locate(d::DArray, I::Int...)
    ntuple(ndims(d), i->searchsortedlast(d.cuts[i], I[i]))
end

chunk{T,N,A}(d::DArray{T,N,A}, i...) = fetch(d.chunks[i...])::A

## convenience constructors ##

dzeros(args...) = DArray(I->zeros(map(length,I)), args...)
dzeros(d::Int...) = dzeros(d)
dones(args...) = DArray(I->ones(map(length,I)), args...)
dones(d::Int...) = dones(d)
dfill(v, args...) = DArray(I->fill(v, map(length,I)), args...)
dfill(v, d::Int...) = dfill(v, d)
drand(args...)  = DArray(I->rand(map(length,I)), args...)
drand(d::Int...) = drand(d)
drandn(args...) = DArray(I->randn(map(length,I)), args...)
drandn(d::Int...) = drandn(d)

## conversions ##

function distribute(a::AbstractArray)
    owner = myid()
    rr = RemoteRef()
    put!(rr, a)
    DArray(size(a)) do I
        remotecall_fetch(owner, ()->fetch(rr)[I...])
    end
end

function convert{S,T,N}(::Type{Array{S,N}}, d::DArray{T,N})
    a = Array(S, size(d))
    @sync begin
        for i = 1:length(d.chunks)
            @async a[d.indexes[i]...] = chunk(d, i)
        end
    end
    a
end

function convert{S,T,N}(::Type{Array{S,N}}, s::SubDArray{T,N})
    I = s.indexes
    d = s.parent
    if isa(I,(UnitRange{Int}...)) && S<:T && T<:S
        l = locate(d, map(first, I)...)
        if isequal(d.indexes[l...], I)
            # SubDArray corresponds to a chunk
            return chunk(d, l...)
        end
    end
    a = Array(S, size(s))
    a[[1:size(a,i) for i=1:N]...] = s
    a
end

function reshape{T,S<:Array}(A::DArray{T,1,S}, d::Dims)
    if prod(d) != length(A)
        error("dimensions must be consistent with array size")
    end
    DArray(d) do I
        sz = map(length,I)
        d1offs = first(I[1])
        nd = length(I)

        B = Array(T,sz)
        nr = size(B,1)
        sztail = size(B)[2:end]

        for i=1:div(length(B),nr)
            i2 = ind2sub(sztail, i)
            globalidx = [ I[j][i2[j-1]] for j=2:nd ]

            a = sub2ind(d, d1offs, globalidx...)

            B[:,i] = A[a:(a+nr-1)]
        end
        B
    end
end

## indexing ##

function getindex(r::RemoteRef, args...)
    if r.where==myid()
        getindex(fetch(r), args...)
    else
        remotecall_fetch(r.where, getindex, r, args...)
    end
end

getindex(d::DArray, i::Int) = getindex(d, ind2sub(size(d), i))
getindex(d::DArray, i::Int...) = getindex(d, sub2ind(size(d), i...))

function getindex{T}(d::DArray{T}, I::(Int...))
    chidx = locate(d, I...)
    chunk = d.chunks[chidx...]
    idxs = d.indexes[chidx...]
    localidx = ntuple(ndims(d), i->(I[i]-first(idxs[i])+1))
    chunk[localidx...]::T
end

getindex(d::DArray) = d[1]
getindex(d::DArray, I::Union(Int,UnitRange{Int})...) = sub(d,I)

copy(d::SubOrDArray) = d

# local copies are obtained by convert(Array, ) or assigning from
# a SubDArray to a local Array.

function setindex!(a::Array, d::DArray, I::UnitRange{Int}...)
    n = length(I)
    @sync begin
        for i = 1:length(d.chunks)
            K = d.indexes[i]
            @async a[[I[j][K[j]] for j=1:n]...] = chunk(d, i)
        end
    end
    a
end

function setindex!(a::Array, s::SubDArray, I::UnitRange{Int}...)
    n = length(I)
    d = s.parent
    J = s.indexes
    if length(J) < n
        a[I...] = convert(Array,s)
        return a
    end
    offs = [isa(J[i],Int) ? J[i]-1 : first(J[i])-1 for i=1:n]
    @sync begin
        for i = 1:length(d.chunks)
            K_c = {d.indexes[i]...}
            K = [ intersect(J[j],K_c[j]) for j=1:n ]
            if !any(isempty, K)
                idxs = [ I[j][K[j]-offs[j]] for j=1:n ]
                if isequal(K, K_c)
                    # whole chunk
                    @async a[idxs...] = chunk(d, i)
                else
                    # partial chunk
                    ch = d.chunks[i]
                    @async a[idxs...] = remotecall_fetch(ch.where, ()->sub(fetch(ch), [K[j]-first(K_c[j])+1 for j=1:n]...))
                end
            end
        end
    end
    a
end

# to disambiguate
setindex!(a::Array{Any}, d::SubOrDArray, i::Int) = arrayset(a, d, i)

setindex!(a::Array, d::SubOrDArray, I::Union(Int,UnitRange{Int})...) =
    setindex!(a, d, [isa(i,Int) ? (i:i) : i for i in I ]...)

## higher-order functions ##

map(f::Callable, d::DArray) = DArray(I->map(f, localpart(d)), d)

reduce(f::Function, d::DArray) =
    mapreduce(fetch, f,
              { @spawnat p reduce(f, localpart(d)) for p in procs(d) })

function map!(f::Callable, d::DArray)
    @sync begin
        for p in procs(d)
            @spawnat p map!(f, localpart(d))
        end
    end
end