text:
include(joinpath(pwd(), "src/direct_policy_optimization/dpo.jl"))
include(joinpath(@__DIR__, "quadrotor_broken_propeller.jl"))

# Additive noise model
model = additive_noise_model(model)

function fd(model::Quadrotor, x⁺, x, u, w, h, t)
    midpoint_implicit(model, x⁺, x, u, w, h) - w
end

# Nominal solution
X̄, Ū = unpack(Z̄, prob)
prob_nom = prob.prob

# DPO
N = 2 * model.n
D = 2 * model.d

β = 1.0
δ = 1.0e-3

# initial samples
x1_sample = resample(x1, Diagonal(ones(model.n)), 1.0e-3)

# mean problem
prob_mean = trajectory_optimization(
				model,
				EmptyObjective(),
				T,
				ul = control_bounds(model, T, [Inf * ones(4); 0.0], [Inf * ones(4); 0.0])[1],
				uu = control_bounds(model, T, [Inf * ones(4); 0.0], [Inf * ones(4); 0.0])[2],
				dynamics = false
				)

# sample problems
uu_sample = []
for i = 1:6
	push!(uu_sample, [uu1; Inf])
end
for i = 1:6
	push!(uu_sample, [uu2; Inf])
end
for i = 1:6
	push!(uu_sample, [uu3; Inf])
end
for i = 1:6
	push!(uu_sample, [uu4; Inf])
end

prob_sample = [trajectory_optimization(
				model,
				EmptyObjective(),
				T,
				xl = state_bounds(model, T, x1 = x1_sample[i])[1],
				xu = state_bounds(model, T, x1 = x1_sample[i])[2],
				ul = ul,
				uu = control_bounds(model, T, Inf * ones(model.m), uu_sample[i]),
				dynamics = false,
				con = con_free_time
				) for i = 1:N]

# sample objective
Q = [(t < T ? Diagonal([10.0; 10.0; 1.0])
	: Diagonal(100.0 * ones(model.n))) for t = 1:T]
R = [Diagonal(1.0e-1 * ones(model.m)) for t = 1:T-1]

obj_sample = sample_objective(Q, R)
policy = linear_feedback(model.n, model.m-1)
dist = disturbances([Diagonal(δ * [1.0; 1.0; 0.1]) for t = 1:T-1])
sample = sample_params(β, T)

prob_dpo = dpo_problem(
	prob_nom, prob_mean, prob_sample,
	obj_sample,
	policy,
	dist,
	sample)

# TVLQR policy
K = tvlqr(model, X̄, Ū, Q, R, h)

z0_dpo = zeros(prob_dpo.num_var)
z0_dpo[prob_dpo.prob.idx.nom] = pack(X̄, Ū, prob_nom)
z0_dpo[prob_dpo.prob.idx.mean] = pack(X̄, Ū, prob_nom)
for i = 1:N
	z0_dpo[prob_dpo.prob.idx.sample[i]] = pack(X̄, Ū, prob_nom)
end
for j = 1:(N + D)
	z0_dpo[prob_dpo.prob.idx.slack[j]] = vcat(X̄[2:end]...)
end
for t = 1:T-1
	z0_dpo[prob_dpo.prob.idx.policy[prob_dpo.prob.idx.θ[t]]] = vec(copy(K[t]))
end

# Solve
Z = solve(prob_dpo, copy(z0_dpo),
	tol = 1.0e-3, c_tol = 1.0e-3,
	max_iter = 1000)

# Simulate TVLQR

# TVLQR cost
Q_lqr = [t < T ? Diagonal(10.0*ones(model.nx)) : Diagonal(100.0*ones(model.nx)) for t = 1:T]
R_lqr = [Diagonal(1.0*ones(model.nu)) for t = 1:T-1]
H_lqr = [1.0 for t = 1:T-1]

K = TVLQR_gains(model,X_nom,U_nom,[H_nom[1] for t = 1:T-1],Q_lqr,R_lqr)

# Sample
N = 2*model.nx
models = [model for i = 1:N]
β = 1.0
w = 1.0e-3*ones(model.nx)
γ = N
x1_sample = resample([x1 for i = 1:N],β=1.0,w=1.0e-3*ones(model.nx))

xl_traj_sample = [[copy(xl) for t = 1:T] for i = 1:N]
xu_traj_sample = [[copy(xu) for t = 1:T] for i = 1:N]

ul_traj_sample = [[copy(ul_nom) for t = 1:T-1] for i = 1:N]
uu_traj_sample = [[copy(uu_nom) for t = 1:T-1] for i = 1:N]
for t = 1:T-1
	for i = 1:6
		uu_traj_sample[i][t][1] *= 0.5
	end
	for i = 6 .+ (1:6)
		uu_traj_sample[i][t][2] *= 0.5
	end
	for i = 12 .+ (1:6)
		uu_traj_sample[i][t][3] *= 0.5
	end
	for i = 18 .+ (1:6)
		uu_traj_sample[i][t][4] *= 0.5
	end
end

for i = 1:N
    xl_traj_sample[i][1] = copy(x1_sample[i])
    xu_traj_sample[i][1] = copy(x1_sample[i])
end

prob_sample = init_sample_problem(prob,models,Q_lqr,R_lqr,H_lqr,
    xl=xl_traj_sample,
    xu=xu_traj_sample,
	ul=ul_traj_sample,
    uu=uu_traj_sample,
    β=β,w=w,γ=γ)

prob_sample_moi = init_MOI_Problem(prob_sample)

Z0_sample = pack(X_nom,U_nom,H_nom[1],K,prob_sample,r=0.001)

# # Solve
# Z_sample_sol = solve(prob_sample_moi,copy(Z0_sample),nlp=:SNOPT7,
#     time_limit=60*60,tol=1.0e-2,c_tol=1.0e-2)

using JLD
# @save joinpath(pwd(),"quadrotor_prop_fail.jld") Z_sample_sol
@load joinpath(pwd(),"quadrotor_prop_fail.jld") Z_sample_sol

# Z_sample_sol = solve(prob_sample_moi,copy(Z_sample_sol),nlp=:SNOPT7,
#     time_limit=60*60,tol=1.0e-2,c_tol=1.0e-2)

# Unpack solutions
X_nom_sample, U_nom_sample, H_nom_sample, X_sample, U_sample, H_sample = unpack(Z_sample_sol,prob_sample)

K_sample = [reshape(Z_sample_sol[prob_sample.idx_K[t]],model.nu,model.nx) for t = 1:T-1]
# Time trajectories
t_nominal = zeros(T)
t_sample = zeros(T)
for t = 2:T
    t_nominal[t] = t_nominal[t-1] + H_nom[t-1]
    t_sample[t] = t_sample[t-1] + H_nom_sample[t-1]
end

display("time (nominal): $(sum(H_nom))s")
display("time (sample): $(sum(H_nom_sample))s")

# # Plots results
using Plots
# # Position trajectory
x_nom_pos = [X_nom[t][1] for t = 1:T]
y_nom_pos = [X_nom[t][2] for t = 1:T]
# pts = Plots.partialcircle(0,2π,100,r)
# cx,cy = Plots.unzip(pts)
# cx1 = [_cx + xc1 for _cx in cx]
# cy1 = [_cy + yc1 for _cy in cy]
# cx2 = [_cx + xc2 for _cx in cx]
# cy2 = [_cy + yc2 for _cy in cy]
# cx3 = [_cx + xc3 for _cx in cx]
# cy3 = [_cy + yc3 for _cy in cy]
# cx4 = [_cx + xc4 for _cx in cx]
# cy4 = [_cy + yc4 for _cy in cy]
# # cx5 = [_cx + xc5 for _cx in cx]
# # cy5 = [_cy + yc5 for _cy in cy]
#
# plt = plot(Shape(cx1,cy1),color=:red,label="",linecolor=:red)
# # plt = plot!(Shape(cx2,cy2),color=:red,label="",linecolor=:red)
# # plt = plot!(Shape(cx3,cy3),color=:red,label="",linecolor=:red)
# # plt = plot!(Shape(cx4,cy4),color=:red,label="",linecolor=:red)
# # # plt = plot(Shape(cx5,cy5),color=:red,label="",linecolor=:red)
#
# for i = 1:N
#     x_sample_pos = [X_sample[i][t][1] for t = 1:T]
#     y_sample_pos = [X_sample[i][t][2] for t = 1:T]
#     plt = plot!(x_sample_pos,y_sample_pos,aspect_ratio=:equal,
#         color=:cyan,label= i != 1 ? "" : "sample")
# end
# display(plt)
plt = plot()
plt = scatter!(x_nom_pos,y_nom_pos,aspect_ratio=:equal,xlabel="x",ylabel="y",width=4.0,label="TO",color=:purple,legend=:topleft)
x_sample_pos = [X_nom_sample[t][1] for t = 1:T]
y_sample_pos = [X_nom_sample[t][2] for t = 1:T]
plt = plot!(x_sample_pos,y_sample_pos,aspect_ratio=:equal,width=4.0,label="DPO",color=:orange,legend=:bottomright)
#
# savefig(plt,joinpath(@__DIR__,"results/quadrotor_trajectory.png"))
#
# Control
plt = plot(t_nominal[1:T-1],Array(hcat(U_nom...))',color=:purple,width=2.0,
    title="quad",xlabel="time (s)",ylabel="control",
    legend=:bottom,linetype=:steppost)
plt = plot!(t_sample[1:T-1],Array(hcat(U_nom_sample...))',color=:orange,
    width=2.0,linetype=:steppost)
savefig(plt,joinpath(@__DIR__,"results/quadrotor_control.png"))

# # Samples
X_sample
# State samples
plt1 = plot(title="Sample states",legend=:bottom,xlabel="time (s)");
for i = 1:N
    # t_sample = zeros(T)
    # for t = 2:T
    #     t_sample[t] = t_sample[t-1] + H_nom_sample[t-1]
    # end
    plt1 = plot!(hcat(X_sample[i]...)',label="");
end
plt1 = plot!(hcat(X_nom_sample...)',color=:red,width=2.0,
    label=["nominal" "" ""])
display(plt1)
# savefig(plt1,joinpath(@__DIR__,"results/quadrotor_sample_states.png"))
#
# # Control samples
# plt2 = plot(title="Sample controls",xlabel="time (s)",legend=:bottom);
# for i = 1:N
#     t_sample = zeros(T)
#     for t = 2:T
#         t_sample[t] = t_sample[t-1] + H_nom_sample[t-1]
#     end
#     plt2 = plot!(t_sample[1:end-1],hcat(U_sample[i]...)',label="",
#         linetype=:steppost);
# end
# plt2 = plot!(t_sample[1:end-1],hcat(U_nom_sample...)',color=:red,width=2.0,
#     label=["nominal" ""],linetype=:steppost)
# display(plt2)
# savefig(plt2,joinpath(@__DIR__,"results/quadrotor_sample_controls.png"))

using Distributions

model_sim = model
x1_sim = copy(x1)
T_sim = 10*T

W = Distributions.MvNormal(zeros(model_sim.nx),
	Diagonal(1.0e-5*ones(model_sim.nx)))
w = rand(W,T_sim)

W0 = Distributions.MvNormal(zeros(model_sim.nx),
	Diagonal(1.0e-5*ones(model_sim.nx)))
w0 = rand(W0,1)

z0_sim = vec(copy(x1_sim) + w0)

t_nom = range(0,stop=sum(H_nom),length=T)
t_sim_nom = range(0,stop=sum(H_nom),length=T_sim)

t_nom_sample = range(0,stop=sum(H_nom_sample),length=T)
t_sim_nom_sample = range(0,stop=sum(H_nom_sample),length=T_sim)

dt_sim_sample = sum(H_nom_sample)/(T_sim-1)
dt_sim_nom = sum(H_nom)/(T_sim-1)

# simulate
z_tvlqr1, u_tvlqr1, J_tvlqr1, Jx_tvlqr1, Ju_tvlqr1 = simulate_linear_controller(K,
    X_nom,U_nom,model_sim,Q_lqr,R_lqr,T_sim,H_nom[1],z0_sim,w,_norm=2,
	ul=ul_nom,uu=uu1)
z_tvlqr2, u_tvlqr2, J_tvlqr2, Jx_tvlqr2, Ju_tvlqr2 = simulate_linear_controller(K,
    X_nom,U_nom,model_sim,Q_lqr,R_lqr,T_sim,H_nom[1],z0_sim,w,_norm=2,
	ul=ul_nom,uu=uu2)
z_tvlqr3, u_tvlqr3, J_tvlqr3, Jx_tvlqr3, Ju_tvlqr3 = simulate_linear_controller(K,
    X_nom,U_nom,model_sim,Q_lqr,R_lqr,T_sim,H_nom[1],z0_sim,w,_norm=2,
	ul=ul_nom,uu=uu3)
z_tvlqr4, u_tvlqr4, J_tvlqr4, Jx_tvlqr4, Ju_tvlqr4 = simulate_linear_controller(K,
    X_nom,U_nom,model_sim,Q_lqr,R_lqr,T_sim,H_nom[1],z0_sim,w,_norm=2,
	ul=ul_nom,uu=uu4)

z_sample1, u_sample1, J_sample1, Jx_sample1, Ju_sample1 = simulate_linear_controller(K_sample,
    X_nom_sample,U_nom_sample,model_sim,Q_lqr,R_lqr,T_sim,H_nom_sample[1],z0_sim,w,_norm=2,
	ul=ul_nom,uu=uu1)
z_sample2, u_sample2, J_sample2, Jx_sample2, Ju_sample2 = simulate_linear_controller(K_sample,
    X_nom_sample,U_nom_sample,model_sim,Q_lqr,R_lqr,T_sim,H_nom_sample[1],z0_sim,w,_norm=2,
	ul=ul_nom,uu=uu2)
z_sample3, u_sample3, J_sample3, Jx_sample3, Ju_sample3 = simulate_linear_controller(K_sample,
    X_nom_sample,U_nom_sample,model_sim,Q_lqr,R_lqr,T_sim,H_nom_sample[1],z0_sim,w,_norm=2,
	ul=ul_nom,uu=uu3)
z_sample4, u_sample4, J_sample4, Jx_sample4, Ju_sample4 = simulate_linear_controller(K_sample,
    X_nom_sample,U_nom_sample,model_sim,Q_lqr,R_lqr,T_sim,H_nom_sample[1],z0_sim,w,_norm=2,
	ul=ul_nom,uu=uu4)

# plot states
plt_x = plot(t_nom,hcat(X_nom...)[1:model.nx,:]',
	legend=:topright,color=:red,
    label="",width=2.0,xlabel="time (s)",
    title="Quadrotor",ylabel="state")
plt_x = plot!(t_sim_nom,hcat(z_tvlqr1...)[1:model.nx,:]',color=:black,label="",
    width=1.0)
plt_x = plot!(t_sim_nom,hcat(z_tvlqr2...)[1:model.nx,:]',color=:black,label="",
    width=1.0)
plt_x = plot!(t_sim_nom,hcat(z_tvlqr3...)[1:model.nx,:]',color=:black,label="",
    width=1.0)
plt_x = plot!(t_sim_nom,hcat(z_tvlqr4...)[1:model.nx,:]',color=:black,label="",
    width=1.0)

plt_x = plot(t_nom_sample,hcat(X_nom_sample...)[1:model.nx,:]',
	legend=:topright,color=:red,
    label="",width=2.0,xlabel="time (s)",
    title="Quadrotor",ylabel="state")
plt_x = plot!(t_sim_nom_sample,hcat(z_sample1...)[1:model.nx,:]',color=:black,label="",
    width=1.0)
plt_x = plot!(t_sim_nom_sample,hcat(z_sample2...)[1:model.nx,:]',color=:black,label="",
	width=1.0)
plt_x = plot!(t_sim_nom_sample,hcat(z_sample3...)[1:model.nx,:]',color=:black,label="",
    width=1.0)
plt_x = plot!(t_sim_nom_sample,hcat(z_sample4...)[1:model.nx,:]',color=:black,label="",
    width=1.0)

# plot COM
plot_traj = plot(hcat(X_nom...)[1,:],hcat(X_nom...)[2,:],
	legend=:topright,color=:red,
    label="",width=2.0,xlabel="y",ylabel="z",
    title="Quadrotor")
plot_traj = plot!(hcat(z_tvlqr1...)[1,:],hcat(z_tvlqr1...)[2,:],
	color=:black,
    label="",width=1.0)
plot_traj = plot!(hcat(z_tvlqr2...)[1,:],hcat(z_tvlqr2...)[2,:],
	color=:black,
	label="",width=1.0)
plot_traj = plot!(hcat(z_tvlqr3...)[1,:],hcat(z_tvlqr3...)[2,:],
	color=:black,
    label="",width=1.0)
plot_traj = plot!(hcat(z_tvlqr4...)[1,:],hcat(z_tvlqr4...)[2,:],
	color=:black,
    label="",width=1.0)

plot_traj = plot(hcat(X_nom_sample...)[1,:],hcat(X_nom_sample...)[2,:],
	legend=:topright,color=:red,
    label="",width=2.0,xlabel="y",ylabel="z",
    title="Quadrotor")
plot_traj = plot!(hcat(z_sample1...)[1,:],hcat(z_sample1...)[2,:],
	color=:black,
    label="",width=1.0)
plot_traj = plot!(hcat(z_sample2...)[1,:],hcat(z_sample2...)[2,:],
	color=:black,
    label="",width=1.0)
plot_traj = plot!(hcat(z_sample3...)[1,:],hcat(z_sample3...)[2,:],
	color=:black,
    label="",width=1.0)
plot_traj = plot!(hcat(z_sample4...)[1,:],hcat(z_sample4...)[2,:],
	color=:black,
    label="",width=1.0)

plot(t_nom[1:end-1],hcat(U_nom...)',
	linetype=:steppost,width=2.0)
plot!(t_sim_nom[1:end-1],hcat(u_tvlqr1...)',
	linetype=:steppost,color=:black,width=1.0)

plot(t_nom[1:end-1],hcat(U_nom_sample...)',
	linetype=:steppost,width=2.0)
plot!(t_sim_nom[1:end-1],hcat(u_sample1...)',
	linetype=:steppost,color=:black,width=1.0)

# state tracking
(Jx_tvlqr1 + Jx_tvlqr2 + Jx_tvlqr3 + Jx_tvlqr4)/4.0
(Jx_sample1 + Jx_sample2 + Jx_sample3 + Jx_sample4)/4.0

# control tracking
(Ju_tvlqr1 + Ju_tvlqr2 + Ju_tvlqr3 + Ju_tvlqr4)/4.0
(Ju_sample1 + Ju_sample2 + Ju_sample3 + Ju_sample4)/4.0

# objective value
(J_tvlqr1 + J_tvlqr2 + J_tvlqr3 + J_tvlqr4)/4.0
(J_sample1 + J_sample2 + J_sample3 + J_sample4)/4.0

using PGFPlots
const PGF = PGFPlots

# TO trajectory
p_u1_nom = PGF.Plots.Linear(t_nom[1:end-1],hcat(U_nom...)[1,:],
    mark="",style="const plot, color=cyan, line width=2pt, solid")
p_u2_nom = PGF.Plots.Linear(t_nom[1:end-1],hcat(U_nom...)[2,:],
    mark="",style="const plot, color=cyan, line width=2pt, solid")
p_u3_nom = PGF.Plots.Linear(t_nom[1:end-1],hcat(U_nom...)[3,:],
    mark="",style="const plot, color=cyan, line width=2pt, solid")
p_u4_nom = PGF.Plots.Linear(t_nom[1:end-1],hcat(U_nom...)[4,:],
    mark="",style="const plot, color=cyan, line width=2pt, solid")

p_u1_dpo = PGF.Plots.Linear(t_nom_sample[1:end-1],hcat(U_nom_sample...)[1,:],
    mark="",style="const plot, color=orange, line width=2pt, solid")
p_u2_dpo = PGF.Plots.Linear(t_nom_sample[1:end-1],hcat(U_nom_sample...)[2,:],
    mark="",style="const plot, color=orange, line width=2pt, solid")
p_u3_dpo = PGF.Plots.Linear(t_nom_sample[1:end-1],hcat(U_nom_sample...)[3,:],
    mark="",style="const plot, color=orange, line width=2pt, solid")
p_u4_dpo = PGF.Plots.Linear(t_nom_sample[1:end-1],hcat(U_nom_sample...)[4,:],
    mark="",style="const plot, color=orange, line width=2pt, solid")

a1 = Axis([p_u1_nom;p_u1_dpo
    ],
    xmin=0, ymin=0,
    hideAxis=false,
	ylabel="u1",
	xlabel="time",
	)
a2 = Axis([p_u2_nom;p_u2_dpo
    ],
    xmin=0, ymin=0,
    hideAxis=false,
	ylabel="u2",
	xlabel="time",
	)
a3 = Axis([p_u3_nom;p_u3_dpo
    ],
    xmin=0, ymin=0,
    hideAxis=false,
	ylabel="u3",
	xlabel="time",
	)
a4 = Axis([p_u4_nom;p_u4_dpo
    ],
    xmin=0, ymin=0,
    hideAxis=false,
	ylabel="u4",
	xlabel="time",
	)

# Save to tikz format
dir = joinpath(@__DIR__,"results")
PGF.save(joinpath(dir,"quad_prop_u1.tikz"), a1, include_preamble=false)
PGF.save(joinpath(dir,"quad_prop_u2.tikz"), a2, include_preamble=false)
PGF.save(joinpath(dir,"quad_prop_u3.tikz"), a3, include_preamble=false)
PGF.save(joinpath(dir,"quad_prop_u4.tikz"), a4, include_preamble=false)

# #
# # visualize
visualize!(vis,model,z_tvlqr,Δt=dt_sim_nom)
visualize!(vis,model,z_sample4,Δt=dt_sim_sample)
#
for (k,z) in enumerate([z_tvlqr1,z_tvlqr2,z_tvlqr3,z_tvlqr4])
	i = k
	q_to = z
	for t = 1:3:T_sim
		setobject!(vis["traj_to$t$i"], Sphere(Point3f0(0),
			convert(Float32,0.025)),
			MeshPhongMaterial(color=RGBA(128.0/255.0,128.0/255.0,128.0/255.0,1.0)))
		settransform!(vis["traj_to$t$i"], Translation((q_to[t][1],q_to[t][2],q_to[t][3])))
		setvisible!(vis["traj_to$t$i"],false)
	end
end
q_to_nom = X_nom
for t = 1:T
	setobject!(vis["traj_to_nom$t"], Sphere(Point3f0(0),
		convert(Float32,0.05)),
		MeshPhongMaterial(color=RGBA(0.0,255.0/255.0,255.0/255.0,1.0)))
	settransform!(vis["traj_to_nom$t"], Translation((q_to_nom[t][1],q_to_nom[t][2],q_to_nom[t][3])))
	setvisible!(vis["traj_to_nom$t"],false)
end

for (k,z) in enumerate([z_sample1,z_sample2,z_sample3,z_sample4])
	i = k
	q_dpo = z
	for t = 1:3:T_sim
		setobject!(vis["traj_dpo$t$i"], Sphere(Point3f0(0),
			convert(Float32,0.025)),
			MeshPhongMaterial(color=RGBA(128.0/255.0,128.0/255.0,128.0/255.0,1.0)))
		settransform!(vis["traj_dpo$t$i"], Translation((q_dpo[t][1],q_dpo[t][2],q_dpo[t][3])))
		setvisible!(vis["traj_dpo$t$i"],true)
	end
end

q_dpo_nom = X_nom_sample
for t = 1:T
	setobject!(vis["traj_dpo_nom$t"], Sphere(Point3f0(0),
		convert(Float32,0.05)),
		MeshPhongMaterial(color=RGBA(255.0/255.0,127.0/255.0,0.0,1.0)))
	settransform!(vis["traj_dpo_nom$t"], Translation((q_dpo_nom[t][1],q_dpo_nom[t][2],q_dpo_nom[t][3])))
	setvisible!(vis["traj_dpo_nom$t"],true)
end

obj_path = joinpath(pwd(),"/home/taylor/Research/direct_policy_optimization/dynamics/quadrotor/drone.obj")
mtl_path = joinpath(pwd(),"/home/taylor/Research/direct_policy_optimization/dynamics/quadrotor/drone.mtl")

ctm = ModifiedMeshFileObject(obj_path,mtl_path,scale=1.0)
setobject!(vis["drone2"],ctm)
settransform!(vis["drone2"], compose(Translation(X_nom[1][1:3]),LinearMap(MRP(X_nom[1][4:6]...)*RotX(pi/2.0))))

max_stars_repo_name: zacmanchester/DirectMotionPlanning

max_issues_repo_name: zacmanchester/DirectMotionPlanning

max_forks_repo_name: zacmanchester/DirectMotionPlanning


text:
import Base: any, all, count, countnz

#############################
# reduce
# functions in base implemented with a direct loop need to be overloaded to use mapreduce
any(pred, A::GPUArray) = Bool(mapreduce(pred, |, Int32(0), A))
all(pred, A::GPUArray) = Bool(mapreduce(pred, &, Int32(1), A))
count(pred, A::GPUArray) = Int(mapreduce(pred, +, UInt32(0), A))
countnz(A::GPUArray) = Int(mapreduce(x-> x != 0, +, UInt32(0), A))
countnz(A::GPUArray, dim) = Int(mapreducedim(x-> x != 0, +, UInt32(0), A, dim))

Base.:(==)(A::GPUArray, B::GPUArray) = Bool(mapreduce(==, &, Int32(1), A, B))

# hack to get around of fetching the first element of the GPUArray
# as a startvalue, which is a bit complicated with the current reduce implementation
function startvalue(f, T)
    error("Please supply a starting value for mapreduce. E.g: mapreduce($f, $op, 1, A)")
end
startvalue(::typeof(+), T) = zero(T)
startvalue(::typeof(*), T) = one(T)
startvalue(::typeof(Base.scalarmin), T) = typemax(T)
startvalue(::typeof(Base.scalarmax), T) = typemin(T)

# TODO widen and support Int64 and use Base.r_promote_type
gpu_promote_type(op, ::Type{T}) where {T} = T
gpu_promote_type(op, ::Type{T}) where {T<:Base.WidenReduceResult} = T
gpu_promote_type(::typeof(+), ::Type{T}) where {T<:Base.WidenReduceResult} = T
gpu_promote_type(::typeof(*), ::Type{T}) where {T<:Base.WidenReduceResult} = T
gpu_promote_type(::typeof(+), ::Type{T}) where {T<:Number} = typeof(zero(T)+zero(T))
gpu_promote_type(::typeof(*), ::Type{T}) where {T<:Number} = typeof(one(T)*one(T))
gpu_promote_type(::typeof(Base.scalarmax), ::Type{T}) where {T<:Base.WidenReduceResult} = T
gpu_promote_type(::typeof(Base.scalarmin), ::Type{T}) where {T<:Base.WidenReduceResult} = T
gpu_promote_type(::typeof(max), ::Type{T}) where {T<:Base.WidenReduceResult} = T
gpu_promote_type(::typeof(min), ::Type{T}) where {T<:Base.WidenReduceResult} = T

function Base.mapreduce(f::Function, op::Function, A::GPUArray{T, N}) where {T, N}
    OT = gpu_promote_type(op, T)
    v0 = startvalue(op, OT) # TODO do this better
    mapreduce(f, op, v0, A)
end
function acc_mapreduce end
function Base.mapreduce(f, op, v0, A::GPUArray, B::GPUArray, C::Number)
    acc_mapreduce(f, op, v0, A, (B, C))
end
function Base.mapreduce(f, op, v0, A::GPUArray, B::GPUArray)
    acc_mapreduce(f, op, v0, A, (B,))
end
function Base.mapreduce(f, op, v0, A::GPUArray)
    acc_mapreduce(f, op, v0, A, ())
end

@generated function mapreducedim_kernel(state, f, op, R, A, range::NTuple{N, Any}) where N
    types = (range.parameters...,)
    indices = ntuple(i-> Symbol("I_$i"), N)
    Iexpr = ntuple(i-> :(I[$i]), N)
    body = :(@inbounds R[$(Iexpr...)] = op(R[$(Iexpr...)], f(A[$(indices...)])))
    for i = N:-1:1
        idxsym = indices[i]
        if types[i] == Void
            body = quote
                $idxsym = I[$i]
                $body
            end
        else
            rsym = Symbol("r_$i")
            body = quote
                $(rsym) = range[$i]
                for $idxsym in UInt32(first($rsym)):UInt32(last($rsym))
                    $body
                end
            end
        end
        body
    end
    quote
        I = @cartesianidx R state
        $body
        return
    end
end

function Base._mapreducedim!(f, op, R::GPUArray, A::GPUArray)
    range = ifelse.(length.(indices(R)) .== 1, indices(A), nothing)
    gpu_call(mapreducedim_kernel, R, (f, op, R, A, range))
    return R
end

for i = 0:10
    args = ntuple(x-> Symbol("arg_", x), i)
    fargs = ntuple(x-> :(broadcast_index($(args[x]), length, global_index)), i)
    @eval begin
        # http://developer.amd.com/resources/articles-whitepapers/opencl-optimization-case-study-simple-reductions/
        function reduce_kernel(state, f, op, v0::T, A, ::Val{LMEM}, result, $(args...)) where {T, LMEM}
            ui0 = UInt32(0); ui1 = UInt32(1); ui2 = UInt32(2)
            tmp_local = @LocalMemory(state, T, LMEM)
            global_index = linear_index(state)
            acc = v0
            # # Loop sequentially over chunks of input vector
            while global_index <= length(A)
                element = f(A[global_index], $(fargs...))
                acc = op(acc, element)
                global_index += global_size(state)
            end
            # Perform parallel reduction
            local_index = threadidx_x(state) - ui1
            tmp_local[local_index + ui1] = acc
            synchronize_threads(state)

            offset = blockdim_x(state) ÷ ui2
            while offset > ui0
                if (local_index < offset)
                    other = tmp_local[local_index + offset + ui1]
                    mine = tmp_local[local_index + ui1]
                    tmp_local[local_index + ui1] = op(mine, other)
                end
                synchronize_threads(state)
                offset = offset ÷ ui2
            end
            if local_index == ui0
                result[blockidx_x(state)] = tmp_local[ui1]
            end
            return
        end
    end

end

to_cpu(x) = x
to_cpu(x::GPUArray) = Array(x)

function acc_mapreduce(
        f, op, v0::OT, A::GPUArray{T, N}, rest::Tuple
    ) where {T, OT, N}
    dev = device(A)
    blocksize = 80
    threads = 256
    if length(A) <= blocksize * threads
        args = zip(Array(A), to_cpu.(rest)...)
        return mapreduce(x-> f(x...), op, v0, args)
    end
    out = similar(A, OT, (blocksize,))
    fill!(out, v0)
    args = (f, op, v0, A, Val{threads}(), out, rest...)
    gpu_call(reduce_kernel, A, args, ((blocksize,), (threads,)))
    reduce(op, Array(out))
end

ext: jl


text:
type Token
    word::Int
    chars::Vector{Int}
end

function setup_dicts()
    h5file = joinpath(dirname(@__FILE__), ".data/wordembeds_nyt100.h5")
    words = h5read(h5file, "s")
    wordvecs = h5read(h5file, "v")
    worddict = Dict(words[i] => i for i=1:length(words))
    chardict = Dict(' ' => 1)
    worddict, chardict
end

function setup_traindata(worddict, chardict)
    data = CoNLL.read(joinpath(dirname(@__FILE__), ".data/wsj_00-18.conll"))
    info("# sentences of train data: $(length(data))")

    data_x, data_y = setup_mldata(data, worddict, chardict, true)
    info("# words: $(length(worddict))")
    info("# chars: $(length(chardict))")
    data_x, data_y
end

function setup_data()
    h5file = joinpath(dirname(@__FILE__), ".data/wordembeds_nyt100.h5")
    words = h5read(h5file, "s")
    worddict = Dict(words[i] => i for i=1:length(words))
    chardict = Dict(' ' => 1)

    #traindata = UD_English.traindata()
    #testdata = UD_English.testdata()
    traindata = CoNLL.read(joinpath(dirname(@__FILE__), ".data/wsj_00-18.conll"))
    testdata = CoNLL.read(joinpath(dirname(@__FILE__), ".data/wsj_22-24.conll"))
    info("# sentences of train data: $(length(traindata))")
    info("# sentences of test data: $(length(testdata))")

    train_x, train_y = setup_mldata(traindata, worddict, chardict, true)
    test_x, test_y = setup_mldata(testdata, worddict, chardict, false)
    info("# words: $(length(worddict))")
    info("# chars: $(length(chardict))")
    train_x, train_y, test_x, test_y
end

function setup_mldata(doc, worddict, chardict, append)
    data_x, data_y = Var[], Var[]
    unkword = worddict["UNKNOWN"]
    for sent in doc
        x = Int[]
        y = Int[]
        for items in sent
            word = items[2]
            word0 = replace(word, r"[0-9]", '0')
            wordid = get(worddict, lowercase(word0), unkword)
            chars = Vector{Char}(word0)
            if append
                charids = map(c -> get!(chardict,c,length(chardict)+1), chars)
            else
                charids = map(c -> get(chardict,c,0), chars)
            end
            head = parse(Int, items[7])
            token = Token(wordid, charids)
            push!(x, wordid)
            push!(y, head)
        end
        push!(data_x, Var(reshape(x,1,length(x))))
        push!(data_y, Var(y))
    end
    data_x, data_y
end

max_stars_repo_name: hshindo/PairwiseMatrix

max_issues_repo_name: hshindo/PairwiseMatrix

max_forks_repo_name: hshindo/PairwiseMatrix


text:
model = Dynamics.Cartpole()
prob = Problems.Cartpole()[1]
bnd = prob.constraints.constraints[1]
n,m = size(prob)

N = prob.N
rollout!(prob)
Z = prob.Z
vals = [@SVector zeros(model.n) for k = 1:N-1]

∇c = [SizedMatrix{n,2n+m}(zeros(n,2n+m)) for k = 1:N-1]
dyn_con = DynamicsConstraint{RK3}(model, N)
@test TO.width(dyn_con) == 2n+m
evaluate!(vals, dyn_con, Z)
jacobian!(∇c, dyn_con, Z)
@test (@allocated evaluate!(vals, dyn_con, Z)) == 0
@test (@allocated jacobian!(∇c, dyn_con, Z)) == 0

con_rk3 = ConstraintVals(dyn_con, 1:N-1)
evaluate!(con_rk3, Z)
jacobian!(con_rk3, Z)
TO.max_violation!(con_rk3)
maximum(con_rk3.c_max)
@test (@allocated evaluate!(con_rk3, Z)) == 0
@test (@allocated jacobian!(con_rk3, Z)) == 0

∇c = [zeros(SizedMatrix{n,2n+2m}) for k = 1:N-1]
dyn_con = DynamicsConstraint{HermiteSimpson}(model, N)
@test TO.width(dyn_con) == 2(n+m)
evaluate!(vals, dyn_con, Z)
jacobian!(∇c, dyn_con, Z)
@test (@allocated evaluate!(vals, dyn_con, Z)) == 0
@test (@allocated jacobian!(∇c, dyn_con, Z)) == 0

con_hs = ConstraintVals(dyn_con, 1:N-1)
evaluate!(con_hs, Z)
jacobian!(con_hs, Z)
@test (@allocated evaluate!(con_hs, Z)) == 0
@test (@allocated jacobian!(con_hs, Z)) == 0

# Test default
dyn_con = DynamicsConstraint(model, N)
@test integration(dyn_con) == RK3 == TO.DEFAULT_Q


text:
@testset "CUFFT" begin

using CuArrays.CUFFT

using CUDAdrv

using FFTW

# notes:
#   plan_bfft does not need separate testing since it is used by plan_ifft

N1 = 8
N2 = 32
N3 = 64
N4 = 8

MYRTOL = 1e-5
MYATOL = 1e-8

## complex

function out_of_place(X::AbstractArray{T,N}) where {T <: Complex,N}
    fftw_X = fft(X)
    d_X = CuArray(X)
    p = plan_fft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_ifft(d_Y)
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)

    pinv2 = inv(p)
    d_Z = pinv2 * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

function in_place(X::AbstractArray{T,N}) where {T <: Complex,N}
    fftw_X = fft(X)
    d_X = CuArray(X)
    p = plan_fft!(d_X)
    p * d_X
    Y = collect(d_X)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_ifft!(d_X)
    pinv * d_X
    Z = collect(d_X)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

function batched(X::AbstractArray{T,N},region) where {T <: Complex,N}
    fftw_X = fft(X,region)
    d_X = CuArray(X)
    p = plan_fft(d_X,region)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test_maybe_broken isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_ifft(d_Y,region)
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

@testset for T in [ComplexF32, ComplexF64]

@testset "1D" begin
    dims = (N1,)
    X = rand(T, dims)
    out_of_place(X)
end
@testset "1D inplace" begin
    dims = (N1,)
    X = rand(T, dims)
    in_place(X)
end

@testset "2D" begin
    dims = (N1,N2)
    X = rand(T, dims)
    out_of_place(X)
end
@testset "2D inplace" begin
    dims = (N1,N2)
    X = rand(T, dims)
    in_place(X)
end

@testset "Batch 1D" begin
    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,1)

    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,2)

    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,(1,2))
end

@testset "3D" begin
    dims = (N1,N2,N3)
    X = rand(T, dims)
    out_of_place(X)
end

@testset "3D inplace" begin
    dims = (N1,N2,N3)
    X = rand(T, dims)
    in_place(X)
end

@testset "Batch 2D (in 3D)" begin
    dims = (N1,N2,N3)
    for region in [(1,2),(2,3),(1,3)]
        X = rand(T, dims)
        batched(X,region)
    end

    X = rand(T, dims)
    @test_throws ArgumentError batched(X,(3,1))
end

@testset "Batch 2D (in 4D)" begin
    dims = (N1,N2,N3,N4)
    for region in [(1,2),(1,4),(3,4)]
        X = rand(T, dims)
        batched(X,region)
    end
    for region in [(1,3),(2,3),(2,4)]
        X = rand(T, dims)
        @test_throws ArgumentError batched(X,region)
    end

end

end

## real

function out_of_place(X::AbstractArray{T,N}) where {T <: Real,N}
    fftw_X = rfft(X)
    d_X = CuArray(X)
    p = plan_rfft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_irfft(d_Y,size(X,1))
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)

    pinv2 = inv(p)
    d_Z = pinv2 * d_Y
    Z = collect(d_Z)
    @test_maybe_broken isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)

    pinv3 = inv(pinv)
    d_W = pinv3 * d_X
    W = collect(d_W)
    @test isapprox(W, Y, rtol = MYRTOL, atol = MYATOL)
end

function batched(X::AbstractArray{T,N},region) where {T <: Real,N}
    fftw_X = rfft(X,region)
    d_X = CuArray(X)
    p = plan_rfft(d_X,region)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_irfft(d_Y,size(X,region[1]),region)
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

@testset for T in [Float32, Float64]

@testset "1D" begin
    X = rand(T, N1)
    out_of_place(X)
end

@testset "Batch 1D" begin
    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,1)

    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,2)

    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,(1,2))
end

@testset "2D" begin
    X = rand(T, N1,N2)
    out_of_place(X)
end

@testset "Batch 2D (in 3D)" begin
    dims = (N1,N2,N3)
    for region in [(1,2),(2,3),(1,3)]
        X = rand(T, dims)
        batched(X,region)
    end

    X = rand(T, dims)
    @test_throws ArgumentError batched(X,(3,1))
end

@testset "Batch 2D (in 4D)" begin
    dims = (N1,N2,N3,N4)
    for region in [(1,2),(1,4),(3,4)]
        X = rand(T, dims)
        batched(X,region)
    end
    for region in [(1,3),(2,3),(2,4)]
        X = rand(T, dims)
        @test_throws ArgumentError batched(X,region)
    end
end

@testset "3D" begin
    X = rand(T, N1, N2, N3)
    out_of_place(X)
end

end

## complex integer

function out_of_place(X::AbstractArray{T,N}) where {T <: Complex{<:Integer},N}
    fftw_X = fft(X)
    d_X = CuArray(X)
    p = plan_fft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    d_Y = fft(d_X)
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)
end

@testset for T in [Complex{Int32}, Complex{Int64}]

@testset "1D" begin
    dims = (N1,)
    X = rand(T, dims)
    out_of_place(X)
end

end

## real integer

function out_of_place(X::AbstractArray{T,N}) where {T <: Integer,N}
    fftw_X = rfft(X)
    d_X = CuArray(X)
    p = plan_rfft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    d_Y = rfft(d_X)
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)
end

@testset for T in [Int32, Int64]

@testset "1D" begin
    X = rand(T, N1)
    out_of_place(X)
end

end

@testset "streams" begin
    X = rand(N1)
    d_X = CuArray(X)
    p = plan_fft(d_X)
    CUFFT.set_stream(p, CUDAdrv.CuDefaultStream())
end

end


text:
using Boards
using Base.Test

# write your own tests here
@test getboard() >= 0


text:
using MPSKit,TensorKit,Test
using OptimKit

println("------------------------------------")
println("|     States                       |")
println("------------------------------------")
@testset "FiniteMPS ($D,$d,$elt)" for (D,d,elt) in [
        (ComplexSpace(10),ComplexSpace(2),ComplexF64),
        (ℂ[SU₂](1=>1,0=>3),ℂ[SU₂](0=>1)*ℂ[SU₂](0=>1),ComplexF32)
        ]

    ts = FiniteMPS(rand,elt,rand(3:20),d,D);

    ovl = dot(ts,ts);
    @test ovl ≈ norm(ts.AC[1])^2

    for i in 1:length(ts)
        @test ts.AC[i] ≈ ts.AL[i]*ts.CR[i]
        @test ts.AC[i] ≈ MPSKit._permute_front(ts.CR[i-1]*MPSKit._permute_tail(ts.AR[i]))
    end

    @test elt == eltype(eltype(ts))

    ts = ts*3
    @test ovl*9 ≈ norm(ts)^2
    ts = 3*ts
    @test ovl*9*9 ≈ norm(ts)^2
end

@testset "InfiniteMPS ($D,$d,$elt)" for (D,d,elt) in [
        (ComplexSpace(10),ComplexSpace(2),ComplexF64),
        (ℂ[U₁](1=>3),ℂ[U₁](0=>1),ComplexF64)
        ]
    tol = Float64(eps(real(elt))*100);

    ts = #=@inferred=# InfiniteMPS([TensorMap(rand,elt,D*d,D),TensorMap(rand,elt,D*d,D)],tol = tol);

    for i in 1:length(ts)
        @tensor difference[-1,-2,-3] := ts.AL[i][-1,-2,1]*ts.CR[i][1,-3]-ts.CR[i-1][-1,1]*ts.AR[i][1,-2,-3];
        @test norm(difference,Inf) < tol*10;

        @test transfer_left(l_LL(ts,i),ts.AL[i],ts.AL[i]) ≈ l_LL(ts,i+1)
        @test transfer_left(l_LR(ts,i),ts.AL[i],ts.AR[i]) ≈ l_LR(ts,i+1)
        @test transfer_left(l_RL(ts,i),ts.AR[i],ts.AL[i]) ≈ l_RL(ts,i+1)
        @test transfer_left(l_RR(ts,i),ts.AR[i],ts.AR[i]) ≈ l_RR(ts,i+1)

        @test transfer_right(r_LL(ts,i),ts.AL[i],ts.AL[i]) ≈ r_LL(ts,i+1)
        @test transfer_right(r_LR(ts,i),ts.AL[i],ts.AR[i]) ≈ r_LR(ts,i+1)
        @test transfer_right(r_RL(ts,i),ts.AR[i],ts.AL[i]) ≈ r_RL(ts,i+1)
        @test transfer_right(r_RR(ts,i),ts.AR[i],ts.AR[i]) ≈ r_RR(ts,i+1)
    end
end

@testset "MPSMultiline ($D,$d,$elt)" for (D,d,elt) in [
        (ComplexSpace(10),ComplexSpace(2),ComplexF64),
        (ℂ[U₁](1=>3),ℂ[U₁](0=>1),ComplexF32)
        ]

    tol = Float64(eps(real(elt))*100);
    ts = MPSMultiline([TensorMap(rand,elt,D*d,D) TensorMap(rand,elt,D*d,D);TensorMap(rand,elt,D*d,D) TensorMap(rand,elt,D*d,D)],tol = tol);

    for (i,j) in Iterators.product(size(ts,1),size(ts,2))
        @tensor difference[-1,-2,-3] := ts.AL[i,j][-1,-2,1]*ts.CR[i,j][1,-3]-ts.CR[i,j-1][-1,1]*ts.AR[i,j][1,-2,-3];
        @test norm(difference,Inf) < tol*10;

        @test transfer_left(l_LL(ts,i,j),ts.AL[i,j],ts.AL[i,j]) ≈ l_LL(ts,i,j+1)
        @test transfer_left(l_LR(ts,i,j),ts.AL[i,j],ts.AR[i,j]) ≈ l_LR(ts,i,j+1)
        @test transfer_left(l_RL(ts,i,j),ts.AR[i,j],ts.AL[i,j]) ≈ l_RL(ts,i,j+1)
        @test transfer_left(l_RR(ts,i,j),ts.AR[i,j],ts.AR[i,j]) ≈ l_RR(ts,i,j+1)

        @test transfer_right(r_LL(ts,i,j),ts.AL[i,j],ts.AL[i,j]) ≈ r_LL(ts,i,j+1)
        @test transfer_right(r_LR(ts,i,j),ts.AL[i,j],ts.AR[i,j]) ≈ r_LR(ts,i,j+1)
        @test transfer_right(r_RL(ts,i,j),ts.AR[i,j],ts.AL[i,j]) ≈ r_RL(ts,i,j+1)
        @test transfer_right(r_RR(ts,i,j),ts.AR[i,j],ts.AR[i,j]) ≈ r_RR(ts,i,j+1)
    end
