text:
using CleanCMB
using Healpix, Libsharp
using PyCall
using Random, Statistics
using Printf
using Plots, LaTeXStrings
using Tables, CSV
# %% Create figures?
showresults = true

# %% Simulation parameters
nrz = 300 # How many realisations?
Random.seed!(5147) # Initial random number seed. Useful if you need reproducible sequence
rsim = 0 # Tensor-to-scalar ratio used for the simulation
Alens = 1 # Lensing power spectrum amplitude (Alens = 1 for the fiducial)
ℓmin, ℓmax = 30, 260 # ℓ range for fitting

# %% Specification of the experiments
# Reference: Simons Observatory Collaboration, JCAP, 02, 056 (2019), Table 1.
# Reference: CCAT-prime Collaboration, J.Low Temp.Phys., 199, 1089 (2020), Table 1.
ν = [27, 39, 93, 145, 225, 280, 350, 410, 850] # in GHz
nν, nνSO = length(ν), 6
FWHM = [91, 63, 30, 17, 11, 9, 0.58, 0.5, 0.23] # in arcmin
σ = FWHM * π / 10800 / √(8 * log(2)) # in radians
uKarcmin = [35, 21, 2.6, 3.3, 6.3, 16, 105, 372, 5.7e5] # in μK arcmin (for temperature; x√2 for pol)
lknee = [30, 30, 50, 50, 70, 100, 700, 700, 700]
αknee = [-2.4, -2.4, -2.5, -3, -3, -3, -1.4, -1.4, -1.4]
# temperature beam
bTl(ℓ, σb) = exp(-ℓ * (ℓ + 1) * σb^2 / 2)
# polarisation beam, Eq.(5.8) of Ng & Liu, Int.J.Mod.Phys.D, 8, 61 (1999)
bPl(ℓ, σb) = ifelse(ℓ >= 2, exp(-(ℓ * (ℓ + 1) - 4) * σb^2 / 2), 0)

# %% Read in the hits map and calculate weights for inhomogeneous noise
nhitsfile = "data/nhits_SAT_r7.FITS"
nhits = readMapFromFITS(nhitsfile, 1, Float64)
nside = nhits.resolution.nside
nhits /= maximum(nhits)
fsky = mean(nhits)^2 / mean(nhits .^ 2)
K = √(mean(nhits) / fsky)
weight = zeros(12 * nside^2)
for ip = 1:12*nside^2
    if nhits[ip] > 0
        weight[ip] = K / √nhits[ip]
    end
end

# %% Compute scalar and tensor power spectra using CLASS
lmax, mmax = 3 * nside - 1, 3 * nside - 1
include("compute_cl_class.jl")

# %% Basic parameters for spherical harmonics transform
geom_info = make_healpix_geom_info(nside, 1)
alm_info = make_triangular_alm_info(lmax, mmax, 1)
elm, blm = Alm{ComplexF64}(lmax, mmax), Alm{ComplexF64}(lmax, mmax)

# %% Read in and smooth foreground maps generated by "d1s1" model of Python Sky Model (PySM), in equatorial coordinates
# Reference: Thorne et al., MNRAS, 469, 2821 (2017), https://github.com/healpy/pysm
f_q, f_u = [], []
for iν = 1:nν
    filename = @sprintf("data/map_d1s1_equ_%03dghz_r7_uKcmb_qu.fits", ν[iν])
    q, u = (readMapFromFITS(filename, ic, Float64) for ic = 1:2)
    sharp_execute!(
        SHARP_MAP2ALM,
        2,
        [elm.alm, blm.alm],
        [q.pixels, u.pixels],
        geom_info,
        alm_info,
        SHARP_DP,
    )
    for l = 0:lmax
        for m = 0:l
            ilm = alm_index(alm_info, l, m) + 1
            elm.alm[ilm] *= bPl(l, σ[iν])
            blm.alm[ilm] *= bPl(l, σ[iν])
        end
    end
    sharp_execute!(
        SHARP_ALM2MAP,
        2,
        [elm.alm, blm.alm],
        [q.pixels, u.pixels],
        geom_info,
        alm_info,
        SHARP_DP,
    )
    push!(f_q, q)
    push!(f_u, u)
end

# %% Setup NaMaster for the power spectrum analysis on a partial sky
Δℓ = 10 # multipole binning size
maskfile = "data/mask_apodized_r7.fits"
include("setup_namaster.jl")
# The theory power spectrum must be binned into bandpowers in the same manner the data has.
# Generate an NmtWorkspace object that we use to compute and store the mode coupling matrix.
# Note that this matrix depends only on the masks of the two fields to correlate, but not on the maps themselves.
w = nmt.NmtWorkspace()
f = nmt.NmtField(mask, [f_q[1], f_u[1]], purify_b = true)
w.compute_coupling_matrix(f, f, b)
cls_th = [cls["ee"], zero(cls["ee"]), zero(cls["bb"]), cls["bb"]] * Tcmb^2
cls_th_binned = w.decouple_cell(w.couple_cell(cls_th))
clt_th = [clt["ee"], zero(clt["ee"]), zero(clt["bb"]), clt["bb"]] * Tcmb^2
clt_th_binned = w.decouple_cell(w.couple_cell(clt_th))

# %% Compute covariance matrices of the foreground EE and BB
# Used to compute the power spectra of residual foreground for checking the results,
# but this information cannot be used for the real data analysis.
f3 = []
for iν = 1:nν
    push!(
        f3,
        nmt.NmtField(
            mask,
            [f_q[iν], f_u[iν]],
            purify_b = true,
            beam = bPl.(0:lmax, σ[iν]),
        ),
    )
end
ce3, cb3 = zeros(nν, nν, nbands), zeros(nν, nν, nbands) # foreground
for iν = 1:nν, jν = iν:nν
    w.compute_coupling_matrix(f3[iν], f3[jν], b)
    cl3 = compute_master(f3[iν], f3[jν], w)
    ce3[iν, jν, :], cb3[iν, jν, :] = cl3[1, :], cl3[4, :]
end

# %% Loop over realisations
ee1, bb1 = zeros(nbands, nrz), zeros(nbands, nrz) # Cleaned power spectra
ee2, bb2 = zeros(nbands, nrz), zeros(nbands, nrz) # Noise power spectra
ee3, bb3 = zeros(nbands, nrz), zeros(nbands, nrz) # Residual foreground power spectra
ee1SO, bb1SO = zeros(nbands, nrz), zeros(nbands, nrz) # Cleaned power spectra
ee2SO, bb2SO = zeros(nbands, nrz), zeros(nbands, nrz) # Noise power spectra
ee3SO, bb3SO = zeros(nbands, nrz), zeros(nbands, nrz) # Residual foreground power spectra
for irz = 1:nrz
    @show irz
    # Create spherical harmonics coefficients of the CMB
    for l = 2:lmax
        ee = cls_th[1][l+1] + (rsim / rclass) * clt_th[1][l+1]
        bb = Alens * cls_th[4][l+1] + (rsim / rclass) * clt_th[4][l+1]
        ilm = alm_index(alm_info, l, 0) + 1
        elm.alm[ilm] = √ee * randn(Float64)
        blm.alm[ilm] = √bb * randn(Float64)
        for m = 1:l
            ilm = alm_index(alm_info, l, m) + 1
            elm.alm[ilm] = √ee * randn(ComplexF64)
            blm.alm[ilm] = √bb * randn(ComplexF64)
        end
    end
    # Smooth CMB and create spherical harmonics coefficients of the noise
    celm, cblm = Alm{ComplexF64}(lmax, mmax), Alm{ComplexF64}(lmax, mmax)
    nelm, nblm = Alm{ComplexF64}(lmax, mmax), Alm{ComplexF64}(lmax, mmax)
    c_q, c_u = HealpixMap{Float64,RingOrder}(nside), HealpixMap{Float64,RingOrder}(nside)
    n_q, n_u = HealpixMap{Float64,RingOrder}(nside), HealpixMap{Float64,RingOrder}(nside)
    f1, f2 = [], [] # List of NaMaster fields
    for iν = 1:nν
        for l = 0:lmax
            ee =
                2 *
                (uKarcmin[iν] * π / 10800)^2 *
                (1 + (l / lknee[iν])^αknee[iν])
            bb =
                2 *
                (uKarcmin[iν] * π / 10800)^2 *
                (1 + (l / lknee[iν])^αknee[iν])
            if l >= 2
                ilm = alm_index(alm_info, l, 0) + 1
                celm.alm[ilm] = elm.alm[ilm] * bPl(l, σ[iν])
                cblm.alm[ilm] = blm.alm[ilm] * bPl(l, σ[iν])
                nelm.alm[ilm] = √ee * randn(Float64)
                nblm.alm[ilm] = √bb * randn(Float64)
                for m = 1:l
                    ilm = alm_index(alm_info, l, m) + 1
                    celm.alm[ilm] = elm.alm[ilm] * bPl(l, σ[iν])
                    cblm.alm[ilm] = blm.alm[ilm] * bPl(l, σ[iν])
                    nelm.alm[ilm] = √ee * randn(ComplexF64)
                    nblm.alm[ilm] = √bb * randn(ComplexF64)
                end
            else
                for m = 0:l
                    ilm = alm_index(alm_info, l, m) + 1
                    celm.alm[ilm], cblm.alm[ilm] = 0, 0
                    nelm.alm[ilm], nblm.alm[ilm] = 0, 0
                end
            end
        end
        sharp_execute!(
            SHARP_ALM2MAP,
            2,
            [celm.alm, cblm.alm],
            [c_q.pixels, c_u.pixels],
            geom_info,
            alm_info,
            SHARP_DP,
        )
        sharp_execute!(
            SHARP_ALM2MAP,
            2,
            [nelm.alm, nblm.alm],
            [n_q.pixels, n_u.pixels],
            geom_info,
            alm_info,
            SHARP_DP,
        )
        # Inhomogenous noise for SO-SAT and homogeneous noise for CCAT-prime
        if iν <= nνSO
            n_q .*= weight
            n_u .*= weight
        end
        ## Create NaMaster fields for the pseudo-Cℓ
        # f1: foreground + CMB + noise
        push!(
            f1,
            nmt.NmtField(
                mask,
                [f_q[iν] + c_q + n_q, f_u[iν] + c_u + n_u],
                purify_b = true,
                beam = bPl.(0:lmax, σ[iν]),
            ),
        )
        # f2: noise
        push!(
            f2,
            nmt.NmtField(
                mask,
                [n_q, n_u],
                purify_b = true,
                beam = bPl.(0:lmax, σ[iν]),
            ),
        )
    end
    ## Compute covariance matrices of EE and BB
    ce1, cb1 = zeros(nν, nν, nbands), zeros(nν, nν, nbands) # total
    ce2, cb2 = zeros(nν, nν, nbands), zeros(nν, nν, nbands) # noise
    for iν = 1:nν, jν = iν:nν
        w.compute_coupling_matrix(f1[iν], f1[jν], b)
        cl1 = compute_master(f1[iν], f1[jν], w)
        cl2 = compute_master(f2[iν], f2[jν], w)
        ce1[iν, jν, :], cb1[iν, jν, :] = cl1[1, :], cl1[4, :]
        ce2[iν, jν, :], cb2[iν, jν, :] = cl2[1, :], cl2[4, :]
    end
    ## Perform ILC to obtain power spectra of clean maps of the CMB
    we, wb = ilc_weights(ce1), ilc_weights(cb1)
    ee1[:, irz], bb1[:, irz] = ilc_clean_cij(ce1, we), ilc_clean_cij(cb1, wb)
    ee2[:, irz], bb2[:, irz] = ilc_clean_cij(ce2, we), ilc_clean_cij(cb2, wb)
    ee3[:, irz], bb3[:, irz] = ilc_clean_cij(ce3, we), ilc_clean_cij(cb3, wb)
    # SO only
    we, wb =
        ilc_weights(ce1[1:nνSO, 1:nνSO, :]), ilc_weights(cb1[1:nνSO, 1:nνSO, :])
    ee1SO[:, irz], bb1SO[:, irz] = ilc_clean_cij(ce1[1:nνSO, 1:nνSO, :], we),
    ilc_clean_cij(cb1[1:nνSO, 1:nνSO, :], wb)
    ee2SO[:, irz], bb2SO[:, irz] = ilc_clean_cij(ce2[1:nνSO, 1:nνSO, :], we),
    ilc_clean_cij(cb2[1:nνSO, 1:nνSO, :], wb)
    ee3SO[:, irz], bb3SO[:, irz] = ilc_clean_cij(ce3[1:nνSO, 1:nνSO, :], we),
    ilc_clean_cij(cb3[1:nνSO, 1:nνSO, :], wb)
    # Show power spectra for visual inspection
    if showresults
        p = plot(
            ell_eff,
            cls_th_binned[1, :] .+ (rsim / rclass) * clt_th_binned[1, :],
            title = @sprintf("Realisation # = %03d", irz),
            label = "True EE",
            xaxis = :log,
            yaxis = :log,
            xlab = L"\ell",
            ylab = L"\ell(\ell+1)C_\ell/2\pi",
            ylims = [2e-5, 20],
            xlims = [ℓmin, ℓmax],
            legend = :topleft,
        )
        p = plot!(
            ell_eff,
            Alens * cls_th_binned[4, :] .+
            (rsim / rclass) * clt_th_binned[4, :],
            label = "True BB",
        )
        p = plot!(ell_eff, ee1[:, irz], m = :circle, lab = "Cleaned EE")
        p = plot!(ell_eff, bb1[:, irz], m = :circle, lab = "Cleaned BB")
        p = plot!(ell_eff, ee2[:, irz], m = :diamond, lab = "Noise EE")
        p = plot!(ell_eff, bb2[:, irz], m = :diamond, lab = "Noise BB")
        p = plot!(ell_eff, ee3[:, irz], m = :star5, lab = "FG EE")
        p = plot!(ell_eff, bb3[:, irz], m = :star5, lab = "FG BB")
        display(p)
    end
end

# %% Calculate the mean power spectra and variance
me1, mb1 = mean(ee1, dims = 2), mean(bb1, dims = 2)
ve1, vb1 = var(ee1, dims = 2), var(bb1, dims = 2)
me2, mb2 = mean(ee2, dims = 2), mean(bb2, dims = 2)
me3, mb3 = mean(ee3, dims = 2), mean(bb3, dims = 2)
me1SO, mb1SO = mean(ee1SO, dims = 2), mean(bb1SO, dims = 2)
ve1SO, vb1SO = var(ee1SO, dims = 2), var(bb1SO, dims = 2)
me2SO, mb2SO = mean(ee2SO, dims = 2), mean(bb2SO, dims = 2)
me3SO, mb3SO = mean(ee3SO, dims = 2), mean(bb3SO, dims = 2)
# Plot and save to "ilc_clbb_sim_sosatccatp.pdf"
ii = findall(x -> x >= ℓmin && x <= ℓmax, ell_eff)
if showresults
    p = scatter(
        ell_eff[ii],
        mb1[ii],
        m = 5,
        lab = "Cleaned power spectrum",
        xlab = L"\ell",
        ylab = L"\ell(\ell+1)C^{BB}_\ell/2\pi~[\mu K^2]",
        legend = :topleft,
    )
    p = scatter!(ell_eff[ii], mb1[ii] - mb2[ii], m = 5, lab = "- noisebias")
    p = scatter!(
        ell_eff[ii],
        mb1[ii] - mb2[ii] - mb3[ii],
        yerr = sqrt.(vb1[ii]),
        m = 5,
        lab = "- foreground",
    )
    p = plot!(ell_eff[ii], Alens * cls_th_binned[4, ii], lab = "Binned scalar")
    p = plot!(
        ell_eff[ii],
        (rsim / rclass) * clt_th_binned[4, ii],
        lab = "Binned tensor",
        ls = :dash,
    )
    savefig("ilc_clbb_sim_sosatccatp.pdf")
    display(p)
end

# %% Calculate the tensor-to-scalar ratio
r, rSO = zeros(nrz), zeros(nrz)
w, wSO = zeros(nrz, 2), zeros(nrz, 2)

# Joint fit for the tensor-to-scalar ratio and the foreground amplitude
y1 = clt_th_binned[4, ii] / rclass
# SO-SAT + CCAT-prime
y2 = mb3[ii]
v = vb1[ii]
Fij = [
    sum(y1 .* y1 ./ v) sum(y1 .* y2 ./ v)
    sum(y2 .* y1 ./ v) sum(y2 .* y2 ./ v)
] # 2x2 Fisher matrix for the tensor-to-scalar ratio and the foreground amplitude
Cij = inv(Fij) # Covariance matrix
for irz = 1:nrz
    x = bb1[ii, irz] .- mb2[ii] .- Alens * cls_th_binned[4, ii]
    r[irz] = sum(x .* y1 ./ v) / sum(y1 .^ 2 ./ v)
    z = [sum(x .* y1 ./ v), sum(x .* y2 ./ v)]
    w[irz, 1:2] = Fij \ z
end
# Report the final results
println("SO-SAT + CCAT-prime results")
println("Fitted ℓs: ", ell_eff[ii])
println("Without foreground marginalisation:")
println("- r = ", mean(r), " ± ", std(r))
println("- Fisher error = ", 1 / √sum(y1 .^ 2 ./ v))
println("With foreground marginalisation:")
println("- r = ", mean(w[:, 1]), " ± ", std(w[:, 1]))
println("- Fisher error = ", √Cij[1, 1])
# println("- FG = ", mean(w[:, 2]), " ± ", std(w[:, 2]))
# println("- Fisher error = ", √Cij[2, 2])
# SO-SAT Only
y2 = mb3SO[ii]
v = vb1SO[ii]
Fij = [
    sum(y1 .* y1 ./ v) sum(y1 .* y2 ./ v)
    sum(y2 .* y1 ./ v) sum(y2 .* y2 ./ v)
] # 2x2 Fisher matrix for the tensor-to-scalar ratio and the foreground amplitude
Cij = inv(Fij) # Covariance matrix
for irz = 1:nrz
    x = bb1SO[ii, irz] .- mb2SO[ii] .- Alens * cls_th_binned[4, ii]
    rSO[irz] = sum(x .* y1 ./ v) / sum(y1 .^ 2 ./ v)
    z = [sum(x .* y1 ./ v), sum(x .* y2 ./ v)]
    wSO[irz, 1:2] = Fij \ z
end
# Report the final results
println("SO-SAT only results")
println("Fitted ℓs: ", ell_eff[ii])
println("Without foreground marginalisation:")
println("- r = ", mean(rSO), " ± ", std(rSO))
println("- Fisher error = ", 1 / √sum(y1 .^ 2 ./ v))
println("With foreground marginalisation:")
println("- r = ", mean(wSO[:, 1]), " ± ", std(wSO[:, 1]))
println("- Fisher error = ", √Cij[1, 1])
# println("- FG = ", mean(wSO[:, 2]), " ± ", std(wSO[:, 2]))
# println("- Fisher error = ", √Cij[2, 2])

# %% Save to the file
t = Tables.table([1:nrz r w[:, 1] rSO wSO[:, 1]])
CSV.write(
    "ilc_results_sosatccatp.csv",
    t,
    header = [
        "irz",
        "r_wo_FGmarg",
        "r_w_FGmarg",
        "r_SOonly_wo_FGmarg",
        "r_SOonly_w_FGmarg",
    ],
)


text:
mutable struct Quadruped{T} <: Model{T}
	nq::Int
	nu::Int
	nw::Int
	nc::Int

	g::T
	μ_world::T
	μ_joint::T

    # torso
    l_torso::T
    d_torso::T
    m_torso::T
    J_torso::T

    # leg 1
        # thigh
    l_thigh1::T
    d_thigh1::T
    m_thigh1::T
    J_thigh1::T
        # calf
    l_calf1::T
    d_calf1::T
    m_calf1::T
    J_calf1::T

    # leg 2
        # thigh
    l_thigh2::T
    d_thigh2::T
    m_thigh2::T
    J_thigh2::T
        # calf
    l_calf2::T
    d_calf2::T
    m_calf2::T
    J_calf2::T

	# leg 3
        # thigh
    l_thigh3::T
    d_thigh3::T
    m_thigh3::T
    J_thigh3::T
        # calf
    l_calf3::T
    d_calf3::T
    m_calf3::T
    J_calf3::T

	# leg 4
        # thigh
    l_thigh4::T
    d_thigh4::T
    m_thigh4::T
    J_thigh4::T
        # calf
    l_calf4::T
    d_calf4::T
    m_calf4::T
    J_calf4::T

	# fast methods
	base
	dyn

	joint_friction
end

# Trunk model
#                             middle   com
#                                267 mm
#             o---------------------|---X-----------------o
#                                   13mm
# BRhip o___________________________|___________________________o FRhip
#                    183 mm                    183mm
# kinematics
function kinematics_1(model::Quadruped, q; body = :torso, mode = :ee)
	x = q[1]
	z = q[2]

	if body == :torso
		l = model.l_torso
		d = model.d_torso
		θ = q[3]
	elseif body == :thigh_1
		l = model.l_thigh1
		d = model.d_thigh1
		θ = q[4]
	elseif body == :thigh_2
		l = model.l_thigh2
		d = model.d_thigh2
		θ = q[6]
	else
		@error "incorrect body specification"
	end

	if mode == :ee
		return [x + l * sin(θ); z - l * cos(θ)]
	elseif mode == :com
		return [x + d * sin(θ); z - d * cos(θ)]
	else
		@error "incorrect mode specification"
	end
end

function jacobian_1(model::Quadruped, q; body = :torso, mode = :ee)
	jac = zeros(eltype(q), 2, model.nq)
	jac[1, 1] = 1.0
	jac[2, 2] = 1.0
	if body == :torso
		r = mode == :ee ? model.l_torso : model.d_torso
		θ = q[3]
		jac[1, 3] = r * cos(θ)
		jac[2, 3] = r * sin(θ)
	elseif body == :thigh_1
		r = mode == :ee ? model.l_thigh1 : model.d_thigh1
		θ = q[4]
		jac[1, 4] = r * cos(θ)
		jac[2, 4] = r * sin(θ)
	elseif body == :thigh_2
		r = mode == :ee ? model.l_thigh2 : model.d_thigh2
		θ = q[6]
		jac[1, 6] = r * cos(θ)
		jac[2, 6] = r * sin(θ)
	else
		@error "incorrect body specification"
	end

	return jac
end

function kinematics_2(model::Quadruped, q; body = :calf_1, mode = :ee)

	if body == :calf_1
		p = kinematics_1(model, q, body = :thigh_1, mode = :ee)
		θb = q[5]

		lb = model.l_calf1
		db = model.d_calf1
	elseif body == :calf_2
		p = kinematics_1(model, q, body = :thigh_2, mode = :ee)
		θb = q[7]

		lb = model.l_calf2
		db = model.d_calf2
	elseif body == :thigh_3
		p = kinematics_1(model, q, body = :torso, mode = :ee)
		θb = q[8]

		lb = model.l_thigh3
		db = model.d_thigh3
	elseif body == :thigh_4
		p = kinematics_1(model, q, body = :torso, mode = :ee)
		θb = q[10]

		lb = model.l_thigh4
		db = model.d_thigh4
	else
		@error "incorrect body specification"
	end

	if mode == :ee
		return p + [lb * sin(θb); -1.0 * lb * cos(θb)]
	elseif mode == :com
		return p + [db * sin(θb); -1.0 * db * cos(θb)]
	else
		@error "incorrect mode specification"
	end
end

function jacobian_2(model::Quadruped, q; body = :calf_1, mode = :ee)

	if body == :calf_1
		jac = jacobian_1(model, q, body = :thigh_1, mode = :ee)

		θb = q[5]

		r = mode == :ee ? model.l_calf1 : model.d_calf1

		jac[1, 5] += r * cos(θb)
		jac[2, 5] += r * sin(θb)
	elseif body == :calf_2
		jac = jacobian_1(model, q, body = :thigh_2, mode = :ee)

		θb = q[7]

		r = mode == :ee ? model.l_calf2 : model.d_calf2

		jac[1, 7] += r * cos(θb)
		jac[2, 7] += r * sin(θb)
	elseif body == :thigh_3
		jac = jacobian_1(model, q, body = :torso, mode = :ee)
		θb = q[8]

		r = mode == :ee ? model.l_thigh3 : model.d_thigh3

		jac[1, 8] += r * cos(θb)
		jac[2, 8] += r * sin(θb)
	elseif body == :thigh_4
		jac = jacobian_1(model, q, body = :torso, mode = :ee)

		θb = q[10]

		r = mode == :ee ? model.l_thigh4 : model.d_thigh4

		jac[1, 10] += r * cos(θb)
		jac[2, 10] += r * sin(θb)
	else
		@error "incorrect body specification"
	end

	return jac
end

function kinematics_3(model::Quadruped, q; body = :calf_3, mode = :ee)

	if body == :calf_3
		p = kinematics_2(model, q, body = :thigh_3, mode = :ee)
		θc = q[9]

		lb = model.l_calf3
		db = model.d_calf3
	elseif body == :calf_4
		p = kinematics_2(model, q, body = :thigh_4, mode = :ee)

		θc = q[11]

		lb = model.l_calf4
		db = model.d_calf4
	else
		@error "incorrect body specification"
	end

	if mode == :ee
		return p + [lb * sin(θc); -1.0 * lb * cos(θc)]
	elseif mode == :com
		return p + [db * sin(θc); -1.0 * db * cos(θc)]
	else
		@error "incorrect mode specification"
	end
end

function jacobian_3(model::Quadruped, q; body = :calf_3, mode = :ee)

	if body == :calf_3
		jac = jacobian_2(model, q, body = :thigh_3, mode = :ee)

		θc = q[9]

		r = mode == :ee ? model.l_calf3 : model.d_calf3

		jac[1, 9] += r * cos(θc)
		jac[2, 9] += r * sin(θc)

	elseif body == :calf_4
		jac = jacobian_2(model, q, body = :thigh_4, mode = :ee)

		θc = q[11]

		r = mode == :ee ? model.l_calf4 : model.d_calf4

		jac[1, 11] += r * cos(θc)
		jac[2, 11] += r * sin(θc)
	else
		@error "incorrect body specification"
	end

	return jac
end

# Lagrangian
function lagrangian(model::Quadruped, q, q̇)
	L = 0.0

	# # torso
	p_torso = kinematics_1(model, q, body = :torso, mode = :com)
	J_torso = jacobian_1(model, q, body = :torso, mode = :com)
	v_torso = J_torso * q̇
	#
	L += 0.5 * model.m_torso * transpose(v_torso) * v_torso
	L += 0.5 * model.J_torso * q̇[3]^2.0
	L -= model.m_torso * model.g * p_torso[2]

	# thigh 1
	p_thigh_1 = kinematics_1(model, q, body = :thigh_1, mode = :com)
	J_thigh_1 = jacobian_1(model, q, body = :thigh_1, mode = :com)
	v_thigh_1 = J_thigh_1 * q̇

	L += 0.5 * model.m_thigh1 * transpose(v_thigh_1) * v_thigh_1
	L += 0.5 * model.J_thigh1 * q̇[4]^2.0
	L -= model.m_thigh1 * model.g * p_thigh_1[2]

	# leg 1
	p_calf_1 = kinematics_2(model, q, body = :calf_1, mode = :com)
	J_calf_1 = jacobian_2(model, q, body = :calf_1, mode = :com)
	v_calf_1 = J_calf_1 * q̇

	L += 0.5 * model.m_calf1 * transpose(v_calf_1) * v_calf_1
	L += 0.5 * model.J_calf1 * q̇[5]^2.0
	L -= model.m_calf1 * model.g * p_calf_1[2]

	# thigh 2
	p_thigh_2 = kinematics_1(model, q, body = :thigh_2, mode = :com)
	J_thigh_2 = jacobian_1(model, q, body = :thigh_2, mode = :com)
	v_thigh_2 = J_thigh_2 * q̇

	L += 0.5 * model.m_thigh2 * transpose(v_thigh_2) * v_thigh_2
	L += 0.5 * model.J_thigh2 * q̇[6]^2.0
	L -= model.m_thigh2 * model.g * p_thigh_2[2]

	# leg 2
	p_calf_2 = kinematics_2(model, q, body = :calf_2, mode = :com)
	J_calf_2 = jacobian_2(model, q, body = :calf_2, mode = :com)
	v_calf_2 = J_calf_2 * q̇

	L += 0.5 * model.m_calf2 * transpose(v_calf_2) * v_calf_2
	L += 0.5 * model.J_calf2 * q̇[7]^2.0
	L -= model.m_calf2 * model.g * p_calf_2[2]

	# thigh 3
	p_thigh_3 = kinematics_2(model, q, body = :thigh_3, mode = :com)
	J_thigh_3 = jacobian_2(model, q, body = :thigh_3, mode = :com)
	v_thigh_3 = J_thigh_3 * q̇

	L += 0.5 * model.m_thigh3 * transpose(v_thigh_3) * v_thigh_3
	L += 0.5 * model.J_thigh3 * q̇[8]^2.0
	L -= model.m_thigh3 * model.g * p_thigh_3[2]

	# leg 3
	p_calf_3 = kinematics_3(model, q, body = :calf_3, mode = :com)
	J_calf_3 = jacobian_3(model, q, body = :calf_3, mode = :com)
	v_calf_3 = J_calf_3 * q̇

	L += 0.5 * model.m_calf3 * transpose(v_calf_3) * v_calf_3
	L += 0.5 * model.J_calf3 * q̇[9]^2.0
	L -= model.m_calf3 * model.g * p_calf_3[2]

	# thigh 4
	p_thigh_4 = kinematics_2(model, q, body = :thigh_4, mode = :com)
	J_thigh_4 = jacobian_2(model, q, body = :thigh_4, mode = :com)
	v_thigh_4 = J_thigh_4 * q̇

	L += 0.5 * model.m_thigh4 * transpose(v_thigh_4) * v_thigh_4
	L += 0.5 * model.J_thigh4 * q̇[10]^2.0
	L -= model.m_thigh4 * model.g * p_thigh_4[2]

	# leg 4
	p_calf_4 = kinematics_3(model, q, body = :calf_4, mode = :com)
	J_calf_4 = jacobian_3(model, q, body = :calf_4, mode = :com)
	v_calf_4 = J_calf_4 * q̇

	L += 0.5 * model.m_calf4 * transpose(v_calf_4) * v_calf_4
	L += 0.5 * model.J_calf4 * q̇[11]^2.0
	L -= model.m_calf4 * model.g * p_calf_4[2]

	return L
end

function _dLdq(model::Quadruped, q, q̇)
	Lq(x) = lagrangian(model, x, q̇)
	ForwardDiff.gradient(Lq, q)
end

function _dLdq̇(model::Quadruped, q, q̇)
	Lq̇(x) = lagrangian(model, q, x)
	ForwardDiff.gradient(Lq̇, q̇)
end

function kinematics(model::Quadruped, q)
	p_calf_1 = kinematics_2(model, q, body = :calf_1, mode = :ee)
	p_calf_2 = kinematics_2(model, q, body = :calf_2, mode = :ee)
	p_calf_3 = kinematics_3(model, q, body = :calf_3, mode = :ee)
	p_calf_4 = kinematics_3(model, q, body = :calf_4, mode = :ee)

	SVector{8}([p_calf_1; p_calf_2; p_calf_3; p_calf_4])
end

# Methods
function M_func(model::Quadruped, q)
	M = Diagonal([0.0, 0.0, model.J_torso, model.J_thigh1, model.J_calf1, model.J_thigh2, model.J_calf2, model.J_thigh3, model.J_calf3, model.J_thigh4, model.J_calf4])

	# torso
	J_torso = jacobian_1(model, q, body = :torso, mode = :com)
	M += model.m_torso * transpose(J_torso) * J_torso

	# thigh 1
	J_thigh_1 = jacobian_1(model, q, body = :thigh_1, mode = :com)
	M += model.m_thigh1 * transpose(J_thigh_1) * J_thigh_1

	# leg 1
	J_calf_1 = jacobian_2(model, q, body = :calf_1, mode = :com)
	M += model.m_calf1 * transpose(J_calf_1) * J_calf_1

	# thigh 2
	J_thigh_2 = jacobian_1(model, q, body = :thigh_2, mode = :com)
	M += model.m_thigh2 * transpose(J_thigh_2) * J_thigh_2

	# leg 2
	J_calf_2 = jacobian_2(model, q, body = :calf_2, mode = :com)
	M += model.m_calf2 * transpose(J_calf_2) * J_calf_2

	# thigh 3
	J_thigh_3 = jacobian_2(model, q, body = :thigh_3, mode = :com)
	M += model.m_thigh3 * transpose(J_thigh_3) * J_thigh_3

	# leg 3
	J_calf_3 = jacobian_3(model, q, body = :calf_3, mode = :com)
	M += model.m_calf3 * transpose(J_calf_3) * J_calf_3

	# thigh 4
	J_thigh_4 = jacobian_2(model, q, body = :thigh_4, mode = :com)
	M += model.m_thigh4 * transpose(J_thigh_4) * J_thigh_4

	# leg 4
	J_calf_4 = jacobian_3(model, q, body = :calf_4, mode = :com)
	M += model.m_calf4 * transpose(J_calf_4) * J_calf_4

	return M
end

function ϕ_func(model::Quadruped, env::Environment, q)
	p_calf_1 = kinematics_2(model, q, body = :calf_1, mode = :ee)
	p_calf_2 = kinematics_2(model, q, body = :calf_2, mode = :ee)
	p_calf_3 = kinematics_3(model, q, body = :calf_3, mode = :ee)
	p_calf_4 = kinematics_3(model, q, body = :calf_4, mode = :ee)

	SVector{4}([p_calf_1[2] - env.surf(p_calf_1[1:1]);
				p_calf_2[2] - env.surf(p_calf_2[1:1]);
				p_calf_3[2] - env.surf(p_calf_3[1:1]);
				p_calf_4[2] - env.surf(p_calf_4[1:1])])
end

function B_func(model::Quadruped, q)
	@SMatrix [0.0  0.0 -1.0  1.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0;
			  0.0  0.0  0.0 -1.0  1.0  0.0  0.0  0.0  0.0  0.0  0.0;
			  0.0  0.0 -1.0  0.0  0.0  1.0  0.0  0.0  0.0  0.0  0.0;
			  0.0  0.0  0.0  0.0  0.0 -1.0  1.0  0.0  0.0  0.0  0.0;
			  0.0  0.0 -1.0  0.0  0.0  0.0  0.0  1.0  0.0  0.0  0.0;
			  0.0  0.0  0.0  0.0  0.0  0.0  0.0 -1.0  1.0  0.0  0.0;
			  0.0  0.0 -1.0  0.0  0.0  0.0  0.0  0.0  0.0  1.0  0.0;
			  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0 -1.0  1.0]
end

function A_func(model::Quadruped, q)
	@SMatrix [1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
			  0.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0]
end

function J_func(model::Quadruped, env::Environment, q)
	J_calf_1 = jacobian_2(model, q, body = :calf_1, mode = :ee)
	J_calf_2 = jacobian_2(model, q, body = :calf_2, mode = :ee)
	J_calf_3 = jacobian_3(model, q, body = :calf_3, mode = :ee)
	J_calf_4 = jacobian_3(model, q, body = :calf_4, mode = :ee)

	return [J_calf_1;
			J_calf_2;
			J_calf_3;
			J_calf_4]
end

function C_func(model::Quadruped, q, q̇)
	tmp_q(z) = _dLdq̇(model, z, q̇)
	tmp_q̇(z) = _dLdq̇(model, q, z)

	ForwardDiff.jacobian(tmp_q, q) * q̇ - _dLdq(model, q, q̇)
end

function contact_forces(model::Quadruped, env::Environment{R2, LinearizedCone}, γ1, b1, q2, k)
	m = friction_mapping(env)

	SVector{8}([transpose(rotation(env, k[1:1])) * [m * b1[1:2]; γ1[1]];
				transpose(rotation(env, k[3:3])) * [m * b1[3:4]; γ1[2]];
				transpose(rotation(env, k[5:5])) * [m * b1[5:6]; γ1[3]];
				transpose(rotation(env, k[7:7])) * [m * b1[7:8]; γ1[4]]])
end

function velocity_stack(model::Quadruped, env::Environment{R2, LinearizedCone}, q1, q2, k, h)
	v = J_func(model, env, q2) * (q2 - q1) / h[1]

	v1_surf = rotation(env, k[1:1]) * v[1:2]
	v2_surf = rotation(env, k[3:3]) * v[3:4]
	v3_surf = rotation(env, k[5:5]) * v[5:6]
	v4_surf = rotation(env, k[7:7]) * v[7:8]

	SVector{8}([transpose(friction_mapping(env)) * v1_surf[1];
				transpose(friction_mapping(env)) * v2_surf[1];
				transpose(friction_mapping(env)) * v3_surf[1];
				transpose(friction_mapping(env)) * v4_surf[1]])
end

################################################################################
# Instantiation
################################################################################
# Dimensions
nq = 2 + 5 + 4            # configuration dimension
nu = 4 + 4                # control dimension
nc = 4                    # number of contact points
nw = 2
nquat = 0

# World parameters
g = 9.81      # gravity
μ_world = 1.0 # coefficient of friction
μ_world_downhill = 0.7 # coefficient of friction
μ_joint = 0.1 # coefficient of torque friction at the joints

# ~Unitree A1
# Model parameters
m_torso = 4.713 + 4 * 0.696
m_thigh = 1.013
m_leg = 0.166

J_torso = 0.01683 + 4 * 0.696 * 0.183^2.0
J_thigh = 0.00552
J_leg = 0.00299

l_torso = 0.183*2
l_thigh = 0.2
l_leg = 0.2

d_torso = 0.5 * l_torso + 0.0127
d_thigh = 0.5 * l_thigh - 0.00323
d_leg = 0.5 * l_leg - 0.006435

m_payload = 3.0
J_payload = 0.03

# Model parameters
# m_torso = 4.713 + 4 * 0.696
# m_thigh = 1.013
# m_leg = 0.166
#
# J_torso = 0.056579028 + 4 * 0.696 * 0.183^2.0
# J_thigh = 0.005139339
# J_leg = 0.003014022
#
# # l_torso = 0.267
# l_torso = 0.267
# l_thigh = 0.2
# l_leg = 0.2
#
# d_torso = 0.5 * l_torso + 0.0127
# d_thigh = 0.5 * l_thigh - 0.00323
# d_leg = 0.5 * l_leg - 0.006435
#
# m_payload = 5.0
# J_payload = 0.05

quadruped = Quadruped(nq,nu,nw,nc,
				g, μ_world, μ_joint,
				l_torso, d_torso, m_torso, J_torso,
				l_thigh, d_thigh, m_thigh, J_thigh,
				l_leg, d_leg, m_leg, J_leg,
				l_thigh, d_thigh, m_thigh, J_thigh,
				l_leg, d_leg, m_leg, J_leg,
				l_thigh, d_thigh, m_thigh, J_thigh,
				l_leg, d_leg, m_leg, J_leg,
				l_thigh, d_thigh, m_thigh, J_thigh,
				l_leg, d_leg, m_leg, J_leg,
				BaseMethods(), DynamicsMethods(),
				SVector{nq}([zeros(3); μ_joint * ones(nq - 3)]))

quadruped_payload = Quadruped(nq,nu,nw,nc,
				g, μ_world, μ_joint,
				l_torso, d_torso,
				m_torso + m_payload,
				J_torso + J_payload,
				l_thigh, d_thigh, m_thigh, J_thigh,
				l_leg, d_leg, m_leg, J_leg,
				l_thigh, d_thigh, m_thigh, J_thigh,
				l_leg, d_leg, m_leg, J_leg,
				l_thigh, d_thigh, m_thigh, J_thigh,
				l_leg, d_leg, m_leg, J_leg,
				l_thigh, d_thigh, m_thigh, J_thigh,
				l_leg, d_leg, m_leg, J_leg,
				BaseMethods(), DynamicsMethods(),
				SVector{nq}([zeros(3); μ_joint * ones(nq - 3)]))

function friction_coefficients(model::Quadruped)
	return [model.μ_world]
end

function initialize_z!(z, model::Quadruped, idx::RoboDojo.IndicesZ, q)
    z .= 1.0
    z[idx.q] .= q
end

ext: jl


text:
export Pitch, ≂, @pitch_str

"""
A pitch (pitch). It is recommended that the `pitch"..."` constructor is used, see
[`parse_pitch`](@ref).

`+` and `-` with integers change the pitch by a given number of semitones.
"""
struct Pitch
    "Semitones relative to C₀."
    semitones::Int
end

Broadcast.broadcastable(pitch::Pitch) = Ref(pitch)

Broadcast.broadcastable(semitones::Integer) = Ref(semitones)

Base.hash(pitch::Pitch, h::UInt) = hash(hash(pitch.semitones, 0x3d948cfcb6fb103d), h)

Base.:+(pitch::Pitch, semitones::Integer) = Pitch(pitch.semitones + semitones)

Base.:-(pitch::Pitch, semitones::Integer) = Pitch(pitch.semitones - semitones)

Base.:-(pitch1::Pitch, pitch2::Pitch) = pitch1.semitones - pitch2.semitones

"""
$(SIGNATURES)

Test if two pitchs *octave equivalent* (ie are separated by an integer number of octaves), ie
are in the same *pitch class*.
"""
≂(pitch1::Pitch, pitch2::Pitch) = mod(pitch1.semitones, 12) == mod(pitch2.semitones, 12)

"""
Pitch names for printing.

We always print sharp notes, eg `C♯` instead of `D♭`.
"""
const PITCH_NAMES = ["C", "C♯", "D", "D♯", "E", "F",
                    "F♯", "G", "G♯", "A", "A♯", "B"]

function Base.string(pitch::Pitch)
    octave, rel_semitone = fldmod(pitch.semitones, 12)
    PITCH_NAMES[rel_semitone + 1] * string(octave)
end

function Base.show(io::IO, pitch::Pitch)
    print(io, "pitch\"", string(pitch), "\"")
end

"""
Major scale in semitones. Used internally, for pitch parsing.
"""
const MAJOR_SCALE_SEMITONES = [0, 2, 4, 5, 7, 9, 11]

