text:
using Plots
include_ddp()

# Model
include_model("rocket3D")

function fd(model::Rocket3D{Midpoint, FixedTime}, x, u, w, h, t)
	return view(x, 1:model.n) + h * f(model, view(x, 1:model.n) + 0.5 * h * f(model, view(x, 1:model.n), view(u, 1:model.m), w), view(u, 1:model.m), w)
end

n = model.n
m = model.m

# Time
T = 101
h = 0.05

# Initial conditions, controls, disturbances
x1 = zeros(model.n)
x1[1] = 2.5
x1[2] = 2.5
x1[3] = 10.0
# mrp = MRP(RotY(-0.45 * π) * RotX(-0.25 * π))
mrp = MRP(RotZ(0.25 * π) * RotY(-0.45 * π))

x1[4:6] = [mrp.x; mrp.y; mrp.z]
x1[9] = -1.0

visualize!(vis, model, [x1], Δt = h)

# visualize!(vis, model, [x1], Δt = h)

xT = zeros(model.n)
# xT[1] = 2.5
# xT[2] = 0.0
xT[3] = model.length
mrpT = MRP(RotZ(0.25 * π) * RotY(0.0))
xT[4:6] = [mrpT.x; mrpT.y; mrpT.z]
visualize!(vis, model, [xT], Δt = h)

u_ref = [0.0; 0.0; 0.0]#model.mass * 9.81]
ū = [u_ref + [1.0e-2; 1.0e-2; 1.0e-2] .* randn(model.m) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)
# x̄ = linear_interpolation(x1, xT, T)
# plot(hcat(x̄...)')

# Objective
Q = h * [(t < T ? 1.0 * Diagonal([1.0e-1 * ones(3); 0.0 * ones(3); 1.0e-1 * ones(3); 1000.0 * ones(3)])
        : 1000.0 * Diagonal(1.0 * ones(model.n))) for t = 1:T]
q = h * [-2.0 * Q[t] * xT for t = 1:T]

R = h * [Diagonal([10000.0; 10000.0; 100.0]) for t = 1:T-1]
r = h * [-2.0 * R[t] * u_ref  for t = 1:T-1]

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

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

# Constraints
x_con = [-0.5; 0.5]
y_con = [-0.75; 0.75]
p = [t < T ? 2 * m + 1 : n - 2 + 4 for t = 1:T]
info_t = Dict(:ul => [-5.0; -5.0; 0.0], :uu => [5.0; 5.0; 12.5], :inequality => (1:2 * m + 1))
info_T = Dict(:xT => xT, :inequality => (1:4))
con_set = [StageConstraint(p[t], t < T ? info_t : info_T) for t = 1:T]
idx_T = collect([3, 4, 5, 6, 7, 8, 9, 10, 11, 12])

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

	if t < T
		ul = cons.con[t].info[:ul]
		uu = cons.con[t].info[:uu]
		c[1:2 * m] .= [ul - u; u - uu]
		c[collect([1, 2])] .= 0.0 # turn off fx, fy lower-bound constraints
		c[collect([4, 5, 6])] .= 0.0 # turn off upper_bound constraints
		c[2 * m + 1] = model.length - x[3]
	elseif t == T
		c[1] = x_con[1] - x[1]
		c[2] = x[1] - x_con[2]
		c[3] = y_con[1] - x[2]
		c[4] = x[2] - y_con[2]
		xT = cons.con[T].info[:xT]
		c[4 .+ (1:(n - 2))] .= (x - xT)[idx_T]
	else
		nothing
	end
end

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

# Solve
@time constrained_ddp_solve!(prob,
    max_iter = 1000,
	grad_tol = 1.0e-3,
	obj_tol = 1.0e-3,
	max_al_iter = 10,
	con_tol = 1.0e-3,
	ρ_init = 1.0, ρ_scale = 10.0)

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

x̄_nominal = x̄
ū_nominal = ū

@save "/home/taylor/Research/motion_planning/examples/implicit_dynamics/examples/trajectories/rocket.jld2" x̄_nominal ū_nominal
@load "/home/taylor/Research/motion_planning/examples/implicit_dynamics/examples/trajectories/rocket.jld2" x̄_nominal ū_nominal