end

@testset "MPSComoving" begin
    ham = nonsym_ising_ham(lambda=4.0);
    (gs,_,_) = find_groundstate(InfiniteMPS([ℂ^2],[ℂ^10]),ham,Vumps(verbose=false));

    window = MPSComoving(gs,copy.([gs.AC[1];[gs.AR[i] for i in 2:10]]),gs);

    for i in 1:length(window)
        @test window.AC[i] ≈ window.AL[i]*window.CR[i]
        @test window.AC[i] ≈ MPSKit._permute_front(window.CR[i-1]*MPSKit._permute_tail(window.AR[i]))
    end

    @test norm(window) ≈ 1
    window = window*3
    @test 9 ≈ norm(window)^2
    window = 3*window
    @test 9*9 ≈ norm(window)^2
    normalize!(window)

    e1 = expectation_value(window,ham);

    (window,pars,_) = find_groundstate(window,ham,Dmrg(verbose=false));

    e2 = expectation_value(window,ham);

    @test e1[1] ≈ e2[1]
    @test e1[2] ≈ e2[2]

    (window,pars) = timestep(window,ham,0.1,Tdvp2(),pars)
    (window,pars) = timestep(window,ham,0.1,Tdvp(),pars)

    e3 = expectation_value(window,ham);

    @test e1[1] ≈ e3[1]
    @test e1[2] ≈ e3[2]
end