"""
$(SIGNATURES)

Parse a pitch in [scientific pitch
notation](https://en.wikipedia.org/wiki/Scientific_pitch_notation), composed of

- an upper- or lowecase letter A-G for the pitch

- an optional accidental `b`, `#` (plain ascii), or `♭`, `♯` (Unicode accidentals),

- an optional octave number (`0` if omitted, can be negative)

Example arguments: "C#2", "D♭-2" `.
"""
function parse_pitch(str::AbstractString)
    m = match(r"^([a-gA-G])([b♭#♯]?)(-?[0-9]*)$", str)
    @argcheck m ≢ nothing "Could not parse $str as a valid pitch."
    _semitone, _accidental, _octave = m.captures
    major_tone = uppercase(first(_semitone)) - 'C' + 1
    if major_tone < 1
        major_tone += 7
    end
    rel_semitone = MAJOR_SCALE_SEMITONES[major_tone]
    if !isempty(_accidental)
        accidental = first(_accidental)
        if accidental == '#' || accidental == '♯'
            rel_semitone += 1
        elseif accidental == 'b' || accidental == '♭'
            rel_semitone -= 1
        end
    end
    octave = isempty(_octave) ? 0 : parse(Int, _octave)
    Pitch(octave * 12 + rel_semitone)
end

"""
pitch"..."

Parse the argument string as a [`Pitch`](@ref), eg

```jldoctest
julia> pitch"C♭2"
pitch"B1"
```

See [`parse_pitch`](@ref) for a full description.
"""
macro pitch_str(str)
    parse_pitch(str)
end

####
#### named intervals
####

"""
Some named intervals, not exported from the main module, use `using Fretboards.Intervals` to
import.
"""
module Intervals

export m3, M3, P4, d5, P5, A5, m7, M7

"Minor 3rd."
const m3 = 3

"Major 3rd."
const M3 = 4

"Perfect 4th."
const P4 = 5

"Diminished 5th."
const d5 = 6

"Perfect 5th."
const P5 = 7

"Augmented 5th."
const A5 = 8

"Minor 7th."
const m7 = 10

"Major 7th."
const M7 = 11

end


text:
xmax(x::Node) = x[:geometry] !== nothing ? xmax(x[:geometry].mesh) : nothing
xmin(x::Node) = x[:geometry] !== nothing ? xmin(x[:geometry].mesh) : nothing
ymax(x::Node) = x[:geometry] !== nothing ? ymax(x[:geometry].mesh) : nothing
ymin(x::Node) = x[:geometry] !== nothing ? ymin(x[:geometry].mesh) : nothing
zmax(x::Node) = x[:geometry] !== nothing ? zmax(x[:geometry].mesh) : nothing
zmin(x::Node) = x[:geometry] !== nothing ? zmin(x[:geometry].mesh) : nothing

xmax(x) = map_coord(maximum, x, 1)
xmin(x) = map_coord(minimum, x, 1)
ymax(x) = map_coord(maximum, x, 2)
ymin(x) = map_coord(minimum, x, 2)
zmax(x) = map_coord(maximum, x, 3)
zmin(x) = map_coord(minimum, x, 3)

"""
    xmax(x)
    ymax(x)
    zmax(x)

Get the maximum x, y or z coordinates of a mesh or a Node.
"""
xmax, ymax, zmax

"""
    xmin(x)
    ymin(x)
    zmin(x)

Get the minimum x, y or z coordinates of a mesh or a Node.
"""
xmin, ymin, zmin

"""
    map_coord(f, mesh, coord)

Apply function `f` over the mesh coordinates `coord`.
Values for `coord` can be 1 for x, 2 for y and 3 for z.
"""
function map_coord(f, mesh, coord)
    f([i.coords[coord] for i in mesh.points])
end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

module APITests

using Pkg, Test

include("utils.jl")

@testset "Pkg.activate" begin
    temp_pkg_dir() do project_path
        cd(mktempdir()) do
            path = pwd()
            Pkg.activate(".")
            mkdir("Foo")
            cd(mkdir("modules")) do
                Pkg.generate("Foo")
            end
            Pkg.develop(Pkg.Types.PackageSpec(path="modules/Foo")) # to avoid issue #542
            Pkg.activate("Foo") # activate path Foo over deps Foo
            @test Base.active_project() == joinpath(path, "Foo", "Project.toml")
            Pkg.activate(".")
            rm("Foo"; force=true, recursive=true)
            Pkg.activate("Foo") # activate path from developed Foo
            @test Base.active_project() == joinpath(path, "modules", "Foo", "Project.toml")
            Pkg.activate(".")
            Pkg.activate("./Foo") # activate empty directory Foo (sidestep the developed Foo)
            @test Base.active_project() == joinpath(path, "Foo", "Project.toml")
            Pkg.activate(".")
            Pkg.activate("Bar") # activate empty directory Bar
            @test Base.active_project() == joinpath(path, "Bar", "Project.toml")
            Pkg.activate(".")
            Pkg.add("Example") # non-deved deps should not be activated
            Pkg.activate("Example")
            @test Base.active_project() == joinpath(path, "Example", "Project.toml")
            Pkg.activate(".")
            cd(mkdir("tests"))
            Pkg.activate("Foo") # activate developed Foo from another directory
            @test Base.active_project() == joinpath(path, "modules", "Foo", "Project.toml")
            Pkg.activate() # activate home project
            @test Base.ACTIVE_PROJECT[] === nothing
        end
    end
end

end # module APITests

max_stars_repo_name: greimel/julia

max_issues_repo_name: greimel/julia

max_forks_repo_name: greimel/julia


text:
module Metatheory

using Base.Meta
using Reexport
using TermInterface

macro log(args...)
  quote
    haskey(ENV, "MT_DEBUG") && @info($(args...))
  end |> esc
end

@inline alwaystrue(x) = true

include("docstrings.jl")
include("utils.jl")
export @timer
export @iftimer
export @timerewrite
export @matchable

include("Patterns.jl")
@reexport using .Patterns

include("ematch_compiler.jl")
@reexport using .EMatchCompiler

include("matchers.jl")
include("Rules.jl")
@reexport using .Rules

include("Syntax.jl")
@reexport using .Syntax
include("EGraphs/EGraphs.jl")
@reexport using .EGraphs

include("Library.jl")
export Library

include("Rewriters.jl")
using .Rewriters
export Rewriters

function rewrite(expr, theory; order = :outer)
  if order == :inner
    Fixpoint(Prewalk(Fixpoint(Chain(theory))))(expr)
  elseif order == :outer
    Fixpoint(Postwalk(Fixpoint(Chain(theory))))(expr)
  end
end
export rewrite

end # module


text:
# ccal
# This is basically a horrible hack around not being able to call C
# varargs functions. Emulates `ccall` but as a function taking a function
# pointer + types, args etc.

Base.eltype{T}(::Type{Type{T}}) = T

# stagedfunction ccal(f::Ptr, t, ts, vals...)
#   R = eltype(t)
#   AS = Expr(:tuple, eltype(ts)...)
#   args = map(x->gensym(), vals)
#   quote
#     let
#       $(Expr(:tuple, args...)) = vals
#       ccall(f, $R, $AS, $(args...))
#     end
#   end
# end

function ccal(f, R, ts, vals...)
  AS = Expr(:tuple, ts...)
  @eval ccall($f, $R, $AS, $(vals...))
end

# Calls

toobject(o::Object) = o
toobject(c::Class) = c
toobject(p::Ptr) = p

ctype(x) = x
ctype(o::Type{Object}) = Ptr{Void}
ctype(s::Type{Selector}) = Ptr{Void}
ctype(a::AbstractArray) = map(ctype, a)

const cmsgsend = cglobal(:objc_msgSend)

function message(obj, sel, args...)
  obj = toobject(obj)
  clarse = class(obj)
  m = method(clarse, sel)
  m == C_NULL && error("$clarse does not respond to $sel")
  types = signature(m)
  ctypes = ctype(types)

  args = Any[args...]
  for i = 1:length(args)
    types[i+3] == Object && (args[i] = toobject(args[i]))
  end

  result = ccal(cmsgsend, ctypes[1], tuple(ctypes[2:end]...),
                obj, sel, args...)
  types[1] in (Object, Selector) && return types[1](result)
  return result
end


text:
using TemporalGPs:
    predict,
    step_marginals,
    step_logpdf,
    step_filter,
    invert_dynamics,
    step_posterior,
    storage_type,
    is_of_storage_type

using Zygote, StaticArrays

println("lgssm:")
@testset "lgssm" begin

    rng = MersenneTwister(123456)

    storages = (
        dense=(name="dense storage Float64", val=ArrayStorage(Float64)),
        static=(name="static storage Float64", val=SArrayStorage(Float64)),
    )
    emission_types = (
        small_output=(name="small output", val=SmallOutputLGC),
        large_output=(name="large output", val=LargeOutputLGC),
        scalar_output=(name="scalar output", val=ScalarOutputLGC),
    )
    settings = [
        (tv=:time_varying, N=1, Dlat=3, Dobs=2, storage=storages.dense),
        (tv=:time_varying, N=49, Dlat=3, Dobs=2, storage=storages.dense),
        (tv=:time_invariant, N=49, Dlat=3, Dobs=2, storage=storages.dense),
        (tv=:time_varying, N=49, Dlat=1, Dobs=1, storage=storages.dense),
        (tv=:time_varying, N=1, Dlat=3, Dobs=2, storage=storages.static),
        (tv=:time_invariant, N=49, Dlat=3, Dobs=2, storage=storages.static),
    ]
    orderings = [
        Forward(),
        Reverse(),
    ]
    Qs = [
        Val(:dense),
        # Val(:diag), diag tests don't work because `FiniteDiffernces.to_vec`.
    ]

    @testset "($tv, $N, $Dlat, $Dobs, $(storage.name), $(emission.name), $order, $Q)" for
        (tv, N, Dlat, Dobs, storage) in settings,
        emission in emission_types,
        order in orderings,
        Q in Qs

        # Print current iteration to prevent CI timing out.
        println(
            "(time_varying=$tv, N=$N, Dlat=$Dlat, Dobs=$Dobs, " *
            "storage=$(storage.name), emissions=$(emission.val), ordering=$order)",
        )

        # Build LGSSM.
        model = if emission.val ∈ (SmallOutputLGC, LargeOutputLGC)
            random_lgssm(rng, order, Val(tv), emission.val, Dlat, Dobs, N, Q, storage.val)
        elseif emission.val == ScalarOutputLGC
            random_lgssm(rng, order, Val(tv), emission.val, Dlat, Dobs, N, storage.val)
        else
            throw(error("Unrecognised storage $(emission.val)"))
        end

        # Verify the correct output types has been obtained.
        @test eltype(model.emissions) <: emission.val

        # Verify that model properties are as requested.
        @test storage_type(model) == storage.val

        @test length(model) == N
        @test getindex(model, N) isa TemporalGPs.ElementOfLGSSM

        y = first(rand(model))
        x = TemporalGPs.x0(model)

        @testset "step_marginals" begin
            @inferred step_marginals(x, model[1])
            adjoint_test(step_marginals, (x, model[1]))
            if storage.val isa SArrayStorage
                check_adjoint_allocations(step_marginals, (x, model[1]))
            end
        end
        @testset "step_logpdf" begin
            args = (ordering(model[1]), x, (model[1], y))
            @inferred step_logpdf(args...)
            adjoint_test(step_logpdf, args)
            if storage.val isa SArrayStorage
                check_adjoint_allocations(step_logpdf, args)
            end
        end
        @testset "step_filter" begin
            args = (ordering(model[1]), x, (model[1], y))
            @inferred step_filter(args...)
            adjoint_test(step_filter, args)
            if storage.val isa SArrayStorage
                check_adjoint_allocations(step_filter, args)
            end
        end
        @testset "invert_dynamics" begin
            args = (x, x, model[1].transition)
            @inferred invert_dynamics(args...)
            adjoint_test(invert_dynamics, args)
            if storage.val isa SArrayStorage
                check_adjoint_allocations(invert_dynamics, args)
            end
        end
        @testset "step_posterior" begin
            args = (ordering(model[1]), x, (model[1], y))
            @inferred step_posterior(args...)
            adjoint_test(step_posterior, args)
            if storage.val isa SArrayStorage
                check_adjoint_allocations(step_posterior, args)
            end
        end

        # Run standard battery of LGSSM tests.
        test_interface(
            rng, model;
            rtol=1e-5,
            atol=1e-5,
            context=NoContext(),
            max_primal_allocs=10,
            max_forward_allocs=35,
            max_backward_allocs=50,
            check_allocs=storage.val isa SArrayStorage,
        )
    end
end


text:
"""
    @sym("x") --> :x
    @sym(x) --> Symbol(<value of x>)

@sym("ab","cd") --> :abcd
@sym(:sym,"_") --> :sym_
@sym(:a,:b,:c) --> :abc
"""
macro sym(x)
    esc(:(Symbol($x)))
end

macro sym(x,y)
    esc(:(Symbol($x,$y)))
end

macro sym(x,y,z)
    esc(:(Symbol($x,$y,$z)))
end

"""
    @maybe(T) --> Union{Nothing, T}
    @maybe(T1,T2) --> Union{Nothing, T1, T2}
"""
macro maybe(T)
    :(Union{Nothing, $T})
end

macro maybe(T1, T2)
    :(Union{Nothing, $T1, $T2})
end


text:
# https://docs.rapids.ai/api/cuml/stable/pickling_cuml_models.html

"""
    persistent(booster)

Private method.

Return a persistent (ie, Julia-serializable) representation of the
RAPIDS.jl model.

Restore the model with [`model`](@ref)

"""
function persistent(model)

    # this implemenation is not ideal; see
    # https://github.com/dmlc/XGBoost.jl/issues/103

    model_file, io = mktemp()
    close(io)

    pickle.dump(model, @py open(model_file, "wb") )
    XGBoost.save(booster, xgb_file)
    persistent_booster = read(xgb_file)
    rm(xgb_file)
    return persistent_booster
end