# Trajectories
plot(hcat(ū...)', linetype = :steppost)
plot(hcat(x̄...)[1:3, :]', linetype = :steppost)
plot(hcat(x̄...)[4:6, :]', linetype = :steppost)

# t = 1
# idx = [3; 1; 2]
# second_order_cone_projection(ū[t][idx])
# ū[t]
# # Visualize
include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
# open(vis)
x_anim = [[x̄[1] for i = 1:100]..., x̄..., [x̄[end] for i = 1:100]...]
visualize!(vis, model,
	x_anim,
	Δt = h, T_off = length(x_anim)-125)
obj_platform = joinpath(pwd(), "models/rocket/space_x_platform.obj")
mtl_platform = joinpath(pwd(), "models/rocket/space_x_platform.mtl")
ctm_platform = ModifiedMeshFileObject(obj_platform,mtl_platform,scale=1.0)

setobject!(vis["platform"],ctm_platform)
settransform!(vis["platform"], compose(Translation(0.0,2.0,-0.6), LinearMap(0.75 * RotZ(pi)*RotX(pi/2))))

settransform!(vis["/Cameras/default"], compose(Translation(0.0, 0.0, 0.0),
	LinearMap(RotZ(1.5 * π + 0.0 * π))))
setvisible!(vis["/Grid"], false)
setprop!(vis["/Cameras/default/rotated/<object>"], "zoom", 1.0)

max_issues_repo_name: thowell/DirectMotionPlanning


text:
include("interpolation.jl")
include("wannierize_utils.jl")

using PyPlot
using LinearAlgebra
#using Interpolations

## Assumptions: the kpoints are contained in a NxNxN cartesian grid, the neighbor list must contain the six cartesian neighbors

## N1xN2xN3 grid, filename.mmn must contain the overlaps
## nbeg and nend specify the window to wannierize
## Input Mmn file is nband x nband x nkpt x nntot
## Output is (nwannier = nend - nbeg + 1) x nband x nkpt, padded with zeros
function make_wannier(p,method)
    t1 = collect(0:p.N1-1)/p.N1
    t2 = collect(0:p.N2-1)/p.N2
    t3 = collect(0:p.N3-1)/p.N3
    Ntot = p.N1*p.N2*p.N3

    if lowercase(method) == "parallel transport"
	p.logMethod = false
    elseif lowercase(method) == "log interpolation"
	p.logMethod = true
    else
	println("method of interpolation '$method' not recognized")
	return
    end

    A0_init = deepcopy(p.A)
    M0 = deepcopy(p.M)
    A = p.A
    M = p.M
    neighbors = p.neighbors

    fill!(A,NaN) #protection: A must be filled by the algorithm
    println("Filling (k,0,0)")
    A[:,:,:,1,1] = propagate(Matrix(1.0I,p.nwannier,p.nwannier), [p.ijk_to_K[i,1,1] for i=1:p.N1], p)

    # compute obstruction matrix
    Obs = normalize_matrix(overlap_A([p.N1,1,1],[1,1,1],p))
    println("Obstruction matrix = ")
    println(Obs)
    dV = eigen(Obs)
    d = dV.values
    V = dV.vectors
    logd = log.(d)
    # Hack to avoid separating eigenvalues at -1. TODO understand that
    for i=1:p.nwannier
        if imag(logd[i]) < -pi+.01
            logd[i] = logd[i] + 2pi*im
        end
    end
    println("log(d) = $(logd)")

    # and pull it back
    for i=1:p.N1
	A[:,:,i,1,1] = A[:,:,i,1,1]*V*diagm(0=>exp.(t1[i]*logd))*V'
    end

    println("Filling (k1,k2,0)")
    for i=1:p.N1
	A[:,:,i,:,1] = propagate(A[:,:,i,1,1], [p.ijk_to_K[i,j,1] for j=1:p.N2],p)
    end
    A0 = deepcopy(A)

    # corner obstruction
    Obs = normalize_matrix(overlap_A([1,p.N2,1],[1,1,1],p))
    dV = eigen(Obs)
    d = dV.values
    V = dV.vectors
    logd = log.(d)
    # Hack to avoid separating eigenvalues at -1. TODO understand that
    for i=1:p.nwannier
        if imag(logd[i]) < -pi+.01
            logd[i] = logd[i] + 2pi*im
        end
    end
    # pull it back
    for i=1:p.N1
        for j=1:p.N2
            A[:,:,i,j,1] = A[:,:,i,j,1]*V*diagm(0=>exp.(t2[j]*logd))*V'
        end
    end

    # Pull back the line obstruction
    phases = zeros(p.N1,p.nwannier)
    Obs_array = zeros(ComplexF64,p.N1,p.nwannier,p.nwannier)
    detObs = zeros(ComplexF64,p.N1)
    eigs = zeros(ComplexF64,p.N1,p.nwannier)
    for i =1:p.N1
        Obs_array[i,:,:] =  normalize_matrix(overlap_A([i,p.N2,1],[i,1,1],p))
    	detObs[i] = det(Obs_array[i,:,:])
    end

    # Find a continuous log of the determinant
    logDet = imag(log.(detObs))
    for i=2:p.N1
    	kmin = argmin([abs(logDet[i]+2*pi*k-logDet[i-1]) for k in -1:1])
    	logDet[i] = logDet[i]+(kmin-2)*2*pi
    end
    for i=1:p.N1
	Obs_array[i,:,:] = exp(-im*logDet[i]/p.nwannier) * Obs_array[i,:,:]
	eigs[i,:] = eigvals(Obs_array[i,:,:])
    end

    # Interpolate the line obstruction
    Uint = zeros(ComplexF64,p.N1,p.N2,p.nwannier,p.nwannier)

    if !p.logMethod
	Uint = matrixTransport(Obs_array,t2)
    end

    for i=1:p.N1
	Obs = Obs_array[i,:,:]
    	dV = eigen(Obs)
    	d = dV.values
    	V = dV.vectors
	if p.logMethod #& (!p.map)
	    for j=1:p.N2
		Uint[i,j,:,:] = powm(Obs,t2[j])
	    end
	end

        for j=1:p.N2
	    A[:,:,i,j,1] = A[:,:,i,j,1] * exp(im*logDet[i]*t2[j]/2)
	    A[:,:,i,j,1] = A[:,:,i,j,1] * Uint[i,j,:,:]
        end
    end

    # Propagate along the third dimension
    println("Filling (k1,k2,k3)")
    for i=1:p.N1,j=1:p.N2
    	A[:,:,i,j,:] = propagate(A[:,:,i,j,1], [p.ijk_to_K[i,j,k] for k=1:p.N3],p)
    end

    # Fix corner
    Obs = normalize_matrix(overlap_A([1,1,p.N3],[1,1,1],p))
    dV = eigen(Obs)
    d = dV.values
    V = dV.vectors
    logd = imag(log.(d))
    #println("Obstruction matrix  - Id = $(norm(Obs-I)))")

    # Hack to avoid separating eigenvalues at -1. TODO understand that
    for i =1:p.nwannier
        if imag(logd[i]) < -pi+.01
            logd[i] = logd[i] + 2pi*im
        end
    end

    for k=1:p.N3
        # fixer = powm(Obs,t3[k])
        fixer = V*diagm(0=>exp.(t3[k]*logd))*V'
        for i=1:p.N1,j=1:p.N2
            A[:,:,i,j,k] = A[:,:,i,j,k]*fixer
        end
    end

    # Fix first edge
    Obs_array_i = zeros(ComplexF64,p.N1,p.nwannier,p.nwannier)
    Uint_ik = zeros(ComplexF64,p.N1,p.N1,p.nwannier,p.nwannier)
    for i=1:p.N1
	Obs_array_i[i,:,:] = normalize_matrix(overlap_A([i,1,p.N3], [i,1,1],p))
    end
    if !p.logMethod
    	Uint_ik = matrixTransport(Obs_array_i,t3)
    end
    for i=1:p.N1
        for k=1:p.N3
	    if p.logMethod
	    	fixer = powm(Obs_array_i[i,:,:], t3[k])
            	for j=1:p.N3
                    A[:,:,i,j,k] = A[:,:,i,j,k]*fixer
		end
	    else
		for j=1:p.N3
		    A[:,:,i,j,k] = A[:,:,i,j,k]*Uint_ik[i,k,:,:]
		end
            end
        end
    end

    # Fix second edge
    Obs_array_j = zeros(ComplexF64,p.N2,p.nwannier,p.nwannier)
    Uint_jk = zeros(ComplexF64,p.N1,p.N1,p.nwannier,p.nwannier)
    for j=1:p.N2
	Obs_array_j[j,:,:] = normalize_matrix(overlap_A([1,j,p.N3], [1,j,1], p))
    end
    if !p.logMethod
    	Uint_jk = matrixTransport(Obs_array_j,t3)
    end
    for j=1:p.N2
        for k=1:p.N3
	    if p.logMethod
	    	fixer = powm(Obs_array_j[j,:,:], t3[k])
            	for i=1:p.N1
                    A[:,:,i,j,k] = A[:,:,i,j,k]*fixer
		end
	    elseif !p.logMethod
		for i=1:p.N1
		    A[:,:,i,j,k] = A[:,:,i,j,k]*Uint_jk[j,k,:,:]
		end
            end
        end
    end

    # Fix whole surface
    for i=1:p.N1,j=1:p.N2
	Obs = normalize_matrix(overlap_A([i,j,p.N3],[i,j,1],p))
        for k=1:p.N3
            A[:,:,i,j,k] = A[:,:,i,j,k]*powm(Obs,t3[k])
        end
    end

    p.A = A
    interp = InterpResults(Obs_array,Obs_array_i,Obs_array_j,Uint,Uint_ik)

    return(p,interp)
end

function plot_results(p,interp)
    #Interpolation method
    if(p.logMethod)
	suffix = "log"
    else
	suffix = "parallel_transport"
    end

    #Initialize arrays
    omegaright = zeros(p.N1,p.N2)
    omegaup = zeros(p.N1,p.N2)
    Aright1 = zeros(ComplexF64,p.N1,p.N2)
    Aright2 = zeros(ComplexF64,p.N1,p.N2)
    detA = zeros(ComplexF64,p.N1,p.N2)

    #Fill arrays
    imax = 1
    jmax = 1
    omegamax = 0.
    for i=1:p.N1,j=1:p.N2
        right = i==p.N1 ? 1 : i+1
        up = j==p.N2 ? 1 : j+1
	Krightj = p.ijk_to_K[right,j,1]
	Kij = p.ijk_to_K[i,j,1]
	Kiup = p.ijk_to_K[i,up,1]
	omegaright[i,j] = (p.N1)*norm(p.A[:,:,right,j,1]-overlap(Krightj,Kij,p)*p.A[:,:,i,j,1])
	omegaup[i,j] = (p.N2)*norm(p.A[:,:,i,up,1]-overlap(Kiup,Kij,p)*p.A[:,:,i,j,1])
        if (omegaright[i,j] + omegaup[i,j] > omegamax)
       	    omegamax = omegaright[i,j] + omegaup[i,j]
       	    imax = i
       	    jmax = j
        end

    end
    println("omegamax = $omegamax, [i,j] = $([imax,jmax])")
    if p.N1>10
    	nb = 10
    else
	nb = 1
    end
    fig_size = (4,3)

    figure(figsize=fig_size)
    ax = PyPlot.axes()
    matshow(omegaright,false,cmap=ColorMap("binary"),origin="lower")
    xticks(0:(div(p.N1,nb)):p.N1,0:(1.0/nb):1)
    yticks(0:(div(p.N1,nb)):p.N1,0:(1.0/nb):1)
    ax[:xaxis][:set_ticks_position]("bottom")
    colorbar()
    xlabel(L"$k_1$")
    ylabel(L"$k_2$")
    #cur_axes = gca()
    #cur_axes[:axis]("off")
    savefig("omega_right_$(p.filename)_$suffix.png")
    title("Regularity of Bloch frame (finite difference with right neighbor)")

    figure(figsize=fig_size)
    ax = PyPlot.axes()
    matshow(omegaup,false,cmap=ColorMap("binary"),origin="lower")
    xticks(0:(div(p.N1,nb)):p.N1,0:(1.0/nb):1)
    yticks(0:(div(p.N1,nb)):p.N1,0:(1.0/nb):1)
    ax[:xaxis][:set_ticks_position]("bottom")
    colorbar()
    xlabel(L"$k_1$")
    ylabel(L"$k_2$")
    #cur_axes = gca()
    #cur_axes[:axis]("off")
    savefig("omega_up_$(p.filename)_$suffix.png")
    title("Regularity of Bloch frame (finite difference with up neighbor)")

    plot_surface_obstructions(p, "_1_none")
    plot_surface_obstructions(p, "_2_corners")
    plot_surface_obstructions(p, "_3_edges")
    plot_surface_obstructions(p, "_4_surface")

    if p.N3 == 1
    	plot_3d = true #false
    else
        plot_3d = false
    end
    if plot_3d

   	Uint = interp.Uint
   	Obs_array = interp.Obs_array

	figure()
	for i in [div(p.N2,3) div(p.N2,3)*2 p.N2] #5:5:p.N2 #[33 66 100]
	    plot3D(real(Uint[[1:p.N2;1],i,1,1]),imag(Uint[[1:p.N2;1],i,1,1]),real(Uint[[1:p.N2;1],i,2,1]))#,"-k")
	end
	xlabel(L"$\mathrm{Re}(\mathbf{U}_{11})$")
	ylabel(L"$\mathrm{Im}(\mathbf{U}_{11})$")
	zlabel(L"$\mathrm{Re}(\mathbf{U}_{21})$")
	title(L"Interpolation of the first column of $s\mapsto\mathbf{U}(s)$ to $\mathbf{e}_1$")

	figure()
	for i in [div(p.N2,3) div(p.N2,3)*2 p.N2] #5:5:p.N2 #[33 66 100]
	    plot3D(real(Uint[[1:p.N2;1],i,1,1]),imag(Uint[[1:p.N2;1],i,1,1]),imag(Uint[[1:p.N2;1],i,2,1]),"-k")
	end
	xlabel(L"$\mathrm{Re}(\mathbf{U}_{11})$")
	ylabel(L"$\mathrm{Im}(\mathbf{U}_{11})$")
	zlabel(L"$\mathrm{Im}(\mathbf{U}_{21})$")
	title(L"Interpolation of the first column of $s\mapsto\mathbf{U}(s)$ to $\mathbf{e}_1$")

	eigs = zeros(ComplexF64,p.N2,2)
	log_eig = zeros(p.N2,2)
	c = zeros(2)

	for i in 2:p.N2
	    eigs[i,:], w = eigen(Obs_array[i,:,:])
	    log_eig[i,:] = imag(log.(eigs[i,:]))
	    log_eig[i,:] += c[:]*2π
	    if log_eig[i-1,1]-log_eig[i,1] > 1.8π
		c[1]+=1
		if log_eig[i-1,1]-log_eig[i,1]>3.8π
		    c[1]+=1
		end
	    else
		c[1]=0
	    end
	    if log_eig[i,2]-log_eig[i-1,2] > 1.8π
		c[2]+=-1
		if log_eig[i,2]-log_eig[i-1,2]>3.8π
		    c[2]+=-1
		end
	    else
		c[2]=0
	    end

	    log_eig[i,:] += c[:]*2π
	end
	figure()
	plot(range(0,step=1,length=p.N2),log_eig[:,1],label=L"$\log(\lambda_1)$")
	plot(range(0,step=1,length=p.N2),log_eig[:,2],label=L"$\log(\lambda_2)$")
	xlabel(L"$k_1$")
	legend()
	title(L"Logarithm of the eigenvalues of $\mathbf{Obs}$")

	figure()
	plot3D(range(0,step=1,length=p.N2),cos.(log_eig[:,1]),sin.(log_eig[:,1]),label=L"$\lambda_1$")
	plot3D(range(0,step=1,length=p.N2),cos.(log_eig[:,2]),sin.(log_eig[:,2]),label=L"$\lambda_2$")
	legend()
	xlabel(L"$k_1$")
	ylabel(L"$\mathrm{Re}(\lambda_1)$")
	zlabel(L"$\mathrm{Im}(\lambda_2)$")
	title(L"Eigenvalues of $\mathbf{Obs}$")
    end

    plot_obs = false
    if(plot_obs)
	eig_obs = zeros(ComplexF64,p.N1,p.nwannier)
	eig_obs_i = zeros(ComplexF64,p.N1,p.nwannier)
	eig_obs_j = zeros(ComplexF64,p.N1,p.nwannier)

	for i=1:p.N1
	    eig_obs[i,:] = eigen(Obs_array[i,:,:])[1]
	    eig_obs_i[i,:] = eigen(Obs_array_i[i,:,:])[1]
	    eig_obs_j[i,:] = eigen(Obs_array_j[i,:,:])[1]
	end

	figure()
	plot(real(eig_obs))
	figure()
	plot(real(eig_obs_i))
	figure()
	plot(real(eig_obs_j))
    end
end

function print_error(p)
    err_before = 0.
    err_after  = 0.
    for i=1:p.N1, j=1:p.N2, k=1:p.N3
	err_before += norm(normalize_matrix(overlap(p.ijk_to_K[i,j,k],p.ijk_to_K[(i%p.N1+1),j,k],p)) - I)^2
	err_after  += norm(normalize_matrix(overlap_A([i,j,k],[(i%p.N1+1),j,k],p))- I)^2
	err_before += norm(normalize_matrix(overlap(p.ijk_to_K[i,j,k],p.ijk_to_K[i,(j%p.N2)+1,k],p)) - I)^2
	err_after  += norm(normalize_matrix(overlap_A([i,j,k],[i,(j%p.N2)+1,k],p)) - I)^2
    end
    err_before = p.N1*sqrt( 1/(p.N1*p.N2*p.N3) * err_before)
    err_after = p.N1*sqrt(1/(p.N1*p.N2*p.N3) * err_after)
    println("err before = $(err_before)")
    println("err after  = $(err_after)")
end

function plot_Bloch_frame_slice(p,A0,A0_init)
    if(p.logMethod & !p.map)
	suffix = "log"
    else
	suffix = "parallel_transport"
    end

    if p.N3 > 2
    	N3_list = [1, div(p.N3,2), p.N3]
    else
	N3_list = [1]
    end
    init_frame = deepcopy(A0)
    for i=1:p.N1,j=1:p.N2,k=1:p.N3
	init_frame[:,:,i,j,k] = A0_init[:,:,i,j,k]'*A0[:,:,i,j,k] #*p.A[:,:,i,j,k]
	#init_frame[:,:,i,j,k] = normalize_matrix(init_frame[:,:,i,j,k])
    end

    for slice in N3_list
	figure(figsize=(6.5,5))
	matshow(real(init_frame[1,1,:,:,slice]+init_frame[2,2,:,:,slice]),false)
	colorbar()
	cur_axes = gca()
	cur_axes[:axis]("off")
	#title("Bloch frame slice before algorithm (A0), N3 = $slice")
	savefig("Bloch_frame_$(p.filename)_sum_$(p.N1)_init.png")

	figure(figsize=(6.5,5))
	matshow(real(init_frame[1,1,:,:,slice]-init_frame[2,2,:,:,slice]),false)
	colorbar()
	cur_axes = gca()
	cur_axes[:axis]("off")
	#title("Bloch frame slice before algorithm (A0), N3 = $slice")
	savefig("Bloch_frame_$(p.filename)_diff_$(p.N1)_init.png")

    end

    smooth_frame = deepcopy(A0)
    for i=1:p.N1,j=1:p.N2,k=1:p.N3
    	smooth_frame[:,:,i,j,k] = p.A[:,:,i,j,k]'*A0[:,:,i,j,k] #*p.A[:,:,i,j,k]
	smooth_frame[:,:,i,j,k] = normalize_matrix(smooth_frame[:,:,i,j,k])
    end

    for slice in N3_list
	figure(figsize=(6.5,5))
	matshow(real(smooth_frame[1,1,:,:,slice]+smooth_frame[2,2,:,:,slice]),false)
	colorbar()
	cur_axes = gca()
	cur_axes[:axis]("off")
	#title("Bloch frame slice after algorithm, N3 = $slice")
	savefig("Bloch_frame_$(p.filename)_sum_$(p.N1)_$suffix.png")

	figure(figsize=(6.5,5))
	matshow(real(smooth_frame[1,1,:,:,slice]-smooth_frame[2,2,:,:,slice]),false)
        colorbar()
	cur_axes = gca()
	cur_axes[:axis]("off")
        #title("Bloch frame slice after algorithm, N3 = $slice")
	savefig("Bloch_frame_$(p.filename)_diff_$(p.N1)_$suffix.png")
    end

end

ext: jl

max_stars_repo_name: antoine-levitt/wannier

max_issues_repo_name: antoine-levitt/wannier

max_forks_repo_name: antoine-levitt/wannier


text:
# TODO: When search is implemented, they'll need their own rate limits.

module GitHub

import ..GitForge: endpoint, into, postprocessor

using ..GitForge
using ..GitForge:
    @json,
    AStr,
    DoSomething,
    Endpoint,
    Forge,
    JSON,
    OnRateLimit,
    RateLimiter,
    HEADERS,
    ORL_THROW,
    @not_implemented

using Dates
using HTTP
using JSON2

export GitHubAPI, NoToken, Token, JWT

const DEFAULT_URL = "https://api.github.com"
const JSON_OPTS = (
    dateformat=dateformat"y-m-d",
    read_datetimeformats=[dateformat"y-m-dTH:M:SZ"],
    write_datetimeformat=dateformat"y-m-dTH:M:SZ",
)

abstract type AbstractToken end

"""
    NoToken()

Represents no authentication.
Only public data will be available.
"""
struct NoToken <: AbstractToken end

"""
    Token(token::$AStr)

An [OAuth2 token](https://developer.github.com/v3/#authentication),
or a [personal access token](https://developer.github.com/v3/#authentication).
"""
struct Token <: AbstractToken
    token::String
end

"""
    JWT(token::$AStr)

A JWT signed by a [private key](https://developer.github.com/apps/building-github-apps/authenticating-with-github-apps/#authenticating-as-a-github-app) for GitHub Apps.
"""
struct JWT <: AbstractToken
    token::String
end

auth_headers(::NoToken) = []
auth_headers(t::Token) = ["Authorization" => "token $(t.token)"]
auth_headers(t::JWT) = ["Authorization" => "Bearer $(t.token)"]

"""
    GitHubAPI(;
        token::AbstractToken=NoToken(),
        url::$AStr="$DEFAULT_URL",
        has_rate_limits::Bool=true,
        on_rate_limit::OnRateLimit=ORL_THROW,
    )

Create a GitHub API client.

## Keywords
- `token::AbstractToken=NoToken()`: Authorization token (or lack thereof).
- `url::$AStr="$DEFAULT_URL"`: Base URL of the target GitHub instance.
- `has_rate_limits::Bool=true`: Whether or not the GitHub server has rate limits.
- `on_rate_limit::OnRateLimit=ORL_THROW`: Behaviour on exceeded rate limits.
"""
struct GitHubAPI <: Forge
    token::AbstractToken
    url::String
    hasrl::Bool
    orl::OnRateLimit
    rl_general::RateLimiter
    rl_search::RateLimiter

    function GitHubAPI(;
        token::AbstractToken=NoToken(),
        url::AStr=DEFAULT_URL,
        has_rate_limits::Bool=true,
        on_rate_limit::OnRateLimit=ORL_THROW,
    )
        return new(token, url, has_rate_limits, on_rate_limit, RateLimiter(), RateLimiter())
    end
end

GitForge.base_url(g::GitHubAPI) = g.url
GitForge.request_headers(g::GitHubAPI, ::Function) = [HEADERS; auth_headers(g.token)]
GitForge.postprocessor(::GitHubAPI, ::Function) = JSON()
GitForge.has_rate_limits(g::GitHubAPI, ::Function) = g.hasrl
GitForge.rate_limit_check(g::GitHubAPI, ::Function) =
    GitForge.rate_limit_check(g.rl_general)
GitForge.on_rate_limit(g::GitHubAPI, ::Function) = g.orl
GitForge.rate_limit_wait(g::GitHubAPI, ::Function) = GitForge.rate_limit_wait(g.rl_general)
GitForge.rate_limit_period(g::GitHubAPI, ::Function) =
    GitForge.rate_limit_period(g.rl_general)
GitForge.rate_limit_update!(g::GitHubAPI, ::Function, r::HTTP.Response) =
    GitForge.rate_limit_update!(g.rl_general, r)

include("users.jl")
include("repositories.jl")
include("pull_requests.jl")
include("organizations.jl")
include("commits.jl")
include("branches.jl")
include("tags.jl")
include("comments.jl")

ismemberorcollaborator(r::HTTP.Response) = !(r.status in [403, 404])

end


text:
#
# derivative()
#

f1(x::Real) = sin(x)
@test norm(derivative(f1, :scalar, :forward)(0.0) - cos(0.0)) < 10e-4
@test norm(derivative(f1, :scalar, :central)(0.0) - cos(0.0)) < 10e-4
@test norm(derivative(f1, :forward)(0.0) - cos(0.0)) < 10e-4
@test norm(derivative(f1, :central)(0.0) - cos(0.0)) < 10e-4
@test norm(derivative(f1)(0.0) - cos(0.0)) < 10e-4

f2(x::Vector) = sin(x[1])
@test norm(derivative(f2, :vector, :forward)([0.0]) .- cos(0.0)) < 10e-4
@test norm(derivative(f2, :vector, :central)([0.0]) .- cos(0.0)) < 10e-4

#
# adjoint overloading
#

f3(x::Real) = sin(x)
for x in range(0.0, stop=0.1, length=11) # seq()
    @test norm(derivative(f3,x) - cos(x)) < 10e-4
end

#
# gradient()
#

f4(x::Vector) = (100.0 - x[1])^2 + (50.0 - x[2])^2
@test norm(Calculus.gradient(f4, :forward)([100.0, 50.0]) - [0.0, 0.0]) < 10e-4
@test norm(Calculus.gradient(f4, :central)([100.0, 50.0]) - [0.0, 0.0]) < 10e-4
@test norm(Calculus.gradient(f4)([100.0, 50.0]) - [0.0, 0.0]) < 10e-4

#
# jacobian()
#

@test norm(Calculus.jacobian(identity, rand(3), :forward) - Matrix(1.0I, 3, 3)) < 10e-4
@test norm(Calculus.jacobian(identity, rand(3), :central) - Matrix(1.0I, 3, 3)) < 10e-4

#
# second_derivative()
#

@test norm(second_derivative(x -> x^2)(0.0) - 2.0) < 10e-4
@test norm(second_derivative(x -> x^2)(1.0) - 2.0) < 10e-4
@test norm(second_derivative(x -> x^2)(10.0) - 2.0) < 10e-4
@test norm(second_derivative(x -> x^2)(100.0) - 2.0) < 10e-4

#
# hessian()
#

f5(x) = sin(x[1]) + cos(x[2])
@test norm(Calculus.gradient(f5)([0.0, 0.0]) - [cos(0.0), -sin(0.0)]) < 10e-4
@test norm(hessian(f5)([0.0, 0.0]) - [-sin(0.0) 0.0; 0.0 -cos(0.0)]) < 10e-4


text:
using JuMP, Combinatorics, Gurobi

function printResults(o, a, b, edges, T, Tf)
    println("Objective value: ", o)

    for i in 1:size(b,1)
        println("bw flow ", i, " = ", getvalue(b[i]))
    end
    print("\n")
    println("-----------------------------------------")

    for i in 1:size(a,1)
        for j in 1:size(a,2)
            println("flow ",i, ", tunnel ", j, " = ", getvalue(a[i,j]))
            print("edges in use: ")
            for e in T[Tf[i][j]]
                print(edges[e])
            end
            println("\n")
        end
    end
end

function FFC(env, edges, capacity, flows, demand, k, T, Tf; minb=0)

    nedges = size(edges,1)
    nflows = size(flows,1)
    ntunnels = size(T,1)

    # CREATE ALL SCENARIOS
    scenarios = []
    for bits in collect(combinations(1:nedges,k))
        s = ones(nedges)
        for bit in bits
            s[bit] = 0
        end
        push!(scenarios, s)
    end

    #CREATE RESIDUAL TUNNELS BY SCENARIO BY FLOW (References Tf)
    Tsf = []
    for s in 1:size(scenarios,1)
        sft = []
        for f in 1:size(Tf,1)
            ft = []
            for t in 1:size(Tf[f],1)
                up = true
                if (length(T[Tf[f][t]]) == 0)
                    up = false
                end
                for e in T[Tf[f][t]]
                    if scenarios[s][e] == 0
                        up = false
                    end
                end
                if up
                    push!(ft, t)
                end
            end
            push!(sft,ft)
        end
        push!(Tsf,sft)
    end

    #CREATE TUNNEL EDGE MATRIX
    L = zeros(ntunnels, nedges)
    for t in 1:ntunnels
        for e in 1:nedges
            if in(e, T[t])
                L[t,e] = 1
            end
        end
    end

    model = Model(solver=GurobiSolver(env, OutputFlag=0))
    @variable(model, b[1:nflows] >= 0, basename="b")
    @variable(model, a[1:nflows,1:size(Tf[1],1)] >= 0, basename="a")
    @variable(model, u >= 0, basename="u")

    for f in 1:nflows
        @constraint(model, sum(a[f,t] for t in 1:size(Tf[f],1)) >= b[f])   #the sum of all allocated bandwidths on every flow must be >= the total bandwidth for that flow
    end

    for e in 1:nedges
        @constraint(model, sum(a[f,t] * L[Tf[f][t],e] for f in 1:nflows for t in 1:size(Tf[f],1)) <= capacity[e])   #overlapping flows cannot add up to the capacity of that link
    end

    for f in 1:nflows
        for s in 1:size(scenarios,1)
            @constraint(model, sum(a[f,t] for t in Tsf[s][f]) >= b[f])   #residual tunnels must be able to carry bandwidth
        end
    end

    for f in 1:nflows
        @constraint(model, b[f] >= minb)
    end

    for f in 1:nflows
        @constraint(model, b[f] <= demand[f])   #all allocated bandwidths must be less than the demand for that flow
        for t in 1:size(Tf[f],1)
            @constraint(model, a[f,t] >= 0)     #each allocated bandwidth on for flow f on tunnel t >= 0
        end
    end

    for f in 1:nflows
        # @constraint(model, u <= b[f]/demand[f])
    end
  #   @objective(model, Max, u)

    @objective(model, Max, sum((b[i] for i in 1:size(b,1))))
    solve(model)
    # printResults(getobjectivevalue(model), a, b, edges, T, Tf)
    return getvalue(a), getvalue(b)
end


text:
module DemoPackageSSD

    print_greeting(io::IO=stdout) = printstyled(io, "Hello, world!"; color=:magenta)

end


text:
import Base: ==, <=, !=, length, eltype, iterate

abstract type LinkedList{T} end

mutable struct Nil{T} <: LinkedList{T}
end

mutable struct Cons{T} <: LinkedList{T}
    head::T
    tail::LinkedList{T}
end

cons(h, t::LinkedList{T}) where {T} = Cons{T}(h, t)

head(x::Cons) = x.head
tail(x::Cons) = x.tail

nil(T) = Nil{T}()
nil() = Nil{Any}()

==(x::Nil, y::Nil) = true
==(x::Cons, y::Cons) = (x.head == y.head) && (x.tail == y.tail)

list() = nil()

function list(elts::T...) where T
    l = nil(T)
    for i=length(elts):-1:1
        l = cons(elts[i],l)
    end
    return l
end

#length(l::Nil{T}) where {T} = 0
#length(l::Cons{T}) where {T} = 1 + length(l.tail)
length(l::Nil) = 0
function length(l::Cons)
    n = 0
    for i in l
        n += 1
    end
    return n
end

iterate(l::LinkedList, ::Nil) = nothing
function iterate(l::LinkedList, state::Cons = l)
    state.head, state.tail
end


text:
# ps3000_get_streaming_last_values.jl

export ps3000_get_streaming_last_values

function ps3000_get_streaming_last_values(handle::Int16, callback::Ptr{Void})
  #=
  c call signature for callback:

  void my_get_overview_buffers (
    int16_t ** overviewBuffers,
    int16_t overflow,
    uint32_t triggeredAt,
    int16_t triggered,
    int16_t auto_stop,
    uint32_t nValues )
  =#
  ps_status = ccall((:ps3000_get_streaming_last_values, ps3000driver),
                  Int16,
                  (Int16, Ptr{Void}),
                  handle, callback)
  if ps_status == 0
    error("ps3000 Error: Parameter Out Of Range (possibly from ps3000_run_streaming_ns()")
  end
  return ps_status  # The actual number of data values returned per channel
end


text:
#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct ReducedOrderPLL <: FrequencyEstimator
        ω_lp::Float64
        kp_pll::Float64
        ki_pll::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
    end

Parameters of a Phase-Locked Loop (PLL) based on Purba, Dhople, Jafarpour, Bullo and Johnson.
"Reduced-order Structure-preserving Model for Parallel-connected Three-phase Grid-tied Inverters."
2017 IEEE 18th Workshop on Control and Modeling for Power Electronics (COMPEL): 1-7.

# Arguments
- `ω_lp::Float64`: PLL low-pass filter frequency (rad/sec), validation range: `(0, nothing)`
- `kp_pll::Float64`: PLL proportional gain, validation range: `(0, nothing)`
- `ki_pll::Float64`: PLL integral gain, validation range: `(0, nothing)`
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: The states of the ReducedOrderPLL model are:
	vq_pll: q-axis of the measured voltage in the PLL synchronous reference frame (SRF),
	ε_pll: Integrator state of the PI controller,
	θ_pll: Phase angle displacement in the PLL SRF
- `n_states::Int`: ReducedOrderPLL has 3 states
"""
mutable struct ReducedOrderPLL <: FrequencyEstimator
    "PLL low-pass filter frequency (rad/sec)"
    ω_lp::Float64
    "PLL proportional gain"
    kp_pll::Float64
    "PLL integral gain"
    ki_pll::Float64
    ext::Dict{String, Any}
    "The states of the ReducedOrderPLL model are:
	vq_pll: q-axis of the measured voltage in the PLL synchronous reference frame (SRF),
	ε_pll: Integrator state of the PI controller,
	θ_pll: Phase angle displacement in the PLL SRF"
    states::Vector{Symbol}
    "ReducedOrderPLL has 3 states"
    n_states::Int
end

function ReducedOrderPLL(ω_lp, kp_pll, ki_pll, ext=Dict{String, Any}(), )
    ReducedOrderPLL(ω_lp, kp_pll, ki_pll, ext, [:vq_pll, :ε_pll, :θ_pll], 3, )
end

function ReducedOrderPLL(; ω_lp, kp_pll, ki_pll, ext=Dict{String, Any}(), states=[:vq_pll, :ε_pll, :θ_pll], n_states=3, )
    ReducedOrderPLL(ω_lp, kp_pll, ki_pll, ext, states, n_states, )
end

# Constructor for demo purposes; non-functional.
function ReducedOrderPLL(::Nothing)
    ReducedOrderPLL(;
        ω_lp=0,
        kp_pll=0,
        ki_pll=0,
        ext=Dict{String, Any}(),
    )
end

"""Get [`ReducedOrderPLL`](@ref) `ω_lp`."""
get_ω_lp(value::ReducedOrderPLL) = value.ω_lp
"""Get [`ReducedOrderPLL`](@ref) `kp_pll`."""
get_kp_pll(value::ReducedOrderPLL) = value.kp_pll
"""Get [`ReducedOrderPLL`](@ref) `ki_pll`."""
get_ki_pll(value::ReducedOrderPLL) = value.ki_pll
"""Get [`ReducedOrderPLL`](@ref) `ext`."""
get_ext(value::ReducedOrderPLL) = value.ext
"""Get [`ReducedOrderPLL`](@ref) `states`."""
get_states(value::ReducedOrderPLL) = value.states
"""Get [`ReducedOrderPLL`](@ref) `n_states`."""
get_n_states(value::ReducedOrderPLL) = value.n_states

"""Set [`ReducedOrderPLL`](@ref) `ω_lp`."""
set_ω_lp!(value::ReducedOrderPLL, val) = value.ω_lp = val
"""Set [`ReducedOrderPLL`](@ref) `kp_pll`."""
set_kp_pll!(value::ReducedOrderPLL, val) = value.kp_pll = val
"""Set [`ReducedOrderPLL`](@ref) `ki_pll`."""
set_ki_pll!(value::ReducedOrderPLL, val) = value.ki_pll = val
"""Set [`ReducedOrderPLL`](@ref) `ext`."""
set_ext!(value::ReducedOrderPLL, val) = value.ext = val


text:
st = state(p=VariableSpec(),t=VariableSpec())
@test mass_volume(water,st(3.0, 300.0)) ≈ 0.100_215_168E-2
@test mass_volume(water,st(3.0u"MPa", 300.0u"K")) ≈ 0.100_215_168E-2u"m^3/kg"

@test mass_volume(water,st(80.0, 300.0)) ≈ 0.971_180_894E-3
@test mass_volume(water,st(80.0u"MPa", 300.0u"K")) ≈ 0.971_180_894E-3u"m^3/kg"

@test mass_volume(water,st(3.0, 500.0)) ≈ 0.120_241_800E-2
@test mass_volume(water,st(3.0u"MPa", 500.0u"K")) ≈ 0.120_241_800E-2u"m^3/kg"

@test mass_enthalpy(water,st(3.0, 300.0)) ≈ 0.115_331_273E3
@test mass_enthalpy(water,st(3.0u"MPa", 300.0u"K")) ≈ 0.115_331_273E3u"kJ/kg"

@test mass_enthalpy(water,st(80.0, 300.0)) ≈ 0.184_142_828E3
@test mass_enthalpy(water,st(80.0u"MPa", 300.0u"K")) ≈ 0.184_142_828E3u"kJ/kg"

@test mass_enthalpy(water,st(3.0, 500.0)) ≈ 0.975_542_239E3
@test mass_enthalpy(water,st(3.0u"MPa", 500.0u"K")) ≈ 0.975_542_239E3u"kJ/kg"

@test mass_internal_energy(water,st(3.0, 300.0)) ≈ 0.112_324_818E3
@test mass_internal_energy(water,st(3.0u"MPa", 300.0u"K")) ≈ 0.112_324_818E3u"kJ/kg"

@test mass_internal_energy(water,st(80.0, 300.0)) ≈ 0.106_448_356E3
@test mass_internal_energy(water,st(80.0u"MPa", 300.0u"K")) ≈ 0.106_448_356E3u"kJ/kg"

@test mass_internal_energy(water,st(3.0, 500.0)) ≈ 0.971_934_985E3
@test mass_internal_energy(water,st(3.0u"MPa", 500.0u"K")) ≈ 0.971_934_985E3u"kJ/kg"

@test mass_entropy(water,st(3.0, 300.0)) ≈ 0.392_294_792
@test mass_entropy(water,st(3.0u"MPa", 300.0u"K")) ≈ 0.392_294_792u"kJ/kg/K"

@test mass_entropy(water,st(80.0, 300.0)) ≈ 0.368_563_852
@test mass_entropy(water,st(80.0u"MPa", 300.0u"K")) ≈ 0.368_563_852u"kJ/kg/K"

@test mass_entropy(water,st(3.0, 500.0)) ≈ 0.258_041_912E1
@test mass_entropy(water,st(3.0u"MPa", 500.0u"K")) ≈ 0.258_041_912E1u"kJ/kg/K"

@test mass_cp(water,st(3.0, 300.0)) ≈ 0.417_301_218E1
@test mass_cp(water,st(3.0u"MPa", 300.0u"K")) ≈ 0.417_301_218E1u"kJ/kg/K"

@test mass_cp(water,st(80.0, 300.0)) ≈ 0.401_008_987E1
@test mass_cp(water,st(80.0u"MPa", 300.0u"K")) ≈ 0.401_008_987E1u"kJ/kg/K"

@test mass_cp(water,st(3.0, 500.0)) ≈ 0.465_580_682E1
@test mass_cp(water,st(3.0u"MPa", 500.0u"K")) ≈ 0.465_580_682E1u"kJ/kg/K"

@test sound_speed(water,st(3.0, 300.0)) ≈ 0.150_773_921E4
@test sound_speed(water,st(3.0u"MPa", 300.0u"K")) ≈ 0.150_773_921E4u"m/s"

@test sound_speed(water,st(80.0, 300.0)) ≈ 0.163_469_054E4
@test sound_speed(water,st(80.0u"MPa", 300.0u"K")) ≈ 0.163_469_054E4u"m/s"

@test sound_speed(water,st(3.0, 500.0)) ≈ 0.124_071_337E4
@test sound_speed(water,st(3.0u"MPa", 500.0u"K")) ≈ 0.124_071_337E4u"m/s"


text:
# Goals: minimize calls to ordered_states (allocates memory)

# needs pomdp for stateindex in pdf(b, s)
# needs list of ordered_states for rand(b)

"""
    DiscreteBelief

A belief specified by a probability vector.

Normalization of `b` is NOT enforced at all times, but the `DiscreteBeleif(pomdp, b)` constructor will warn, and `update(...)` always returns a belief with normalized `b`.

Constructor:

`DiscreteBelief(pomdp, b::Vector{Float64}; check::Bool=true)`

# Fields
    - `pomdp` : the POMDP problem
    - `state_list` : a vector of ordered states
    - `b` : the probability vector
"""
struct DiscreteBelief{P<:POMDP, S}
    pomdp::P
    state_list::Vector{S}       # vector of ordered states
    b::Vector{Float64}
end

function DiscreteBelief(pomdp, b::Vector{Float64}; check::Bool=true)
    if check
        if !isapprox(sum(b), 1.0, atol=0.001)
            @warn """
                  b in DiscreteBelief(pomdp, b) does not sum to 1.

                  To suppress this warning use `DiscreteBelief(pomdp, b, check=false)`
                  """
        end
        if !all(0.0 <= p <= 1.0 for p in b)
            @warn """
                  b in DiscreteBelief(pomdp, b) contains entries outside [0,1].

                  To suppress this warning use `DiscreteBelief(pomdp, b, check=false)`
                  """
        end
    end
    return DiscreteBelief(pomdp, ordered_states(pomdp), b)
end

"""
     uniform_belief(pomdp)

Return a DiscreteBelief with equal probability for each state.
"""
function uniform_belief(pomdp)
    state_list = ordered_states(pomdp)
    ns = length(state_list)
    return DiscreteBelief(pomdp, state_list, ones(ns) / ns)
end

pdf(b::DiscreteBelief, s) = b.b[stateindex(b.pomdp, s)]

function Random.rand(rng::Random.AbstractRNG, b::DiscreteBelief)
    i = sample(rng, Weights(b.b))
    return b.state_list[i]
end

function Base.fill!(b::DiscreteBelief, x::Float64)
    fill!(b.b, x)
    return b
end

Base.length(b::DiscreteBelief) = length(b.b)

support(b::DiscreteBelief) = b.state_list

==(b1::DiscreteBelief, b2::DiscreteBelief) = b1.state_list == b2.state_list && b1.b == b2.b

Base.hash(b::DiscreteBelief, h::UInt) = hash(b.b, hash(b.state_list, h))

"""
    DiscreteUpdater

An updater type to update discrete belief using the discrete Bayesian filter.

# Fields
    - pomdp <: POMDP

"""
mutable struct DiscreteUpdater{P<:POMDP} <: Updater
    pomdp::P
end

uniform_belief(up::DiscreteUpdater) = uniform_belief(up.pomdp)

function initialize_belief(bu::DiscreteUpdater, dist::Any)
    state_list = ordered_states(bu.pomdp)
    ns = length(state_list)
    b = zeros(ns)
    belief = DiscreteBelief(bu.pomdp, state_list, b)
    for s in support(dist)
        sidx = stateindex(bu.pomdp, s)
        belief.b[sidx] = pdf(dist, s)
    end
    return belief
end

function update(bu::DiscreteUpdater, b::DiscreteBelief, a, o)
    pomdp = b.pomdp
    state_space = b.state_list
    bp = zeros(length(state_space))

    bp_sum = 0.0   # to normalize the distribution

    for (spi, sp) in enumerate(state_space)

        # po = O(a, sp, o)
        od = observation(pomdp, a, sp)
        po = pdf(od, o)

        if po == 0.0
            continue
        end

        b_sum = 0.0
        for (si, s) in enumerate(state_space)
            td = transition(pomdp, s, a)
            pp = pdf(td, sp)
            b_sum += pp * b.b[si]
        end

        bp[spi] = po * b_sum
        bp_sum += bp[spi]
    end

    if bp_sum == 0.0
        error("""
              Failed discrete belief update: new probabilities sum to zero.

              b = $b
              a = $a
              o = $o

              Failed discrete belief update: new probabilities sum to zero.
              """)
    else
        for i = 1:length(bp); bp[i] /= bp_sum; end
    end

    return DiscreteBelief(pomdp, b.state_list, bp)
end

update(bu::DiscreteUpdater, b::Any, a, o) = update(bu, initialize_belief(bu, b), a, o)


text:
using FlightMechanicsUtils
using FlightMechanicsSimulator
using Test

@testset "LHDownGravity" begin
    # Check constructor with default value
    grav_1 = LHDownGravity()
    # Check constructor with custom value
    grav_2 = LHDownGravity(F16Stevens.GD * FT2M)

    # Check return type
    @test isa(grav_1, LHDownGravity)
    # Check inheritance
    @test isa(grav_1, FlightMechanicsSimulator.Gravity)

    # Check getters
    @test isapprox(get_gravity_accel(grav_1), 9.80665)
    @test isapprox(get_gravity_accel(grav_2), 9.805416)

    @test isapprox(get_gravity_horizon(grav_1), [0, 0, 9.80665])
    @test isapprox(get_gravity_horizon(grav_2), [0, 0, 9.805416])
end


text:
function XC_x_slater( Rhoe::Float64 )
    third = 1.0/3.0
    pi34 = 0.6203504908994  # pi34=(3/4pi)^(1/3)
    rs = pi34/Rhoe^third
    f = -0.687247939924714 # -9/8*(3/2pi)^(2/3)
    alpha = 2.0/3.0
    # f*alpha = -3/4*(3/2pi)^(2/3)
    ex = f * alpha / rs
    vx = 4.0 / 3.0 * f * alpha / rs
    return ex, vx
end


text:
module ConvergenceTestWorkflow

using Express: Scf
using QuantumESPRESSO.Outputs.PWscf: parse_electrons_energies
using Unitful: @u_str
using UnitfulAtomic

import Express.ConvergenceTestWorkflow: parseoutput

include("Config.jl")
include("actions.jl")

function parseoutput(file)
    str = read(file, String)
    return parse_electrons_energies(str, :converged)[end, :ε]
end

end


text:
using Logging

"""
jlfile = @__FILE__
"""
module Experiments
using Distributed
using ExSup
using JLD2
using LarSurf
# using ViewerGL
function report_init_row(jlfile)
    data = Dict()
    data["jlfile"] = jlfile
    data["nprocs"] = nprocs()
    data["nworkers"] = nworkers()
    data["hostname"] = gethostname()
    data["ncores"] = length(Sys.cpu_info())
	ExSup.datetime_to_dict!(data)
    return data
end

function report_add_data_info(data, segmentation, voxelsize_mm)

	data["data size 1"] = size(segmentation, 1)
	data["data size 2"] = size(segmentation, 2)
	data["data size 3"] = size(segmentation, 3)
	nvoxels = sum(segmentation)
	@info "nvoxels=$(nvoxels)"
	data["nvoxels"] = nvoxels
	data["voxelsize_mm 1"] = voxelsize_mm[1]
	data["voxelsize_mm 2"] = voxelsize_mm[2]
	data["voxelsize_mm 3"] = voxelsize_mm[3]
	return data
end

"""
Get experiment and record all time and machine data to CSV file.

include_core_time_data parameter allow to get time data from every core but
it causes some problems with saving the data into CSV file
"""
function experiment_get_surface(
	data3d_full, voxelsize_mm;
	output_path=".",
	threshold=1,
	mask_label="data",
	stepxy=1, stepz=1, do_crop=false, cropx=1, cropy=1, cropz=1,
	block_size_scalar=64, data=nothing, time_start=nothing,
	output_csv_file = "exp_surface_extraction_times.csv",
	# show=False
	# include_core_time_data=false
	include_core_time_data=true
	)
	if time_start == nothing
		time_start = time()
	end
	if data == nothing
		data = Dict()
	end
	# fn = output_csv_file
	block_size = [block_size_scalar, block_size_scalar, block_size_scalar]
	data["block size"] = block_size[1]
	@info "raw data size=$(size(data3d_full))"
	data3d_full = data3d_full[1:stepz:end, 1:stepxy:end, 1:stepxy:end]
	if do_crop
		data3d_full = data3d_full[1:cropx, 1:cropy, 1:cropz]
	end
	    # round(size(data3d_full, 1) / target_size1)
	    # return data3d_full
	# segmentation = convert(Array{Int8, 3}, data3d_full .> threshold)
	segmentation = Int8.(data3d_full .> threshold)

	sz = [
		size(data3d_full, 1),
	 	size(data3d_full, 2),
	 	size(data3d_full, 3)
		]
	voxelsize_mm[1] = voxelsize_mm[1] * stepz
	voxelsize_mm[2] = voxelsize_mm[2] * stepxy
	voxelsize_mm[3] = voxelsize_mm[3] * stepxy
	@info "voxelsize mm = $(voxelsize_mm), size = $(sz)"
	data = report_add_data_info(data, segmentation, voxelsize_mm)
	# segmentation = LarSurf.generate_cube(data_size1; remove_one_pixel=true)
	@info "==== using done, data generated time from start: $(time() - time_start) [s]"
	data["data generated"] = time() - time_start

	@info "Setup..."
	# setup_time = @elapsed LarSurf.lsp_setup(block_size;reference_time=time_start)
	setup_time = @elapsed LarSurf.lsp_setup(block_size;reference_time=time_start)
	@info "setup time: $setup_time"
	@info "==== setup done, time from start: $(time() - time_start) [s]"
	data["setup done"] = time() - setup_time
	# for wid in workers()
	#     # println("testing on $wid")
	#     ftr = @spawnat wid LarSurf._single_boundary3
	#     @test fetch(ftr) != nothing
	# end

	# @debug "Setup done"
	tmd = @timed larmodel = LarSurf.lsp_get_surface(segmentation; voxelsize=voxelsize_mm)
	val, tm, mem, gc = tmd
	data["surface extraction time [s]"] = tm
	@info "Total time per surface extraction $(mask_label): $tm"
	@info "==== time from start: $(time() - time_start) [s]"
	data["finished"] = time() - time_start
	ExSup.datetime_to_dict!(data)
	@info "csv filename" output_csv_file
	# data["smoothing time [s]"] = 0.0
	data["smoothing time [s]"] = missing
	data["operation"] = "surface extraction"

	if include_core_time_data
		# the time data are updated with individual cores statistics by this command
    	LarSurf.get_time_data()
	end

	V, FV = larmodel

	# println("sz $(size(V)), $(size(FV))")
	data["vertex #"] = size(V, 2)
	data["quads #"] = size(FV, 1)
	FVtri = LarSurf.triangulate_quads(FV)

	data["triangles #"] = size(FVtri, 1)

	# println("sz FVtri $(size(FVtri))")

	ExSup.add_to_csv(data, output_csv_file)
	@info "csv export done"

	@JLD2.save "$(mask_label)_V_FV.jld2" V FV
	# if show
	# 	@info "Starting ViewerGL..."
	# 	ViewerGL.VIEW([
	# 	    ViewerGL.GLGrid(V, FVtri, ViewerGL.Point4d(1,1,1,0.1))
	# 		ViewerGL.GLAxis(ViewerGL.Point3d(-1,-1,-1),ViewerGL.Point3d(1,1,1))
	# 	])
	# end
	# FV = FVtri
	@JLD2.save "$output_path/$(mask_label)_V_FVtri.jld2" V FVtri
	objlines = LarSurf.Lar.lar2obj(V, FVtri, "$output_path/$(mask_label)_tri.obj")
	return V, FV, FVtri
end

function experiment_make_smoothing(V, FV, FVtri;
	output_path=".",
	mask_label="data",
	taubin=true, taubin_lambda=0.33, taubin_mu=-0.34, taubin_n=5,
	data=nothing,
	output_csv_file = "exp_surface_extraction_times.csv",
	# show=false
	)
	@info "smoothing"
	if data == nothing
		data = Dict()
	end
	if taubin
		t = @elapsed Vs = LarSurf.Smoothing.smoothing_FV_taubin(
		V, FV, taubin_lambda, taubin_mu, taubin_n)
		@info "smoothing taubin time", t
		# t = @elapsed FVtri = LarSurf.triangulate_quads(FV)

	else
		# t = @elapsed Vs = LarSurf.Smoothing.smoothing_FV(V, FVtri, 0.6, 3)
		t = @elapsed Vs = LarSurf.Smoothing.smoothing_FV(V, FVtri, taubin_lambda, taubin_n)
		@info "smoothing time", t
	end
	data["operation"] = "smoothing"
	data["smoothing time [s]"] = t
	@info "Smoothing numer of Vs: $(size(Vs))"
	@JLD2.save "$output_path/$(mask_label)_Vs_FVtri.jld2" Vs FVtri
	# @JLD2.save "liver01tri.jld2" V FVtri
	t = @elapsed objlines = LarSurf.Lar.lar2obj(Vs, FVtri, "$output_path/$(mask_label)_tri_sm.obj")
	data["LAR to .obj time [s]"] = t
	ExSup.add_to_csv(data, output_csv_file)
	# ExSup.add_to_csv(data, output_csv_file)
	# if show
	# 	ViewerGL.VIEW([
	# 	    ViewerGL.GLGrid(Vs, FVtri, ViewerGL.Point4d(1, 0, 1, 0.1))
	# 		ViewerGL.GLAxis(ViewerGL.Point3d(-1, -1, -1),ViewerGL.Point3d(1, 1, 1))
	# 	])
	# end
	return Vs
end
# end
# Plasm.view(val)
"""
Read data from file, make surface extraction and smoothing.
Surface models are stored into .obj files and all statistics are stored into
CSV file.

:param pth: filename with 3D data(dcm, tif, pklz, ...) or
.jld file generated with V and FV or
.jld file with datap
generated in previous run
"""
function experiment_make_surf_extraction_and_smoothing(
	data3d_full, voxelsize_mm;
	output_path=".",
	threshold=1, mask_label="data",
	stepxy=1, stepz=1, do_crop=false, cropx=1, cropy=1, cropz=1,
	block_size_scalar=64, data=nothing, time_start=nothing,
	output_csv_file = "exp_surface_extraction_times.csv",
	taubin::Bool=true, taubin_lambda=0.33, taubin_mu=-0.34, taubin_n=5,
	smoothing::Bool=true
	)
	# @info "pth"
	# println(pth)
	if data == nothing
		data = Dict()
	end
	datap_readed = true
		# FVtri = LarSurf.triangulate_quads(FV)
	# @load "ircad_$(mask_label).jld2" V FV
	V, FV, FVtri = experiment_get_surface(
		data3d_full, voxelsize_mm;
		output_path=output_path,
		threshold=threshold,
		mask_label=mask_label,
		stepxy=stepxy, stepz=stepz, do_crop=do_crop, cropx=cropx, cropy=cropy, cropz=cropz,
		block_size_scalar=block_size_scalar, data=data, time_start=time_start,
		output_csv_file = output_csv_file,
		# show=show
	)
	if smoothing
		Vs = experiment_make_smoothing(V, FV, FVtri;
				output_path=output_path,
		mask_label=mask_label,
		taubin=taubin,
		taubin_lambda = taubin_lambda,
		taubin_mu = taubin_mu,
		taubin_n = taubin_n,
		data=data,
		output_csv_file=output_csv_file,
		# show=show
		)
		return Vs, FVtri
	else
		return V, FVtri
	end
end

end


text:
using SafeTestsets

@safetestset "Day06" begin
    using AdventOfCode.TestUtils
    using AdventOfCode.Day06

    @testset "orbits_count" begin
        using AdventOfCode.Day06: Tree, count!

        t = Tree([("COM", "B"),
                  ("B", "C"),
                  ("C", "D"),
                  ("D", "E"),
                  ("E", "F"),
                  ("B", "G"),
                  ("G", "H"),
                  ("D", "I"),
                  ("E", "J"),
                  ("J", "K"),
                  ("K", "L")])

        count!(t)
        @test t["COM"].orbits_count == 0
        @test t["D"].orbits_count == 3
        @test t["L"].orbits_count == 7
        @test t.orbits_count == 42
    end

    @testset "orbital_transfers" begin
        using AdventOfCode.Day06: Tree, orbital_transfers
        t = Tree([("COM", "B"),
                  ("B", "C"),
                  ("C", "D"),
                  ("D", "E"),
                  ("E", "F"),
                  ("B", "G"),
                  ("G", "H"),
                  ("D", "I"),
                  ("E", "J"),
                  ("J", "K"),
                  ("K", "L"),
                  ("K", "YOU"),
                  ("I", "SAN")])
        @test orbital_transfers(t, "SAN", "YOU") == 4
    end

    btest("6") do
         both_parts() == (140608, 337)
    end
end


text:
# I'm aware of Reexport.jl, but it doesn't work for QuantumControl.jl
#
# For one thing, `@reexport using QuantumControlBase` re-exports not just the
# members of `QuantumControlBase`, but also `QuantumControlBase` itself. Also,
# as far as I can tell `@reexport using QuantumControlBase.Shapes` does not
# work when it's inside a module also called `Shapes`, as we're doing.
#
# Besides, the macro below is pretty trivial

macro reexport_members(modname::Symbol)
    mod = getfield(__module__, modname)
    member_names = _exported_names(mod)
    Expr(:export, member_names...)
end

function _exported_names(m::Module)
    return filter!(
        x -> (Base.isexported(m, x) && (x != nameof(m))),
        names(m; all=true, imported=true)
    )
end


text:
using BinaryBuilder

# Collection of sources required to build HDF5
name = "HDF5"
version = v"1.12.1"

sources = [
    # 32-bit Windows from https://packages.msys2.org/package/mingw-w64-i686-hdf5
    ArchiveSource("https://mirror.msys2.org/mingw/mingw32/mingw-w64-i686-hdf5-1.12.1-2-any.pkg.tar.zst", "ab7e4568c99e2a7a9b96ad06867ef9e829d286539e0735a93b9295ad8778a818"; unpack_target="i686-w64-mingw32"),
    ArchiveSource("https://mirror.msys2.org/mingw/mingw32/mingw-w64-i686-libaec-1.0.6-1-any.pkg.tar.zst", "43ef1aee2be7cc192d14c6641627ed0a1eca50c21234eba3405fce77302517a8"; unpack_target="i686-w64-mingw32"),
    ArchiveSource("https://mirror.msys2.org/mingw/mingw32/mingw-w64-i686-zlib-1.2.11-9-any.pkg.tar.zst", "8d594fc14497d41a66415bfdd200d7d1d56a6ecd3fe81b9190bec0c4841a5eff"; unpack_target="i686-w64-mingw32"),
    # We need some special compiler support libraries from mingw for i686 (libgcc_s_dw2)
    ArchiveSource("https://mirror.msys2.org/mingw/mingw32/mingw-w64-i686-gcc-libs-11.2.0-6-any.pkg.tar.zst", "bdc359047f61c8e96401ba25b17c80f5f8039c25a063c622e3680123bb0de9d1"; unpack_target="i686-w64-mingw32"),

    # 64-bit Windows from https://packages.msys2.org/package/mingw-w64-x86_64-hdf5
    ArchiveSource("https://mirror.msys2.org/mingw/mingw64/mingw-w64-x86_64-hdf5-1.12.1-2-any.pkg.tar.zst", "540d8b9fd71d0d848e28723f531948dea497c7cb8f95c53e1a3b06f750f12c40"; unpack_target="x86_64-w64-mingw32"),
    ArchiveSource("https://mirror.msys2.org/mingw/mingw64/mingw-w64-x86_64-libaec-1.0.6-1-any.pkg.tar.zst", "839b271b1313f013227c13944876651429b8e0edb40760d64f817a36c8b9435c"; unpack_target="x86_64-w64-mingw32"),
    ArchiveSource("https://mirror.msys2.org/mingw/mingw64/mingw-w64-x86_64-zlib-1.2.11-9-any.pkg.tar.zst", "9da9ebafaef832dba2f442ad44d9ae8759784b86478dcbe326500195f8ea6339"; unpack_target="x86_64-w64-mingw32"),

    # x86_64 and aarch64 for Linux and macOS from https://anaconda.org/conda-forge/hdf5/files
    # NOTE: make sure to select those compatible with OpenSSL 1.1.1
    ArchiveSource("https://anaconda.org/conda-forge/hdf5/1.12.1/download/linux-64/hdf5-1.12.1-nompi_h2750804_103.tar.bz2", "bd7bb0657d63acf52c9d30d1b89276356c6da4ff8a90dd5fcbd0cfde6578f317"; unpack_target="x86_64-linux-gnu"),
    ArchiveSource("https://anaconda.org/conda-forge/hdf5/1.12.1/download/linux-aarch64/hdf5-1.12.1-nompi_h774d4d8_103.tar.bz2", "8688cfc983962bf7a59a97becb0d67ee64eb4a7dd5793b915cf50dccd90bfa2d"; unpack_target="aarch64-linux-gnu"),
    ArchiveSource("https://anaconda.org/conda-forge/hdf5/1.12.1/download/osx-64/hdf5-1.12.1-nompi_h2f0ef1a_102.tar.bz2", "4a4640e44adea33833e7efb6ac3070dd4c80a8a156c2fd3aa7cfcac8865f5a26"; unpack_target="x86_64-apple-darwin14"),
    ArchiveSource("https://anaconda.org/conda-forge/hdf5/1.12.1/download/osx-arm64/hdf5-1.12.1-nompi_had0e5e0_103.tar.bz2", "ef48b684b22c6b0077bc9836e0cc6d15abb88868d7a6c842226666ebb8bbd449"; unpack_target="aarch64-apple-darwin20"),
]

# Bash recipe for building across all platforms
script = raw"""
cd ${WORKSPACE}/srcdir/
mkdir -p ${libdir} ${includedir}

if [[ ${target} == *mingw* ]]; then
    cd ${target}/mingw${nbits}

    rm -f bin/{*_cpp*,*fortran*,*f90*} # we do not need these
    mv -v bin/*.dll ${libdir}
    mv -v include/* ${includedir}

    install_license share/doc/hdf5/COPYING
else
    cd ${target}

    rm -f lib/{*_cpp*,*_fortran*} # we do not need these
    mv -v lib/* ${libdir}
    mv -v include/* ${includedir}

    install_license info/licenses/COPYING
fi
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = [
    Platform("x86_64", "linux"),
    Platform("aarch64", "linux"; libc="glibc"),
    Platform("x86_64", "macos"),
    Platform("x86_64", "windows"),
    Platform("i686", "windows"),
    Platform("aarch64", "macos"),
]

# The products that we will ensure are always built
products = [
    LibraryProduct("libhdf5", :libhdf5),
    LibraryProduct("libhdf5_hl", :libhdf5_hl),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency("Zlib_jll"),
    Dependency("OpenSSL_jll"; compat="1.1.10"),
    Dependency("LibCURL_jll"),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat="1.6")

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil