################################## Platform ###########################################
"""
    struct Platform{M,C}
      models::Dict{Symbol,M}     # dictionary storing Models
      scenarios::Dict{Symbol,C} # dictionary storing Scenarios
    end

The main storage representing a modeling platform.
Typically HetaSimulator works with one platform object which can include several models and scenarios.

Usually a `Platform` is created based on Heta formatted files using [`load_platform`]{@ref}.

To get the platform content use methods: `models(platform)`, `scenarios(platform).
"""
struct Platform{M,C}
  models::OrderedDict{Symbol,M}
  scenarios::OrderedDict{Symbol,C}
end

models(p::Platform) = p.models
scenarios(p::Platform) = p.scenarios

function Base.show(io::IO, mime::MIME"text/plain", p::Platform)
  models_names = join(keys(p.models), ", ")
  scn_names = join(keys(p.scenarios), ", ")

  measurements_count = 0
  for (x, y) in scenarios(p)
    measurements_count += length(measurements(y))
  end

  println(io, "Platform with $(length(models(p))) model(s), $(length(scenarios(p))) scenario(s), $measurements_count measurement(s)")
  println(io, "   Models: $models_names")
  println(io, "   Scenarios: $scn_names")
end

################################## Model ###########################################

abstract type AbstractModel end

"""
    struct Model{IF,OF,EV,SG,C,EA} <: AbstractModel
      init_func::IF
      ode_func::OF
      events::EV
      saving_generator::SG
      records_output::AbstractVector{Pair{Symbol,Bool}}
      constants::C
      events_active::EA
    end

Structure storing core properties of ODE model.
This represent the content of one namespace from a Heta platform.

To get list of model content use methods: constants(model), records(model), switchers(model).

To get the default model options use methods: `parameters(model)`,
`events_active(model)`, `events_save(model)`, `observables(model)`.
These options can be rewritten by a [`Scenario`]{@ref}.
"""
struct Model{IF,OF,EV,SG,C,EA} <: AbstractModel
  init_func::IF
  ode_func::OF
  events::EV # IDEA: use (:TimeEvent, ...) instead of TimeEvent(...)
  saving_generator::SG
  records_output::AbstractVector{Pair{Symbol,Bool}}
  constants::C # LArray{Float64,1,Array{Float64,1},(:a, :b)}
  events_active::EA
end

constants(m::Model) = [keys(m.constants)...]
records(m::Model) = first.(m.records_output)
switchers(m::Model) = [keys(m.events)...]

#parameters(m::Model) = collect(Pair{Symbol, Real}, pairs(m.constants))
events_active(m::Model) = collect(Pair{Symbol, Bool}, pairs(m.events_active))
events_save(m::Model) = [first(x) => (true,true) for x in pairs(m.events)]
observables(m::Model) = begin # observables
  only_true = filter((p) -> last(p), m.records_output)
  first.(only_true)
end

# auxilary function to display first n components of vector
function print_lim(x::Union{Vector, Tuple}, n::Int)
  first_n = ["$y" for y in first(x, n)]
  if length(x) > n
    push!(first_n, "...")
  end
  return join(first_n, ", ")
end
function print_lim(::Nothing, n::Int)
  return "-"
end

function Base.show(io::IO, mime::MIME"text/plain", m::AbstractModel)
  const_str = print_lim(constants(m), 10)
  record_str = print_lim(records(m), 10)
  switchers_str = print_lim(switchers(m), 10)

  println(io, "Model contains $(length(m.constants)) constant(s), $(length(m.records_output)) record(s), $(length(m.events)) switcher(s).")
  println(io, "   Constants: $const_str")
  println(io, "   Records: $record_str")
  println(io, "   Switchers (events): $switchers_str")
end

################################## Params ###########################################

struct Params{C,S}
  constants::C
  static::S
end

################################## Measurement ###########################################

# abstract type AbstractMeasurement end # not used
abstract type AbstractMeasurementPoint end

struct NormalMeasurementPoint{M,SD} <: AbstractMeasurementPoint
  t::Float64
  val::Float64
  scope::Symbol
  μ::M
  σ::SD
end

NormalMeasurementPoint(t,val,scope::Missing,μ,σ) = NormalMeasurementPoint(t,val,:ode_,μ,σ)

struct LogNormalMeasurementPoint{M,SD} <: AbstractMeasurementPoint
  t::Float64
  val::Float64
  scope::Symbol
  μ::M
  σ::SD
end

LogNormalMeasurementPoint(t,val,scope::Missing,μ,σ) = LogNormalMeasurementPoint(t,val,:ode_,μ,σ)

const MeasurementVector{P} = AbstractVector{P} where P<:AbstractMeasurementPoint

################################## Scenario ###########################################
abstract type AbstractScenario end

"""
  struct Scenario{F,P,M} <: AbstractScenario
    init_func::F
    prob::P
    measurements::M
  end

  Type representing simulation conditions, i.e. model variant with updated constants and outputs.

  To get the internal properties use methods: `saveat(scenario)`, `tspan(scenario)`, `parameters(scenario)`, `measurements(scenario)`
"""
struct Scenario{F,P,M} <: AbstractScenario
  init_func::F
  prob::P
  measurements::M
end

saveat(scn::Scenario) = scn.prob.kwargs[:callback].discrete_callbacks[1].affect!.saveat_cache
tspan(scn::Scenario) = scn.prob.tspan
parameters(scn::Scenario) = scn.prob.p.constants
measurements(scn::Scenario) = scn.measurements
# events_active(scn::Scenario)
# events_save(scn::Scenario)
# observables(scn::Scenario)

function Base.show(io::IO, mime::MIME"text/plain", scn::Scenario)
  saveat_str = print_lim(saveat(scn), 10)
  parameters_str = print_lim(
    keys(parameters(scn)),
    10
    )
  measurements_count = length(measurements(scn))

  if length(saveat(scn)) == 0
    time_points_str = "for tspan=$(tspan(scn))"
  else
    time_points_str = "for saveat=$(saveat_str)"
  end

  println(io, "Scenario $time_points_str")
  println(io, "   Time range (tspan): $(tspan(scn))")
  println(io, "   Time points to save the solution at (saveat): $(saveat_str)")
  println(io, "   Parameters: $(parameters_str)")
  println(io, "   Number of measurement points: $(measurements_count)")
end

################################## SimResult ###########################################
abstract type AbstractResult end

struct SavedValues{uType,tType,scopeType}
  u::uType
  t::tType
  scope::scopeType
end

function clear_savings(sv::SavedValues)
  !isempty(sv.u) && deleteat!(sv.u,1:length(sv.u))
  !isempty(sv.t) && deleteat!(sv.t,1:length(sv.t))
  !isnothing(sv.scope) && !isempty(sv.scope) && deleteat!(sv.scope,1:length(sv.scope))
  return nothing
end

struct Simulation{V,scopeType,P}
  vals::V
  scope::scopeType
  params::P
  status::Symbol
end

# copy fix is tmp needed not to rewrite SavedValues with new simulation
Simulation(sv::SavedValues, params, status::Symbol) = Simulation(
  DiffEqBase.SciMLBase.DiffEqArray(copy(sv.u),copy(sv.t)),
  sv.scope,
  params,
  status
)

status(s::Simulation) = s.status
times(s::Simulation) = s.vals.t
vals(s::Simulation) = s.vals.u
parameters(s::Simulation) = s.params

@inline Base.length(S::Simulation) = length(S.vals.t)

# tmp fix to support https://github.com/SciML/SciMLBase.jl/blob/ccaba96f4d7e29e9980cd4cd7270086fc5e542d6/src/ensemble/ensemble_analysis.jl#L56
function Base.getproperty(s::Simulation, sym::Symbol)
  if sym == :t
    return s.vals.t
  else
    return getfield(s, sym)
  end
end

"""
    struct SimResult{S, C<:Scenario} <: AbstractResult
      sim::S
      scenario::C
    end

Structure storing results from [`sim`]{@ref} method applied for one [`Scenario`]{@ref}.

To get the content use methods: `status(results)`, `times(results)`, `vals(results)`, `parameters(results)`, observables(results).

The results can be transformed using `DataFrame` method or visualized using `plot` method.
"""
struct SimResult{S, C<:Scenario} <: AbstractResult
  sim::S
  scenario::C
end

status(sr::SimResult) = status(sr.sim)
times(sr::SimResult) = times(sr.sim)
vals(sr::SimResult) = vals(sr.sim)
parameters(sr::SimResult) = parameters(sr.sim)
measurements(sr::SimResult) = sr.scenario.measurements

@inline Base.length(sr::SimResult) = length(sr.sim)

function Base.show(io::IO, mime::MIME"text/plain", sr::SimResult)
  dim2 = length(keys(sr.sim[1])) # number of observables
  dimentions_str = "$(length(sr))x$dim2"
  times_str = print_lim(times(sr), 10)
  outputs_str = print_lim(observables(sr), 10)
  parameters_str = print_lim(keys(parameters(sr)), 10)

  println(io, "$dimentions_str SimResult with status :$(status(sr)).")
  println(io, "    Solution status: $(status(sr))")
  println(io, "    Time points (times): $(times_str)")
  println(io, "    Observables (outputs): $(outputs_str)")
  println(io, "    Parameters: $(parameters_str)")
end

function Base.show(io::IO, mime::MIME"text/plain", srp::Pair{Symbol, S}, short::Bool = false) where S<:SimResult
  sr = last(srp)
  dim2 = length(keys(sr.sim[1])) # number of observables
  dimentions_str = "$(length(sr))x$dim2"

  short || println(io, "Pair{Symbol, SimResult}")
  println(io, "    :$(first(srp)) => $dimentions_str SimResult with status :$(status(sr)).")
end

function Base.show(io::IO, mime::MIME"text/plain", vector::Vector{Pair{Symbol, S}}) where S<:SimResult
  println(io, "$(length(vector))-element Vector{Pair{Symbol, SimResult}}")

  for x in vector
    show(io, mime, x, true)
  end
end

function Base.getindex(vector::Vector{Pair{Symbol, S}}, id::Symbol) where S<:SimResult
  ind = findfirst((x) -> first(x)===id, vector)
  if ind === nothing
    throw("Index :$id is not found.")
  else
    return vector[ind]
  end
end

################################## Monte-Carlo Simulation ##############################

"""
    struct MCResult{S,C} <: AbstractResult
      sim::S
      saveat::Bool
      scenario::C
    end

Structure storing results of [`mc`]{@ref} method applied for one `Scenario`.

To convert into tabular format, use `DataFrame` method.

The base visulaization can be done with `plot` method.
"""
struct MCResult{S,C} <: AbstractResult
  sim::S
  saveat::Bool
  scenario::C
  # converged
  # elapsed_time
end

@inline Base.length(mcr::MCResult) = length(mcr.sim)
parameters(mcr::MCResult) = [parameters(mcr[i]) for i in 1:length(mcr)]
vals(mcr::MCResult) = [vals(mcr[i]) for i in 1:length(mcr)]
status_summary(mcr::MCResult) = counter([s.status for s in mcr.sim])

function Base.show(io::IO, mime::MIME"text/plain", mcr::MCResult, short::Bool = false)
  # dimentions
  dim0 = length(mcr)
  if mcr.saveat
    dim1 = length(times(mcr.sim[1]))
  else
    dim1 = "?"
  end
  dim2 = length(mcr.sim[1]) # test only first Simulation
  dimentions_str = "$(dim0)x$(dim1)x$(dim2)"

  # other
  status_pairs = [":$(first(x)) x $(last(x))" for x in status_summary(mcr)]
  status = join(status_pairs, ", ")
  outputs_str = print_lim(observables(mcr.sim[1]), 10)
  parameters_1 = parameters(mcr.sim[1])
  parameters_str = isnothing(parameters_1) ? "-" : print_lim(keys(parameters_1), 10)

  println(io, "$(dimentions_str) MCResult with status $(status)" )
  if !short
  println(io, "    Solution status: $(status)")
  println(io, "    Observables (outputs): $(outputs_str)")
  println(io, "    Parameters: $(parameters_str)")
  end
end

function Base.show(io::IO, mime::MIME"text/plain", mcrp::Pair{Symbol, S}, short=false) where S<:MCResult
  short || println(io, "Pair{Symbol, MCResult}")
  print(io, "    :$(first(mcrp)) => ")
  Base.show(io, mime, last(mcrp), true)
end

function Base.show(io::IO, mime::MIME"text/plain", vector::Vector{Pair{Symbol, S}}) where S<:MCResult
  println(io, "$(length(vector))-element Vector{Pair{Symbol, MCResult}}")

  for x in vector
    show(io, mime, x, true)
  end
end

function Base.getindex(V::Vector{Pair{Symbol, S}}, id::Symbol) where S<:MCResult
  ind = findfirst((x) -> first(x)===id, V)
  if ind === nothing
    throw("Index :$id is not found.")
  else
    return V[ind]
  end
end

################################## Events ##############################################
abstract type AbstractEvent end

struct TimeEvent{F1,F2} <: AbstractEvent
  condition_func::F1
  affect_func::F2
  atStart::Bool
end

struct CEvent{F1,F2} <: AbstractEvent
  condition_func::F1
  affect_func::F2
  atStart::Bool
end

struct StopEvent{F1} <: AbstractEvent
  condition_func::F1
  atStart::Bool
end

################################## Fitting ###########################################
"""
    struct FitResult{L<:Real, I}
      obj::L
      optim::Vector{Pair{Symbol,Float64}}
      status::Symbol
      numevals::I
    end

Result of [`fit`](@ref).

Use `optim` method to get optimal values.

Use `object` method to get optimal objective function.

Use `status` to get status.

The optimal parameters can be saved in heta file, see [`save_as_heta`](@ref) method.
"""
struct FitResult{L<:Real, I}
  obj::L
  optim::Vector{Pair{Symbol,Float64}}
  status::Symbol
  numevals::I