"""
    booster(persistent)
Private method.
Return the XGBoost.jl model which has `persistent` as its persistent
(Julia-serializable) representation. See [`persistent`](@ref) method.
"""
function booster(persistent)

    model_file, io = mktemp()
    write(io, persistent)
    close(io)
    model = pickle.load( @py open(model_file, "rb") )
    rm(model_file)

    return model
end

MMI.save(::CUML_MODELS, fitresult; kwargs...) = persistent(fitresult)

MMI.restore(::CUML_MODELS, serializable_fitresult) = booster(serializable_fitresult)


text:
using LinearAlgebra, LinearOperators, SparseArrays, Printf, Test

using MPCC
using MPCCSolver

using Ipopt, JuMP, NLPModels, NLPModelsIpopt, NLPModelsJuMP, Stopping

include("problems.jl")
include("rosenbrock.jl")

printstyled("MPCCSolver param tests... ")
include("test-mpccsolvers-param.jl")
printstyled("passed ✓ \n", color = :green)
if false
printstyled("NLMPCC tests for unconstrained nlps... ")
include("test-nlp-nlmpcc.jl")
printstyled("passed ✓ \n", color = :green)
printstyled("NLMPCC tests for constrained nlps...[TODO] ")
#include("test-nlp-nlmpcc.jl")
printstyled("passed ✓ \n", color = :green)
printstyled("NLMPCC tests for mpccs... ")
include("test-mpcc-nlmpcc.jl")
printstyled("passed ✓ \n", color = :green)
end

printstyled("Theta functions tests... ")
include("test_thetaFct.jl")
printstyled("passed ✓ \n", color = :green)
printstyled("Regularization functions tests... ")
include("test_relaxation_map.jl")
printstyled("passed ✓ \n", color = :green)
printstyled("RlxMPCC tests for mpccs...[TODO] ")
include("test-rlxmpcc.jl")
printstyled("passed ✓ \n", color = :green)
printstyled("Regularization method tests... ")
include("test_relaxation.jl")
printstyled("passed ✓ \n", color = :green)


text:
using VoronoiDelaunay
using VoronoiCells

# Routines for printing points, edges, triangles, and cells to text files for plotting
# Much of this may not work anymore

# I think getx,gety not implemented for abstract Points
print_pt(str::IOStream, p::GeometricalPredicates.Point2D) = println(str,getx(p), " ", gety(p))

# For plotting triangle, print 4 points
# TODO: check if we can use abstract triangles
function print_triangle(str::IOStream, t::VoronoiDelaunay.DelaunayTriangle)
    print_pt(str,geta(t))
    print_pt(str,getb(t))
    print_pt(str,getc(t))
    print_pt(str,geta(t))
end

function print_delaunay_edges(str::IOStream, t::DelaunayTessellation2D)
    for edge in delaunayedges(t)
        print_pt(str,geta(edge))
        print_pt(str,getb(edge))
        println(str)
    end
end

function print_delaunay_edges(fname::String, t::DelaunayTessellation2D)
    str = open(fname,"w")
    print_delaunay_edges(str, t)
    close(str)
end

# Print all edges in voronoi cell, one point per line, for plotting
function print_cell(str::IOStream, cell::VoronoiCell)
    verts = cell._verts
    for i in 1:length(verts)
        print_pt(str,verts[i])
    end
    print_pt(str,verts[1])
end

# Print all cells
function print_cells(str::IOStream, cell_array)
    i = 0
    for cell in cell_array
        println(str,"# $i")
        print_cell(str,cell)
        println(str)
        i += 1
    end
end

function print_cells(fname::String,cell_array)
    str = open(fname,"w")
    print_cells(str, cell_array)
    close(str)
end

###### These are for testing code

# prints pairs of common vertices in two triangles
#  a : x
#  b : y
#  c : z
# Where x,y,z are one of a,b,c or nothing.
function print_common_points(tr1::VoronoiDelaunay.DelaunayTriangle,tr2::VoronoiDelaunay.DelaunayTriangle)
    print(" a : ")
    haspt = print_match_point(tr2,geta(tr1))
    if ! haspt println(" x ") end
    print(" b : ")
    haspt = print_match_point(tr2,getb(tr1))
    if ! haspt println(" x ") end
    print(" c : ")
    haspt = print_match_point(tr2,getc(tr1))
    if ! haspt println(" x ") end
end

function print_match_point(tr::VoronoiDelaunay.DelaunayTriangle, p0::Point2D)
    eps0 = eps(min_coord) # same as for max_coord
    has_vertex_f = false
    pa = geta(tr)
    if distanceL1(pa,p0) <= eps0
        println(" a ")
        has_vertex_f = true
    else
        pb = getb(tr)
        if distanceL1(pb,p0) <= eps0
            println(" b ")
            has_vertex_f = true
        else
            pc = getc(tr)
            if distanceL1(pc,p0) <= eps0
                println(" c ")
                has_vertex_f = true
            end
        end
    end
    has_vertex_f
end


text:
module Controller

# include("./modules/PreProcess.jl")
include("./modules/Filtering.jl")
include("./modules/Scoring.jl")
include("./modules/Detection.jl")
include("./modules/PostProcess.jl")

export controller

SCORING_WINDOW_SIZE = 30
DETECTION_TIMETHRESHOLD = 0.5
POSTPROCESS_TIMETHRESHOLD = 200

function controller(rawdata:: Channel)
    return function (output:: Channel)
        # ppdata = Channel(PreProcess.run(rawdata))
        smoothdata = Channel(Filtering.run(rawdata))
        peakScoreData = Channel(Scoring.run(smoothdata, SCORING_WINDOW_SIZE))
        peakData = Channel(Detection.run(peakScoreData, DETECTION_TIMETHRESHOLD))

        for x in Channel(PostProcess.run(peakData, POSTPROCESS_TIMETHRESHOLD))
            put!(output, x)
        end
    end
end

end

max_stars_repo_name: eliabiraschi/pedometer

max_issues_repo_name: eliabiraschi/pedometer

max_forks_repo_name: eliabiraschi/pedometer


text:
using Images, ImageFeatures, ImageDraw

img1 = Gray.(load("cat-3417184_640.jpg"))
img2 = Gray.(load("cat-3417184_640_watermarked.jpg"))

keypoints_1 = Keypoints(fastcorners(img1, 12, 0.5));
keypoints_2 = Keypoints(fastcorners(img2, 12, 0.5));

size(keypoints_1)
size(keypoints_2)

brief_params = BRIEF()
desc_1, ret_features_1 = create_descriptor(img1, keypoints_1, brief_params);
desc_2, ret_features_2 = create_descriptor(img2, keypoints_2, brief_params);
matches = match_keypoints(ret_features_1, ret_features_2, desc_1, desc_2, 0.5)

grid = Gray.(hcat(img1, img2))
offset = CartesianIndex(0, size(img1, 2))
map(m -> draw!(grid, LineSegment(m[1], m[2] + offset)), matches)
grid

# Creating panorama

using Images, ImageFeatures
img = load("sample-images/cat-3418815_640.jpg")

img_width = size(img, 2)
img_left_width = 400
img_right_width = 340

img_left = view(img, :, 1:img_left_width)
img_left_gray = Gray.(img_left)
img_right = view(img, :, (img_width - img_right_width):img_width)
img_right_gray = Gray.(img_right)

keypoints_1 = Keypoints(fastcorners(img_left_gray, 12, 0.3));
keypoints_2 = Keypoints(fastcorners(img_right_gray, 12, 0.3));

brief_params = BRIEF()
desc_1, ret_features_1 = create_descriptor(img_left_gray, keypoints_1, brief_params);
desc_2, ret_features_2 = create_descriptor(img_right_gray, keypoints_2, brief_params);
matches = match_keypoints(ret_features_1, ret_features_2, desc_1, desc_2, 0.15)

grid = hcat(img_left, img_right)
offset = CartesianIndex(0, size(img_left_gray, 2))
map(m -> draw!(grid, LineSegment(m[1], m[2] + offset)), matches)
imshow(grid)

offset_x = median(map(m -> (img_left_width - m[1][2]) + m[2][2], matches))
offset_x_half = Int(trunc(offset_x / 2))
img_output = hcat(
    img_left[:, 1:(img_left_width-offset_x_half)],
    img_right[:, offset_x_half:img_right_width]
)

imshow(hcat(restrict(img), restrict(img_output)))


text:
function deref(pointer::Int, program::Array{Int,1}, mode::Int, base::Int)::Int
    if mode == 0
        return program[pointer+1]
    elseif mode == 1
        return pointer
    elseif mode == 2
        return program[pointer+base+1]
    else
        throw(ErrorException("Invalid mode: " * mode))
    end
end

function write_pointer(pointer::Int, mode::Int, base::Int)::Int
    if mode == 0
        return pointer+1
    elseif mode == 2
        return pointer+base+1
    else
        throw(ErrorException("Invalid mode for writing: " * mode))
    end
end

function run_intcode(program::Vector{Int}, input::Channel, output::Channel)::Int
    program = vcat(program, zeros(Int, 1000))
    pointer = 1
    ret = nothing
    base = 0
    while pointer <= length(program)
        # println(program)
        instruction = reverse(digits(program[pointer], pad=5))
        # println(instruction)
        opcode = instruction[end] + instruction[end-1]*10
        modes = reverse(instruction[1:end-2])
        # println(join([opcode,": ",modes]))
        if opcode == 1 # add
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            pt3 = program[pointer+3]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            program[write_pointer(pt3, modes[3], base)] = val1 + val2
            pointer = pointer + 4
        elseif opcode == 2 # multiply
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            pt3 = program[pointer+3]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            program[write_pointer(pt3, modes[3], base)] = val1 * val2
            pointer = pointer + 4
        elseif opcode == 3 # store input
            pt1 = program[pointer+1]
            program[write_pointer(pt1, modes[1], base)] = take!(input)
            pointer = pointer + 2
        elseif opcode == 4 # output
            pt1 = program[pointer+1]
            val1 = deref(pt1, program, modes[1], base)
            pointer = pointer + 2
            ret = val1
            put!(output, val1)
        elseif opcode == 5 # jump if true
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            if val1 != 0
                pointer = val2 + 1
            else
                pointer = pointer + 3
            end
        elseif opcode == 6 # jump if false
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            if val1 == 0
                pointer = val2 + 1
            else
                pointer = pointer + 3
            end
        elseif opcode == 7 # less than
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            pt3 = program[pointer+3]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            program[write_pointer(pt3, modes[3], base)] = val1 < val2 ? 1 : 0
            pointer = pointer + 4
        elseif opcode == 8 # equal
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            pt3 = program[pointer+3]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            program[write_pointer(pt3, modes[3], base)] = val1 == val2 ? 1 : 0
            pointer = pointer + 4
        elseif opcode == 9 # modify base
            pt1 = program[pointer+1]
            val1 = deref(pt1, program, modes[1], base)
            base = base + val1
            pointer = pointer + 2
        elseif opcode == 99
            close(output)
            return ret
        else
            throw(ErrorException("Invalid opcode: " * string(opcode)))
        end
    end
end

function process_output(in::Channel, out::Channel)::Int
    score = 0
    screen = zeros(100, 100)
    player, x,y,tile = [0,0,0,0]

    while true
        try
            x = take!(in)
            y = take!(in)
            tile = take!(in)
        catch e
            if isa(e, InvalidStateException)
                break
            else
                rethrow
            end
        end
        println("(", x,", ",y,"): ",tile)
        if x == -1 && y == 0
            score = tile
        else
            screen[x+1, y+1] = tile
        end
        if tile == 3
            player = x
        end
        if tile == 4
            if player == x
                put!(out, 0)
            elseif player < x
                put!(out, 1)
            else
                put!(out, -1)
            end
        end
    end
    return sum(x -> x == 2, screen)
end

program = open("arcade.csv") do f
    parse.(Int, split(readlines(f)[1], ","))
end

in = Channel(1)
out = Channel(2)

program[1] = 2
@async run_intcode(deepcopy(program), in, out)
println(process_output(out, in))


text:
# # Code execution
#
# For documentation sites and blog articles, including code snippets
# in your documents is a great way to connect the theoretical concepts
# with the application. Here we implement [`ExecuteCode`](#), a [`Rewriter`](#)
# that executes marked code blocks and includes the output in the document.
#
# You can use the rewriter by passing it to the [`Project`](#) constructor:
# ```julia
# project = Project(dir, [ExecuteCode()])
# ```
#
# Executing code blocks takes a lot of time, so to give a pleasant interactive
# document creation experience we should make code execution incremental. If a document
# changes, we'll only rerun the code blocks that changed and the ones following them
# (since they may depend on the state of the previous blocks).
#
# Another consideration is that groups of code blocks should be isolated so that they
# don't affect the execution of other groups. We'll assume by default that all code
# blocks marked executable in one document are a group but give the ability to specify
# a different group where needed.
#
#
# ## Cached code execution
#
# Let's focus on the cached code execution before integrating it into the document
# processing system.
#
# We'll implement [`runblockscached`](#), a function that runs a collection of code blocks
# inside a scope and captures the printed output as well as the result of the last code
# line. It caches the results so that only blocks that change have to be rerun.
#
# The cache holds the code, outputs and results for a group of code blocks, as well as
# a module that the code is evaluated in.

mutable struct RunCache
    blocks::Vector{String}
    outputs::Vector{Any}
    results::Vector{Any}
    module_::Module
end

RunCache() = RunCache(String[], [], [], Module())
RunCache(s::Symbol) = RunCache(String[], [], [], Module(s))
Base.show(io::IO, cache::RunCache) = print(io, "RunCache($(nameof(cache.module_)))")

# In case we don't have the result of a block cached, we capture output and result with
# [`runblock`](#).

function runblock(m::Module, codeblock)
    c = IOCapture.capture(rethrow=InterruptException) do
        Base.include_string(m, codeblock)
    end
    return c.output, c.value
end

# We put these two together for [`runblockscached`](#):

function runblockscached(cache::RunCache, blocks)
    outputs = []
    results = []
    updated = false
    for (i, block) in enumerate(blocks)
        oldblock = length(cache.blocks) >= i ? cache.blocks[i] : nothing
        if !updated && oldblock == block
            output, result = cache.outputs[i], cache.results[i]
        else
            output, result = runblock(cache.module_, block)
            updated = true
        end
        push!(outputs, output)
        push!(results, result)
    end
    return RunCache(blocks, outputs, results, cache.module_)