println("------------------------------------")
println("|     Operators                    |")
println("------------------------------------")
@testset "mpoham $(i)" for (i,(th,Dspaces)) in enumerate([
        (nonsym_ising_ham(),[ℂ^1]),
        (u1_xxz_ham(),[ℂ[U₁](1//2=>1)]),
        (repeat(su2_xxx_ham(),2),[ℂ[SU₂](0=>1),ℂ[SU₂](1//2=>1)])
        ])

    ts = InfiniteMPS(th.pspaces,Dspaces); # generate a product state

    (ts,_) = changebonds(ts,th,OptimalExpand()) # optimal expand a la vumps paper
    ndim = dim(space(ts.AC[1],1))
    (ts,_) = changebonds(ts,th,VumpsSvdCut()) # idmrg2 step to expand the bond dimension
    @test dim(space(ts.AC[1],1)) > ndim;

    e1 = expectation_value(ts,th);

    e2 = expectation_value(ts,2*th); #multiplication with a constant
    @test 2*e1≈e2;

    e2 = expectation_value(ts,0.5*th+th); #addition
    @test 1.5*e1≈e2;

    th -= expectation_value(ts,th);
    v = expectation_value(ts,th*th);
    @test real(v[1])>=0;
end

@testset "comact $(i)" for (i,th) in enumerate([
        nonsym_ising_ham(),
        u1_xxz_ham(),
        su2_xxx_ham(),
        su2u1_grossneveu()
        ])

    len = 20;

    inftemp = repeat(infinite_temperature(th),len);
    ndat = collect(map(x-> TensorMap(rand,ComplexF64,codomain(x),domain(x)),inftemp));
    ts = FiniteMPS(ndat);

    (ts,_) = changebonds(ts,anticommutator(th),OptimalExpand());

    e1 = expectation_value(ts,anticommutator(th));
    e2 = expectation_value(ts,2*anticommutator(th));

    @test 2*e1≈e2;

    e3 = expectation_value(ts,anticommutator(th)+commutator(th));
    e4 = expectation_value(ts,anticommutator(th)-commutator(th));

    @test e3+e4≈e2;

    diff = [rand() for i in th.pspaces];
    e5 = expectation_value(ts,anticommutator(th)-diff);
    @test sum([e1[j]-diff[mod1(j,end)] for j in 1:len])≈sum(e5);
end

println("------------------------------------")
println("|     Algorithms                   |")
println("------------------------------------")

@testset "find_groundstate $(ind)" for (ind,(state,alg,ham)) in enumerate([
        (InfiniteMPS([ℂ^2],[ℂ^10]),Vumps(tol_galerkin=1e-8,verbose=false),nonsym_ising_ham(lambda=2.0)),
        (InfiniteMPS([ℂ^2],[ℂ^10]), GradientGrassmann(tol=1e-8, verbosity=0), nonsym_ising_ham(lambda=2.0)),
        (InfiniteMPS([ℂ^2],[ℂ^10]), GradientGrassmann(method=LBFGS(6; gradtol=1e-8, verbosity=0)), nonsym_ising_ham(lambda=2.0)),
        (InfiniteMPS([ℂ^2],[ℂ^10]),Idmrg1(tol_galerkin=1e-8,maxiter=400,verbose=false),nonsym_ising_ham(lambda=2.0)),
        (FiniteMPS(rand,ComplexF64,10,ℂ^2,ℂ^10),Dmrg2(verbose=false),nonsym_ising_ham(lambda=2.0)),
        (FiniteMPS(rand,ComplexF64,10,ℂ^2,ℂ^10),Dmrg(verbose=false),nonsym_ising_ham(lambda=2.0))
        ])

    #vumps type inferrence got broken by @threads, so worth it?
    (ts,pars,delta) =  #=@inferred=# find_groundstate(state,ham,alg)

    @test sum(delta) < 1e-6

    evals = expectation_value(ts,ham);
    th = repeat(ham,length(evals))-evals

    @test real(sum(expectation_value(ts,th*th))) < 1e-2 #is the ground state variance relatively low?
end

@testset "timestep $(ind)" for (ind,(state,alg,opp)) in enumerate([
    (FiniteMPS(fill(TensorMap(rand,ComplexF64,ℂ^1*ℂ^2*adjoint(ℂ^2),ℂ^1),5)),Tdvp(),commutator(nonsym_xxz_ham(spin=1//2))),
    (FiniteMPS(fill(TensorMap(rand,ComplexF64,ℂ^1*ℂ^2*adjoint(ℂ^2),ℂ^1),7)),Tdvp2(),anticommutator(nonsym_xxz_ham(spin=1//2))),
    (InfiniteMPS([ℂ^3,ℂ^3],[ℂ^50,ℂ^50]),Tdvp(),repeat(nonsym_xxz_ham(spin=1),2))
    ])

    edens = expectation_value(state,opp);

    (state,_) = timestep(state,opp,rand()/10,alg)

    @test sum(expectation_value(state,opp)) ≈ sum(edens) atol = 1e-2
end

@testset "leading_boundary $(ind)" for (ind,alg) in enumerate([
        Vumps(tol_galerkin=1e-5,verbose=false),
        PowerMethod(tol_galerkin=1e-5,verbose=false,maxiter=1000)])

    mpo = #=@inferred=# nonsym_ising_mpo();
    state = InfiniteMPS([ℂ^2],[ℂ^10]);
    (state,pars,_) = leading_boundary(state,mpo,alg);

    @test expectation_value(state,mpo,pars)[1,1] ≈ 2.5337 atol=1e-3
end

@testset "quasiparticle_excitation" begin
    th = nonsym_xxz_ham()
    ts = InfiniteMPS([ℂ^3],[ℂ^48]);
    (ts,pars,_) = find_groundstate(ts,th,Vumps(maxiter=400,verbose=false));
    (energies,Bs) = quasiparticle_excitation(th,Float64(pi),ts,pars);
    @test energies[1] ≈ 0.41047925 atol=1e-4
end

@testset "dynamicaldmrg" begin
    ham = nonsym_ising_ham(lambda=4.0);
    (gs,_,_) = find_groundstate(InfiniteMPS([ℂ^2],[ℂ^10]),ham,Vumps(verbose=false));
    window = MPSComoving(gs,copy.([gs.AC[1];[gs.AR[i] for i in 2:10]]),gs);

    szd = TensorMap([1 0;0 -1],ℂ^2,ℂ^2);
    @test expectation_value(gs,szd)[1] ≈ expectation_value(window,szd)[1] atol=1e-10

    polepos = expectation_value(gs,ham,10)
    @test polepos ≈ expectation_value(window,ham)[2]

    vals = (-0.5:0.2:0.5).+polepos
    eta = 0.3im;

    predicted = [1/(v+eta-polepos) for v in vals];

    data = similar(predicted);
    for (i,v) in enumerate(vals)
        (data[i],_) = dynamicaldmrg(window,v+eta,ham,verbose=false)
    end

    @test data ≈ predicted atol=1e-8
end


text:
export run!, collect_agent_data!, collect_model_data!,
       init_agent_dataframe, init_model_dataframe, aggname

###################################################
# Definition of the data collection API
###################################################
get_data(a, s::Symbol) = getproperty(a, s)
get_data(a, f::Function) = f(a)

should_we_collect(s, model, when::AbstractVector) = s ∈ when
should_we_collect(s, model, when::Bool) = when
should_we_collect(s, model, when) = when(model, s)

"""
    run!(model, agent_step! [, model_step!], n; kwargs...) → agent_df, model_df

Run the model (step it with the input arguments propagated into `step!`) and collect
data specified by the keywords, explained one by one below. Return the data as
two `DataFrame`s, one for agent-level data and one for model-level data.

## Data-deciding keywords
* `adata::Vector` means "agent data to collect". If an entry is a `Symbol`, e.g. `:weight`,
  then the data for this entry is agent's field `weight`. If an entry is a `Function`, e.g.
  `f`, then the data for this entry is just `f(a)` for each agent `a`.
  The resulting dataframe columns are named with the input symbol (here `:weight, :f`).

* `adata::Vector{<:Tuple}`: if `adata` is a vector of 2-tuples instead,
  data aggregation is done over the agent properties. For each 2-tuple, the first entry
  is the "key" (any entry like the ones mentioned above, e.g. `:weight, f`). The second
  entry is an aggregating function that aggregates the key, e.g. `mean, maximum`. So,
  continuing from the above example, we would have
  `adata = [(:weight, mean), (f, maximum)]`. The resulting data name columns
  use the function [`aggname`](@ref), and create something like `:mean_weight` or
  `:maximum_f`. This name doesn't play well with anonymous functions!

  **Notice:** Aggregating only works if there are agents to be aggregated over.
  If you remove agents during model run, you should modify the aggregating functions.
  *E.g.* instead of passing `mean`, pass `mymean(a) = isempty(a) ? 0.0 : mean(a)`.

* `mdata::Vector` means "model data to collect" and works exactly like `adata`.
  For the model, no aggregation is possible (nothing to aggregate over).

By default both keywords are `nothing`, i.e. nothing is collected/aggregated.

### Other keywords
* `when=true` : at which steps `s` to perform the data collection and processing.
  A lot of flexibility is offered based on the type of `when`. If `when::Vector`,
  then data are collect if `s ∈ when`. Otherwise data are collected if `when(model, s)`
  returns `true`. By default data are collected in every step.
* `when_model = when` : same as `when` but for model data.
* `replicates=0` : Run `replicates` replicates of the simulation.
* `parallel=false` : Only when `replicates>0`. Run replicate simulations in parallel.
"""
function run! end

run!(model::ABM, agent_step!, n; kwargs...) =
run!(model::ABM, agent_step!, dummystep, n; kwargs...)

function run!(model::ABM, agent_step!, model_step!, n;
  replicates::Int=0, parallel::Bool=false, kwargs...)

  r = replicates
  if r > 0
    if parallel
      return parallel_replicates(model, agent_step!, model_step!, n, r; kwargs...)
    else
      return series_replicates(model, agent_step!, model_step!, n, r; kwargs...)
    end
  else
    return _run!(model, agent_step!, model_step!, n; kwargs...)
  end
end

###################################################
# Core data collection loop
###################################################
"""
  _run!(model, agent_step!, model_step!, n; kwargs...)
Core function that loops over stepping a model and collecting data at each step.
"""
function _run!(model, agent_step!, model_step!, n;
               when = true, when_model = when,
               agent_properties=nothing, model_properties=nothing,
               mdata=model_properties, adata=agent_properties)

    agent_properties ≠ nothing && @warn "use `adata` instead of `agent_properties`"
    model_properties ≠ nothing && @warn "use `mdata` instead of `model_properties`"
    df_agent = init_agent_dataframe(model, adata)
    df_model = init_model_dataframe(model, mdata)
    if n isa Integer
        if when == true; for c in eachcol(df_agent); sizehint!(c, n); end; end
        if when_model == true; for c in eachcol(df_model); sizehint!(c, n); end; end
    end

    s = 0
    while until(s, n, model)
        if should_we_collect(s, model, when)
            collect_agent_data!(df_agent, model, adata, s)
        end
        if should_we_collect(s, model, when_model)
            collect_model_data!(df_model, model, mdata, s)
        end
        step!(model, agent_step!, model_step!, 1)
        s += 1
    end
    if should_we_collect(s, model, when)
        collect_agent_data!(df_agent, model, adata, s)
    end
    if should_we_collect(s, model, when_model)
        collect_model_data!(df_model, model, mdata, s)
    end
    return df_agent, df_model
end

###################################################
# core data collection functions per step
###################################################
"""
    init_agent_dataframe(model, properties) → agent_df
Initialize a dataframe to add data later with [`collect_agent_data!`](@ref).
"""
init_agent_dataframe(model, properties::Nothing) = DataFrame()

"""
    collect_agent_data!(df, model, properties, step = 0)
Collect and add agent data into `df` (see [`run!`](@ref) for the dispatch rules
of `properties`). `step` is given because the step number information is not known.
"""
collect_agent_data!(df, model, properties::Nothing, step::Int=0) = df

function init_agent_dataframe(model::ABM, properties::AbstractArray)
    nagents(model) < 1 && throw(ArgumentError("Model must have at least one agent to "*
    "initialize data collection"))

    headers = Vector{Symbol}(undef, 2+length(properties))
    headers[1] = :step
    headers[2] = :id
    for i in 1:length(properties); headers[i+2] = Symbol(properties[i]); end

    types = Vector{Vector}(undef, 2+length(properties))
    types[1] = Int[]
    types[2] = Int[]
    a = random_agent(model)
    for (i, k) in enumerate(properties)
        current_type = typeof(get_data(a, k))
        isconcretetype(current_type) || warn("Type is not concrete when using $(k)"*
        "on agents. Consider narrowning the type signature of $(k).")
        types[i+2] = current_type[]
    end
    DataFrame(types, headers)
end

function collect_agent_data!(df, model, properties::Vector, step::Int=0)
    alla = sort!(collect(values(model.agents)), by=a->a.id)
    dd = DataFrame()
    dd[!, :step] = fill(step, length(alla))
    dd[!, :id] = map(a->a.id, alla)
    for fn in properties
        dd[!, Symbol(fn)] = collect(get_data(a, fn) for a in alla)
    end
    append!(df, dd)
    return df
end

# Aggregating version
function init_agent_dataframe(model::ABM, properties::Vector{<:Tuple})
    nagents(model) < 1 && throw(ArgumentError("Model must have at least one agent to "*
    "initialize data collection"))
    headers = Vector{Symbol}(undef, 1+length(properties))
    types = Vector{Vector}(undef, 1+length(properties))
    alla = allagents(model)

    headers[1] = :step
    types[1] = Int[]

    for (i, (k, agg)) in enumerate(properties)
        headers[i+1] = aggname(k, agg)
        # This line assumes that `agg` can work with iterators directly
        current_type = typeof( agg( get_data(a, k) for a in Iterators.take(alla,1) ) )
        isconcretetype(current_type) || warn("Type is not concrete when using function $(agg) "*
        "on key $(k). Consider using type annotation, e.g. $(agg)(a)::Float64 = ...")
        types[i+1] = current_type[]
    end
    DataFrame(types, headers)
end

"""
    aggname(k, agg) → name

Return the name of the column of the aggregated data with key `k` and aggregating
function `agg`.
"""
function aggname(k, agg)
    @static if VERSION >= v"1.1"
        Symbol(join([string(agg), string(k)], "_"))
    else
        Symbol(join([split(string(agg), ".")[end], string(k)], "_"))
    end
end

function collect_agent_data!(df, model, properties::Vector{<:Tuple}, step::Int=0)
    alla = allagents(model)
    push!(df[!, 1], step)
    for (i, (k, agg)) in enumerate(properties)
        _add_col_data!(df[!, i+1], agg, k, alla)
    end
    return df
end
# Function barrier
function _add_col_data!(col::AbstractVector{T}, agg, k, alla) where {T}
    res::T = agg(get_data(a, k) for a in alla)
    push!(col, res)
end

# Model data
"""
    init_model_dataframe(model, properties) → model_df
Initialize a dataframe to add data later with [`collect_model_data!`](@ref).
"""
function init_model_dataframe(model::ABM, properties::Vector)
    headers = Vector{Symbol}(undef, 1+length(properties))
    headers[1] = :step
    for i in 1:length(properties); headers[i+1] = Symbol(properties[i]); end

    types = Vector{Vector}(undef, 1+length(properties))
    types[1] = Int[]
    for (i,k) in enumerate(properties)
        types[i+1] =
            if typeof(k) <: Symbol
                typeof(model.properties[k])[]
            else
                current_type = typeof(k(model))
                isconcretetype(current_type) || warn("Type is not concrete when using $(k)"*
                "on the model. Considering narrowing the type signature of $(k).")
                current_type[]
            end
    end
    DataFrame(types, headers)
end

init_model_dataframe(model::ABM, properties::Nothing) = DataFrame()

"""
    collect_model_data!(df, model, properties, step = 0)
Same as [`collect_agent_data!`](@ref) but for model data instead.
"""
function collect_model_data!(df, model, properties::Vector, step::Int=0)
  push!(df[!, :step], step)
  for fn in properties
    push!(df[!, Symbol(fn)], get_data(model, fn))
  end
  return df
end

collect_model_data!(df, model, properties::Nothing, step::Int=0) = df

###################################################
# Parallel / replicates
###################################################
function replicate_col!(df, rep)
  df[!, :replicate] = [rep for i in 1:size(df, 1)]
end

"Run replicates of the same simulation"
function series_replicates(model, agent_step!, model_step!, n, replicates; kwargs...)

  df_agent, df_model = _run!(deepcopy(model), agent_step!, model_step!, n; kwargs...)
  replicate_col!(df_agent, 1)
  replicate_col!(df_model, 1)

  for rep in 2:replicates
    df_agentTemp, df_modelTemp = _run!(deepcopy(model), agent_step!, model_step!, n; kwargs...)
    replicate_col!(df_agentTemp, rep)
    replicate_col!(df_modelTemp, rep)

    append!(df_agent, df_agentTemp)
    append!(df_model, df_modelTemp)
  end
  return df_agent, df_model
end

"Run replicates of the same simulation in parallel"
function parallel_replicates(model::ABM, agent_step!, model_step!, n, replicates; kwargs...)

  all_data = pmap(j -> _run!(deepcopy(model), agent_step!, model_step!, n; kwargs...),
                  1:replicates)

  df_agent = DataFrame()
  df_model = DataFrame()
  for (rep, d) in enumerate(all_data)
    replicate_col!(d[1], rep)
    replicate_col!(d[2], rep)
    append!(df_agent, d[1])
    append!(df_model, d[2])
  end

  return df_agent, df_model
end


text:
"""
# Population genetics simulations
Repository:    https://www.github.com/pdimens/PopGenSims.jl/
Documentation: https://biojulia.net/PopGen.jl/

\nA few things things you can do to get started:

## Import Data using PopGenCore
- `PopGenCore.read(filename; kwargs...)`
- `genepop(infile; kwargs...)`  or similar file-specific importer
- use available `@gulfsharks` or `@nancycats` datasets

## Simulate Samples within Populations
- `simulate(popdata, n = 100)` to simulate samples using population-specific allele frequencies

## Simulate breeding crosses
- `cross(popdata, parent1, parent2, ...)` to cross two parents from the same PopData
- `cross(popdata => parent1, popdata => parent2, ...) to cross two parents from different PopData (e.g. backcross)

## Simulate siblingship
- `simulatekin(popdata, fullsib = , halfsib = , unrelated = , ...)` to simulate breeding events generating pairs of individuals of known kinship.
"""
module PopGenSims

using DataFrames, PooledArrays, StaticArrays
using StatsBase: sample, Weights
import PopGenCore: read, write
using PopGenCore:
    allelefreq,
    copy,
    Genotype,
    GenoArray,
    genotypes,
    PopData,
    sort

include("Cross.jl")
export cross, backcross

include("Samples.jl")
export simulate

include("Sibship.jl")
export simulatekin

include("Utils.jl")
export append, append!

end


text:
"""
    QR()
Finds a vector vanishing 'gamma' such that all of its entries sum to 1.

Input arguments
---------------
gamma::Array{Float64, 2}    This matrix verifies properties necessary for QR() to work. We know that the rank of 'Gamma' n-1 (otherwise we wouldn't reach this function).
tolerance::Float64          Same as always.

Returns
-------
lambda::Vec{Float64}        A vector vanishing 'gamma'. lambda is a column vector containing the components of a vector generating Nul(M) with all the components summing up to 1.
"""
function QR(Gamma, tolerance)
   # We know that rank(Gamma)=n-1
   #and rank(Gamma0)=n
   # In addition all columns of M are non vanishing
   ## Outcome
   # lambda is a column vector containing the components of a vector generating
   # Nul(M) with all the components summing up to 1
   ##

   n = size(Gamma, 2)
   m = size(Gamma, 1) # m>=n-1
   if m == 1
      #then n<=2 but n=1 is excluded since it corresponds to a 1-vertex boundary, thus n=2 and Gamma=[a b]
      #also notice that since Gamma0=[[a b];[1 1]] has full rank it is ruled
      #out that a=b
      lambda = [Gamma[2]; -Gamma[1]] / (Gamma[2] - Gamma[1])
   else

      # M=Q*R
      #Q is an orthogonal matrix of dimension m x m
      qr_decomposition = qr(Gamma).R # access the second element of the tuple
      # R is an upper triangular matrix of dimension m x n
      R = triu(qr_decomposition)

      # diagonal is a column vector with dimension either n-1 (if m=n-1) or n (if m>=n)
      # Set entries that are too small relative to `tolerance` to zero.
      diagonal = heaviside(tolerance .- abs.(diag(R)))
      index = round.(Int, transpose(collect(1:min(m, n))) * diagonal)[1]

      if index == 0
         # the first n-1 columns are linearly independent
         index = n - 1
      else
         #the index-th entry in diag(R) is zero
         index = index - 1
      end

      lambda = zeros(n, 1)
      lambda[1:index] = - R[1:index, 1:index] \ R[1:index, index + 1] # A\b means inv(A)*b
      lambda[index + 1] = 1
      lambda = lambda / (ones(1, n) * lambda)
   end
   lambda = lambda .* heaviside(abs.(lambda) .- tolerance)

   return(lambda)
end


text:
module TropicalModule

using Hecke
import AbstractAlgebra

function slope_eigenspace(M::MatElem{T}) where T <: Hecke.NonArchLocalFieldElem
  f = charpoly(M)
  lf = Hecke.slope_factorization(f)
#  @assert all(x->x==1, values(lf))

  se = Dict{typeof(f), typeof(M)}()
  k = base_ring(M)
  zk = maximal_order(k)

  for f = keys(lf)
    se[f] = kernel(f(M))[2] #hopefully, this is in rref
  end
  @assert sum(ncols(x) for x = values(se)) == nrows(M)
  return se
end

function _intersect(M::MatElem{T}, N::MatElem{T}) where T <: Hecke.FieldElem
  k = base_ring(M)
  I = [M N]
  PR = maximum(precision, I)
  pr = minimum(precision, I)
  if pr != PR
    for i = eachindex(I)
      I[i] = setprecision(I[i], pr)
    end
  end

  r, v = kernel(I) #precision issues...
  l = M*v[1:ncols(M), 1:r]
  return transpose(rref(transpose(l))[2])
end

function valuation_of_roots(f::PolyElem{<:Hecke.NonArchLocalFieldElem})
  iszero(f) && error("polynomial must not be zero")
  return (valuation(constant_coefficient(f)) - valuation(leading_coefficient(f)))//degree(f)
end

function simultaneous_diagonalization(v::Vector{<:MatElem{T}}) where T <: Hecke.NonArchLocalFieldElem

  k = base_ring(v[1])
  @assert all(x->base_ring(x) == k, v)
  n = nrows(v[1])
  @assert all(x->ncols(x) == nrows(x) == n, v)

  vv = map(slope_eigenspace, v)

  d = Dict(v => [valuation_of_roots(k)] for (k,v) = vv[1])
  @assert sum(ncols(x) for x = keys(d)) == n
  for i=2:length(vv)
    dd = typeof(d)()
    for (mat, pol_vec) = d
      for (p, m) = vv[i]
        j = _intersect(mat, m)
        if ncols(j) > 0
          dd[j] = push!(copy(pol_vec), valuation_of_roots(p))
        end
      end
    end
    d = dd
    @assert sum(ncols(x) for x = keys(d)) == n
  end

  return d
end

end #TropicalModule

#=
Hecke.example("Tropics.jl")
Hecke.revise("Tropics.jl")

Zx, x = ZZ["x"]
k, (a,b) = number_field([x^2-2, x^2-7])
zk = maximal_order(k)
prime_decomposition(zk, 7)
lp = [x[1] for x = ans]
lp[1].gen_two*lp[2].gen_two^2
ma = representation_matrix(a)
mb = representation_matrix(k(ans))
@assert iszero(ma*mb - mb*ma)
Qp = PadicField(7, 10)
Main.TropicalModule.simultaneous_diagonalization([map_entries(Qp, ma), map_entries(Qp, mb)])

=#


text:
"""
Iterates a function `f` on `datum` until a fixed point is reached where `f(x) == x`
"""
function normalize(f, datum, fargs...; callback=()->())
    old = datum
    new = f(old, fargs...)
    while new != old
        old = new
        new = f(old, fargs...)
        callback()
    end
    new
end
export normalize

"""
Like [`normalize`](@ref) but keeps a vector of hashes to detect cycles,
returns the current datum when a cycle is detected
"""
function normalize_nocycle(f, datum, fargs...; callback=()->())
    hist = UInt[]
    push!(hist, hash(datum))
    x = f(datum, fargs...)
    while hash(x) ∉ hist
        push!(hist, hash(x))
        x = f(x, fargs...)
        callback()
    end
    x
end
export normalize_nocycle


text:
## Documentation for wrapped functions
## Kept separate, since the sources are generated

"""
    mfb_open(title, width, height)

Create a window that is used to display the buffer sent into the mfb_update function, returns 0 if fails""" mfb_open

"""
    mfb_open_ex(title, width, height, flags)

Create a window, with flags""" mfb_open_ex
"""
    mfb_update(window, buffer)

Update the display.
Input buffer is assumed to be a 32-bit buffer of the size given in the open call.
Will return a negative status if something went wrong or the user want to exit.
Also updates the window event""" mfb_update

"""
    mfb_update(window)

Only updates the window events. Either `mfb_update` or `mfb_update_events` must be called in a loop to clear the event queue.""" mfb_update_events
"""
    mfb_close(window)

Close the window""" mfb_close

"""
    mfb_set_user_data(window , user_data::Ptr{Cvoid})

Set a per-window user data. Data is supplied as a pointer. The data at the end of that pointer should be prevented from garbage collection.

# Example
```julia
x=Ref(2)
mfb_set_user_data(w, x)
```
See also: [`mfb_get_user_data`](@ref)
""" mfb_set_user_data

"""
    mfb_get_user_data(window)::Ptr{Cvoid}

Retrive the user data associated with this window. The return value is a void pointer that should be derferenced

# Example
```julia
g=mfb_get_user_data(w)
unsafe_load(convert(Ptr{Int}, g)) # 2
```
See also: [`mfb_set_user_data`](@ref)
""" mfb_get_user_data

"""
    mfb_set_viewport(window, offset_x, offset_y, width, height)
Set viewport (useful when resize)""" mfb_set_viewport

"""
    mfb_set_active_callback(window, mfb_active_func callback)

""" mfb_set_active_callback

"""
    mfb_set_resize_callback(window, mfb_resize_func callback)

""" mfb_set_resize_callback

"""
    mfb_set_keyboard_callback(window, mfb_keyboard_func callback)

""" mfb_set_keyboard_callback

"""
    mfb_set_char_input_callback(window, mfb_char_input_func callback)

""" mfb_set_char_input_callback

"""
    mfb_set_mouse_button_callback(window, mfb_mouse_button_func callback)

""" mfb_set_mouse_button_callback

"""
    mfb_set_mouse_move_callback(window, mfb_mouse_move_func callback)

""" mfb_set_mouse_move_callback

"""
    mfb_set_mouse_scroll_callback(window, mfb_mouse_scroll_func callback)

""" mfb_set_mouse_scroll_callback

"""
    mfb_get_key_name(mfb_key key)

""" mfb_get_key_name

"""
    mfb_is_window_active(window)

""" mfb_is_window_active

"""
    mfb_get_window_width(window)

""" mfb_get_window_width

"""
    mfb_get_window_height(window)

""" mfb_get_window_height

"""
    mfb_get_mouse_x(window)

Last mouse pos X
""" mfb_get_mouse_x

"""
    mfb_get_mouse_y(window)

Last mouse pos Y
""" mfb_get_mouse_y

"""
    mfb_get_mouse_scroll_x(window)

Mouse wheel X as a sum. When you call this function it resets.
""" mfb_get_mouse_scroll_x

"""
    mfb_get_mouse_scroll_y(window)

Mouse wheel Y as a sum. When you call this function it resets.
""" mfb_get_mouse_scroll_y

"""
    mfb_get_mouse_button_buffer(window)

One byte for every button. Press (1), Release 0. (up to 8 buttons)
""" mfb_get_mouse_button_buffer

"""
    mfb_get_key_buffer(window)

One byte for every key. Press (1), Release 0.
""" mfb_get_key_buffer

"""
    mfb_set_target_fps(fps)

""" mfb_set_target_fps

"""
    mfb_wait_sync(window)

""" mfb_wait_sync

"""
    mfb_timer_create()

Create a timer. Returns an opaque pointer that should be stored, and passed to subsequent timer related calls.
""" mfb_timer_create

"""
    mfb_timer_destroy(tmr)

""" mfb_timer_destroy

"""
    mfb_timer_reset(tmr)

""" mfb_timer_reset
"""
    mfb_timer_now(tmr)

""" mfb_timer_now
"""
    mfb_timer_delta(tmr)

""" mfb_timer_delta
"""
    mfb_timer_get_frequency()

""" mfb_timer_get_frequency
"""
    mfb_timer_get_resolution()

""" mfb_timer_get_resolution


text:
let
problem = BabyPOMDP()
policy = RandomPolicy(problem, rng=MersenneTwister(2))
steps=10
sim = HistoryRecorder(max_steps=steps, rng=MersenneTwister(3))
POMDPLinter.@show_requirements simulate(sim, problem, policy, updater(policy), initialstate_distribution(problem))
r1 = simulate(sim, problem, policy, updater(policy), initialstate_distribution(problem))
policy.rng = MersenneTwister(2)
Random.seed!(sim.rng, 3)
r2 = simulate(sim, problem, policy)

@test length(state_hist(r1)) == steps+1
@test length(action_hist(r1)) == steps
@test length(observation_hist(r1)) == steps
@test length(belief_hist(r1)) == steps+1
@test length(state_hist(r2)) == steps+1
@test length(action_hist(r2)) == steps
@test length(observation_hist(r2)) == steps
@test length(belief_hist(r2)) == steps+1
@test length(info_hist(r2)) == steps
@test length(ainfo_hist(r2)) == steps
@test length(uinfo_hist(r2)) == steps

@test exception(r1) == nothing
@test exception(r2) == nothing
@test backtrace(r1) == nothing
@test backtrace(r2) == nothing

@test n_steps(r1) == n_steps(r2)
@test undiscounted_reward(r1) == undiscounted_reward(r2)
@test discounted_reward(r1) == discounted_reward(r2)

@test length(collect(r1)) == n_steps(r1)
@test length(collect(r2)) == n_steps(r2)

display(r1)
println()

for tuple in r1
    @test length(tuple) == length(POMDPSimulators.default_spec(problem))
end

for ui in eachstep(r2, "update_info")
    @test ui == nothing
end

@test r1[1] == first(r1)
@test r1[end] == last(r1)

# test that complete step is returned
step = first(eachstep(r2))
for key in POMDPSimulators.default_spec(problem)
    @test haskey(step, key)
    @test first(r2[key]) == step[key]
end

problem = LegacyGridWorld()
policy = RandomPolicy(problem, rng=MersenneTwister(2))
steps=10
sim = HistoryRecorder(max_steps=steps, rng=MersenneTwister(3))
POMDPLinter.@show_requirements simulate(sim, problem, policy, initialstate(problem, sim.rng))
r1 = simulate(sim, problem, policy, initialstate(problem, sim.rng))

@test length(state_hist(r1)) <= steps + 1 # less than or equal because it may reach the goal too fast
@test length(action_hist(r1)) <= steps
@test length(reward_hist(r1)) <= steps

@test r1[end] == last(r1)
@test r1[1] == first(r1)
for tuple in r1
    @test length(tuple) == length(POMDPSimulators.default_spec(problem))
    @test isa(tuple.s, statetype(problem))
    @test isa(tuple.a, actiontype(problem))
    @test isa(tuple.r, Float64)
    @test isa(tuple.sp, statetype(problem))
end

display(r1)
println()

step = first(eachstep(r1))
for key in POMDPSimulators.default_spec(problem)
    @test haskey(step, key)
    @test first(r1[key]) == step[key]
end

@test length(collect(r1)) == n_steps(r1)

hv = view(r1, 2:length(r1))
@test n_steps(hv) == n_steps(r1)-1
@test undiscounted_reward(r1) == undiscounted_reward(hv) + first(reward_hist(r1))

# iterators
rsum = 0.0
len = 0
for (s, a, r, sp, ai, t) in eachstep(hv, (:s,:a,:r,:sp,:action_info,:t))
    @test isa(s, statetype(problem))
    @test isa(a, actiontype(problem))
    @test isa(r, Float64)
    @test isa(sp, statetype(problem))
    @test isa(ai, Nothing)
    @test isa(t, Int)
    rsum += r
    len += 1
end
@test len == length(hv)
@test rsum == undiscounted_reward(hv)

# it = eachstep(hv, "(r,sp,s,a)")
# @test eltype(collect(it)) == Tuple{Float64, statetype(problem), statetype(problem), actiontype(problem)}
tuples = collect(eachstep(hv, "(r, sp, s, a)"))
@test sum(first(t) for t in tuples) == undiscounted_reward(hv)
@test sum(t.r for t in tuples) == undiscounted_reward(hv)
tuples = collect(eachstep(hv, "r,sp,s,a,t"))
@test sum(first(t) for t in tuples) == undiscounted_reward(hv)
@test sum(t.r for t in tuples) == undiscounted_reward(hv)

hi = HistoryIterator(hv, :r)
@inferred POMDPSimulators.step_tuple(hi, 1)
@test collect(hi) == collect(reward_hist(hv))
@test collect(eachstep(hv, "r")) == collect(reward_hist(hv)) # why isn't collect able to infer the type here??

# test show_progress
gw = SimpleGridWorld()
hr = HistoryRecorder(show_progress=true, max_steps=100)
println("Should be a progress bar below:")
@test length(simulate(hr, gw, FunctionPolicy(s->:left), rand(hr.rng, initialstate(gw)))) <= 100

# test capture_exception
gw = SimpleGridWorld()
hr = HistoryRecorder(show_progress=true, capture_exception=true, max_steps=100, rng=MersenneTwister(2))
counter = []
error_policy = FunctionPolicy(function (s)
                                  push!(counter, true)
                                  if length(counter) <= 3
                                      sleep(0.1) # so that the progress bar gets shown
                                      return :left
                                  else
                                      error("Policy Error")
                                  end
                              end)
println("Should be a progress bar below:")
exhist = simulate(hr, gw, error_policy, rand(hr.rng, initialstate(gw)))
@test 2 <= length(exhist) <= 10
@test exception(exhist) !== nothing

# test showprogress without max_steps
gw = SimpleGridWorld()
hr = HistoryRecorder(show_progress=true)
@test_throws ErrorException simulate(hr, gw, FunctionPolicy(s->:left), rand(hr.rng, initialstate(gw)))

#=
function f(hv)
    rs = 0.0
    for (r,a) in HistoryIterator{typeof(hv), (:r,:a)}(hv)
        rs += r
    end
    return rs
end
@code_warntype f(hv)
hi = HistoryIterator{typeof(r1), (:r,)}(r1)
t = step_tuple(hi, 1)
@code_warntype step_tuple(hi, 1)
=#
end


text:
# Info object return after a SLOPE solve.
type SLOPESolveInfo
  runtime       # Runtime
  Aprods        # Number of products with A
  ATprods       # Number of products with A^T
  objPrimal     # Primal objective
  objDual  # Dual objective (possibly for infeasible dual point)
  infeas        # Dual infeasibility
  status        # Status: 1 = optimal, 2 = iterations
end

# Constants for exit status
const STATUS_RUNNING   = 0
const STATUS_OPTIMAL   = 1
const STATUS_ITERATIONS = 2
const STATUS_MSG = ["Optimal", "Iteration limit reached"]

# The default choice for lambdas is Benjamini-Hochberg sequence where q is the FDR:
using Distributions
NormalDistr = Normal()
bh_lambdafunc(i, p, q = 0.10) = quantile(NormalDistr, 1 - i*q/(2*p))

# Lets implement the stack-based FastProx SL1 alg of Bogdan et al to see if that helps...
# This is Alg 4 in the Bogdan2014 paper.
function fast_prox_sl1(yproxvals, lambdas)
  p = length(yproxvals)
  tuples = Any[]
  t = 0

  # Find optimal group levels
  for k in 1:p
    t += 1
    i, j, s = k, k, (yproxvals[k] - lambdas[k])
    w = abs(s)
    push!(tuples, (i, j, s, w))
    while (t > 1) && (tuples[t-1][4] .<= w)
      itm1, jtm1, stm1, wtm1 = tuples[t-1]
      it, jt, st, wt = tuples[t]
      wnew = (jtm1 - itm1 + 1)/(jt - itm1 + 1) * stm1 .+ (jt - it + 1)/(jt - itm1 + 1) * st
      tuples[t-1] = (itm1, jt, stm1 .+ st, abs(wnew))
      pop!(tuples) # Delete last entry
      t -= 1
    end
  end

  # Set entries in x for each block
  x = zeros(p)
  for l in 1:t
    for k in tuples[l][1]:tuples[l][2]
      x[k] = tuples[l][4]
    end
  end
  x
end

# The prox problem we solve assumes y are sorted and positive. We keep the permutation vector
# around so we can revert later.
sort_and_normalize(y) = begin
  absy = abs(y)
  return sortperm(absy, rev = true), sign(y), absy
end

# Go back from a sortperm to orig order.
function revert_sortperm(y, yperm)
  p = length(y)
  yr = zeros(p)
  for i in 1:p
    yi = yperm[i]
    yr[yi] = y[i]
  end
  yr
end

solve_prox_problem_with_fastproxsl1(yproxvals, lambdas) = begin
  yperm, sngy, absy = sort_and_normalize(yproxvals)
  ynormalized = absy[yperm]

  xvals = fast_prox_sl1(ynormalized, lambdas)

  # Reorder and put back the signs...
  sngy .* revert_sortperm(xvals, yperm)
end

# Test values for params
#iterations = 10000      # Maximum number of iterations
#q = 0.05
#verbosity = 1         # 0 = nothing, 1 = major, 2 = every
#optimIter = 1         # Iterations between optimality-condition checks
#gradIter = 20         # Iterations between full gradient computations
#tolInfeas = 1e-6       # Maximum allowed dual infeasibility
#tolRelGap = 1e-6       # Stopping criterion for relative primal-dual gap
#xInit = zeros(size(A,2))  # Initial value of x
#lambdas = nothing

# solve_SLOPE  Sorted L1 parameter estimation solver
#
# x, info = solve_SLOPE(A, b, lambda; options) solves the Slope problem
#
#     Minimize 1/2*||Ax-b||_2^2 + sum_i (lambda_i * |x|_[i])
#
# where |x|_[i] denotes the i-th largest entry in |x|. The entries in
# lambda must be nonnegative and in non-increasing order. When lambda is a
# scalar, the above formulation is reduced to the Lasso:
#
#     Minimize 1/2*||Ax-b||_2^2 + lambda * ||x||_1.
#
# Original MatLab version is:
# Copyright 2013, M. Bogdan, E. van den Berg, W. Su, and E.J. Candes
# Julia version by Robert Feldt, robert.feldt@gmail.com
# was based on Adlas.m in SLOPE Toolbox version 1.0 while the FastProxSL1 algorithm
# was implemented from scratch in Julia by Robert based on the Bogdan et al 2014 paper
# on SLOPE. We also excluded the Lasso from the implementation since there are already
# LASSO implementations for Julia.
#
#   The SLOPE Toolbox is free software: you can redistribute it
#   and/or  modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation, either version 3 of
#   the License, or (at your option) any later version.
#
#   The SLOPE Toolbox is distributed in the hope that it will
#   be useful, but WITHOUT ANY WARRANTY; without even the implied
#   warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#   See the GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with the SLOPE Toolbox. If not, see
#   <http://www.gnu.org/licenses/>.
#
function solve_SLOPE(A, b, lambdas = nothing;
  iterations = 10000,      # Maximum number of iterations
  q = 0.05,                 # Expected FDR, False Discovery Rate
  verbosity = 1,         # 0 = nothing, 1 = major, 2 = every
  optimIter = 1,         # Iterations between optimality-condition checks
  gradIter = 20,         # Iterations between full gradient computations
  tolInfeas = 1e-6,       # Maximum allowed dual infeasibility
  tolRelGap = 1e-6,       # Stopping criterion for relative primal-dual gap
  xInit = zeros(size(A,2))  # Initial value of x
  )

# Start timer
t0 = tic()

# Get problem dimension
n = size(A, 2)

# Create lambda sequence unless already done.
if lambdas == nothing
  lambdas = map(k -> bh_lambdafunc(k, n, q), 1:n)
end

# Ensure that lambda is non-increasing, nonnegative and has at least one non-zero entry.
if ((length(lambdas) > 1) && any(lambdas[2:end] .> lambdas[1:end-1]))
  throw("Lambdas must be non-increasing.")
end
if (lambdas[end] < 0)
  throw("Lambdas must be nonnegative")
elseif (lambdas[1] == 0)
  throw("Lambdas must have at least one nonnegative entry.")
end

# Get initial lower bound on the Lipschitz constant
x = randn(n,1)
x = x / norm(x,2)
x = A'*(A*x)
L = norm(x,2)

# Initialize parameters and iterates
t     = 1
eta    = 2
lambdas  = lambdas[:]
b     = b[:]
x     = xInit
y     = x
Ax    = A*x
fPrev  = Inf
iter   = 0

info = SLOPESolveInfo(
  0.0, # runtime is set after finishing
  2,   # Aprods
  1,   # ATprods
  0.0, # objPrimal is set later
  0.0, # info.objDual   is set later
  0.0, # Dual infeasibility
  STATUS_RUNNING
)

solve_prox = solve_prox_problem_with_fastproxsl1

if (verbosity > 0)
  @printf("%5s  %9s  %9s  %9s  %9s", "Iter", "||r||_2", "Gap", "Infeas.", "Rel. gap")
end

# Main loop
while (true)

  # Compute the gradient at f(y)
  if (mod(iter, gradIter) == 0) # Includes first iterations
   r = A*y - b
  else
   r = (Ax + ((tPrev - 1) / t) * (Ax - AxPrev)) - b
  end
  g = A' * r
  f = dot(r, r) / 2

  # Increment iteration count
  iter += 1

  # Check optimality conditions
  if ((mod(iter, optimIter) == 0))
    # Compute 'dual', check infeasibility and gap
    gs    = sort(abs(g), rev = true)
    ys    = sort(abs(y), rev = true)
    info.infeas = max(maximum(cumsum(gs-lambdas)),0)

    # Compute primal and dual objective
    info.objPrimal = ( f + lambdas'*ys)[1]
    info.objDual   = (-f - r'*b       )[1]

    if (verbosity > 0)
      str = @sprintf("  %9.2e  %9.2e  %9.2e", info.objPrimal - info.objDual, info.infeas/lambdas[1],
        abs(info.objPrimal - info.objDual) / max(1,info.objPrimal))
    end

    # Check primal-dual gap
    if ((abs(info.objPrimal - info.objDual)/max(1,info.objPrimal) < tolRelGap) &&
        (info.infeas < tolInfeas * lambdas[1]))
      info.status = STATUS_OPTIMAL;
    end
  else
    str = ""
  end

  if ((verbosity == 2) || ((verbosity == 1) && (mod(iter,optimIter) == 0)))
    @printf("%5d  %9.2e%s\n", iter, f, str)
  end

  # Stopping criteria
  if (info.status == STATUS_RUNNING)
    if (iter >= iterations)
      info.status = STATUS_ITERATIONS
    end
  end

  if (info.status != 0)
    if (verbosity > 0)
      @printf("Exiting with status %d -- %s\n", info.status, STATUS_MSG[info.status])
    end
    break;
  end

  # Keep copies of previous values
  AxPrev = Ax
  xPrev  = x
  fPrev  = f
  tPrev  = t

  # Lipschitz search
  while (true)
    # Compute prox mapping
    x = solve_prox(y - (1/L)*g, lambdas/L)
    d = x - y

    Ax = A*x
    r  = Ax-b
    f  = dot(r, r)/2
    q  = fPrev + dot(d, g) + (L/2)*dot(d, d)

    info.Aprods += 1

    if (q >= f*(1-1e-12))
      break
    else
      L = L * eta
    end
  end

  # Update
  t = (1 + sqrt(1 + 4*t^2)) / 2
  y = x + ((tPrev - 1) / t) * (x - xPrev)
end

# Final touchup of info values
info.runtime = toc()
info.Aprods = info.Aprods + ceil(iter / gradIter)
info.ATprods  = info.ATprods + iter

# ...and we are finished!
return y, info

end # function solve_SLOPE

function rand_sparse_problem(p; numactive = iceil(log(p)),
  n = int(p/2), amplitude = 4.1, sigma = 1.0, shuffle = true)

  beta = vcat(amplitude*ones(numactive)+randn(numactive), zeros(p-numactive))
  indices = collect(1:p)
  if shuffle
    shuffle!(indices)
    beta = beta[indices]
  end
  X = randn(n, p)
  errors = sigma*randn(n)
  y = X * beta .+ errors
  return beta, X, y, indices, n, p, numactive, sigma, amplitude, errors
end

# Find the sharpest difference between the found values and use that as the estimate
# of the num active values.
function estimate_active(xhat)
  relsize = abs(xhat) / maximum(abs(xhat))
  relsizeperm = sortperm(relsize, rev=true)
  srelsize = relsize[relsizeperm]
  relfactors = srelsize[1:end-1] ./ srelsize[2:end]
  numactive = find(rf -> rf == maximum(relfactors), relfactors)[1]
  return relsizeperm[1:numactive]
end

# First estimate the number of active coefficients by a SLOPE solve then
# do a linreg on only those factors.
function sparse_regression_SLOPE(A, b; iterations = 250, options...)
  xhat, info = solve_SLOPE(A, b; iterations = iterations, options...)
  active = estimate_active(xhat)
  Asub = A[:, active]
  res = linreg(Asub, b)
  xhatsub = zeros(size(A, 2))
  xhatsub[active] = res[2:end]
  return xhatsub, res[2:end], active
end

if length(ARGS) >= 1
  numcoefs = int(ARGS[1])
else
  numcoefs = 1000
end
if length(ARGS) >= 2
  numact = int(ARGS[2])
else
  numact = 10
end
if length(ARGS) >= 3
  maxiterations = int(ARGS[3])
else
  maxiterations = 500
end

x, A, b, indices, n, p, numactive, sigma, amplitude, errors = rand_sparse_problem(numcoefs;
  shuffle = false, numactive = numact);

# Split in train and test set
idxs = shuffle(collect(1:n))
trainsize = ifloor(0.80*n)
trainidxs = idxs[1:trainsize]
testidxs = idxs[(trainsize+1):end]
Atrain = A[trainidxs, :]
btrain = b[trainidxs]
Atest = A[testidxs, :]
btest = b[testidxs]

xhat, xselected, selected = sparse_regression_SLOPE(Atrain, btrain; iterations = maxiterations)

println(x[1:numactive])
println(xhat[1:numactive])
println("Norm(x .- xhat) = ", norm(x .- xhat))
mape(y, yhat) = mean(100.0 * abs((yhat - y)/y))
@printf("Num active = %d, Num active selected = %d\n", numactive, length(selected))
@printf("Train MAPE = %.2f%%\n", mape(btrain, Atrain*xhat))
@printf("Test MAPE = %.2f%%\n", mape(btest, Atest*xhat))


text:
## Tuning
# Relative weight of acceleration (vs null space accuracy)
alpha = 1e3; # Accel
beta = 1e1; # Vel
gamma = 1e1 # Loss

# Out
#c_tight = 0*1e-3;
#c_usage = 2*1e-3;

# Tight
c_tight = 2*1e-3;
c_usage = 0*1e-3;

# Closeness to previous trajectory
epsilon = deg2rad(10);
# Constraint the joints to be close on how many iterations...
# More skips makes the formulation of the problem easier
# Only works with proper acceleration weight
nSkip = 0; # Default on all
#nSkip = 10; # One constraint per second
#nSkip = 3;
nSkip = max(floor(nSkip), 0) + 1;
# TODO: Truncate the path if possible, once the difference is small
# Then further optimization steps just make the path smaller

# Optimization tools
#using ECOS; mysolver = ECOSSolver()
using SCS;mysolver = SCSSolver()
#using Gurobi; mysolver = GurobiSolver()
using Convex

# Set the listener properly
using ZMQ
using MAT

ctx=Context()
s1=Socket(ctx, REP)

ZMQ.bind(s1, "ipc:///tmp/armopt")

while true

  println("Waiting for a message")
  msg = ZMQ.recv(s1)
  println("Received a message!")
  matfile = bytestring(msg);
  println("I received: ", matfile)

  file = matopen(matfile)
  varnames = names(file)
  println(varnames);
  close(file);

  # This works fine
  vars = matread(matfile);

  # wh: ctight,cusage,csimilar
  wh = vars["wh"]
  println("wh", wh)
  c_similar = wh[3] # * 1e-3
  c_usage = wh[2] # * 1e-3
  c_tight = wh[1] # * 1e-3

  ## Convert the variables

  # Task Space target
  vw0 = vars["vwPath"]
  vwa = Array{Float64}(size(vw0, 1), size(vw0[1], 1))
  for iter in eachindex(vw0)
  #    @show qw0[iter]
  #    @show iter
      vwa[iter, :] = vw0[iter]
  end
  vw0 = vwa
  vwa = 0;
  # Original Path
  qw0 = vars["qwPath"]
  qw0a = Array{Float64}(size(qw0, 1), size(qw0[1], 1))
  for iter in eachindex(qw0)
  #    @show qw0[iter]
  #    @show iter
      qw0a[iter, :] = qw0[iter]
  end
  qwPath0 = qw0a
  qw0a = 0;qw0 = 0;

  # Task Space effective
  vw = vars["vwEffective"]
  vwa = Array{Float64}(size(vw, 1), size(vw[1], 1))
  for iter in eachindex(vw)
  #    @show qw0[iter]
  #    @show iter
      vwa[iter, :] = vw[iter]
  end
  vw = vwa
  vwa = 0;

  # Number of trajectory points
  np = size(qwPath0, 1)
  # Number of joints
  nq = size(qwPath0, 2)
  # Number of optimization points
  n = length(qwPath0)
  println(n,'=',np, 'x', nq)

  ## Velocity matrix
  v0 = ones(n-nq);
  v0[1:nq, :] = 2;
  v1 = zeros(n);
  v1[1:nq] = -2;
  v1[end-nq+1:end] = 2;
  V = diagm(v0, nq) + diagm(v1);
  V = V + diagm(-flipdim(v0, 1), -nq);
  V = 0.5*V;

  ## Acceleration matrix
  d2 = 2 * ones(n);
  # Proper doundary condition (Central difference 2nd order):
  # http://www.mathematik.uni-dortmund.de/~kuzmin/cfdintro/lecture4.pdf
  d2[1:nq] = 1;
  d2[end-nq+1:end] = 1;
  d1 = ones(n - nq);

  A0 = diagm(-d2);
  A1 = diagm(d1, nq);
  A = A0 + A1 + A1';
  A = sparse(A);
  ATA = A' * A;

  # Null Space effective
  nulls = vars["nulls"]
  nulls_sparse = Array{Any}(size(vw, 1),)
  for iter in eachindex(nulls)
      nulls_sparse[iter] = sparse(nulls[iter]);
  end

  nulls = 0;
  N = blkdiag(nulls_sparse...)';
  nulls_sparse = 0;
  #NTN = N' * N;

  # Jacobian effective
  Js = vars["Js"]
  Js_sparse = Array{Any}(size(vw, 1),)
  for iter in eachindex(Js)
      Js_sparse[iter] = sparse(Js[iter]);
  end

  Js = 0;
  J = blkdiag(Js_sparse...)';
  Js_sparse = 0;
  JV = J * V;

  # Constant Constraints
  qMid = [0, 0.759218, 0, -1.39626, 0, 0, 0];
  qMid = repmat(qMid[:], np, 1);
  elbow_diag = ones(n);
  elbow_diag[3:7:n] = 1;
  N_elbow = diagm(elbow_diag) * N;
  #NTN_elbow = N_elbow' * N_elbow;

  qGravity = deg2rad([0, 60, 90, -120, -90, -15, 0]);
  qGravity = repmat(qGravity[:], np, 1);

  interactions = getkey(vars, "interactions", false)
  sumInteractions = zeros(n);
  ni = 0
  if interactions != false
    interactions = vars["interactions"]
    ni = length(interactions);
    for iter in eachindex(interactions)
      @show interactions[iter]
      qI = interactions[iter]
      sumInteractions = sumInteractions + repmat(qI[:], np, 1)
    end
    sumInteractions = 2 * sumInteractions / ni
    gamma = 1e2;
  else
    gamma = 0;
  end
  # Normalize (the optimization has ni divided out of q' q)

  println("ni", ni)

  # Solve the problem
  q0 = reshape(qwPath0', n, 1)
  q = Convex.Variable(n)
  p = minimize(
    alpha * sumsquares(A * q)
    + beta * sumsquares(JV * q)
    + gamma * (sumsquares(q) - sumInteractions' * q) # Augmented loss
    + c_tight * sumsquares(N_elbow * q)
    + c_similar * sumsquares(N * (q - qGravity))
    + c_usage * sumsquares(N * (q-qMid))
  )

  p.constraints += q[1:nq]==q0[1:nq]
  if ni==0
    p.constraints += q[n-nq+1:n]==q0[n-nq+1:n]
    for k = nq+1 : nSkip*nq : n-nq
        p.constraints += norm(q[k:k+nq-1] - q0[k:k+nq-1], 2) <= epsilon;
    end
  else
    for k = nq+1 : nSkip*nq : n
        p.constraints += norm(q[k:k+nq-1] - q0[k:k+nq-1], 2) <= epsilon;
    end
  end
  #@time Convex.solve!(p, mysolver);
  Convex.solve!(p, mysolver);
  #println("Optimal objective: ", p.optval);
  #println(q.value);

  q_opt = reshape(q.value, nq, np)';

  println("Solution size: ", size(q_opt));
  file = matopen(matfile, "w")
  write(file, "qw", q_opt)
  close(file)

  ZMQ.send(s1, Message(matfile))

end

# Cleanup
ZMQ.close(s1)
ZMQ.close(ctx)

max_stars_repo_name: ToyotaResearchInstitute/rad-robot

max_issues_repo_name: ToyotaResearchInstitute/rad-robot

max_forks_repo_name: ToyotaResearchInstitute/rad-robot


text:
#using Distributed
#@everywhere using MambaModels
using MambaModels

globe_toss = Dict{Symbol, Any}(
  :w => [6, 7, 5, 6, 6],
  :n => [9, 9, 9, 9, 9]
)
globe_toss[:N] = length(globe_toss[:w]);

model = Model(
  w = Stochastic(1,
    (n, p, N) ->
      UnivariateDistribution[Binomial(n[i], p) for i in 1:N],
    false
  ),
  p = Stochastic(() -> Beta(1, 1))
);

inits = [
  Dict(:w => globe_toss[:w], :n => globe_toss[:n], :p => 0.5),
  Dict(:w => globe_toss[:w], :n => globe_toss[:n], :p => rand(Beta(1, 1)))
];

scheme = [NUTS(:p)]
setsamplers!(model, scheme);

chn = mcmc(model, globe_toss, inits, 10000, burnin=2500, thin=1, chains=2);

describe(chn)

chn2 = MCMCChains.Chains(chn.value, String.(chn.names))

MCMCChains.describe(chn2)

MCMCChains.plot(chn2)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


text:
#Function to find local maximums and minimums of a function
# assuming the function is smooth, we just simply check
# y(i-1) < y(i) > y(i+1) for max and y(i-1)> y(i) <y(i+1) for min
#

function findExtrema(y::Vector{Float64}, t::Vector{Float64})

	maxAr  = Float64[]	#array for storing local maxima
	minAr  = Float64[]	#array for storing local minima
	tMax = Float64[]	#array for storing location of maxima
	tMin = Float64[]	#array for storing location of minima
	n = length(y)

	#add initial points to all arrays
	push!(maxAr, y[1])
	push!(minAr, y[1])
	push!(tMax, t[1])
	push!(tMin, t[1])

	for i = 2:length(y)-1
		if(y[i-1] < y[i] && y[i] > y[i+1])
			push!(maxAr, y[i])
			push!(tMax, t[i])
		end

		if(y[i-1] > y[i] && y[i] < y[i+1])
			push!(minAr, y[i])
			push!(tMin, t[i])
		end
	end

	#add final points to all arrays
	push!(maxAr, y[n])
	push!(minAr, y[n])
	push!(tMax, t[n])
	push!(tMin, t[n])

	#fix value to extrapolate
	maxAr[end] = max(maxAr[end], maxAr[end-1])
	minAr[end] = min(minAr[end], minAr[end-1])
	maxAr[1] = max(maxAr[2], maxAr[1])
	minAr[1] = min(minAr[2], minAr[1])

	return maxAr, minAr, tMax, tMin
end


text:
strs = ["mut_pattern" "ref_pattern" "captured_pattern" "reference_pattern" "slice_pattern" "tuple_pattern" "or_pattern" "tuple_struct_pattern" "struct_pattern > field_pattern"]

bases = ["parameter" "closure_parameters"]

ident = "identifier"

final = []

for base ∈ bases
    string = base * " > " * ident * ","
    push!(final, string)
end

for base ∈ bases, str ∈ strs
    if str == "or_pattern"
        # string = base * " > "
        # for i ∈ 1:5
        #     string = string * str * " > "
        #     push!(final, string * ident * ",")
        # end
    else
        string = base * " > " * str * " > " * ident * ","
        push!(final, string)
    end
end

for base ∈ bases, str ∈ strs, str2 ∈ strs
    if str == "or_pattern"
        # string = base * " > "
        # aux = str2 * " > " * ident * ","
        # for i ∈ 1:5
        #     string = string * str * " > "
        #     push!(final, string * aux)
        # end
    elseif str2 == "or_pattern"
        string = base * " > " * str * " > "
        for i ∈ 1:5
            string = string * str2 * " > "
            push!(final, string * ident * ",")
        end
    else
        string = base * " > " * str * " > " * str2 * " > " * ident * ","
        push!(final, string)
    end
end

for base ∈ bases, str ∈ strs, str2 ∈ strs, str3 ∈ strs
    if str == "or_pattern"
        # string = base * " > "
        # aux = str2 * " > " * str3 * " > " * ident * ","
        # for i ∈ 1:5
        #     string = string * str * " > "
        #     push!(final, string * aux)
        # end
    elseif str2 == "or_pattern"
        string = base * " > " * str * " > "
        aux = str3 * " > " * ident * ","
        for i ∈ 1:5
            string = string * str2 * " > "
            push!(final, string * aux)
        end
    elseif str3 == "or_pattern"
        string = base * " > " * str * " > " * str2 * " > "
        for i ∈ 1:5
            string = string * str3 * " > "
            push!(final, string * ident * ",")
        end
    else
        string = base * " > " * str * " > " * str2 * " > " * str3 * " > " * ident * ","
        push!(final, string)
    end
end

# for base ∈ bases, str ∈ strs, str2 ∈ strs, str3 ∈ strs, str4 ∈ strs
#     if str == "or_pattern"
#         string = base * " > "
#         aux = str2 * " > " * str3 * " > " * str4 * " > " * ident * ","
#         for i ∈ 1:5
#             string = string * str * " > "
#             push!(final, string * aux)
#         end
#     elseif str2 == "or_pattern"
#         string = base * " > " * str * " > "
#         aux = str3 * " > " * str4 * " > " * ident * ","
#         for i ∈ 1:5
#             string = string * str2 * " > "
#             push!(final, string * aux)
#         end
#     elseif str3 == "or_pattern"
#         string = base * " > " * str * " > " * str2 * " > "
#         aux = str4 * " > " * ident * ","
#         for i ∈ 1:5
#             string = string * str3 * " > "
#             push!(final, string * aux)
#         end
#     elseif str4 == "or_pattern"
#         string = base * " > " * str * " > " * str2 * " > " * str3 * " > "
#         aux = ident * ","
#         for i ∈ 1:5
#             string = string * str4 * " > "
#             push!(final, string * aux)
#         end
#     else
#         string = base * " > " * str * " > " * str2 * " > " * str3 * " > " * str4 * " > " * ident * ","
#         push!(final, string)
#     end
# end
#
# for base ∈ bases, str ∈ strs, str2 ∈ strs, str3 ∈ strs, str4 ∈ strs, str5 ∈ strs
#     if str == "or_pattern"
#         string = base * " > "
#         aux = str2 * " > " * str3 * " > " * str4 * " > " * str5 * " > " * ident * ","
#         for i ∈ 1:5
#             string = string * str * " > "
#             push!(final, string * aux)
#         end
#     elseif str2 == "or_pattern"
#         string = base * " > " * str * " > "
#         aux = str3 * " > " * str4 * " > " * str5 * " > " * ident * ","
#         for i ∈ 1:5
#             string = string * str2 * " > "
#             push!(final, string * aux)
#         end
#     elseif str3 == "or_pattern"
#         string = base * " > " * str * " > " * str2 * " > "
#         aux = str4 * " > " * str5 * " > " * ident * ","
#         for i ∈ 1:5
#             string = string * str3 * " > "
#             push!(final, string * aux)
#         end
#     elseif str4 == "or_pattern"
#         string = base * " > " * str * " > " * str2 * " > " * str3 * " > "
#         aux = str5 * " > " * ident * ","
#         for i ∈ 1:5
#             string = string * str4 * " > "
#             push!(final, string * aux)
#         end
#     elseif str5 == "or_pattern"
#         string = base * " > " * str * " > " * str2 * " > " * str3 * " > " * str4 * " > "
#         aux = ident * ","
#         for i ∈ 1:5
#             string = string * str5 * " > "
#             push!(final, string * aux)
#         end
#     else
#         string = base * " > " * str * " > " * str2 * " > " * str3 * " > " * str4 * " > " * str5 * " > " * ident * ","
#         push!(final, string)
#     end
# end

sort!(final)

for f in final
    println(f)
end

max_stars_repo_name: GrayJack/language-rust-treesitter

max_issues_repo_name: GrayJack/language-rust-treesitter

max_forks_repo_name: GrayJack/language-rust-treesitter


text:
using MLBase
using Base.Test

## Kfold

ss = collect(Kfold(12, 3))
@test length(ss) == 3
for i = 1:3
	@test isa(ss[i], Vector{Int})
	@test issorted(ss[i])
end
x = vcat(ss...)
@test sort(x) == [1:12]

## LOOCV

ss = collect(LOOCV(4))
@test length(ss) == 4

@test ss[1] == [2, 3, 4]
@test ss[2] == [1, 3, 4]
@test ss[3] == [1, 2, 4]
@test ss[4] == [1, 2, 3]

## RandomSub

# temporary disable until StatsBase passes travis

# ss = collect(RandomSub(10, 5, 6))
# @test length(ss) == 6
# for i = 1:6
# 	@test length(ss[i]) == 5
# 	@test 1 <= minimum(ss[i]) <= maximum(ss[i]) <= 10
# 	@test length(unique(ss[i])) == 5
# end


text:
# Cosine similarity baseline

# using JLD
# isdefined(:data) || @load("groupedMelSegData.jld")

cosine(x,y) = dot(x,y)/(vecnorm(x)*vecnorm(y))

normal(x; mx=mean(x), sx=std(x)) = (x .- mx) ./ sx

cosnorm(x,y; o...) = cosine(normal(x; o...), normal(y; o...))

function allcos(xx,yy; o...)
    a = Float32[]
    for x in xx
        for y in yy
            push!(a, cosnorm(x,y; o...))
        end
    end
    return a
end

function printcos(data)
    q = [0,.25,.5,.75,1]
    for d in data
        isa(d["PerSegsTrue"],Void) && continue
        r1 = vcat(d["RefSegsTrue"]...)
        mx = mean(r1)
        sx = std(r1)
        println(quantile(allcos(d["RefSegsTrue"],d["RefSegsTrue"],mx=mx,sx=sx),q))
        println(quantile(allcos(d["RefSegsTrue"],d["PerSegsTrue"],mx=mx,sx=sx),q))
        println(quantile(allcos(d["RefSegsTrue"],d["PerSegsFalse"],mx=mx,sx=sx),q))
        println()
    end
end

# Output of printcos(data):

# Float32[0.51783,0.843342,0.898515,0.936384,1.0]
# Float32[0.281831,0.601221,0.722581,0.808675,0.955657]
# Float32[-0.528476,0.52987,0.683308,0.779417,0.927576]

# Float32[-0.417565,0.420223,0.641019,0.82493,1.0]
# Float32[-0.248075,0.108873,0.162213,0.335487,0.83023]
# Float32[-0.48789,-0.0693233,0.0767633,0.30544,0.889108]

# Float32[-0.274607,0.600598,0.778746,0.870678,1.0]
# Float32[-0.470661,0.524548,0.7144,0.815901,0.943585]
# Float32[-0.512623,0.0904965,0.32138,0.577408,0.908277]

# Float32[0.608947,0.866173,0.918163,0.946268,1.0]
# Float32[0.428327,0.78044,0.852204,0.898739,0.97381]
# Float32[-0.305821,0.394973,0.71889,0.808471,0.949947]

# Float32[0.430234,0.825169,0.883162,0.924848,1.0]
# Float32[0.0956748,0.554943,0.69163,0.81978,0.944336]
# Float32[-0.308146,0.438314,0.711019,0.793845,0.954038]

# Float32[0.321806,0.750683,0.84002,0.900515,1.0]
# Float32[0.100405,0.614332,0.735065,0.819915,0.962736]
# Float32[0.0353075,0.514158,0.649768,0.770142,0.938516]

# Float32[0.250692,0.802514,0.882671,0.9294,1.0]
# Float32[-0.140958,0.315888,0.760035,0.856603,0.963303]
# Float32[-0.628668,0.142457,0.505574,0.753477,0.947517]

# Float32[0.145144,0.629469,0.769253,0.871738,1.0]
# Float32[-0.196362,0.314905,0.529402,0.690695,0.940649]
# Float32[-0.604712,-0.0881434,0.371628,0.602275,0.917258]

# Float32[0.294274,0.70889,0.81317,0.88929,1.0]
# Float32[0.329316,0.709471,0.772187,0.833477,0.955523]
# Float32[-0.325076,0.389826,0.561363,0.675829,0.909722]

# Float32[0.286949,0.765829,0.853494,0.910613,1.0]
# Float32[0.17036,0.660064,0.7828,0.858932,0.953863]
# Float32[-0.674776,0.0588915,0.357284,0.719501,0.903627]

# Float32[0.360445,0.728305,0.82139,0.887689,1.0]
# Float32[0.451495,0.713871,0.804782,0.874524,0.949332]
# Float32[-0.0533012,0.421255,0.584092,0.723032,0.947801]

# Float32[0.0841178,0.583282,0.725749,0.826258,1.0]
# Float32[-0.0518345,0.0845716,0.143931,0.291672,0.683634]
# Float32[-0.380692,0.213723,0.343918,0.509993,0.922219]

# Float32[0.308682,0.746849,0.820015,0.879318,1.0]
# Float32[0.378033,0.685838,0.78525,0.858608,0.963286]
# Float32[-0.414857,0.149977,0.311626,0.541122,0.920682]

# Float32[0.533278,0.821617,0.881512,0.921837,1.0]
# Float32[0.392578,0.698664,0.790452,0.844808,0.952638]
# Float32[-0.5963,0.0910082,0.400182,0.794856,0.928753]

# Float32[0.643091,0.845797,0.894519,0.935809,1.0]
# Float32[-0.0705272,-0.0264076,-0.0110646,0.0237362,0.0863015]
# Float32[-0.605406,-0.0665198,0.294374,0.619856,0.88169]

# Float32[-0.0810161,0.429112,0.699983,0.827616,1.0]
# Float32[-0.0189678,0.0587066,0.281121,0.730419,0.943307]
# Float32[-0.7917,-0.0391697,0.139415,0.350606,0.868362]

# Float32[0.208546,0.671326,0.789261,0.862174,1.0]
# Float32[0.0571182,0.341359,0.433556,0.822178,0.929826]
# Float32[-0.399101,0.120124,0.282105,0.520031,0.803401]

# Float32[0.399047,0.797122,0.866454,0.915409,1.0]
# Float32[0.34736,0.562841,0.676371,0.796767,0.934049]
# Float32[-0.306154,0.190556,0.495839,0.659734,0.915149]

# Float32[0.633719,0.852629,0.894732,0.928784,1.0]
# Float32[-0.127642,-0.0774944,-0.0511268,-0.0208007,0.0557352]
# Float32[-0.52172,-0.0498513,0.079309,0.551165,0.90503]

# Float32[0.209224,0.670564,0.793241,0.883707,1.0]
# Float32[0.0703688,0.501442,0.6363,0.748235,0.965291]
# Float32[-0.253693,0.14136,0.398783,0.528235,0.899706]

# Float32[0.352818,0.866856,0.910514,0.946583,1.0]
# Float32[0.269033,0.68343,0.777905,0.821568,0.93462]
# Float32[-0.639206,0.254124,0.550059,0.738388,0.93474]

# Float32[0.256018,0.584986,0.705985,0.811057,1.0]
# Float32[0.110211,0.480432,0.62279,0.732447,0.943336]
# Float32[-0.407086,0.154326,0.396867,0.515015,0.859458]

# Float32[-0.177102,0.540115,0.72831,0.846759,1.0]
# Float32[-0.325715,0.270696,0.468592,0.687751,0.932782]
# Float32[-0.694876,-0.0509394,0.21477,0.508979,0.837127]

# Float32[0.413238,0.754692,0.842189,0.907146,1.0]
# Float32[-0.134084,0.523553,0.657909,0.751509,0.957077]
# Float32[-0.234611,0.43265,0.576692,0.694131,0.941075]

# Float32[0.16522,0.733814,0.827267,0.895219,1.0]
# Float32[0.095797,0.534112,0.724223,0.8293,0.949398]
# Float32[-0.510714,0.31645,0.45511,0.604286,0.891072]

# Float32[0.125046,0.558113,0.720568,0.84261,1.0]
# Float32[0.0820815,0.413827,0.577948,0.712663,0.951903]
# Float32[-0.397263,0.0877374,0.339798,0.483134,0.921224]

# Float32[0.597537,0.77237,0.849037,0.911701,1.0]
# Float32[0.270412,0.620787,0.681177,0.757054,0.950191]
# Float32[-0.216431,0.381294,0.596779,0.650016,0.902907]

# Float32[0.289887,0.754108,0.854132,0.916598,1.0]
# Float32[0.0517729,0.595549,0.705233,0.825866,0.958469]
# Float32[-0.350919,0.392623,0.594915,0.713161,0.94204]

# Float32[0.622863,0.854053,0.906569,0.942579,1.0]
# Float32[0.584877,0.730615,0.815238,0.879687,0.982877]
# Float32[0.243019,0.550825,0.636882,0.788584,0.966161]

# Float32[0.0864844,0.678707,0.813338,0.907985,1.0]
# Float32[0.462249,0.790146,0.865742,0.904563,0.95724]
# Float32[-0.199216,0.338543,0.600747,0.758231,0.926958]

# Float32[0.212164,0.656818,0.807896,0.901439,1.0]
# Float32[0.246509,0.706877,0.822019,0.878912,0.957814]
# Float32[-0.150007,0.159957,0.494455,0.649792,0.911291]

# Float32[0.415903,0.66202,0.778553,0.877506,1.0]
# Float32[0.369278,0.592971,0.716744,0.840488,0.94605]
# Float32[0.0135101,0.351836,0.481143,0.636292,0.916417]

# Float32[0.555317,0.850503,0.90208,0.935524,1.0]
# Float32[-0.16492,0.750758,0.842921,0.897972,0.966333]
# Float32[-0.260984,0.476758,0.743343,0.808493,0.95907]

# Float32[0.210637,0.75971,0.857757,0.926666,1.0]
# Float32[-0.021836,0.707973,0.776229,0.864274,0.974328]
# Float32[0.257579,0.572664,0.712575,0.787581,0.893648]

# Float32[0.493405,0.811113,0.879992,0.922542,1.0]
# Float32[0.343702,0.678016,0.749672,0.825126,0.954717]
# Float32[0.0793594,0.31623,0.504254,0.714104,0.940022]

# Float32[0.0768921,0.63624,0.781214,0.866409,1.0]
# Float32[0.158472,0.619118,0.723303,0.801102,0.91985]
# Float32[-0.412722,0.0460212,0.330416,0.524561,0.800956]

# Float32[0.531199,0.817854,0.874247,0.91973,1.0]
# Float32[0.613064,0.791686,0.854856,0.900375,0.971927]
# Float32[-0.0387285,0.489283,0.654501,0.766355,0.951595]

# Float32[0.361937,0.709528,0.799786,0.891958,1.0]
# Float32[0.193897,0.584114,0.698858,0.786687,0.944286]
# Float32[-0.145143,0.267834,0.64085,0.759461,0.884658]

# Float32[0.252463,0.642773,0.77884,0.884377,1.0]
# Float32[0.115067,0.584067,0.733906,0.823004,0.942661]
# Float32[-0.44058,0.0329857,0.483374,0.705563,0.913877]

max_stars_repo_name: JuliaTagBot/melseg

max_issues_repo_name: JuliaTagBot/melseg

max_forks_repo_name: JuliaTagBot/melseg