end

function Base.show(io::IO, mime::MIME"text/plain", fr::FitResult)

  println(io, "FitResult with status :$(fr.status)")
  println(io, "   Status: $(fr.status)")
  println(io, "   Optimal values: $(fr.optim)")
  println(io, "   OF value: $(fr.obj)")
  println(io, "   OF count: $(fr.numevals)")
end

optim(f::FitResult) = f.optim
obj(f::FitResult) = f.obj
status(f::FitResult) = f.status


using Plots
using Random
Random.seed!(1)

include_ddp()

contact_control_path = "/home/taylor/Research/ContactControl.jl/src"

using Parameters

# Utilities
include(joinpath(contact_control_path, "utils.jl"))

# Solver
include(joinpath(contact_control_path, "solver/cones.jl"))
include(joinpath(contact_control_path, "solver/interior_point.jl"))
# include(joinpath(contact_control_path, "solver/lu.jl"))

# Environment
include(joinpath(contact_control_path, "simulator/environment.jl"))
include(joinpath(contact_control_path, "simulation/environments/flat.jl"))

# Dynamics
include(joinpath(contact_control_path, "dynamics/model.jl"))

# Simulator
include(joinpath(contact_control_path, "simulation/contact_methods.jl"))
include(joinpath(contact_control_path, "simulation/simulation.jl"))
include(joinpath(contact_control_path, "simulator/trajectory.jl"))

include(joinpath(contact_control_path, "dynamics/code_gen_dynamics.jl"))
include(joinpath(contact_control_path, "dynamics/fast_methods_dynamics.jl"))

# Models
include(joinpath(contact_control_path, "dynamics/quaternions.jl"))
include(joinpath(contact_control_path, "dynamics/mrp.jl"))
include(joinpath(contact_control_path, "dynamics/euler.jl"))

# include("dynamics/particle_2D/model.jl")
# include("dynamics/particle/model.jl")
# include(joinpath(contact_control_path, "dynamics/cartpole/model.jl"))
include(joinpath(contact_control_path, "dynamics/cartpole/model_soc.jl"))

include(joinpath(contact_control_path, "dynamics/cartpole/visuals.jl"))

# include("dynamics/hopper_3D/model.jl")
# include("dynamics/hopper_3D_quaternion/model.jl")
# include("dynamics/quadruped/model.jl")
# include("dynamics/quadruped_simple/model.jl")
# include("dynamics/biped/model.jl")
# include("dynamics/flamingo/model.jl")
# include("dynamics/pushbot/model.jl")
# include("dynamics/planarpush/model.jl")
# include("dynamics/planarpush_2D/model.jl")
# include("dynamics/rigidbody/model.jl")
# include("dynamics/box/model.jl")

# Simulation
include(joinpath(contact_control_path, "simulation/environments/flat.jl"))
# include("simulation/environments/piecewise.jl")
# include("simulation/environments/quadratic.jl")
# include("simulation/environments/slope.jl")
# include("simulation/environments/sinusoidal.jl")
# include("simulation/environments/stairs.jl")

include(joinpath(contact_control_path, "simulation/residual_approx.jl"))
include(joinpath(contact_control_path, "simulation/code_gen_simulation.jl"))
#
# # Visuals
# using MeshCatMechanisms
# include(joinpath(contact_control_path, "dynamics/visuals.jl"))
# include(joinpath(contact_control_path, "dynamics/visual_utils.jl"))

# s = get_simulation("hopper_2D", "flat_2D_lc", "flat")

nq = s.model.dim.q
m = s.model.dim.u

T = 26
h = 0.1

q0 = [0.0; 0.0]
q1 = [0.0; 0.0]
qT = [0.0; π]
q_ref = [0.0; π]

x1 = [q1; q1]
xT = [qT; qT]

struct Dynamics{T}
	s::Simulation
	ip_dyn::InteriorPoint
	ip_jac::InteriorPoint
	h::T
end

function gen_dynamics(s::Simulation, h;
		dyn_opts =  InteriorPointOptions{Float64}(
						r_tol = 1.0e-8,
						κ_tol = 1.0e-4,
						κ_init = 0.1,
						diff_sol = true),
		jac_opts =  InteriorPointOptions{Float64}(
						r_tol = 1.0e-8,
						κ_tol = 1.0e-2,
						κ_init = 0.1,
						diff_sol = true))

	z = zeros(nz)
	θ = zeros(nθ)

	ip_dyn = interior_point(z, θ,
		idx_ineq = idx_ineq,
		# idx_soc = idx_soc,
		r! = s.res.r!,
		rz! = s.res.rz!,
		rθ! = s.res.rθ!,
		rz = s.rz,
		rθ = s.rθ,
		opts = dyn_opts)

	ip_dyn.opts.diff_sol = false

	ip_jac = interior_point(z, θ,
		idx_ineq = idx_ineq,
		# idx_soc = idx_soc,
		r! = s.res.r!,
		rz! = s.res.rz!,
		rθ! = s.res.rθ!,
		rz = s.rz,
		rθ = s.rθ,
		opts = jac_opts)

	ip_jac.opts.diff_sol = true

	Dynamics(s, ip_dyn, ip_jac, h)
end

d = gen_dynamics(s, h,
	dyn_opts = InteriorPointOptions{Float64}(κ_tol = 1.0e-4, κ_init = 0.1),
	jac_opts = InteriorPointOptions{Float64}(κ_tol = 1.0e-2, κ_init = 0.1))

function f!(d::Dynamics, q0, q1, u1, mode = :dynamics)
	s = d.s
	ip = (mode == :dynamics ? d.ip_dyn : d.ip_jac)
	h = d.h

	# ip.z .= copy([q1; 1.0; 0.1; 1.0; 0.1; 1.0; 0.1; 1.0; 0.1])
	ip.z .= copy(q1)
	ip.θ .= copy([q0; q1; u1; h])

	status = interior_point_solve!(ip)

	!status && (@warn "dynamics failure (res norm: $(norm(ip.r, Inf)))")
end

function f(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :dynamics)
	return copy(d.ip_dyn.z[1:d.s.model.dim.q])
end

f(d, q0, q1, zeros(m))

function fq0(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 1:d.s.model.dim.q])
end

fq0(d, q0, q1, zeros(m))

function fq1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, d.s.model.dim.q .+ (1:d.s.model.dim.q)])
end

fq1(d, q0, q1, zeros(m))

function fx1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 1:(2 * d.s.model.dim.q)])
end

fx1(d, q0, q1, zeros(m))

function fu1(d::Dynamics, q0, q1, u1)
	f!(d, q0, q1, u1, :jacobian)
	return copy(d.ip_jac.δz[1:d.s.model.dim.q, 2 * d.s.model.dim.q .+ (1:d.s.model.dim.u)])
end

fu1(d, q0, q1, zeros(m))

struct CartpoleI{I, T} <: Model{I, T}
    n::Int
    m::Int
    d::Int
	dynamics::Dynamics
end

model = CartpoleI{Midpoint, FixedTime}(2 * s.model.dim.q, s.model.dim.u, 0, d)

function fd(model::CartpoleI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]

	q2 = f(model.dynamics, q0, q1, u)

	return [q1; q2]
end

fd(model, [q0; q1], zeros(model.m), zeros(model.d), h, 1)

function fdx(model::CartpoleI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]
	dq2dx1 = fx1(model.dynamics, q0, q1, u)

	return [zeros(nq, nq) I; dq2dx1]
end

fdx(model, [q0; q1], zeros(model.m), zeros(model.d), h, 1)

function fdu(model::CartpoleI{Midpoint, FixedTime}, x, u, w, h, t)
	nq = model.dynamics.s.model.dim.q
	q0 = x[1:nq]
	q1 = x[nq .+ (1:nq)]
	dq2du1 = fu1(model.dynamics, q0, q1, u)
	return [zeros(nq, model.m); dq2du1]
end

fdu(model, [q0; q1], zeros(model.m), zeros(model.d), h, 1)

n = model.n
m = model.m

# Objective
V = 1.0 * Diagonal(ones(s.model.dim.q))
Q_velocity = [V -V; -V V] ./ h^2.0
Q_track = 1.0 * Diagonal(ones(2 * s.model.dim.q))

Q = [t < T ? Q_velocity + Q_track : Q_velocity + 1.0 * Q_track for t = 1:T]
q = [-2.0 * (t == T ? 1.0 : 1.0) * Q_track * xT for t = 1:T]
R = [Diagonal(1.0 * ones(model.m)) for t = 1:T-1]
r = [zeros(model.m) for t = 1:T-1]

obj = StageCosts([QuadraticCost(Q[t], q[t],
	t < T ? R[t] : nothing, t < T ? r[t] : nothing) for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
    if t < T
		Q = obj.cost[t].Q
		q = obj.cost[t].q
	    R = obj.cost[t].R
		r = obj.cost[t].r
        return x' * Q * x + q' * x + u' * R * u + r' * u
    elseif t == T
		Q = obj.cost[T].Q
		q = obj.cost[T].q
        return x' * Q * x + q' * x
    else
        return 0.0
    end
end

# Constraints
ul = [-10.0]
uu = [10.0]
p = [t < T ? 2 * m : n for t = 1:T]
info_t = Dict(:ul => ul, :uu => uu, :inequality => (1:2 * m))
info_T = Dict(:xT => xT)
con_set = [StageConstraint(p[t], t < T ? info_t : info_T) for t = 1:T]

function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	p = cons.con[t].p

	if t < T
		ul = cons.con[t].info[:ul]
		uu = cons.con[t].info[:uu]
		c .= [ul - u; u - uu]
	else
		c .= x - cons.con[T].info[:xT]
	end
end

ū = [(t == 1 ? -1.0 : 0.0) * ones(model.m) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)

q̄ = state_to_configuration(x̄)
# visualize!(vis, s.model, q̄, Δt = h)

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T,
	analytical_dynamics_derivatives = true)

# Solve
@time constrained_ddp_solve!(prob,
	max_iter = 1000, max_al_iter = 10,
	ρ_init = 1.0, ρ_scale = 10.0,
	con_tol = 0.001)

x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)

q̄ = state_to_configuration(x̄)
v̄ = [(q̄[t+1] - q̄[t]) ./ h for t = 1:length(q̄)-1]

# q̄_friction = q̄
# v̄_friction = v̄
# ū_friction = ū
# q̄_smooth = q̄
# v̄_smooth = v̄
# ū_smooth = ū
# @save "/home/taylor/Research/motion_planning/examples/differential_dynamic_programming/implicit_dynamics/cartpole_friction.jld2" q̄_friction v̄_friction ū_friction
# @load "/home/taylor/Research/motion_planning/examples/differential_dynamic_programming/implicit_dynamics/cartpole_friction.jld2"
# @save "/home/taylor/Research/motion_planning/examples/differential_dynamic_programming/implicit_dynamics/cartpole_smooth.jld2" q̄_smooth v̄_smooth ū_smooth
# @load "/home/taylor/Research/motion_planning/examples/differential_dynamic_programming/implicit_dynamics/cartpole_smooth.jld2"

t = range(0, stop = h * (length(q̄) - 1), length = length(q̄))
plt = plot();
plt = plot!(t, hcat(q̄...)', width = 2.0,
	color = [:magenta :orange],
	labels = ["q1" "q2"],
	legend = :topleft,
	xlabel = "time (s)",
	ylabel = "configuration",
	# title = "cartpole (w / o friction)")

	title = "cartpole (w/ friction)")

plt = plot();
plt = plot!(t, hcat(v̄..., v̄[end])', width = 2.0,
	color = [:magenta :orange],
	labels = ["q1" "q2"],
	legend = :topleft,
	xlabel = "time (s)",
	ylabel = "velocity",
	linetype = :steppost,
	# title = "cartpole (w / o friction)")

	title = "cartpole (w/ friction)")
	# title = "acrobot (w/ joint limits)")

# show(plt)
# savefig(plt, "/home/taylor/Research/implicit_dynamics_manuscript/figures/cartpole_friction.png")
# savefig(plt, "/home/taylor/Research/implicit_dynamics_manuscript/figures/cartpole_no_friction.png")