end

## Rewriter
#
# The rewriter for code execution finds code blocks matching a selector,
# groups them, executes them and includes the output and result in a document.
# We also want to warn when an error is thrown which is useful when changes in
# source code break examples.

# The default selector mimicks the behavior of Publish.jl. "pre[lang="julia" cell]
const PUBLISH_CODEBLOCK_SELECTOR = SelectTag(:pre) & SelectAttrEq(:lang, "julia") & SelectHasAttr(:cell)

Base.@kwdef struct ExecuteCode <: Rewriter
    # Execution caches for each group of code blocks
    caches::Dict{Symbol, RunCache} = Dict{Symbol, RunCache}()
    # Whether to log when a code block throws an error
    warnonerror::Bool = true
    # Selector for code blocks to execute
    codeblocksel::Selector = PUBLISH_CODEBLOCK_SELECTOR
    # Function to get a group from a selected block
    groupfn = x -> get(attributes(x), :cell, "main")
    # Lock to avoid crashes
    lock::ReentrantLock = ReentrantLock()
end

# The rewriter acts on the document tree with [´rewritedoc`](#).

function rewritedoc(executecode::ExecuteCode, p, doc)
    blocks = collect(select(doc, executecode.codeblocksel))
    codes = gettext.(blocks)

    groupids = [creategroupid(p, executecode.groupfn(block)) for block in blocks]
    outputs, results = lock(executecode.lock) do
        executegrouped!(executecode.caches, codes, groupids)
    end

    if executecode.warnonerror
        for (i, result) in enumerate(results)
            if result isa LoadError
                @warn "Got evaluation error in code block:\n\n$(codes[i])\n" error=result.error doc=p block=i line=result.line
            end
        end
    end

    newblocks = [
        XNode(
            :div,
            merge(attributes(block), Dict(:class => "cellcontainer")),
            XTree[
                withattributes(block, merge(attributes(block), Dict(:class => "codecell"))),
                get(attributes(block), :output, "true") == "true" ? viewcodeoutput(outputs[i]) : XLeaf(""),
                get(attributes(block), :result, "true") == "true" ? viewcoderesult(results[i]) : XLeaf(""),

            ],
        )
        for (i, block) in enumerate(blocks)
    ]

    return replacemany(doc, newblocks, executecode.codeblocksel)
end

# We use a helper for executing the code blocks in groups and restoring the ordering
# afterwards:

function executegrouped!(caches, codes, groupids)
    is = []
    codesgrouped = Dict{Symbol, Vector{String}}()

    for (i, (gid, code)) in enumerate(zip(groupids, codes))
        groupcodes = get!(codesgrouped, gid, String[])
        push!(groupcodes, code)
        push!(is, (gid, length(groupcodes)))
    end

    for gid in keys(codesgrouped)
        cache = get!(caches, gid, RunCache(gid))
        caches[gid] = runblockscached(cache, codes)
    end

    outputs = [caches[gid].outputs[i] for (gid, i) in is]
    results = [caches[gid].results[i] for (gid, i) in is]

    return outputs, results
end

creategroupid(path, groupname) = Symbol("$(CommonMark.slugify(string(path)))_$groupname")

# The code output and result are tagged with classes "codeoutput" and
# "coderesult" in the output document.

function viewcodeoutput(output::AbstractString)
    if isempty(output)
        return XLeaf("")
    else
        return XNode(
            :pre,
            Dict(:class => "codeoutput"),
            [XNode(:code, [XLeaf(output)])])
    end
end

function viewcoderesult(result::AbstractString)
    return XNode(
        :pre,
        Dict(:class => "coderesult"),
        [XNode(:code, [XLeaf(result)])],
    )
end

viewcoderesult(result::Nothing) = XLeaf("")

function viewcoderesult(result)
    if any([showable(m, result) for m in HTML_MIMES[1:end-1]])
        return XNode(
            :div,
            Dict(:class => "coderesult"),
            [XLeaf(result)],
        )
    else
        return XNode(
            :pre,
            Dict(:class => "coderesult"),
            [XNode(:code, [XLeaf(result)])],
        )
    end
end

# Resetting the rewriter clears the caches:

function reset!(executecode::ExecuteCode)
    foreach(k -> delete!(executecode.caches, k), keys(executecode.caches))
    return
end


text:
"""
Computes the path of consumption and debt for the previously described
complete markets model where exogenous income follows a linear
state space
"""
function complete_ss(β::AbstractFloat,
                     b0::Union{AbstractFloat, Array},
                     x0::Union{AbstractFloat, Array},
                     A::Union{AbstractFloat, Array},
                     C::Union{AbstractFloat, Array},
                     S_y::Union{AbstractFloat, Array},
                     T::Integer=12)

    # Create a linear state space for simulation purposes
    # This adds "b" as a state to the linear state space system
    # so that setting the seed places shocks in same place for
    # both the complete and incomplete markets economy
    # Atilde = vcat(hcat(A, zeros(size(A,1), 1)),
    #               zeros(1, size(A,2) + 1))
    # Ctilde = vcat(C, zeros(1, 1))
    # S_ytilde = hcat(S_y, zeros(1, 1))

    lss = LSS(A, C, S_y, mu_0=x0)

    # Add extra state to initial condition
    # x0 = hcat(x0, 0)

    # Compute the (I - β*A)^{-1}
    rm = inv(eye(size(A, 1)) - β * A)

    # Constant level of consumption
    cbar = (1 - β) * (S_y * rm * x0 - b0)
    c_hist = ones(T) * cbar[1]

    # Debt
    x_hist, y_hist = simulate(lss, T)
    b_hist = (S_y * rm * x_hist - cbar[1] / (1.0 - β))

    return c_hist, vec(b_hist), vec(y_hist), x_hist
end

using PyPlot

N_simul = 150

# Define parameters
α, ρ1, ρ2 = 10.0, 0.9, 0.0
σ = 1.0
# N_simul = 1
# T = N_simul
A = [1.0 0.0 0.0;
       α  ρ1  ρ2;
     0.0 1.0 0.0]
C = [0.0, σ, 0.0]
S_y = [1.0 1.0 0.0]
β, b0 = 0.95, -10.0
x0 = [1.0, α / (1 - ρ1), α / (1 - ρ1)]

# Do simulation for complete markets
s = rand(1:10000)
srand(s)  # Seeds get set the same for both economies
out = complete_ss(β, b0, x0, A, C, S_y, 150)
c_hist_com, b_hist_com, y_hist_com, x_hist_com = out

fig, ax = subplots(1, 2, figsize = (15, 5))

# Consumption plots
ax[1][:set_title]("Cons and income", fontsize = 17)
ax[1][:plot](1:N_simul, c_hist_com, label = "consumption", lw = 3)
ax[1][:plot](1:N_simul, y_hist_com, label = "income",
            lw = 2, alpha = 0.6, linestyle = "--")
ax[1][:legend](loc = "best", fontsize = 15)
ax[1][:set_xlabel]("Periods", fontsize = 13)
ax[1][:set_ylim]([-5.0, 110])

# Debt plots
ax[2][:set_title]("Debt and income", fontsize = 17)
ax[2][:plot](1:N_simul, b_hist_com, label = "debt", lw = 2)
ax[2][:plot](1:N_simul, y_hist_com, label = "Income",
            lw = 2, alpha = 0.6, linestyle = "--")
ax[2][:legend](loc = "best", fontsize = 15)
ax[2][:axhline](0, color = "k", lw = 1)
ax[2][:set_xlabel]("Periods", fontsize = 13)


text:
using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
Random.seed!(233)

sk = load_op_and_grad("./build/libSinkhornKnopp","sinkhorn_knopp", multiple=true)

# TODO: specify your input parameters
a = constant(ones(10)/10)
b = constant(ones(20)/20)
m = constant(rand(10,20))
reg = constant(1.0)
iter = constant(10000, dtype=Int64)
tol = constant(1e-10)
method = constant(1);
u = sk(a,b,m,reg,iter,tol,method)
sess = tf.Session()
init(sess)
M, _ = run(sess, u)
@show sum(M, dims=1), sum(M, dims=2)
# error()

# TODO: change your test parameter to `m`
# gradient check -- v
function scalar_function(m)
    return sk(a,b,m,reg,iter,tol,method)[2]
end

# TODO: change `m_` and `v_` to appropriate values
G = rand(10,20)
m_ = constant(G)
v_ = rand(10,20)
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session()
init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


text:
"""
    Named(name, block)

Wrapper `Block` to attach a name to a block. Can be used in conjunction
with [`Only`](#) to apply encodings to specific blocks only.
"""
struct Named{Name,B<:AbstractBlock} <: WrapperBlock
    block::B
end
Named(name::Symbol, block::B) where {B<:AbstractBlock} = Named{name,B}(block)

setwrapped(named::Named{name}, block) where name = Named(name, block)

# The name is preserved through encodings and decodings, which is the
# behavior of `propagatewrapper(::W) = PropagateAlways()` so it does
# not need to be overwritten

"""
    Only(fn, encoding)
    Only(BlockType, encoding)
    Only(name, encoding)

Wrapper that applies the wrapped `encoding` to a `block` if
`fn(block) === true`. Instead of a function you can also pass in
a type of block `BlockType` or the `name` of a `Named` block.
"""
struct Only{E<:Encoding} <: StatefulEncoding
    fn::Any
    encoding::E
end

function Only(name::Symbol, encoding::Encoding)
    return Only(Named{name}, encoding)
end

function Only(B::Type{<:AbstractBlock}, encoding::Encoding)
    return Only(block -> block isa B, encoding)
end

encodedblock(only::Only, block::Block) =
    only.fn(block) ? encodedblock(only.encoding, block) : nothing
encodedblock(only::Only, block::WrapperBlock) =
    only.fn(block) ? encodedblock(only.encoding, block) : nothing
encodedblock(only::Only, block::Named) =
    only.fn(block) ? encodedblock(only.encoding, block) : nothing

function decodedblock(only::Only, block::Block)
    inblock = decodedblock(only.encoding, block)
    only.fn(inblock) || return nothing
    return inblock
end
function decodedblock(only::Only, block::WrapperBlock)
    inblock = decodedblock(only.encoding, block)
    only.fn(inblock) || return nothing
    return inblock
end
function decodedblock(only::Only, block::Named)
    inblock = decodedblock(only.encoding, block)
    only.fn(inblock) || return nothing
    return inblock
end

encodestate(only::Only, args...) = encodestate(only.encoding, args...)
decodestate(only::Only, args...) = decodestate(only.encoding, args...)

function encode(only::Only, ctx, block::Block, obs; kwargs...)
    _encode(only, ctx, block, obs; kwargs...)
end
function encode(only::Only, ctx, block::WrapperBlock, obs; kwargs...)
    _encode(only, ctx, block, obs; kwargs...)
end
function _encode(only, ctx, block, obs; kwargs...)
    return only.fn(block) ? encode(only.encoding, ctx, block, obs; kwargs...) : obs
end

function decode(only::Only, ctx, block::Block, obs; kwargs...)
    _decode(only, ctx, block, obs; kwargs...)
end
function decode(only::Only, ctx, block::WrapperBlock, obs; kwargs...)
    _decode(only, ctx, block, obs; kwargs...)
end
function _decode(only, ctx, block, obs; kwargs...)
    return only.fn(decodedblock(only.encoding, block)) ?
           decode(only.encoding, ctx, block, obs; kwargs...) : obs
end

# ## Test

InlineTest.@testset "Only [block]" begin
    encx = Only(:x, OneHot())
    inblock = Label(1:100)
    inblocknamed = Named(:x, inblock)
    obs = mockblock(inblock)
    encobs = encode(OneHot(), Training(), inblock, obs)

    @test encodedblock(encx, inblock) === nothing
    @test encodedblock(encx, inblocknamed) isa Named{:x}
    @test encodedblock(Only(Named, OneHot()), inblock) === nothing
    @test encodedblock(Only(Named, OneHot()), inblocknamed) isa Named{:x}

    outblock = encodedblock(OneHot(), inblock)
    outblocknamed = encodedblock(OneHot(), inblocknamed)
    @test decodedblock(encx, outblock) === nothing
    @test decodedblock(encx, outblocknamed) isa Named{:x}
    @test decodedblock(Only(Named, OneHot()), outblock) === nothing
    @test decodedblock(Only(Named, OneHot()), outblocknamed) isa Named{:x}

    @test encode(encx, Training(), inblock, obs) == obs
    @test encode(encx, Training(), inblocknamed, obs) != obs

    @test decode(encx, Training(), outblock, encobs) == encobs
    @test decode(encx, Training(), outblocknamed, encobs) != encobs

    tfm = OneHot()
    only = Only(:name, tfm)
    block = Named(:name, Label(["cat", "dog"]))
    obs = mockblock(block)
    testencoding(only, block, obs)
    testencoding(only, (block, wrapped(block)), (obs, obs))
    @test encodedblock(only, wrapped(block)) isa Nothing
end


text:
# checking the soundness of overapproximations for the benchmark problems
using OVERT

##################### MODE 3: BUILD UP PROOFS OF EACH LEVEL OF THE APPROX

## works but not needed
# function special_compare(v1, v2)
#     # returns TRUE is v1 <= v2
#     # for variables of the form: :v_1, :v_2, etc.
#     @debug("special_compare: ", v1, v2)
#     v1num = Meta.parse(match(r"(?<=_)(.*)", string(v1)).captures[1])
#     v2num = Meta.parse(match(r"(?<=_)(.*)", string(v2)).captures[1])
#     return v1num <= v2num
# end

function get_free_vars(e::Expr)
    return Symbol.(free_symbols(Basic(e)))
end

function run_soundness_query(ϕ, ϕ̂, defs, domain, problem_name, delta_sat, jobs)
    println("ϕ = $ϕ, ϕ̂ = $ϕ̂, all defs: $defs")
    sq = SoundnessQuery(ϕ, # ϕ
                            defs,
                            ϕ̂, # definitions for ϕ̂
                            domain)
        # println("Soundness Query: $sq")

        result = check("dreal", sq, problem_name*"_soundness_query.smt2", δ=delta_sat, jobs=jobs) # TODO: pass dreal delta
        println("result for var ",problem_name," is: ", result)
        return occursin("unsat", result)
end

import Base.setdiff
function setdiff(s::Set{Symbol}, e::Symbol)
    elements = [s...]
    return Set(filter(i -> i != e, elements))
end

function get_approx_equality_def(var, approx_equality_list)
    defs = filter(c -> c.args[2] == var, approx_equality_list)
    @assert(length(defs) == 1)
    return defs[1]
end

function assert_ineq_bounded(cur_var_bounds, current_var)
    println("current_var=$(current_var)")
    println("cur_var_bounds = $(cur_var_bounds)")
    # assert that the bounds sandwhich the current_var e.g. v1 <= current_var <= v2
    @assert length(cur_var_bounds) == 2
    b1 = cur_var_bounds[1]
    b2 = cur_var_bounds[2]
    mixed_set = Set([:(>=),:(<=)])
    if Set([b1.args[1], b2.args[1]]) == mixed_set # if opposite signs
        # current_var must be on same side of ineqs to be sandwhiched
        @assert (b1.args[2] == current_var && b2.args[2] == current_var) || (b1.args[3] == current_var && b2.args[3] == current_var)
    elseif (b1.args[1] ∈ mixed_set) && (b2.args[1] ∈ mixed_set) && (b1.args[1] == b2.args[1]) # share same ineq sense
        # current var must be on dif sides of ineq to be sandwhiched
        @assert (b1.args[2] == current_var && b2.args[3] == current_var) || (b1.args[3] == current_var && b2.args[2] == current_var)
    else
        @assert false
    end
end

# Overall plan: do this in a more structured way.
# recurse down the fun_eq tree
# and then collect the approx pieces one by one
# e.g. first check if a variable definition exists in approx_eq. if not, then look in approx_ineq
# make sure all variables are defined before generating smt2 file
    # begin the recursion from the output variable
    # output = oa.output
function __check_overapprox(current_var, defs, oa, domain, input_vars, problem_name; jobs=1, delta_sat=0.001)
    #println("\n current_var = $(current_var), \n defs = $defs")
    # base case: IS an input var, not v_i variable
    if current_var ∈ input_vars
        return [], true
    # if not base case
    else
        # check_overapprox of free_vars in RHS expression --> should return defs for each one, merge these into existing defs
        free_vars = Set(get_free_vars(oa.fun_eq[current_var])) # get free vars in exact expression
        #println("free_vars = $(free_vars)")
        # get dependencies of each free var (and also check proof goals of dependencies)
        data = [__check_overapprox(v, [], oa, domain, input_vars, problem_name; jobs=jobs, delta_sat=delta_sat) for v in free_vars]
        # check that each free var is either bounded by two inequalities of the right signs or one equality
        #println("current_var again= $(current_var)")
        dependencies = vcat([d[1] for d in data]...)
        defs = vcat(defs, dependencies)
        results = [d[2] for d in data]
        is_unsat = all(results)
        # if v_2 affine, add to defs and return
        if is_affine(oa.fun_eq[current_var])
            defs = vcat(defs, :($(current_var) == $(oa.fun_eq[current_var])))
        else
        # if v_2 not affine, check for bounds for v2 in oa.approx_ineq, gen proof goal...
            cur_var_bounds = filter(e -> current_var ∈ get_free_vars(e), oa.approx_ineq) # get the two inequalities: v2 <= v3 and v3 <= v4
            assert_ineq_bounded(cur_var_bounds, current_var)
            #println("cur_var_bounds = $(cur_var_bounds)")
            bound_vars = setdiff(Set(vcat(get_free_vars.(cur_var_bounds)...)), current_var) # get the new vars, e.g. v2, v4
            #println("bound_vars = $(bound_vars)")
            bound_var_defs = [get_approx_equality_def(v, oa.approx_eq) for v in bound_vars] # get the definitions for the new vars v2 = ... and v4 = ...
            defs = vcat(defs, bound_var_defs) # add to defs
            # construct proof goal and run
            ϕ = [:($(current_var) == $(oa.fun_eq[current_var]))]
            ϕ̂ = cur_var_bounds
            problem_specific_name = problem_name*"_$(current_var)_"
            is_unsat &= run_soundness_query(ϕ, ϕ̂, defs, domain, problem_specific_name, delta_sat, jobs)
            # then move ineq bounds to defs and return
            defs = vcat(defs, ϕ̂)
        end
    end
    return defs, is_unsat
end

function replace_unicode(oa::OverApproximation)
    # why did we use ≦ in the oa ??? This removes it and replaces with <= or >=
    oa_approx_ineq_fixed = [( (c.args[1] == :≤) || (c.args[1] == :≦) ) ? :($(c.args[2]) <= $(c.args[3]))  : c for c in oa.approx_ineq ]
    oa_approx_ineq_double_fixed = [( (c.args[1] == :≥) || (c.args[1] == :≧) ) ? :($(c.args[2]) >= $(c.args[3]))  : c for c in oa_approx_ineq_fixed]
    oa.approx_ineq = oa_approx_ineq_double_fixed
    return oa
end

function check_overapprox(oa, domain, input_vars, problem_name; jobs=1, delta_sat=0.001)
    defs = []
    oa_clean = replace_unicode(oa)
    defs, result = __check_overapprox(oa.output, defs, oa_clean, domain, input_vars, problem_name; jobs=jobs, delta_sat=delta_sat)
    result ? println("All checks pass!") : println("Some checks fail :( ")
    return result
end


text:
function simulate!(s::Simulator{T}; verbose=false) where T
    status = false

    p = s.policy
	w = s.dist
	traj = s.traj
	N = length(traj.u)

	# last_computed_control = p.ref_traj.u[1] / N_sample
	buffer_time = 0.0
    for t = 1:N
		# println("timestep $t / $N, time $time")
		policy_time = @elapsed traj.u[t] = exec_policy(p, traj, t, buffer_time)
		s.opts.record && (s.stats.policy_time[t] = policy_time)
		(buffer_time < 0.0) && (buffer_time = policy_time)
		buffer_time -= s.h # decrement the buffer time of the sim step

		# traj.u[t] .= last_computed_control
		# if buffer_time <= 0.0
			# policy
			# buffer_time = @elapsed last_computed_control = exec_policy(p, traj, t, buffer_time)
			# s.opts.record && (s.stats.policy_time[t] = buffer_time)
        # end
		# buffer_time -= s.h # decrement the buffer time of the sim step

        # disturbances
        traj.w[t] .= disturbances(w, traj.q[t+1], t)

        # step
        status = RoboDojo.step!(s, t, verbose=verbose)
        !status && break
    end

    return status
end

function exec_policy(p::CIMPC{T,NQ,NU,NW,NC}, traj::Trajectory{T}, t::Int, buffer_time::T) where {T,NQ,NU,NW,NC}
	# reset
	if t == 1
		p.cnt[1] = p.N_sample
		p.q0 .= p.ref_traj.q[1]
		p.altitude .= 0.0
		set_trajectory!(p.traj, p.ref_traj)
		set_implicit_trajectory!(p.im_traj, p.im_traj_cache)
		reset_window!(p.window)
	end

    if p.cnt[1] == p.N_sample
		# altitude_update
		(p.opts.altitude_update && t > 1) && (update_altitude!(p.altitude, p.ϕ, p.s,
									traj, t, NC, p.N_sample,
									threshold = p.opts.altitude_impact_threshold,
									verbose = p.opts.altitude_verbose))
		set_altitude!(p.im_traj, p.altitude)

		# # # optimize
		# q1 = traj.q[t+1]
		# newton_solve!(p.newton, p.s, p.q0, q1,
		# 	p.window, p.im_traj, p.traj, warm_start = t > 1)

		update!(p.im_traj, p.traj, p.s, p.altitude, p.κ[1], p.traj.H)

		# visualize
		p.opts.live_plotting && live_plotting(p.s.model, p.traj, traj, p.newton, p.q0, traj.q[t+1], t)

		# shift trajectory
		rot_n_stride!(p.traj, p.traj_cache, p.stride, p.window)

		# update
		update_window!(p.window, p.ref_traj.H)
		# p.q0 .= q1

		# reset count
		p.cnt[1] = 0
    end
	p.cnt[1] += 1

	if buffer_time <= 0.0
		@show t
		# optimize
		q0 = traj.q[max(1,t+1-N_sample)]
		q1 = traj.q[t+1]
		newton_solve!(p.newton, p.s, q0, q1,
			p.window, p.im_traj, p.traj, warm_start = t > 1)

		update!(p.im_traj, p.traj, p.s, p.altitude, p.κ[1], p.traj.H)
	end

	# scale control
	if p.newton_mode == :direct
		p.u .= p.newton.traj.u[1]
		p.u ./= p.N_sample
	elseif p.newton_mode == :structure
		p.u .= p.newton.u[1]
		p.u ./= p.N_sample
	else
		println("newton mode specified not available")
	end

	return p.u
end


text:
function copy_nthreads(C, A; nthreads, thread_indices, kwargs...)
    @threads :static for tid in 1:nthreads
        @inbounds for i in thread_indices[tid]
            C[i] = A[i]
        end
    end
    return nothing
end

function scale_nthreads(B, C, s; nthreads, thread_indices, kwargs...)
    @threads :static for tid in 1:nthreads
        @inbounds for i in thread_indices[tid]
            B[i] = s * C[i]
        end
    end
    return nothing
end

function add_nthreads(C, A, B; nthreads, thread_indices, kwargs...)
    @threads :static for tid in 1:nthreads
        @inbounds for i in thread_indices[tid]
            C[i] = A[i] + B[i]
        end
    end
    return nothing
end

function triad_nthreads(A, B, C, s; nthreads, thread_indices, kwargs...)
    @threads :static for tid in 1:nthreads
        @inbounds for i in thread_indices[tid]
            A[i] = B[i] + s * C[i]
        end
    end
    return nothing
end


text:
"""
    timecons(data)

Compute a sparse matrix indicating schedule conflicts between courses in a specified program.

# Arguments

- `data` : a dataframe containing the names of the program, list of courses in those programs and
the time slot at which those courses take place.

# Example

The CSV files used for this example are stored in the "docs" directory of the repo.
We assume that there are 3 programs, Master in Economics ("MIE"), Law degree ("Law")
and Sociology degree ("Socio").

The Students.csv files contains the list of students and indicates the program they are
currently enrolled in and if they are in first or second year.

The Programs.csv files contains the detailed list of classes offer in a program. Different types
of classes are available in each program. The "FC" prefix stands for "Formation Commune". The "TC"
prefix stands for "Tronc Commun" and the "OP" prefix stands for "Optional class".
The first number in the name of the course indicates to which students the course is proposed. The
second number is the course ID.
For instance "TC11" is the first "Tronc Commun" offered to first year student, and "OP23" is the third
"Optional course" offeref to second year students.
I assumed that all formations communes were open to all students (which is the case in Sciences Po), which
justifies the use of only one number in the course ID. For instance "FC2" is the second "Formation Commune"
offered to all students.

The "slot" column indicates the day and slot number at which the course takes place. The days are indicated
by the letter and the time slot by the number. I assumed that there were 5 slots in a day.
Note that Sciences Po schedule is made such that no course slot overlaps another. All slots are independant.
For instance "L1" stands for the first slot on Monday, and "R5" for the fifth slot on Tuesday.

"""
function timecons(data)

    d = Dict()
    programs = unique(data[:Program])

    for p in programs, s in 1:2
        df = data[findin(data[:Program],[p]),:]
        df1 = df[findin(df[:ClassS], [s, 12]),:]
        P = fill(0, nrow(df1), nrow(df1))
        for i in 1:(length(df1)-1)
            if df1[:slot][i] == df1[:slot][i+1]
                P[i, i+1] = 1
            end
        end
        d["$p$s"] = sparse(P)
    end
    return d
end

"""
# constraints

* `time`: a sparse matrix C x C that flags with 1 if the ith element HAS a clash with jth element. 0 otherwise.
* `mandatory`: a vector C x 1. the i-th element is 1 if the course must be taken, 0 otherwise
* `notTCprogram`: a vector C x 1. the i-th element is 1 if the course is NOT in the TC of the students's program
* `notTCsemester`: a vector C x 1. the i-th element is 1 if the course is NOT in the TC of the students's semester
* `isTC`: a vector C x 1. the i-th element is 1 if the course IS in the TC of the students
* `notFCprogram`: a vector C x 1. the i-th element is 1 if the course is NOT in the FC of the students's program
* `notFCsemester`: a vector C x 1. the i-th element is 1 if the course is NOT in the FC of the students's semester
* `isFC`: a vector C x 1. the i-th element is 1 if the course IS in the FC of the students
* `notELprogram`: a vector C x 1. the i-th element is 1 if the course is NOT in the EL (electives) of the students's program
* `notELsemester`: a vector C x 1. the i-th element is 1 if the course is NOT in the EL of the students's semester
* `isEL`: a vector C x 1. the i-th element is 1 if the course IS in the EL of the students

"""
struct Constraints
    clash         :: SparseMatrixCSC{Int,Int}
    mandatory     :: Vector{Bool}
    notTCprogram  :: Vector{Bool}
    notTCsemester :: Vector{Bool}
    isTC          :: Vector{Bool}
    notFCprogram  :: Vector{Bool}
    notFCsemester :: Vector{Bool}
    isFC          :: Vector{Bool}
    notELprogram  :: Vector{Bool}
    notELsemester :: Vector{Bool}
    isEL          :: Vector{Bool}
end

"random constraints"
function rand_constraints(N)
    if N<2
        error("N must be > 2")
    end
    clash = sparse(collect(1:2),collect(1:2),ones(Int,2),N,N)
    mandatory = rand([true false],N)
    notTCprogram = rand([true false false false false false],N)
    notTCsemester = rand([true false false false false false],N)
    isTC = rand([true false],N)
    notFCprogram = rand([true false false false false false],N)
    notFCsemester = rand([true false false false false false],N)
    isFC = rand([true false],N)
    notELprogram = rand([true false false false false false],N)
    notELsemester = rand([true false false false false false],N)
    isEL = rand([true false],N)
    return Constraints(clash,mandatory,notTCprogram,notTCsemester,isTC,notFCprogram,notFCsemester,isFC,notELprogram,notELsemester,isEL)