plot(hcat(ū..., ū[end])', linetype = :steppost)

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
open(vis)
# default_background!(vis)
# settransform!(vis["/Cameras/default"],
#         compose(Translation(0.0, -95.0, -1.0), LinearMap(RotY(0.0 * π) * RotZ(-π / 2.0))))
# setprop!(vis["/Cameras/default/rotated/<object>"], "zoom", 1)
q̄ = state_to_configuration(x̄)
q_anim = [[q̄[1] for t = 1:20]..., q̄..., [q̄[end] for t = 1:20]...]
visualize!(vis, s.model, q_anim, Δt = h)

using PGFPlots
const PGF = PGFPlots

plt_q1_smooth = PGF.Plots.Linear(t, hcat(q̄_smooth...)[1,:],
	mark="none",style="color=cyan, line width = 2pt, dashed",legendentry="q1")

plt_q2_smooth = PGF.Plots.Linear(t, hcat(q̄_smooth...)[2,:],
	mark="none",style="color=orange, line width = 2pt, dashed",legendentry="q2")

plt_qd1_smooth = PGF.Plots.Linear(t, hcat(v̄_smooth..., v̄_smooth[end])[1,:],
	mark="none",style="const plot, color=magenta, line width = 2pt, dashed",legendentry="q1")

plt_qd2_smooth = PGF.Plots.Linear(t, hcat(v̄_smooth..., v̄_smooth[end])[2,:],
	mark="none",style="const plot, color=green, line width = 2pt, dashed",legendentry="q2")

plt_q1_friction = PGF.Plots.Linear(t, hcat(q̄_friction...)[1,:],
	mark="none",style="color=cyan, line width = 2pt",legendentry="q1 (friction)")

plt_q2_friction = PGF.Plots.Linear(t, hcat(q̄_friction...)[2,:],
	mark="none",style="color=orange, line width = 2pt",legendentry="q2 (friction)")

plt_qd1_friction = PGF.Plots.Linear(t, hcat(v̄_friction..., v̄_friction[end])[1,:],
	mark="none",style="const plot, color=magenta, line width = 2pt",legendentry="q1 (friction)")

plt_qd2_friction = PGF.Plots.Linear(t, hcat(v̄_friction..., v̄_friction[end])[2,:],
	mark="none",style="const plot, color=green, line width = 2pt",legendentry="q2 (friction)")

aq = Axis([plt_q1_friction; plt_q2_friction; plt_q1_smooth; plt_q2_smooth],#; plt_qd1_smooth; plt_qd1_friction; plt_qd2_smooth; plt_qd2_friction],
    axisEqualImage=false,
    hideAxis=false,
	ylabel="configuration",
	xlabel="time (s)",
	xlims=(0.0, 2.5),
	legendStyle="{at={(0.01,0.99)},anchor=north west}")

av = Axis([plt_qd1_friction; plt_qd2_friction; plt_qd1_smooth; plt_qd2_smooth],
    axisEqualImage=false,
    hideAxis=false,
	ylabel="velocity",
	xlabel="time (s)",
	legendStyle="{at={(0.01,0.99)},anchor=north west}")

# Save to tikz format
PGF.save("/home/taylor/Research/implicit_dynamics_manuscript/figures/cartpole_friction_configuration.tikz", aq, include_preamble=false)
PGF.save("/home/taylor/Research/implicit_dynamics_manuscript/figures/cartpole_friction_velocity.tikz", av, include_preamble=false)


addIntegratorMedium_! =
Dict{String,Function}(
    "FTCS"=>addIntegratorMediumFTCS_!,
    "Lax"=>addIntegratorMediumFTCS_!,
    "Leapfrog"=>addIntegratorMediumFTCS_!,
)

δMedium_(i1,i2) = if i1==i2 1. else 0. end


#not available in Julia 0.3.8, but is in 0.4dev
#@enum(KernelType,"Gaussian", "Epanetchnikov", "Laplacian")
#@enum(Gradient,WRTMean, WRTVariance, WRTWeight)

type GaussianKer
    val::Float64
end

DirectSize = 100;        # if N*M is less than this, just compute.

type BallTreeDensity
  bt::BallTree

  KernelType
  multibandwidth::Int64               # flag: is bandwidth uniform?

  means::Array{Float64,1}                  # Weighted mean of points from this level down
  bandwidth::Array{Float64,1}              # Variance or other multiscale bandwidth
  bandwidthMin::Array{Float64,1}
  bandwidthMax::Array{Float64,1}           # Bounds on BW in non-uniform case

  calcStatsHandle::Function
  swapHandle::Function
end

getType(bd::BallTreeDensity) = bd.KernelType

#function Ndim(bd::BallTreeDensity)
#    return Ndim(bd.bt)
#end
Ndim(bd::BallTreeDensity) = Ndim(bd.bt)
Npts(bd::BallTreeDensity) = Npts(bd.bt)

Npts(bd::BallTreeDensity, i::Int64) = Npts(bd.bt, i)

## todo make sure these two are working properly
#function center(bd::BallTreeDensity, i::Int64)
#    return center(bd.bt, i)
#end
center(bd::BallTreeDensity, i::Int64) = center(bd.bt, i)

#function rangeB(bd::BallTreeDensity, i::Int64)
#    return rangeB(bd.bt, i)
#end
rangeB(bd::BallTreeDensity, i::Int64) = rangeB(bd.bt, i)

weight(bd::BallTreeDensity, i::Int64) = weight(bd.bt,i)

isLeaf(bd::BallTreeDensity, ind::Int64) = isLeaf(bd.bt, ind)

validIndex(bd::BallTreeDensity, ind::Int64) = validIndex(bd.bt, ind)

left(bd::BallTreeDensity, i::Int64) = left(bd.bt, i)

right(bd::BallTreeDensity, i::Int64) = right(bd.bt, i)

leafFirst(bd::BallTreeDensity, i::Int64) = leafFirst(bd.bt, i)

leafLast(bd::BallTreeDensity, i::Int64) = leafLast(bd.bt, i)

getIndexOf(bd::BallTreeDensity, i::Int64) = getIndexOf(bd.bt, i)

mean(bd::BallTreeDensity, i::Int64) = bd.means[((i-1)*bd.bt.dims+1):end]
mean(bd::BallTreeDensity, i::Int64, k::Int) = bd.means[((i-1)*bd.bt.dims+k)]

variance(bd::BallTreeDensity, i::Int64) = bd.bandwidth[((i-1)*bd.bt.dims+1):end] # !!! only works for Gaussian
variance(bd::BallTreeDensity, i::Int64, k::Int) = bd.bandwidth[((i-1)*bd.bt.dims+k)]

bw(bd::BallTreeDensity, i::Int64) = bd.bandwidth[((i-1)*bd.bt.dims+1):end]
bw(bd::BallTreeDensity, i::Int64, k::Int) = bd.bandwidth[((i-1)*bd.bt.dims+k)]

bwMax(bd::BallTreeDensity, i::Int64) = bd.bandwidthMax[((i-1)*bd.bt.dims*bd.multibandwidth+1):end]
bwMax(bd::BallTreeDensity, i::Int64, k::Int) = bd.bandwidthMax[((i-1)*bd.bt.dims*bd.multibandwidth+k)]

bwMin(bd::BallTreeDensity, i::Int64) = bd.bandwidthMin[((i-1)*bd.bt.dims*bd.multibandwidth+1):end]
bwMin(bd::BallTreeDensity, i::Int64, k::Int) = bd.bandwidthMin[((i-1)*bd.bt.dims*bd.multibandwidth+k)]

bwUniform(bd::BallTreeDensity) = bd.multibandwidth==0

root(bd::BallTreeDensity) = root()

function buildTree!(bd::BallTreeDensity)
  #println("buildTree!(::BallTreeDensity)")
  buildTree!(bd.bt)
  nothing
end

# Swap the ith leaf with the jth leaf.
function swapDensity!(bd::BallTreeDensity, i::Int64, j::Int64)
  if (i==j)
      return Union{}
  end

  swapBall!(bd.bt,i,j);

  i = (i-1)*bd.bt.dims+1;
  j = (j-1)*bd.bt.dims+1;
  for k in 1:bd.bt.dims
    tmp = bd.means[i];
    bd.means[i] = bd.means[j];
    bd.means[j] = tmp;
    tmp = bd.bandwidth[i];
    bd.bandwidth[i]  = bd.bandwidth[j];
    bd.bandwidth[j]  = tmp;
    if (!bwUniform(bd))
      tmp = bd.bandwidthMax[i];bd.bandwidthMax[i]=bd.bandwidthMax[j];bd.bandwidthMax[j]=tmp;
      tmp = bd.bandwidthMin[i];bd.bandwidthMin[i]=bd.bandwidthMin[j];bd.bandwidthMin[j]=tmp;
    end
    i+=1
    j+=1
  end
end

function calcStatsDensity!(bd, root::Int64)
  #println("calcStatsDensity! -- is running")
  calcStatsBall!(bd.bt, root)

  k = 0
  wtL=0.
  wtR = 0.
  wtT = 0.

  leftI = left(bd.bt, root);
  rightI=right(bd.bt, root);        # get children indices
  if (!validIndex(bd.bt,leftI) || !validIndex(bd.bt, rightI))
    return Union{}   # nothing to do if this
  end    #   isn't a parent node

  Ni  = bd.bt.dims*(root-1)
  NiL = bd.bt.dims*(leftI-1)
  NiR = bd.bt.dims*(rightI-1)
  wtL = bd.bt.weights[leftI]
  wtR = bd.bt.weights[rightI]
  wtT = wtL + wtR + eps(Float64);
  wtL /= wtT
  wtR /= wtT
  if (!bwUniform(bd))
    for k in 1:bd.bt.dims
      if (bd.bandwidthMax[NiL+k] > bd.bandwidthMax[NiR+k])
        bd.bandwidthMax[Ni+k] =  bd.bandwidthMax[NiL+k]
      else
        bd.bandwidthMax[Ni+k] =   bd.bandwidthMax[NiR+k]
      end
      if (bd.bandwidthMin[NiL+k] < bd.bandwidthMin[NiR+k])
        bd.bandwidthMin[Ni+k] = bd.bandwidthMin[NiL+k]
      else
        bd.bandwidthMin[Ni+k] = bd.bandwidthMin[NiR+k]
      end
    end
  end
  #switch(type) {
  #  case Gaussian:
    for k in 1:bd.bt.dims
      #@show k, Ni, NiL, NiR, round([wtL, wtR, wtT],2), round(bd.means[NiL+k],2), round(bd.means[NiR+k],2)
      bd.means[Ni+k]     = wtL * bd.means[NiL+k] + wtR * bd.means[NiR+k];
      bd.bandwidth[Ni+k] = wtL* (bd.bandwidth[NiL+k] + bd.means[NiL+k]*bd.means[NiL+k]) +
                        wtR* (bd.bandwidth[NiR+k] + bd.means[NiR+k]*bd.means[NiR+k]) -
                        bd.means[Ni+k]*bd.means[Ni+k];
      #@show round([bd.means[Ni+k], bd.bandwidth[Ni+k]],2)
    end
  #  break;
  #case Laplacian:
  #  for(unsigned int k=0; k < dims; k++) {
  #    means[Ni+k]     = wtL * means[NiL+k] + wtR * means[NiR+k];
  #    bandwidth[Ni+k] = wtL* (2*bandwidth[NiL+k]*bandwidth[NiL+k] + means[NiL+k]*means[NiL+k]) +
  #                      wtR* (2*bandwidth[NiR+k]*bandwidth[NiR+k] + means[NiR+k]*means[NiR+k]) -
  #                      means[Ni+k]*means[Ni+k];     // compute in terms of variance
  #    bandwidth[Ni+k] = sqrt(.5*bandwidth[Ni+k]);    //  then convert back to normal BW rep.
  #  }; break;
  #case Epanetchnikov:
  #  for(unsigned int k=0; k < dims; k++) {
  #    means[Ni+k]     = wtL * means[NiL+k] + wtR * means[NiR+k];
  #    bandwidth[Ni+k] = wtL* (.2*bandwidth[NiL+k]*bandwidth[NiL+k] + means[NiL+k]*means[NiL+k]) +
  #                      wtR* (.2*bandwidth[NiR+k]*bandwidth[NiR+k] + means[NiR+k]*means[NiR+k]) -
  #                     means[Ni+k]*means[Ni+k];     // compute in terms of variance
  #   bandwidth[Ni+k] = sqrt(5*bandwidth[Ni+k]);     //  then convert back to normal BW rep.
  #  }; break;
 #}
  return Union{}
end

# Create new matlab arrays and put them in the given structure
function makeBallTreeDensity(_pointsMatrix::Array{Float64,2}, _weightsMatrix::Array{Float64,1}, _bwMatrix::Array{Float64,1}, _type::DataType=GaussianKer)

  bt = makeBallTree(_pointsMatrix, _weightsMatrix, true)
  bt.calcStatsHandle = calcStatsDensity!
  bt.swapHandle = swapDensity!

  Np::Int64 = bt.num_points
  Nd::Int64 = bt.dims

  means = zeros(Nd*2*Np)
  means[(Np*Nd+1):end] = reshape(_pointsMatrix, Nd*Np, 1)

  if (size(_bwMatrix,2) == 1)
    multibw = 0
    bandwidth = zeros(Nd*2*Np)
    # @show Nd, Np, size(_bwMatrix), size(repmat(_bwMatrix,Np))
    bandwidth[(Np*Nd+1):end] = repmat(_bwMatrix,Np)
    bandwidthMin = bandwidth[(Np*Nd+1):end]
    bandwidthMax = bandwidth[(Np*Nd+1):end]
  else
    multibw = 1
    bandwidth = zeros(Nd*2*Np)
    bandwidthMin = zeros(Nd*2*Np)
    bandwidthMax = zeros(Nd*2*Np)
    bandwidth[(Np*Nd+1):end] = _bwMatrix
    bandwidthMin[(Np*Nd+1):end] = _bwMatrix
    bandwidthMax[(Np*Nd+1):end] = _bwMatrix
  end
  bd = BallTreeDensity(bt, _type, multibw, means, bandwidth, bandwidthMin, bandwidthMax, bt.calcStatsHandle, bt.swapHandle)
  bd.bt.data = bd # we need the circular reference for emulating polymorphism

  buildTree!(bd)

  return bd
end

function makeBallTreeDensity(_pointsMatrix::Array{Float64,2}, _weightsMatrix::Array{Float64,1}, _type::DataType=GaussianKer)
  makeBallTreeDensity(_pointsMatrix, _weightsMatrix, ones(size(_pointsMatrix,1)), _type)
end

# Figure out which of two children in this tree is closest to a given
# ball in another tree.  Returns the index in this tree of the closer
# child.
function closer!(bd::BallTreeDensity, myLeft::Int64, myRight::Int64, otherTree::BallTreeDensity, otherRoot::Int64)
  if (myRight==NO_CHILD || otherRoot==NO_CHILD)
    return myLeft
  end
  dist_sq_l = 0.0
  dist_sq_r = 0.0
  #@fastmath @inbounds begin
    for i in 1:bd.bt.dims
      dist_sq_l += (center(otherTree.bt, otherRoot, i) - center(bd.bt, myLeft, i)) *
        (center(otherTree.bt, otherRoot, i) - center(bd.bt, myLeft, i));
      dist_sq_r += (center(otherTree.bt, otherRoot, i) - center(bd.bt, myRight, i)) *
        (center(otherTree.bt, otherRoot, i) - center(bd.bt, myRight, i));
    end
  #end

  if (dist_sq_l < dist_sq_r)
    return myLeft;
  else
    return myRight;
  end
end

function closer!(bd::BallTreeDensity, i::Int64, j::Int64, other_tree::BallTreeDensity)
  return closer!(bd,i,j,other_tree,root()) #othertree.root() = 1 anyway
end

# Perform a *slight* adjustment of the tree: move the points by delta, but
#   don't reform the whole tree; just fix up the statistics.
#
function movePoints!(bd::BallTreeDensity, delta::Array{Float64,1})
  for i in leafFirst(bd,root()):leafLast(bd,root())
    for k in 1:bt.dims                      # first adjust locations by delta
      bd.bt.centers[bd.bt.dims*(i-1)+k] += delta[ (getIndexOf(bd,i)-1)*bd.bt.dims + k ];
    end
  end
  for i in bd.bt.num_points:-1:1               # then recompute stats of
    calcStats!(bd.bt.data, i)                       #   parent nodes
  end
  calcStats!(bd.bt.data, root())                    #   and finally root node
  return Union{}
end

# Assumes newWeights is the right size (num_points)
function changeWeights!(bd::BallTreeDensity, newWeights::Array{Float64,1})
  for i in (bd.bt.num_points+1):bd.bt.num_points*2
    bd.bt.weights[i] = newWeights[ getIndexOf(bt, i) ];
  end

  for i in bd.bt.num_points:-1:1
    calcStats!(bd.bt.data,i)
  end
  calcStats!(bd.bt.data, root());
  return Union{}
end

function resample(p::BallTreeDensity, Np::Int64=-1, ksType::ASCIIString="lcv")
# resample(p,Np,KSType) -- construct a new estimate of the KDE p by sampling
#                      Np new points; determines a bandwidth by ksize(pNew,KSType)
#                      NOTE: KStype = 'discrete' resamples points by weight &
#                            preserves original kernel size
  if (Np==-1)
    Np = getNpts(p)
  end
  if (ksType == "discrete")
    q = kde!(getPoints(p),zeros(getDim(p),1),getWeights(p))
    samplePts,ind = sample(q,Np)
    if (size(p.bandwidth,2)>2*p.bt.num_points)
      ks = getBW(p,ind)
    else
      ks = getBW(p,1)
    end
    p2 = kde!(samplePts,ks);
  else
    samplePts, = sample(p,Np)
    p2 = kde!(samplePts, ksType)
  end
  return p2
end

function printBallTree(bd::BallTreeDensity, rnd=15)
  printBallTree(bd.bt);
  @show round(bd.means,rnd);
  @show round(bd.bandwidth,rnd);
  @show round(bd.bandwidthMin,rnd);
  @show round(bd.bandwidthMax,rnd);

  nothing
end

function test02()
  ## matlab results
  #centers=[30]=4.4, 3.6, 5, 3.2, 0.18, 2, 5.5, 6.5, 7.5, 2.6, 0.18, 2.5, 0, 0, 0, 1.9, 1.2, 3, 3.3, -0.81, 2, 4.6, -0.56, 1, 4, 5, 6, 7, 8, 9,
  #weights=[10]=1, 0.6, 0.4, 0.4, 0, 0.2, 0.2, 0.2, 0.2, 0.2,
  #ranges=[30]=2.6, 4.4, 4, 1.4, 0.98, 1, 1.5, 1.5, 1.5, 0.7, 0.98, 0.5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
  #highest_leaf=[10]=9, 7, 9, 6, 0, 5, 6, 7, 8, 9,
  #lowest_leaf=[10]=5, 5, 8, 5, 0, 5, 6, 7, 8, 9,
  #permutation=[10]=0, 0, 0, 0, 0, 2, 1, 0, 3, 4,

  mus = [[4.6173,    3.2641,    1.8729, 4, 7 ]',
     [-0.5592,   -0.8088,    1.1610, 5, 8]',
     [1.,2.,3, 6, 9]']

  bd = makeBallTreeDensity(mus, 0.2*ones(5), 0.04*ones(3))
  printBallTree(bd)
end


###############################################
############  FUNCTIONS   #####################
###############################################

using Dates # to provide correct Julian time slices

export get_gdirs_with_climate

"""
    get_gdirs_with_climate(gdirs; overwrite=false, plot=true)

Gets the climate data for a given number of gdirs. Returns a tuple of `(gdirs, climate`.
"""
function get_gdirs_with_climate(gdirs, tspan; overwrite=false, plot=true)
    climate_raw = get_climate(gdirs, overwrite)
    climate = filter_climate(climate_raw, tspan)
    if plot
        plot_avg_longterm_temps(climate, gdirs)
    end
    gdirs_climate = get_gdir_climate_tuple(gdirs, climate)
    return gdirs_climate
end

"""
    get_climate(gdirs, overwrite)

Gets the climate data for multiple gdirs in parallel.
"""
function get_climate(gdirs, overwrite)
    println("Getting climate data...")
    # Retrieve and compute climate data in parallel
    # /!\ Keep batch size small in order to avoid memory problems
    climate = pmap(gdir -> get_climate_glacier(gdir, overwrite), gdirs; batch_size=10)
    # climate = map(gdir -> get_climate_glacier(gdir, overwrite), gdirs)
    return climate
end

"""
    get_climate_glacier(gdir::PyObject, overwrite; mb=true, period_y=(1979,2019))

Gets the W5E5 climate data for a single gdir. Returns a dictionary with the following format:

    Dict("longterm_temps"=>longterm_temps, "annual_temps"=>annual_temps,
    "MB_dataset"=>MB_ds, "climate_dataset"=>climate_ds)
"""
function get_climate_glacier(gdir::PyObject, overwrite; mb=true, period_y=(1979,2019))
    # Generate downscaled climate data
    if !isfile(joinpath(gdir.dir, "annual_temps.nc")) || overwrite # Retrieve unless overwrite
        mb_type = "mb_real_daily"
        grad_type = "var_an_cycle" # could use here as well 'cte'
        # fs = "_daily_".*climate
        fs = "_daily_W5E5"
        MB_ds, climate_ds = get_MB_climate_datasets(gdir, mb, period_y, fs)
        # Convert to annual values to force UDE
        annual_temps = climate_ds.annual.temp.groupby("time.year").mean(dim=["time", "x", "y"])
        # yb,ye = annual_temps.year.data[1], annual_temps.year.data[end]
        println("Storing climate data in: ", joinpath(gdir.dir, "annual_temps.nc"))
        annual_temps.to_netcdf(joinpath(gdir.dir, "annual_temps.nc"))
        annual_temps = xr.open_dataset(joinpath(gdir.dir, "annual_temps.nc"))
    else
        annual_temps = xr.open_dataset(joinpath(gdir.dir, "annual_temps.nc"))
        MB_ds, climate_ds = nothing, nothing # dummy empty variables
    end
    # Compute a 20-year rolling mean for long-term air temperature variability
    longterm_temps = annual_temps.rolling(year=20).mean().dropna("year")
    # return longterm_temps, annual_temps, MB_ds, climate_ds
    climate = Dict("longterm_temps"=>longterm_temps, "annual_temps"=>annual_temps,
                    "MB_dataset"=>MB_ds, "climate_dataset"=>climate_ds)
    return climate
end

"""
    get_MB_climate_datasets(gdir, mb, period_y, fs)

Downloads the climate and mass balance data for a single gdir.
"""
function get_MB_climate_datasets(gdir, mb, period_y, fs)
    # Retrieve reference MB data
    rgi_id = gdir.rgi_id
    println("Downloading climate data for $rgi_id...")
    climate = get_raw_climate_data(gdir)

    if mb
        println("Fetching mass balance reference data...")
        try
            # Building MB and climate datasets
            mb_glaciological = gdir.get_ref_mb_data(input_filesuffix=fs)
            # Get hydrological period matching the reference MB data
            hydro_period = to_hydro_period(mb_glaciological)
            climate_ds, MB_ds = build_MB_climate_ds(hydro_period, climate, gdir, mb_glaciological)
            return MB_ds, climate_ds
        catch
            @warn "Error retrieving reference mass balance data. Retrieving only climate data."
            # Building only climate dataset
            hydro_period = collect(Date(period_y[1],10,1):Day(1):Date(period_y[2],09,30))
            climate_ds, MB_ds = build_climate_ds(hydro_period, climate, gdir)
            return MB_ds, climate_ds
        end
    else
        # Building only climate dataset
        hydro_period = collect(Date(period_y[1],10,1):Day(1):Date(period_y[2],09,30))
        climate_ds, MB_ds = build_climate_ds(hydro_period, climate, gdir)
        return MB_ds, climate_ds
    end
end

"""
    build_MB_climate_ds(hydro_period, climate, gdir, mb_glaciological)

Builds a Dataset with the mass balance and climate data for a glacier. It stores "annual", "accumulation" and/or "ablation" series.
"""
function build_MB_climate_ds(hydro_period, climate, gdir, mb_glaciological)
    println("Processing data for reference period ", hydro_period[1], " - ", hydro_period[end])
    # Build the training and reference dataset including annual and seasonal data
    # Get the climate forcings for the MB data
    seasons = ["annual", "accumulation", "ablation"]
    balances = ["ANNUAL_BALANCE", "WINTER_BALANCE", "SUMMER_BALANCE"]
    ds_clim_buffer, ds_MB_buffer = [],[]
    for (season, balance) in zip(seasons, balances)
        season_MB = mb_glaciological[balance]
        MB_clim = get_climate_forcing(gdir, climate, hydro_period, season)
        push!(ds_MB_buffer, season_MB)
        push!(ds_clim_buffer, MB_clim)
    end
    # Store everything in a global dataset
    println("Storing data in metadataset")
    MB_ds = Dataset(ds_MB_buffer...) # splat it!!!
    climate_ds = Dataset(ds_clim_buffer...)
    return climate_ds, MB_ds
end

"""
    build_climate_ds(hydro_period, climate, gdir)

Builds a Dataset with only the climate data for a glacier. It stores "annual" series.
"""
function build_climate_ds(hydro_period, climate, gdir)
    ds_clim_buffer = []
    println("Processing data for reference period ", hydro_period[1], " - ", hydro_period[end])
    MB_clim = get_climate_forcing(gdir, climate, hydro_period, "annual")
    push!(ds_clim_buffer, MB_clim)
    push!(ds_clim_buffer, nothing)
    push!(ds_clim_buffer, nothing)
    # Store everything in a global dataset
    println("Storing data in metadataset")
    climate_ds = Dataset(ds_clim_buffer...)
    MB_ds = nothing
    return climate_ds, MB_ds
end

"""
    get_climate_forcing(gdir, climate, period, season)

Processes raw climate data to seasonal data, with PDDs, snowfall and rainfall.
It projects the climate data into the glacier matrix by applying temperature gradients and precipitation lapse rates.
"""
function get_climate_forcing(gdir, climate, period, season)
    @assert any(season .== ["annual","accumulation", "ablation"]) "Wrong season type, must be `annual`, `accumulation` or `ablation`"

    # Make sure the desired period is covered by the climate data
    period = trim_period(period, climate)
    @assert any((climate.time[1].dt.date.data[1] <= period[1]) & any(climate.time[end].dt.date.data[1] >= period[end])) "No overlapping period available between climate and MB data!"
    climate = climate.sel(time=period) # Crop desired time period

    if season == "accumulation"
        climate = climate.sel(time=is_acc.(clim_period.time.dt.month.data))
    elseif season == "ablation"
        climate = climate.sel(time=is_abl.(clim_period.time.dt.month.data))
    end

    # Transform climate to cumulative PDDs, snowfall and rainfall
    climate = get_cumulative_climate(climate)

    # Convert climate dataset to 2D based on the glacier's DEM
    g_dem = xr.open_rasterio(gdir.get_filepath("dem"))
    climate = create_2D_climate_data(climate, g_dem)

    # Apply temperature gradients and compute snow/rain fraction for the selected period
    apply_t_grad!(climate, g_dem)
    climate = climate.drop("gradient")

    return climate

end

"""
    get_cumulative_climate(climate, gradient_bounds=[-0.009, -0.003], default_grad=-0.0065)

Computes Positive Degree Days (PDDs) and cumulative rainfall and snowfall from climate data.
"""
function get_cumulative_climate(climate, gradient_bounds=[-0.009, -0.003], default_grad=-0.0065)
    avg_temp = climate.temp.resample(time="1M").mean()
    avg_gradients = climate.gradient.resample(time="1M").mean()
    climate.temp.data = climate.temp.where(climate.temp > 0, 0).data # get PDDs
    climate.gradient.data = utils.clip_array(climate.gradient.data, gradient_bounds[1], gradient_bounds[2]) # Clip gradients within plausible values
    attributes = climate.attrs
    climate_sum = climate.resample(time="1M").sum() # get monthly cumulative values
    climate_sum_avg1 = climate_sum.assign(Dict("avg_temp"=>avg_temp))
    climate_sum_avg2 = climate_sum_avg1.assign(Dict("avg_gradient"=>avg_gradients))
    climate_sum_avg2.attrs = attributes
    return climate_sum_avg2
end

"""
    get_raw_climate_data(gdir, temp_resolution="daily", climate="W5E5")

Downloads the raw W5E5 climate data with a given resolution (daily by default). Returns an xarray Dataset.
"""
function get_raw_climate_data(gdir, temp_resolution="daily", climate="W5E5")
    MBsandbox.process_w5e5_data(gdir, climate_type=climate, temporal_resol=temp_resolution)
    fpath = gdir.get_filepath("climate_historical", filesuffix="_daily_W5E5")
    climate = xr.open_dataset(fpath)
    return climate
end

"""
    apply_t_grad!(climate, g_dem)

Applies temperature gradients to the glacier 2D climate data based on a DEM.
"""
function apply_t_grad!(climate, g_dem)
    # We apply the gradients to the temperature
    climate.temp.data = climate.temp.data .+ climate.avg_gradient.data .* (g_dem.data .- climate.ref_hgt)
    climate.PDD.data = climate.PDD.data .+ climate.gradient.data .* (g_dem.data .- climate.ref_hgt)
    # We adjust the rain/snow fractions with the updated temperature
    climate.snow.data = climate.snow.where(climate.temp < 0, 0).data
    climate.rain.data = climate.rain.where(climate.temp > 0, 0).data
end

"""
    create_2D_climate_data(climate, g_dem)

Projects climate data to the glacier matrix by simply copying the closest gridpoint to all matrix gridpoints.
Generates a new xarray Dataset which is returned.
"""
function create_2D_climate_data(climate, g_dem)
    # Create dummy 2D arrays to have a base to apply gradients afterwards
    dummy_grid = ones(size(permutedims(g_dem.data, (1,2,3))))
    temp_2D = climate.avg_temp.data .* dummy_grid
    PDD_2D = climate.temp.data .* dummy_grid
    snow_2D = climate.prcp.data .* dummy_grid
    rain_2D = climate.prcp.data .* dummy_grid

    # We generate a new dataset with the scaled data
    climate_2D = xr.Dataset(
        data_vars=Dict([
            ("temp", (["time","y","x"], temp_2D)),
            ("PDD", (["time","y","x"], PDD_2D)),
            ("snow", (["time","y","x"], snow_2D)),
            ("rain", (["time","y","x"], rain_2D)),
            ("gradient", (["time"], climate.gradient.data)),
            ("avg_gradient", (["time"], climate.avg_gradient.data))
            ]),
        coords=Dict([
            ("time", climate.time),
            ("x", g_dem.x),
            ("y", g_dem.y)
        ]),
        attrs=climate.attrs
    )

    return climate_2D

end

"""
    to_hydro_period(mass_balance::PyObject, trim_edges=true)

Converts the dates of mass balance series to hydrological periods.
"""
function to_hydro_period(mass_balance::PyObject, trim_edges=true)
    # Select hydro period between October 1st and 30th of September
    if trim_edges
        hydro_period = collect(Date(mass_balance.index[1],10,1):Day(1):Date(mass_balance.index[end]-1,09,30))
    else
        hydro_period = collect(Date(mass_balance.index[1]-1,10,1):Day(1):Date(mass_balance.index[end],09,30))
    end

    return hydro_period
end

"""
    to_hydro_period(mass_balance::PyObject, trim_edges=true)

Generates a Date range based the hydrological period for an Array of years.
"""
function to_hydro_period(years::Array)
    # Select hydro period between October 1st and 30th of September
    hydro_period = collect(Date(years[1]-1,10,1):Day(1):Date(years[end],09,30))

    return hydro_period
end

"""
    trim_period(period, climate)

Trims a time period based on the time range of a climate series.
"""
function trim_period(period, climate)
    if any(climate.time[1].dt.date.data[1] > period[1])
        head = jldate(climate.time[1])
        period = Date(year(head), 10, 1):Day(1):period[end] # make it a hydrological year
    end
    if any(climate.time[end].dt.date.data[1] < period[end])
        tail = jldate(climate.time[end])
        period = period[1]:Day(1):Date(year(tail), 9, 30) # make it a hydrological year
    end

    return period
end

"""
    is_abl(month)

Determines if the month belongs the the ablation season or not.
"""
function is_abl(month)
    return (month >= 4) & (month <= 10)
end

"""
    is_acc(month)

Determines if the month belongs the the accumulation season or not.
"""
function is_acc(month)
    return (month <= 4) | (month >= 10)
end

# Metastructure to store xarray dataset with MB and climate data for training
struct Dataset
    annual
    accumulation
    ablation
end

"""
    fake_temp_series(t, means=[0,-2.0,-3.0,-5.0,-10.0,-12.0,-14.0,-15.0,-20.0])

Generates fake random temperature series based on an array of mean temperature values.
"""
function fake_temp_series(t, means=[0,-2.0,-3.0,-5.0,-10.0,-12.0,-14.0,-15.0,-20.0])
    temps, norm_temps, norm_temps_flat = [],[],[]
    for mean in means
       push!(temps, mean .+ rand(t).*1e-1) # static
       append!(norm_temps_flat, mean .+ rand(t).*1e-1) # static
    #    push!(temps, (mean .+ rand(t)) .+ 0.5.*ts) # with trend
    #    push!(temps, (mean .+ rand(t)) .+ -0.5.*ts) # with trend
    end

    # Normalise temperature series
    norm_temps_flat = Flux.normalise([norm_temps_flat...]) # requires splatting

    # Re-create array of arrays
    for i in 1:t₁:length(norm_temps_flat)
        push!(norm_temps, norm_temps_flat[i:i+(t₁-1)])
    end

    return temps, norm_temps
end

"""
    filter_climate(climate)

Filters pairs of Glacier Directory/Climate series which have climate series shorter than the simulation period.
"""
function filter_climate(climate, tspan)
    updated_climate = []
    for climate_glacier in climate
        if length(climate_glacier["longterm_temps"].temp.data) >= tspan[end]
            push!(updated_climate, climate_glacier)
        else
            @warn "Filtered glacier due to short climate series"
        end
    end
    return updated_climate
end

"""
    get_gdir_climate_tuple(gdirs, climate)

Generates a tuple with all the Glacier directories and climate series for the simulations.
"""
function get_gdir_climate_tuple(gdirs, climate)
    dates, longterm_temps, annual_temps = [],[],[]
    for climate_batch in climate
        push!(dates, climate_batch["longterm_temps"].year.data)
        push!(longterm_temps, climate_batch["longterm_temps"].temp.data)
        push!(annual_temps, climate_batch["annual_temps"].temp.data)
    end
    gdirs_climate = (dates, gdirs, longterm_temps, annual_temps)
    return gdirs_climate
end


const XIntV = Union{XInt, Nothing}

_vec(n::Integer) = Vector{Limb}(undef, n)
@inline vec(x::XInt) = x.v::Vector{Limb}

lenvec(x::XSigned) = len(x), vec(x)

vec!(x::XInt, n::Integer) =
    if is_short(x)
        _vec(n)
    else
        vec!(vec(x), n)
    end

function vec!(xv::Vector, n::Integer)
    len = length(xv)
    if n > len
        Base._growend!(xv, n-len)
    end
    xv
end

vec!(::Nothing, n::Integer) = _vec(n)

# copy x into r, allocating for r when r===nothing
# assumes !is_short(x)
_copy!(r::XIntV, x::XSigned) =
    if r === x
        x
    else
        xl, xv = lenvec(x)
        _XInt(flipsign(xl, _sign(x)), _copy!(vec!(r, xl), xv, 1:xl))
    end

function _copy!(r::Vector, xv::Vector, span::UnitRange)
    a, b = first(span), last(span)
    if r !== xv
        unsafe_copyto!(r, a, xv, a, b-a+1)
    end
    r
end

function _copy!(r::Vector, xv::SVec, span::UnitRange)
    for i=span
        @inbounds r[i] = xv[i]
    end
    r
end

@inline function assign!(r::XIntV, vals::T...) where T
    v = vec!(r, length(vals))
    for i = eachindex(vals)
        @inbounds v[i] = vals[i] % Limb
    end
    v
end

@inline XInt!(r::XIntV, z::Short) =
    fits(z) ? _XInt(z % SLimb) :
              _XInt(flipsign(one(SLimb), z), assign!(r, z))

# make a XInt into r, with value -z
@inline XInt_neg!(r::XIntV, z::Limb) =
    fits(z) ? _XInt(-z % SLimb) :
              _XInt(-one(SLimb), assign!(r, z))

@inline XInt!(r::XIntV, z::Limb, flip::SLimb) =
    flip < 0 ?
        XInt_neg!(r, z) :
        XInt!(r, z)

@inline XInt!(r::XIntV, z::SLimbW, extra::SLimb=0) =
    # @assert extra >= 0
    if slimbmin < z <= slimbmax
        _XInt(z % SLimb)
    else
        XInt_big!(r, z, extra)
    end

@noinline function XInt_big!(r::XIntV, z::SLimbW, extra::SLimb)
    # @assert extra >= 0
    zz = abs(z)
    z1 = zz % Limb
    z2 = (zz >>> BITS_PER_LIMB) % Limb
    if iszero(z2)
        rv = vec!(r, 1+extra)
        @inbounds rv[1] = z1
        _XInt(sign(z) % SLimb, rv)
    else
        rv = vec!(r, 2+extra)
        @inbounds rv[1] = z1
        @inbounds rv[2] = z2
        _XInt(flipsign(SLimb(2), z), rv)
    end
end

XInt!(r::XIntV, z::LimbW, extra::SLimb=0) =
    # @assert extra >= 0
    if z <= slimbmax
        _XInt(z % SLimb)
    else
        z1 = z % Limb
        z2 = (z >>> BITS_PER_LIMB) % Limb
        if iszero(z2)
            rv = vec!(r, 1+extra)
            @inbounds rv[1] = z1
            _XInt(SLimb(1), rv)
        else
            rv = vec!(r, 2+extra)
            @inbounds rv[1] = z1
            @inbounds rv[2] = z2
            _XInt(SLimb(2), rv)
        end
    end

@inline function normalize(v::Vector, l::Integer)
    # @assert l <= length(v)
    while l > 0
        iszero(@inbounds v[l]) || break
        l -= 1
    end
    l
end

function add1!(r::XIntV, x::XSigned, y::Short)
    # @assert !is_short(x)
    samesign = signbit(x) == signbit(y)
    if samesign
        _add1!(r, x, abs(y) % Limb)
    else
        _sub1!(r, x, abs(y) % Limb)
    end
end

function _add1!(r::XIntV, x::XSigned, y::Limb)
    xl = len(x)
    rv = vec!(r, xl+1)
    c = MPN.add_1(rv, x=>xl, y)
    @inbounds rv[xl+1] = c
    rl = xl + !iszero(c)
    _XInt(flipsign(rl, _sign(x)), rv)
end

function _sub1!(r::XIntV, x::XSigned, y::Limb)
    xl = len(x)
    # we know that abs(x) > y, as y comes from a SLimb, whose absolute value
    # is < |typemmin(SLimb)| == limb1min (upper bit of y is not set)
    if xl == 1
        xv = vec(x)
        # xv[1] >= limb1min by specification
        rv1 = @inbounds xv[1] - y
        XInt!(r, rv1, _sign(x))
    else
        # if xl > 1, we see that x-y can't be an immediate integer
        rv = vec!(r, xl)
        MPN.sub_1(rv, x=>xl, y)
        rl = xl-iszero(@inbounds rv[xl])
        _XInt(flipsign(rl, _sign(x)), rv)
    end
end

_widen(x::XSigned) = widen(short(x))
_promote(x::XInt) = x
_promote(x::BitInteger) = Static(x)

@inline add!(r::XIntV, x, y) = add!(r, _promote(x), _promote(y))

@inline function add!(r::XIntV, x::XSigned, y::XSigned)
    xshort = is_short(x)
    yshort = is_short(y)
    if xshort & yshort
        XInt!(r, _widen(x) + _widen(y))
    elseif xshort
        iszero(x) ? XInt(y) : add1!(r, y, short(x))
    elseif yshort
        iszero(y) ? XInt(x) : add1!(r, x, short(y))
    else
        addbig!(r, x, y)
    end
end

# fast path for sum
# this allocate in x even if x == 0
# (i.e. this doesn't return the other arg y in this case, so no aliasing)
@inline function addbig_positive!(x::XInt, y::XInt)
    if iszero(x.x)
        yl, yv = lenvec(y)
        xv = _vec(yl+2) # +1 for later invocation which requests it, and +1 for margin
        unsafe_copyto!(xv, 1, yv, 1, yl)
        return _XInt(yl, xv)
    end
    xl, xv = lenvec(x)
    yl, yv = lenvec(y)
    rv = xv # result in x, even if we have to swap the arguments
    if xl < yl
        xl, yl = yl, xl
        xv, yv = yv, xv
    end
    rv = vec!(rv, xl+1)
    c = MPN.add(rv, xv=>xl, yv=>yl) # c ∈ (0, 1)
    @inbounds rv[xl+1] = c
    _XInt(xl + c % SLimb, rv)
end

# NOTE: this is still unfortunately roughly 2x slower than MPZ.add! for BigInt
# we recover a lot of perfs if we inline instead, but that's a big function...
@inline function addbig!(r::XIntV, x::XSigned, y::XSigned)
    xl, xv = lenvec(x)
    yl, yv = lenvec(y)
    xz, yz = _sign(x), _sign(y)
    if xl < yl
        # we don't just exchange the variables within the main function to
        # avoid type-unstability
        addbig!(r, yl, yv, yz, xl, xv, xz)
    else
        addbig!(r, xl, xv, xz, yl, yv, yz)
    end
end

@noinline function addbig!(r::XIntV, xl, xv, xz, yl, yv, yz)
    if xz == yz
        rv = vec!(r, xl+1)
        c = MPN.add(rv, xv=>xl, yv=>yl) # c ∈ (0, 1)
        @inbounds rv[xl+1] = c
        _XInt(flipsign(xl + c % SLimb, xz), rv)
    elseif xl >= yl+2 # >= 3
        # we are sure that the result will have length xl or xl-1
        rv = vec!(r, xl)
        MPN.sub(rv, xv=>xl, yv=>yl)
        rl = xl - iszero(@inbounds rv[xl])
        _XInt(flipsign(rl, xz), rv)
    else
        local xvl, yvl
        flip = false
        if xl == yl
            # same length, need to compare the content; we could use:
            # MPN.cmp(x=>xl, y=>yl) % Int
            # but we can as well do it manually, to record an upper bound on rl
            while xl > 0
                xvl = @inbounds xv[xl]
                yvl = @inbounds yv[xl]
                if xvl != yvl
                    if xvl < yvl
                        flip = true
                    end
                    break
                end
                xl -= 1
            end
            if xl == 0
                return _XInt(0)
            elseif xl == 1
                if flip
                    return XInt!(r, yvl-xvl, yz)
                else
                    return XInt!(r, xvl-yvl, xz)
                end
            end
            yl = xl
        else # xl == yl+1
            xvl = @inbounds xv[xl]
            yvl = zero(Limb)
        end
        if flip
            _addbig!(r, yl, yv, yz, yvl, xl, xv, xz, xvl)
        else
            _addbig!(r, xl, xv, xz, xvl, yl, yv, yz, yvl)
        end
    end
end

@inline function _addbig!(r, xl, xv, xz, xvl, yl, yv, yz, yvl)
    # here, xvl > yvl && xl >= 2
    rvl = xvl - yvl # value of rv[xl] "so far"
    @inbounds if rvl == 1
        rv1, c = Base.sub_with_overflow(xv[1], yv[1])
        if c && fits(rv1) # above limbs could cancel out, check if result is immediate
            i = 2
            local ri::Limb
            while i < xl
                ri, c2 = Base.sub_with_overflow(xv[i], yv[i])
                ri, c = ri-c,
                        c2 | c & iszero(ri)
                iszero(ri) || break
                # at this point (ri==0), we can see that c has the same value as c2,
                # but letting c = c2 above is not enough, as we re-use c later for
                # the "real" computation, so it must be set according to the general
                # case
                i += 1
            end
            # i == xl || ri != 0
            if i == xl && c
                return _XInt(flipsign(rv1 % SLimb, xz)) # we know that rv1 fits
            else
                rv = vec!(r, xl)
                rv[1] = rv1
                fill!(@view(rv[2:i-1]), zero(Limb))
                if i == xl # !c, as the case c==true was handled above
                    rv[xl] = rvl
                    rl = xl
                else
                    rv[i] = ri
                    # we would like to use native MPN.sub_n, but we can't pass to it
                    # a pre-existing carry
                    c = MPN.sub_n(rv, xv, yv, i+1:xl-1, c)
                    rv[xl] = rvl-c
                    rl = normalize(rv, xl)
                end
                return _XInt(flipsign(rl, xz), rv)
            end
        end
    end
    @assert xl == yl || xl == yl+1
    rv = vec!(r, xl)
    c = MPN.sub_n(rv, xv, yv, yl)
    if xl > yl
        rv[xl] = @inbounds(xv[xl]-c)
    end
    rl = normalize(rv, xl)
    _XInt(flipsign(rl, xz), rv)
end

add!(x::XInt, y) = add!(x, x, _promote(y))

neg!(x::XInt) = _XInt(-x.x, x.v)

sub!(r::XIntV, x, y::XInt) = add!(r, _promote(x), neg!(y))
# we don't implement neg!(s::Static), as s might contain unsigned, or typemin(SLimb)
sub!(r::XIntV, x::XInt, y::BitInteger) = neg!(add!(r, neg!(x), Static(y)))

sub!(x::XInt, y) = sub!(x, x, _promote(y))

is_limb(x::XSigned) = is_short(x) # we don't bother for non-short with one limb
is_limb(x::ShortMax) = true
limb(x::XSigned) = short(x)
limb(x::SLimbMax) = x % SLimb
limb(x::ULimbMax) = x % Limb
_widen(x::ShortMax) = SLimbW(x)
_sign(x::Limb) = SLimb(1)
_sign(x::SLimb) = sign(x)

mul!(x::XInt, y::Union{ShortMax,XInt}) = mul!(x, x, y)

@inline function mul!(r::XIntV, x::XInt, y::Union{ShortMax,XInt})
    xshort = is_short(x)
    yshort = is_limb(y)
    if xshort & yshort
        XInt!(r, _widen(x) * _widen(y))
    elseif xshort
        mul1!(r, y, short(x))
    elseif yshort
        mul1!(r, x, limb(y))
    else
        mulbig!(r, x, y)
    end
end

function mul1!(r::XIntV, x::XInt, y::Short)
    # @assert !is_short(x) # y can be slimbmin
    iszero(y) && return XInt(0)
    isone(y) && return x
    y == -1 && return neg!(x)
    sz = flipsign(_sign(x), _sign(y))
    xl = len(x)
    rv = vec!(r, xl + 1)
    high = MPN.mul_1(rv, x, abs(y) % Limb) # rv === vec(x) is supported
    @inbounds rv[xl+1] = high
    _XInt(flipsign(xl + !iszero(high), sz), rv)
end

function mulbig!(r::XIntV, x::XInt, y::XInt)
    # @assert !is_short(x) && !is_short(y)
    xl, xv = lenvec(x)
    yl, yv = lenvec(y)
    sz = flipsign(_sign(x), _sign(y))
    if xl < yl
        xl, yl = yl, xl
        xv, yv = yv, xv
    end

    tv = rv = vec!(r, xl+yl)
    # check for aliasing
    if rv === xv
        # MPN can't do mpn_mul inplace; rv is already big enough to hold the result,
        # so copy xv to a temporary, which can be discarded afterwards
        tv = tmpvec(xl)
        unsafe_copyto!(tv, 1, xv, 1, xl)
        xv = tv
        if rv === yv
            # assume xl == yl; make yv egal again to xv, the square will be computed
            yv = xv
        end
    elseif rv === yv
        tv = tmpvec(yl)
        unsafe_copyto!(tv, 1, yv, 1, yl)
        yv = tv
    end
    high = MPN.mul(rv, xv => xl, yv => yl)
    tv !== rv && freetmpvec(tv)
    _XInt(flipsign(xl+yl-iszero(high), sz), rv)
end

# set r to ~x == -(x+1)
com!(r::XIntV, x::XInt=r) =
    if is_short(x)
        XInt!(r, ~x.x) # could be non-short for x.x == slimbmax
    else
        neg!(add1!(r, x, one(SLimb)))
    end

lshift!(x::XInt, c::Integer) = lshift!(x, x, c)

# from base/operators.jl
@inline function lshift!(r::XIntV, x::XInt, c::Integer)
    typemin(Cint) <= c <= typemax(Cint) && return lshift!(r, x, c % Cint)
    (x >= 0 || c >= 0) && return zero(XInt)
    XInt(-1)
end

@inline lshift!(r::XIntV, x::XInt, c::Unsigned) =
    c <= typemax(Cuint) ? lshift!(r, x, c % Cuint) : zero(XInt)

@inline lshift!(r::XIntV, x::XInt, c::Cint) =
    c >= 0 ? lshift!(r, x, unsigned(c)) : rshift!(r, x, unsigned(-c))

@inline function lshift!(r::XIntV, x::XInt, c::Cuint)
    z = x.x
    iszero(z) && return x # NOTE: needs updating if a reduce argument is added
    if is_short(x)
        iszero(c) && return x
        if z > 0
            leading_zeros(z) > c &&
                return _XInt(z << c)
        else
            if leading_ones(z) > c
                s = z << c
                s !== slimbmin &&
                    return _XInt(z << c)
                # NOTE: it seems to be almost 2x faster for the general case to not
                # return the (known) result here for s==slimbmin, but rather to let
                # this special case be handled by lshift_small!
            end
        end
        lshift_small!(r, z, c)
    else
        lshift_big!(r, x, c)
    end
end

@noinline function lshift_small!(r::XIntV, z::SLimb, c::Cuint)
    zu = abs(z) % Limb
    offset = _divLimb(c) % SLimb
    cnt = _modLimb(c) % SLimb
    rl = 1+offset
    if iszero(cnt)
        rv = vec!(r, rl)
        @inbounds rv[rl] = zu
    else
        rv = vec!(r, rl+1)
        w = widen(zu) << cnt
        @inbounds rv[rl] = w % Limb
        high = @inbounds rv[rl+1] = (w >> BITS_PER_LIMB) % Limb
        rl += !iszero(high)
    end
    fill!(@view(rv[1:offset]), Limb(0))
    _XInt(flipsign(rl, z), rv)
end

@noinline function lshift_big!(r::XIntV, x::XInt, c::Cuint)
    r === x && c == 0 && return x
    xl = len(x)
    xv = vec(x)
    offset = _divLimb(c) % SLimb
    cnt = _modLimb(c) % SLimb
    rl = xl+offset
    if iszero(cnt)
        rv = vec!(r, rl)
        # TODO: when r===x and rv requires allocation anyway in vec! above,
        # maybe it would be more efficient to prepend! zeros ?
        unsafe_copyto!(rv, offset+1, xv, 1, xl)
    else
        rv = vec!(r, rl+1)
        high = MPN.lshift(@view(rv[offset+1:rl]), xv=>xl, cnt)
        @inbounds rv[rl+1] = high
        rl += !iszero(high)
    end
    fill!(@view(rv[1:offset]), Limb(0))
    _XInt(flipsign(rl, x.x), rv)
end

rshift!(x::XInt, c::Integer) = rshift!(x, x, c)

# from base/operators.jl
@inline function rshift!(r::XIntV, x::XInt, c::Integer)
    typemin(Cint) <= c <= typemax(Cint) && return rshift!(r, x, c % Cint)
    (x >= 0 || c < 0) && return zero(XInt)
    XInt(-1)
end

@inline rshift!(r::XIntV, x::XInt, c::Cint) =
    c >= 0 ? rshift!(r, x, unsigned(c)) : lshift!(r, x, unsigned(-c))

@inline function rshift!(r::XIntV, x::XInt, c::Cuint)
    z = x.x
    is_short(x) && return _XInt(z >> c)
    (iszero(z) || iszero(c)) && return x
    rshift_big!(r, x, c)
end

# For non-negative integers, right-shift is as obvious as it gets.
# For negative numbers, we use as usual the formula ~y = -(y+1) to emulate 2-complement
# integers; assuming x < 0 and c > 0, letting u = -x and v = u - 1 >= 0, we have
# x == -(v + 1) == ~v, from which we infer
# x >> c == ~v >> c == ~(v >> c) == -(v >> c + 1)
# Now, we don't want to compute v >> c, only u >> c. There are two cases:
# 1. u is a multiple of 2^c, i.e. u = k*2^c and u >> c == k
#    then v == u-1 == (k-1)*2^c + (2^c-1), therefore v >> c == k-1
#    hence: x >> c == -(k-1 + 1) == -k == -(u >> c) == -(-x >> c)
# 2. otherwise, u = k*2^c + r where 1 <= r < 2^c
#    then v == u-1 == k*2^c + r-1 where 0 <= r-1 < 2^c, therefore v >> c == u >> c
#    hence: x >> c == -(-x >> c + 1)
# Also, this can be derived more direcly from x >> c == fld(x, 2^c) where x = -k*2^c - r
@noinline function rshift_big!(r::XIntV, x::XInt, c::Cuint)
    xl = len(x)
    xv = vec(x)
    offset = _divLimb(c) % SLimb
    cnt = _modLimb(c) % SLimb
    rl = xl-offset
    neg = signbit(x.x)
    rl <= 0 && return XInt(-neg)

    rest = false # records whether there is any set bit among the first c in -x (when x < 0)
    if neg
        rest = any(!=(zero(Limb)), @view(xv[1:offset]))
    end

    if  rl <= 2 # special case, to avoid allocating for nothing
        xvend = @inbounds xv[xl] % LimbW
        if rl == 2
            xvend = xvend << BITS_PER_LIMB | @inbounds xv[xl-1] % LimbW
        end
        u = xvend >> cnt
        if neg & !rest
            rest |= trailing_zeros(xvend) < cnt
        end
        if rest && u == typemax(LimbW)
            # happens e.g. with x = -(XInt(2)^192-1) and c == 64 (with sizeof(Limb) == 8)
            _XInt(SLimb(-3), assign!(r, 0, 0, 1))
        else
            u += rest
            if u <= slimbmax
                _XInt(flipsign(u % SLimb, x.x))
            elseif u <= typemax(Limb)
                _XInt(flipsign(one(SLimb), x.x),
                      assign!(r, u % Limb))
            else
                _XInt(flipsign(SLimb(2), x.x),
                      assign!(r, u % Limb , (u >> BITS_PER_LIMB) % Limb))
            end
        end
    else
        rv = vec!(r, rl + neg)
        if iszero(cnt)
            unsafe_copyto!(rv, 1, xv, offset+1, rl)
        else
            discarded = MPN.rshift(rv=>rl, @view(xv[offset+1:xl]), cnt)
            if neg
                rest |= !iszero(discarded)
            end
            rl -= iszero(@inbounds rv[rl])
        end
        if rest
            d = MPN.add_1(rv, rv=>rl, Limb(1))
            @inbounds rv[rl+1] = d
            rl += !iszero(d)
        end
        _XInt(flipsign(rl, x.x), rv)
    end
end

@inline and!(r::XIntV, x, y) = and!(r, _promote(x), _promote(y))

@inline and!(x::XInt, y::Union{ShortMax,XSigned}) = and!(x, x, y)

@inline function and!(r::XIntV, x::Union{ShortMax,XSigned}, y::Union{ShortMax,XSigned})
    xshort = is_limb(x)
    yshort = is_limb(y)
    if xshort & yshort
        XInt!(r, limb(x) & limb(y)) # if y is unsigned, result of & is also unsigned and
                                    # XInt! does the right thing
    elseif xshort
        iszero(x) ? _XInt(0) : and_small!(r, y, limb(x))
    elseif yshort
        iszero(y) ? _XInt(0) : and_small!(r, x, limb(y))
    else
        and_big!(r, x, y)
    end
end

function and_small!(r::XIntV, x::XInt, y::Limb)
    xv = vec(x)
    x1 = @inbounds xv[1]
    if ispos(x)
        XInt!(r, x1 & y)
    else
        # x & y = ~(-x - 1) & y
        XInt!(r, ~(x1-1) & y)
    end
end

@noinline function and_small!(r::XIntV, x::XSigned, y::SLimb)
    # @assert !iszero(x) && !iszero(y) && !is_short(x) && is_short(y)
    # y is short, but can also be typemin(SLimb)
    xl, xv = lenvec(x)
    yy = y % Limb
    x1 = @inbounds xv[1]
    if ispos(x)
        if y > 0
            # x & y <= min(x, y) == y, so the result is short
            _XInt(x1 & yy)
        else # x > 0, y < 0
            u = x1 & yy
            res = _copy!(r, x)
            rv = vec(res)
            @inbounds rv[1] = u
            # @assert_reduced x # if x is not reduced, the result could be immediate
            res
        end
    elseif y > 0 # x < 0
        # x & y = ~(-x - 1) & y; result >= 0, and <= y, so is immediate
        _XInt((~(x1-1) & yy) % SLimb)
    else # x < 0, y < 0
        # let u = -x: -u & y = ~(u-1) & y = ~(u-1 | ~y) = -(u-1 | ~y) - 1
        # the result is < 0, and it's easy to see that it's also <= min(x, y), so
        # it can't be an immediate integer (provided x isn't)

        if iszero(x1)
            # if x1 == 0, we see that a 2-complement representation of x would also
            # have a null first limb, so x & y == x
            XInt(x)
        else
            # x1 > 0, so x-1 differs from x only in the 1st limb
            r1 = (x1-1) | ~y
            if r1 == typemax(Limb)
                # +1 overflows, so we need to allocate one more limb, just in case
                rv = vec!(r, xl+1)
                @inbounds rv[1] = zero(Limb)
                carry = true
                @inbounds for i=2:xl
                    u = xv[i]
                    if carry && u == typemax(Limb)
                        rv[i] = zero(Limb)
                    else
                        rv[i] = u+carry
                        carry = false
                    end
                end
                @inbounds rv[xl+1] = carry
                rl = xl + carry
                _XInt(-rl, rv)
            else
                rv = vec!(r, xl)
                _copy!(rv, xv, 2:xl)
                @inbounds rv[1] = r1+1
                _XInt(flipsign(xl, _sign(x)), rv)
            end
        end
    end
end

@noinline function and_big!(r::XIntV, x::XSigned, y::XSigned)
    # @assert !iszero(x) && !iszero(y) && !is_short(x) && !is_short(y)
    if ispos(x)
        if ispos(y)
            xl, xv = lenvec(x)
            yl, yv = lenvec(y)
            if xl < yl
                xv, yv = yv, xv
                xl, yl = yl, xl
            end
            rl = yl
            local rvl # stores the last limb in r
            @inbounds while rl > 0
                rvl = xv[rl] & yv[rl]
                iszero(rvl) || break
                rl -= 1
            end
            if rl <= 1
                XInt!(r, rvl)
            else
                rv = vec!(r, rl)
                # MPN.and_n(rv, xv, yv, rl)
                for idx = 1:rl
                    # we already computed the value of the last iteration, but
                    # it seems no more expensive to just recompute it in this loop
                    @inbounds rv[idx] = xv[idx] & yv[idx]
                end
                _XInt(rl, rv)
            end
        else # x > 0, y < 0
            x, y = y, x
            @goto mixed
        end
    elseif ispos(y) # x < 0
        @label mixed
        # result is positive
        # x & y = ~(-x - 1) & y
        xl, xv = lenvec(x)
        yl, yv = lenvec(y)
        if yl <= xl
            # we don't know the resulting rl in advance; on the high limbs of -x,
            # we don't know yet the effect of `-1`. So we first scan from the low
            # end: as long as x's limbs are 0, they become 0xF...F after `-1`, and
            # 0 after `~`, so r's corresponding limbs are 0.
            i = 1
            local xi
            while i <= yl
                xi = @inbounds xv[i]
                iszero(xi) || break
                i += 1
            end
            if yl < i
                _XInt(0)
            else
                rl = yl
                local rvl
                # we scan from the high end to compute rl: above i, the `-1` had no effect,
                # and we search for the hightest non-zero limb in r
                while rl > i
                    rvl = @inbounds ~xv[rl] & yv[rl]
                    iszero(rvl) || break
                    rl -= 1
                end
                if rl == i
                    # all the limbs of r above and below i are 0
                    rvl = @inbounds ~(xi-1) & yv[i]
                    if iszero(rvl)
                        rl = 0
                    end
                end
                if rl <= 1
                    XInt!(r, rvl)
                else
                    rv = vec!(r, rl)
                    for j = 1:i-1
                        @inbounds rv[j] = zero(Limb)
                    end
                    if i < rl
                        @inbounds rv[i] = ~(xv[i]-1) & yv[i] # rvl if i == rl
                        for j = i+1:rl-1
                            @inbounds rv[j] = ~xv[j] & yv[j]
                        end
                    end
                    @inbounds rv[rl] = rvl
                    _XInt(rl, rv)
                end
            end
        else # yl > xl
            # no need to scan, we know that rl == yl
            rv = vec!(r, yl)
            xok = false
            @inbounds for i = 1:xl
                xi = xv[i]
                yi = yv[i]
                if xok
                    rv[i] = ~xi & yi
                else
                    xok = !iszero(xi)
                    rv[i] = ~(xi-1) & yi
                end
            end
            for i = xl+1:yl
                @inbounds rv[i] = yv[i]
            end
            _XInt(yl, rv)
        end
    else # x < 0, y < 0
        # result is negative
        # -u & -v = ~(u-1) & ~(v-1) = ~(u-1 | v-1) = -(u-1 | v-1) - 1
        xl, xv = lenvec(x)
        yl, yv = lenvec(y)
        if xl < yl
            xv, yv = yv, xv
            xl, yl = yl, xl
        end
        rl = xl
        rv = vec!(r, rl+1)
        xok, yok, oneok = false, false, false
        for i=1:xl
            xi = @inbounds xv[i]
            yi = i <= yl ? @inbounds(yv[i]) : zero(Limb)
            if !xok
                xok = !iszero(xi) # carry won't propagate anymore
                xi -= 1
            end
            if !yok
                yok = !iszero(yi)
                yi -= 1
            end
            u = xi | yi
            if !oneok
                oneok = u != typemax(Limb)
                u += 1
            end
            @inbounds rv[i] = u
        end
        if !oneok
            rl += 1
            @inbounds rv[rl] = one(Limb)
        end
        _XInt(-rl, rv)
    end
end

## ior (|)

@inline ior!(r::XIntV, x, y) = ior!(r, _promote(x), _promote(y))
@inline ior!(x::XInt, y::Union{ShortMax,XSigned}) = ior!(x, x, y)

@inline function ior!(r::XIntV, x::Union{ShortMax,XSigned}, y::Union{ShortMax,XSigned})
    xshort = is_limb(x)
    yshort = is_limb(y)
    if xshort & yshort
        x1 = limb(x)
        y1 = limb(y)
        r1 = x1 | y1
        if x1 >= 0 && y1 >= 0 # result is >= 0
            XInt!(r, unsigned(r1))
        else # result is < 0
            XInt!(r, signed(r1))
        end
    elseif xshort
        iszero(x) ? XInt(y) : ior_small!(r, y, limb(x))
    elseif yshort
        iszero(y) ? XInt(x) : ior_small!(r, x, limb(y))
    else
        ior_big!(r, x, y)
    end
end

@noinline function ior_small!(r::XIntV, x::XSigned, y::Short)
    # @assert !iszero(x) && !iszero(y) && !is_short(x) && is_short(y)
    # y is small
    xl, xv = lenvec(x)
    yy = y % Limb
    x1 = @inbounds xv[1]
    if ispos(x)
        if y > 0
            res = _copy!(r, x)
            rv = vec(res)
            @inbounds rv[1] = x1 | yy
            res
        else # x > 0, y < 0
            # when fits(y), y is short, so more than its sign bit is set,
            # and x | y >= y (result is short); but when y doesn't come
            # from a XInt, y==slimbmin is possible, so result is not
            # guarranteed to be short
            # _XInt(x1 | yy) # when fits(y)
            XInt!(r, (x1 % SLimb) | y)
        end
    elseif y > 0 # x < 0
        # x | y = ~(-x - 1) | y = ~((-x - 1) & ~y) = -((-x - 1) & ~y) - 1
        carry = iszero(x1)
        r1 = (x1-1) & ~yy + 1
        # a small analysis shows that, given y > 0, the `+1` can't carry over upper limbs

        if xl == 1
            # then !carry as -x > 0
            XInt_neg!(r, r1)
        else
            x2 = @inbounds xv[2]
            r2 = x2 - carry
            carry &= iszero(x2)
            if xl == 2 && iszero(r2)
                XInt_neg!(r, r1)
            else
                rv = vec!(r, xl)
                @inbounds rv[1] = r1
                @inbounds rv[2] = r2
                @inbounds for i=3:xl
                    xi = xv[i]
                    rv[i] = xi - carry
                    carry &= iszero(xi)
                end
                rl = xl - iszero(@inbounds rv[xl])
                _XInt(-rl, rv)
            end
        end
    else # x < 0, y < 0
        # x | y = ~(-x-1) | y # similar to x > 0, y < 0
        # _XInt(~(x1-1) | yy)
        XInt!(r, (~(x1-1) | yy) % SLimb)
    end
end

@noinline function ior_big!(r::XIntV, x::XSigned, y::XSigned)
    # @assert !iszero(x) && !iszero(y) && !is_short(x) && !is_short(y)
    # we keep here (and in and_big!) some type instabilities if
    # typeof(x) != typeof(y), but the overhead is not that bad, and might
    # not be concerning enough to warrant type-stabilizing this function
    # (like we did for addbig!)
    if ispos(x)
        if ispos(y) # r > 0
            xl, xv = lenvec(x)
            yl, yv = lenvec(y)
            if xl < yl
                xv, yv = yv, xv
                xl, yl = yl, xl
            end
            rv = vec!(r, xl)
            for i=1:yl
                @inbounds rv[i] = xv[i] | yv[i]
            end
            for i=yl+1:xl
                @inbounds rv[i] = xv[i]
            end
            _XInt(xl, rv)
        else # x > 0, y < 0 : r < 0
            x, y = y, x
            @goto mixed
        end
    elseif ispos(y) # x < 0 : r < 0
        @label mixed
        # x | y = ~(-x - 1) | y = ~((-x - 1) & ~y) = -((-x - 1) & ~y) - 1
        xl, xv = lenvec(x)
        yl, yv = lenvec(y)
        if yl + 2 <= xl
            # then -x-1 still has at least one more limb than y, which is not cancelled
            # when &-ed with ~y; then rl could be xl, or xl-1 when xv[l] == 1 and
            # xv[l-1] == 0
            rv = vec!(r, xl)
            cm = cp = true # carries for -1 and +1
            @inbounds for i=1:yl
                xi = xv[i]
                yi = yv[i]
                ri = (xi - cm) & ~yi + cp
                cm &= iszero(xi)
                cp &= iszero(ri) # note: !cp ==> !cm, i.e. cm can't last longer than cp
                rv[i] = ri
            end
            # at this point, we know for sure that !cp
            # bewond y's limbs, ~y is all ones, so we just write x limbs, carrying out cm
            for i=yl+1:xl
                xi = xv[i]
                rv[i] = xi - cm
                cm &= iszero(xi)
            end
            rl = xl - iszero(@inbounds rv[xl])
            _XInt(-rl, rv)
        else # yl >= xl - 1
            # rl could be anything, e.g.  -(2^128) | (2^128-1) == -1
            # so the idea is to scan from the end to find the highest non-zero limb
            # but for this, we need to know where the carry from `-1` stops: this is the
            # first non-zero position in xv, call it k; we therefore need to first scan
            # from the beginning to find k, and then from the end with the
            # formula xi & ~yi; if it happens that this is all zero beyond k, we would
            # then need to scan again from the beginning to find rl, before allocating;
            # this is a rare case, so we might not be concerned with having to scan
            # twice from the beginning; but for simplicity, in the first scan from the
            # beginning, we also record the highest non-zero ri upto k (this is a bit
            # more expensive than just finding the first non-zero xi, but in general
            # this loop won't have many iterations)
            cm = cp = true # carries for -1 and +1
            i = 0 # the first non-zero position in xv
            rl = 0
            @inbounds while cm
                i += 1
                xi = xv[i]
                yi = i > yl ? zero(Limb) : yv[i]
                ri = (xi - cm) & ~yi + cp
                cm &= iszero(xi)
                cp &= iszero(ri)
                if !iszero(ri)
                    rl = i
                end
            end
            k = i
            # beyond i, the formula is xi & ~yi: we scan from the end to find the first non-0
            i = xl
            @inbounds while i > k
                yi = i > yl ? zero(Limb) : yv[i]
                iszero(xv[i] & ~yi) || break
                i -= 1
            end
            # so far, rl points at the highest non-zero limb upto k in the result
            if i != k # not all zero above k
                rl = i
            end
            # we are now ready to store the result
            if rl == 1
                r1 = @inbounds (xv[1]-1) & ~yv[1] + 1
                XInt_neg!(r, r1)
            else
                rv = vec!(r, rl)
                cm = cp = true # carries for -1 and +1
                @inbounds for i=1:rl
                    xi = xv[i]
                    yi = i > yl ? zero(Limb) : yv[i]
                    ri = (xi - cm) & ~yi + cp
                    cm &= iszero(xi)
                    cp &= iszero(ri)
                    rv[i] = ri
                end
                _XInt(-rl, rv)
            end
        end
    else # x < 0, y < 0 : r < 0
        # -u | -v = ~(u-1) | ~(v-1) = ~(u-1 & v-1) = -(u-1 & v-1) - 1
        # we do a scan, similar to what we do for x < 0, y > 0 and yl >= xl - 1
        xl, xv = lenvec(x)
        yl, yv = lenvec(y)
        if xl < yl
            xv, yv = yv, xv
            xl, yl = yl, xl
        end

        cx, cy, cp = true, true, true
        i = 0
        rl = 0
        while (cx | cy) && i < yl
            i += 1
            xi = @inbounds xv[i]
            yi = @inbounds yv[i]
            u = (xi-cx) & (yi-cy) + cp
            if !iszero(u)
                rl = i
            end
            cx &= iszero(xi)
            cy &= iszero(yi)
            cp &= iszero(u) # cp = cx & cy would also work
        end
        k = i
        i = yl
        @inbounds while i > k
            iszero(xv[i] & yv[i]) || break
            i -= 1
        end
        if i > k
            rl = i
        end
        if rl == 1
            r1 = @inbounds (xv[1]-1) & (yv[1]-1) + 1
            XInt_neg!(r, r1)
        else
            rv = vec!(r, rl)
            cx, cy, cp = true, true, true
            @inbounds for i=1:rl
                xi = xv[i]
                yi = yv[i]
                rv[i] = u = (xi-cx) & (yi-cy) + cp
                cx &= iszero(xi)
                cy &= iszero(yi)
                cp &= iszero(u)
            end
            _XInt(-rl, rv)
        end
    end
end

@inline xor!(r::XIntV, x, y) = xor!(r, _promote(x), _promote(y))
@inline xor!(x::XInt, y::Union{ShortMax,XSigned}) = xor!(x, x, y)

@inline function xor!(r::XIntV, x::Union{ShortMax,XSigned}, y::Union{ShortMax,XSigned})
    xshort = is_limb(x)
    yshort = is_limb(y)
    if xshort & yshort
        x1 = limb(x)
        y1 = limb(y)
        XInt!(r, _widen(x1) ⊻ _widen(y1))
    elseif xshort
        iszero(x) ? XInt(y) : xor_small!(r, y, limb(x))
    elseif yshort
        iszero(y) ? XInt(x) : xor_small!(r, x, limb(y))
    else
        xor_big!(r, x, y)
    end
end

@noinline function xor_small!(r::XIntV, x::XSigned, y::Limb)
    xl, xv = lenvec(x)
    x1 = @inbounds xv[1]
    if ispos(x)
        if xl == 1
            XInt!(r, x1 ⊻ y)
        else
            res = _copy!(r, x)
            rv = vec(res)
            @inbounds rv[1] = x1 ⊻ y
            res
        end
    else
        # x ⊻ y = ~(~x ⊻ y) = ~((-x-1) ⊻ y) = -((-x-1) ⊻ y) - 1
        cm = iszero(x1)
        r1 = (x1-1) ⊻ y + 1
        cp = iszero(r1) # if cm, then !cp (as y != 0)
        if !cp && xl == 1
            # xl == 1 ==> !cm
            XInt_neg!(r, r1)
        elseif xl == 2
            x2 = @inbounds xv[2]
            r2 = x2 - cm + cp
            if iszero(r2)
                if cp
                    _XInt(SLimb(-3), assign!(r, r1, r2, one(Limb)))
                else
                    XInt_neg!(r, r1)
                end
            else
                _XInt(SLimb(-2), assign!(r, r1, r2))
            end
        else
            rv = vec!(r, xl+1)
            @inbounds rv[1] = r1
            @inbounds for i=2:xl
                xi = xv[i]
                rv[i] = xi - cm + cp
                cm &= iszero(xi)
                cp &= xi == typemax(Limb)
            end
            @inbounds rv[xl+1] = cp
            _XInt(-(xl + !iszero(cp)), rv)
        end
    end
end

@noinline function xor_small!(r::XIntV, x::XSigned, y::SLimb)
    # @assert !iszero(x) && !iszero(y) && !is_short(x) && is_short(y)
    xl, xv = lenvec(x)
    yy = y % Limb
    x1 = @inbounds xv[1]
    if ispos(x)
        if y > 0
            res = _copy!(r, x)
            rv = vec(res)
            @inbounds rv[1] = x1 ⊻ yy
            res
        else # x > 0, y < 0
            # x ⊻ y = ~(x ⊻ ~y) = -(x ⊻ ~y) - 1
            # if x is one limb, then the high bit of its limb is set,
            # but the high bit of ~y is not, so the result can't be immediate
            rv = vec!(r, xl+1) # ?TODO: avoid +1 in the majority of cases where its not needed?
            @inbounds rv[1] = r1 = x1 ⊻ ~yy + 1
            c = iszero(r1)
            for i=2:xl
                xi = @inbounds xv[i]
                @inbounds rv[i] = xi + c
                c &= xi == typemax(Limb)
            end
            @inbounds rv[xl+1] = c
            _XInt(-(xl + !iszero(c)), rv)
        end
    elseif y > 0 # x < 0
        # x ⊻ y = ~(~x ⊻ y) = ~((-x-1) ⊻ y) = -((-x-1) ⊻ y) - 1
        cm = iszero(x1)
        r1 = (x1-1) ⊻ yy + 1
        cp = iszero(r1) # if cm, then !cp
        if !cp && xl == 1
            # xl == 1 ==> !cm
            XInt_neg!(r, r1)
        else
            rv = vec!(r, xl+1)
            @inbounds rv[1] = r1
            @inbounds for i=2:xl
                xi = xv[i]
                rv[i] = xi - cm + cp
                cm &= iszero(xi)
                cp &= xi == typemax(Limb)
            end
            @inbounds rv[xl+1] = cp
            rl =
                if cp
                    xl + 1
                elseif iszero(@inbounds rv[xl])
                    # can happen because of cm, but then !iszero(rv[xl-2])
                    xl - 1
                else
                    xl
                end
            _XInt(-rl, rv)
        end
    else # x < 0, y < 0 (result >= 0)
        # x ⊻ y = ~x ⊻ ~y = (-x - 1) ⊻ ~y
        r1 = (x1-one(Limb)) ⊻ ~yy
        if xl == 1
            # res is immediate if x1 == limb1min
            XInt!(r, r1)
        else
            rv = vec!(r, xl)
            @inbounds rv[1] = r1
            c = iszero(x1)
            local ri
            for i=2:xl
                xi = @inbounds xv[i]
                @inbounds ri = rv[i] = xi-c
                c &= iszero(xi)
            end
            _XInt(xl - iszero(ri), rv)
        end
    end
end

@noinline function xor_big!(r::XIntV, x::XSigned, y::XSigned)
    # @assert !iszero(x) && !iszero(y) && !is_short(x) && !is_short(y)
    xl, xv = lenvec(x)
    yl, yv = lenvec(y)
    if ispos(x)
        if ispos(y)
            if xl < yl
                xl, yl = yl, xl
                xv, yv = yv, xv
            end
            if xl == yl
                local i = xl
                local rvl
                while i > 0
                    xi = @inbounds xv[i]
                    yi = @inbounds yv[i]
                    rvl = xi ⊻ yi
                    iszero(rvl) || break
                    i -= 1
                end
                if i == 0
                    XInt(0)
                elseif i == 1
                    XInt!(r, rvl)
                else
                    rv = vec!(r, i)
                    rl = i
                    @inbounds rv[rl] = rvl
                    i -= 1
                    while i > 0
                        xi = @inbounds xv[i]
                        yi = @inbounds yv[i]
                        @inbounds rv[i] = xi ⊻ yi
                        i -= 1
                    end
                    _XInt(rl, rv)
                end
            else
                rv = vec!(r, xl)
                for i=1:yl
                    @inbounds rv[i] = xv[i] ⊻ yv[i]
                end
                for i=yl+1:xl
                    @inbounds rv[i] = xv[i]
                end
                _XInt(xl, rv)
            end
        else # x > 0, y < 0
            xl, yl = yl, xl
            xv, yv = yv, xv
            @goto mixed
        end
    elseif ispos(y) # x < 0
        @label mixed
        # x ⊻ y = ~(~x ⊻ y) = ~((-x-1) ⊻ y) = -((-x-1) ⊻ y) - 1
        if yl > xl
            rv = vec!(r, yl + 1)
            cp = cm = true
            @inbounds for i=1:xl
                xi = xv[i]
                yi = yv[i]
                ri = rv[i] = (xi-cm) ⊻ yi + cp
                cm &= iszero(xi)
                cp &= iszero(ri)
            end
            @inbounds for i=xl+1:yl
                ri = rv[i] = yv[i] + cp
                cp &= iszero(ri)
            end
            @inbounds rv[yl+1] = cp
            _XInt(-(yl + !iszero(cp)), rv)
        elseif yl + 2 <= xl
            rv = vec!(r, xl + 1)
            cp = cm = true
            @inbounds for i=1:yl
                xi = xv[i]
                yi = yv[i]
                ri = rv[i] = (xi-cm) ⊻ yi + cp
                cm &= iszero(xi)
                cp &= iszero(ri)
            end
            @inbounds for i=yl+1:xl
                xi = xv[i]
                rv[i] = xi-cm+cp
                cm &= iszero(xi)
                cp &= iszero(ri)
            end
            @inbounds rv[xl+1] = cp
            _XInt(-(xl + !iszero(cp)), rv)
        else # xl == yl || xl == yl+1
            x1 = @inbounds xv[1]
            y1 = @inbounds yv[1]
            r1 = (x1-1) ⊻ y1 + 1
            cm = iszero(x1)
            cp = iszero(r1)
            # when fits(r1), all upper limbs could be 0 and result could be immediate,
            # so we scan from the low end to check that out
            i = 2
            local ri
            @inbounds while i <= yl
                xi = xv[i]
                yi = yv[i]
                ri = (xi-cm) ⊻ yi + cp
                cm &= iszero(xi)
                cp &= iszero(ri)
                iszero(ri) || break
                i += 1
            end
            if i > yl && xl > yl
                ri = @inbounds xv[xl] - cm + cp
                cp &= iszero(ri)
                if iszero(ri)
                    i += 1
                else
                    cp = false
                end
            end
            if !cp && i > xl
                XInt_neg!(r, r1)
            else
                rv = vec!(r, xl + (i > xl))
                @inbounds rv[1] = r1
                fill!(@view(rv[2:i-1]), zero(Limb))
                if i <= xl # then !cp
                    @inbounds rv[i] = ri
                    rl = i
                    i += 1
                    @inbounds while i <= yl
                        xi = xv[i]
                        yi = yv[i]
                        ri = rv[i] = (xi-cm) ⊻ yi
                        cm &= iszero(xi)
                        if !iszero(ri)
                            rl = i
                        end
                        i += 1
                    end
                    if xl > yl && i <= xl
                        @inbounds ri = rv[i] = xv[i]-cm
                        if !iszero(ri)
                            rl = i
                        end
                    end
                    _XInt(-rl, rv)
                else
                    @inbounds rv[xl+1] = cp
                    _XInt(-SLimb(xl+1), rv)
                end
            end
        end
    else # x < 0, y < 0 (result >= 0)
        # x ⊻ y = ~x ⊻ ~y = (-x - 1) ⊻ (-y - 1)
        if xl < yl
            xl, yl = yl, xl
            xv, yv = yv, xv
        end
        @inbounds x1 = xv[1]
        @inbounds y1 = yv[1]
        r1 = (x1-1) ⊻ (y1-1)
        cx = iszero(x1)
        cy = iszero(y1)
        i = 2
        local ri
        while i <= xl
            @inbounds xi = xv[i]
            @inbounds yi = i <= yl ? yv[i] : zero(Limb)
            ri = (xi-cx) ⊻ (yi-cy)
            cx &= iszero(xi)
            cy &= iszero(yi)
            iszero(ri) || break
            i += 1
        end
        if i > xl
            XInt!(r, r1)
        else
            rv = vec!(r, xl)
            @inbounds rv[1] = r1
            fill!(@view(rv[2:i-1]), zero(Limb))
            @inbounds rv[i] = ri
            rl = i
            @inbounds for i=i+1:xl
                xi = xv[i]
                yi = i <= yl ? yv[i] : zero(Limb)
                ri = rv[i] = (xi-cx) ⊻ (yi-cy)
                cx &= iszero(xi)
                cy &= iszero(yi)
                if !iszero(ri)
                    rl = i
                end
            end
            _XInt(rl, rv)
        end
    end
end