end

"""
# Student

## Fields

- `price` : a column vector of dimension C x 1, which i-th element is the price that was assigned to class i.
- `preferences` : an array containing S elements, and such that its n-th element is the (sparse) matrix representing student n preferences. Each of the matrices contained in that array should be a squared matrix of dimension C.
- `budget` : a column vector of dimension N x 1, which n-th element is the budget that was allocated to the n-th student.
- `constraints`: an instance of a [`Constraints`](@ref)
- `req_TC` : number of TC courses that the student is required to take
- `req_FC` : number of FC courses that the student is required to take
- `req_EL` : number of EL courses that the student is required to take
"""
mutable struct Student
    preferences :: SparseMatrixCSC{Int64, Int64}
    program :: String
    year :: Int64
    constraints :: Constraints
    req_FC :: Int64
    req_TC :: Int64
    req_EL :: Int64
    budget :: Float64
    allocation :: Array{Int64,1}

    function Student(N=10;preferences=spzeros(N,N),
                     program="MIE",
                     year::Int=1,
                     constraints=Constraints(N),
                     req_FC::Int=1,
                     req_TC::Int=2,
                     req_EL::Int=1,
                     budget = rand()*100)
        new(
            preferences,
            program,
            year,
            constraints,
            req_FC,
            req_TC,
            req_EL,
            budget,
            zeros(N))
    end
end

"""
    read_students()

Construct `Student`s from artificial time sheet data.

# Example

julia> CourseMatch.read_students()
Dict{Int64,CourseMatch.Student} with 120 entries:
  68  => CourseMatch.Student("Law", 2, …
  2   => CourseMatch.Student("MIE", 1, …
  89  => CourseMatch.Student("Socio", 1, …
  11  => CourseMatch.Student("MIE", 1, …
  39  => CourseMatch.Student("MIE", 2, …
  46  => CourseMatch.Student("Law", 1, …
  85  => CourseMatch.Student("Socio", 1, …
  25  => CourseMatch.Student("MIE", 2, …
  55  => CourseMatch.Student("Law", 1, …
  42  => CourseMatch.Student("Law", 1, …
  29  => CourseMatch.Student("MIE", 2, …
  58  => CourseMatch.Student("Law", 1, …
  66  => CourseMatch.Student("Law", 2, …
  59  => CourseMatch.Student("Law", 1, …
  8   => CourseMatch.Student("MIE", 1, …
  ⋮   => ⋮
"""
function read_students()
	data = CSV.read(joinpath(dirname(@__FILE__),"..","docs","Students.csv") ; delim = ";")
	ptable = CSV.read(joinpath(dirname(@__FILE__),"..","docs","Programs.csv") ; delim = ";")
	dictp = timecons(ptable)
    d = Dict{Int,Student}()
    for i in 1:nrow(data)
        s = data[:ID][i]
        d[s] = Student(sparse(collect(1:10), collect(1:10), rand(0:100, 10)),
                    data[:ProgSem][i],
                    data[:ProgSem][i][end],
                    dictp[p],
                    #progcons,
                    1,
                    rand(1:2),
                    3,
                    150+rand(),
                    fill(0, size(dictp[p])[1])
                    )
    end
    return d
end

max_stars_repo_name: ScPo-CompEcon/CourseMatch

max_issues_repo_name: ScPo-CompEcon/CourseMatch

max_forks_repo_name: ScPo-CompEcon/CourseMatch


text:
using Documenter, SphericalGeometry

makedocs(
    sitename = "SphericalGeometry.jl",
    modules = [SphericalGeometry],
    pages = Any[
        "Home" => "index.md",
    ],
)

deploydocs(
    repo = "github.com/rjdverbeek-tud/SphericalGeometry.jl.git",
)


text:
#****************************************************************************
#                     Molecular Dynamics Potentials (MDP)
#                            CESMIX-MIT Project
#
#  Contributing authors: Ngoc-Cuong Nguyen (cuongng@mit.edu, exapde@gmail.com)
#****************************************************************************

__precompile__()

module Potential

using Revise, LinearAlgebra

export empiricalpotential, snappotential, lammpspotential, addpotential, getpotential, initsna, PotentialStruct
export empiricaldescriptors, snapdescriptors, lammpssnapdescriptors, emldescriptors
export boundingbox, domainmaps, atomlist, neighborlist, fullneighborlist, neighsingles, accumarray
export neighpairlist, neighpairs, neightripletlist, neightriplets, neighquadletlist, neighquadlets
export tallysingle, tallypair, tallytriplet, tallyquadlet, findatomtype, periodicimages

function latticevolume(a::Vector{Float64}, b::Vector{Float64}, c::Vector{Float64})
    vol = dot(a, cross(b, c))
    return vol
end

# function getpotentials(potentials)

#     for i = length(potentials):-1:1
#         if potentials[i] === nothing
#             deleteat!(potentials, i)
#         end
#     end

#     n = length(potentials)
#     empotentials = Array{Any}(nothing, n)
#     mlpotentials = nothing;

#     n = 0
#     for i = 1:length(potentials)
#         if typeof(potentials[i]) == Potential.PotentialStruct
#             n = n + 1
#             empotentials[n] = potentials[i]
#         else
#             mlpotentials = potentials[i]
#         end
#     end

#     if n>0
#         if length(potentials) > n
#             for i = length(potentials):-1:(n+1)
#                 deleteat!(empotentials, i)
#             end
#         end
#     else
#         empotentials = []
#     end

#     return empotentials, mlpotentials
# end

include("accumarray.jl");
include("atomlist.jl");
include("boundingbox.jl");
include("domainmaps.jl");
include("findatomtype.jl");
include("periodicimages.jl");
include("prebox.jl");
include("potentialstruct.jl");
include("latticecoords.jl");
include("neighborlist.jl");
include("fullneighborlist.jl");
include("neighsingles.jl");
include("neighpairlist.jl");
include("neighpairs.jl");
include("neightripletlist.jl");
include("neightriplets.jl");
include("neighquadletlist.jl");
include("neighquadlets.jl");
include("empiricalpotential.jl");
include("empiricaldescriptors.jl");
include("tally.jl");
include("snastruct.jl");
include("snapcpp.jl");
include("lammpspotential.jl");
include("lammpssnapdescriptors.jl");
include("emlpotential.jl");
include("emldescriptors.jl");

end

max_stars_repo_name: cesmix-mit/MDP

max_issues_repo_name: cesmix-mit/MDP

max_forks_repo_name: cesmix-mit/MDP


text:
function _lll_gram(A::NfOrdIdl)
  K = nf(order(A))
  @assert istotally_real(K)
  g = trace_matrix(A)
  @assert det(g) != 0
  @assert isposdef(g)
  l, t = lll_gram_with_transform(g)
  return FakeFmpqMat(l, fmpz(1)), t::fmpz_mat
end

function lll_basis(A::NfOrdIdl, v::fmpz_mat = zero_matrix(FlintZZ, 1, 1); prec::Int = 100)
  L, T = lll(A, v, prec=prec)
  S = FakeFmpqMat(T)*basis_mat(A)*basis_mat(order(A))
  K = nf(order(A))
  q = nf_elem[elem_from_mat_row(K, numerator(S), i, denominator(S)) for i=1:degree(K)]
  return q
end

function cmpindex(A::fmpz_mat, i::Int, j::Int, b::fmpz)
  a = ccall((:fmpz_mat_entry, :libflint), Ptr{fmpz}, (Ref{fmpz_mat}, Int, Int), A, i-1, j-1)
  return ccall((:fmpz_cmp, :libflint), Int32, (Ptr{fmpz}, Ref{fmpz}), a, b)
end

function prod_diag(A::fmpz_mat)
  a = fmpz()
  for i=1:rows(A)
    b = ccall((:fmpz_mat_entry, :libflint), Ptr{fmpz}, (Ref{fmpz_mat}, Int, Int), A, i-1, i-1)
    ccall((:fmpz_mul, :libflint), Void, (Ref{fmpz}, Ref{fmpz}, Ptr{fmpz}), a, a, b)
  end
  return a
end

function lll(A::NfOrdIdl, v::fmpz_mat = zero_matrix(FlintZZ, 1, 1); prec::Int = 100)

  K = nf(order(A))
  if iszero(v) && istotally_real(K)
    #in this case the gram-matrix of the minkowski lattice is the trace-matrix
    #which is exact.
    return _lll_gram(A)
  end

  c = minkowski_mat(nf(order(A)), prec) ## careful: current iteration
                                        ## c is NOT a copy, so don't change.
  l, t1 = lll_with_transform(basis_mat(A, Val{false}))
  b = FakeFmpqMat(l)*basis_mat(order(A), Val{false})

  n = degree(order(A))

  rt_c = roots_ctx(K)
  if !isdefined(rt_c, :cache_z1)
    rt_c.cache_z1 = zero_matrix(FlintZZ, n, n)
    rt_c.cache_z2 = zero_matrix(FlintZZ, n, n)
  end

  d::fmpz_mat = rt_c.cache_z1
  g::fmpz_mat = rt_c.cache_z2

  sv = fmpz(0)
  if !iszero(v)
    @v_do :ClassGroup 2 println("using inf val", v)
    c = deepcopy(c)
    mult_by_2pow_diag!(c, v)
    sv = max(fmpz(0), sum(v[1,i] for i=1:cols(l)))
  end

  round_scale!(d, c, prec)
  ccall((:fmpz_mat_mul, :libflint), Void, (Ref{fmpz_mat}, Ref{fmpz_mat},  Ref{fmpz_mat}), g, (b.num), d)
  den = b.den

  ccall((:fmpz_mat_gram, :libflint), Void, (Ref{fmpz_mat}, Ref{fmpz_mat}), d, g)
  shift!(d, -prec)
  prec = div(prec, 2)
  shift!(d, -prec)  #TODO: remove?

  for i=1:n
    fmpz_mat_entry_add_ui!(d, i, i, UInt(rows(d)))
  end

  ctx=Nemo.lll_ctx(0.99, 0.51, :gram)

  ccall((:fmpz_mat_one, :libflint), Void, (Ref{fmpz_mat}, ), g)
  ccall((:fmpz_lll, :libflint), Void, (Ref{fmpz_mat}, Ref{fmpz_mat}, Ref{Nemo.lll_ctx}), d, g, ctx)

  l, t = d, g
  ## test if entries in l are small enough, if not: increase precision
  ## or signal that prec was too low
  @v_do :ClassGroup 3 print_with_color(:green, "lll basis length profile\n");
  @v_do :ClassGroup 3 for i=1:rows(l)
    print(div(l[i,i], fmpz(2)^prec), " : ")
  end
  @v_do :ClassGroup 3 println("")
  if nbits(maximum(t)) >  div(prec, 2)
    @v_do :ClassGroup 2 print_with_color(:red, "lll trafo too large\n");
    throw(LowPrecisionLLL())
  end
  ## lattice has lattice disc = order_disc * norm^2
  ## lll needs to yield a basis sth
  ## l[1,1] = |b_i|^2 <= 2^((n-1)/2) disc^(1/n)
  ## and prod(l[i,i]) <= 2^(n(n-1)/2) disc
  n = rows(l)
  disc = abs(discriminant(order(A)))*norm(A)^2 * den^(2*n) * fmpz(2)^(2*sv)
  di = root(disc, n)+1
  di *= fmpz(2)^(div(n+1,2)) * fmpz(2)^prec
  pr = fmpz(1)
  if cmpindex(l, 1, 1, di) > 0
    @v_do :ClassGroup 2 print_with_color(:red, "LLL basis too large\n");
    @v_do :ClassGroup 3 println("bound is ", di, " value at ", 1, " is ", l[1,1]);
    throw(LowPrecisionLLL())
  end
  pr = prod_diag(l)
  if pr > fmpz(2)^(div(n*(n-1), 2)) * disc * fmpz(2)^(n*prec)
    @v_do :ClassGroup 2 print_with_color(:red, "LLL basis too large\n");
    @v_do :ClassGroup 2 println("prod too large: ", pr, " > 2^(n(n-1)/2) disc = ", fmpz(2)^(div(n*(n-1), 2)) * disc * fmpz(2)^(n*prec));
    throw(LowPrecisionLLL())
  end

  return FakeFmpqMat(deepcopy(l), fmpz(2)^prec), t*t1
end

function short_elem(A::NfOrdFracIdl,
                v::fmpz_mat=zero_matrix(FlintZZ, 1,1); prec::Int = 100)
  return divexact(short_elem(A.num, v, prec = prec), A.den)
end
function _short_elem(A::NfOrdFracIdl,
                v::fmpz_mat=zero_matrix(FlintZZ, 1, 1))
  return divexact(_short_elem(A.num, v), A.den)
end

function short_elem(A::NfOrdIdl,
                v::fmpz_mat = zero_matrix(FlintZZ, 1,1); prec::Int = 100)
  K = nf(order(A))
  temp = FakeFmpqMat(basis_mat(A))*basis_mat(order(A))
  b = temp.num
  b_den = temp.den
  l, t = lll(A, v, prec = prec)
  w = view(t, 1,1, 1, cols(t))
  c = w*b
  q = elem_from_mat_row(K, c, 1, b_den)
  return q
end

function _short_elem(A::NfOrdIdl,
                v::fmpz_mat = zero_matrix(FlintZZ, 1, 1))
  p = 64
  while true
    if p > 40000
      error("Something wrong in reduce_ideal")
    end
    try
      b = short_elem(A, v, prec = p)
      return b
    catch e
      if e isa LowPrecisionLLL
        p = 2*p
      else
        rethrow(e)
      end
    end
  end
end

function reduce_ideal(A::NfOrdIdl)
  B = inv(A)
  success = false
  b = _short_elem(B)
  C = b*A
  simplify(C)
  @assert C.den == 1
  return C.num
end

function reduce_ideal(A::NfOrdFracIdl)
  B = inv(A)
  b = _short_elem(B.num)
  C = divexact(b, B.den)*A
  simplify(C)
  @assert C.den == 1
  return C.num
end

function reduce_ideal2(A::NfOrdIdl)
  B = inv(A)
  b = _short_elem(B)
  C = b*A
  simplify(C)
  @assert C.den == 1
  return C.num, b
end