text:
using wgregseq, FASTX, DataFrames, CSV, BioSequences, CairoMakie, Statistics, StatsBase, Dates

# Set path
dir = @__DIR__
home_dir = joinpath(split(dir, "/")[1:end-2])

# Import genome
re = open(FASTA.Reader, "/$home_dir/data/ecocyc/mg1655_genome.fasta")
wt_sequence = [sequence(record) for record in re][1]

##
println("Importing Data...")
# Import gene list to generate sequences
gene_table = CSV.read(
    "../data/100_genes.csv",
    DataFrames.DataFrame,
    delim=",",
    comment="#",
    missingstring="none",
)

# Give IDs to groups for adding primers later
group_dict = Dict{String, Int64}(filter(x -> occursin("Antibiotic/toxin", x), gene_table.group) .=> 1)
group_dict["Gold Standard"] = 2
group_dict["Heinemann dataset"] = 3
group_dict["uncharacterized protein"] = 4
group_dict["Heinemann dataset uncharacterized"] = 7
group_dict["YmfT_modulon"] = 5
group_dict["YgeV_modulon"] = 6
groups = zeros(nrow(gene_table))
for i in 1:nrow(gene_table)
    if gene_table[i, "group"] in keys(group_dict)
        groups[i] = group_dict[gene_table[i, "group"]]
    else
        groups[i] = 8
    end
end

# Add IDs to table
insertcols!(gene_table, 4, :group_ID => groups)

# Print number of genes per group
println(combine(groupby(gene_table, "group"), nrow))

## Import promoters
# Import promoter list and infer types that can be infered automatically
promoter_list = CSV.read(
    "/$home_dir/data/promoter_list_processed.csv",
    DataFrames.DataFrame,
    types=Dict(
        "promoter"=>String,
        "tss"=>Float64,
        "direction"=>String
    )
)

operons_without_promoters = CSV.read(
    "/$home_dir/data/operons_without_promoters.csv",
    DataFrames.DataFrame,
    types=Dict(
        "direction"=>String
    )
)

# Replace columns by nicer types
promoter_list.genes = parse.(Vector{String}, promoter_list.genes)
promoter_list.gene_position = parse.(Vector{Float64}, promoter_list.gene_position)
promoter_list.evidence = parse.(Vector{String}, promoter_list.evidence)

operons_without_promoters.genes = parse.(Vector{String}, operons_without_promoters.genes)
operons_without_promoters.gene_position = parse.(Vector{Float64}, operons_without_promoters.gene_position)

## Some genes may have the wrong synomym
all_gene_list = CSV.read(
    "/$home_dir/data/all_genes_table.csv",
    DataFrame,
    types=Dict(
        "ID"=>String,
        "gene"=>String,
        "gene_position"=>Float64,
        "direction"=>String
    )
)

all_gene_list.synonyms = parse.(Vector{String}, all_gene_list.synonyms)
all_gene_list.transcription_units = parse.(Vector{String}, all_gene_list.transcription_units)

# Replace gene names if they are synonyms
for i in 1:nrow(gene_table)
    name = String(gene_table[i, "name"])
    if name ∉ all_gene_list.gene
        syn = all_gene_list[map(x -> name in x, all_gene_list.synonyms), :gene][1]
        gene_table[i, "name"] = syn
    end
end
gene_group_ID_dict = Dict(gene_table.name .=> gene_table.group_ID)
println("Done!")
println()
## Get promoters for genes

println("Finding Promoters...")
df_list = DataFrame[]
genes_no_tss = []
for gene in gene_table.name
    _df = promoter_list[map(x -> gene in x, promoter_list.genes), [:tss, :direction, :gene_position, :genes, :promoter, :evidence]]
    if nrow(_df) > 0
        push!(df_list, _df)
    else
        push!(genes_no_tss, gene)
    end
end
df = vcat(df_list...) |> unique
println("Done!")
println()
## Missing TSS
println("Looking for TSS in Urtecho data set...")
# Group genes into operons / transcription units
df_no_prom = DataFrame()
for gene in genes_no_tss
    append!(df_no_prom, operons_without_promoters[map(x -> gene in x, operons_without_promoters.genes), :])
end
if nrow(df_no_prom) > 0
    unique!(df_no_prom)
end

# Look for these units in Urtecho et al. 2020 dataset
# import Urtecho data
urtecho_tss = CSV.read(
    "/$home_dir/data/urtecho_2020/tss_operon_regulation.txt",
    DataFrame
)

function occursin_operon(gene, operon)
    split_operon = split(operon, "-")
    a, b = gene[1:3], gene[4]
    return prod(occursin.(a, split_operon) .* occursin.(b, split_operon))
end

# Find operon in Urtecho data
df_tss_urtecho = DataFrame()

# Make array to store indeces of genes to be removed from list
delete_index_list = Int64[]
for i in 1:nrow(df_no_prom)
    genes = df_no_prom[i, "genes"]
    gene_position = df_no_prom[i, "gene_position"]
    for gene in genes
        operons = filter(x -> prod(occursin_operon.(gene, x)), urtecho_tss.operon)
        if (operons |> unique |> length) > 1
            throw(ErrorException("More than one operon for genes: $(genes)"))
        elseif  (operons |> unique |> length) == 0
            #println("No operons for genes $(genes) in Urtecho dataset.")
            #println()
        else
            operon = unique(operons)[1]
            temp = urtecho_tss[(urtecho_tss.operon .== operon) .& ((urtecho_tss.active .== "active")), :]
            if nrow(temp) != 0
                insertcols!(temp, 2, :genes => fill(genes, nrow(temp)))
                insertcols!(temp, 2, :gene_position => fill(gene_position, nrow(temp)))
                insertcols!(temp, 2, :evidence => fill(["EXP"], nrow(temp)))
                rename!(temp, "tss_strand" => "direction", "tss_position"=> "tss", "tss_name"=>"promoter")
                append!(df_tss_urtecho, temp)
                push!(delete_index_list, i)
            end
        end
    end
end

# Add found promoters to list
if nrow(df_tss_urtecho) > 0
    append!(df, df_tss_urtecho[:, ["genes", "tss", "direction", "gene_position", "promoter", "evidence"]])
end

# Remove genes with identified promoters from list of genes without promoters
df_no_prom = df_no_prom[Not(delete_index_list), :]
println("Done!")
println()

## Now we tackle the genes for which we have not found a TSS yet
# We look for the site with the strongest predicted promoter within
# 500 bases upstream of the start of the coding region
println("Predicting TSS using La Fleur model...")
p = wgregseq.promoter_finder.Promoter_Calculator()
tss_list = Int64[]
name_list = String[]

for i in 1:nrow(df_no_prom)
    if df_no_prom[i, "direction"] == "+"
        ind = Int64(minimum(df_no_prom[i, "gene_position"]))
        sequence = wt_sequence[ind-500:ind]
        r = p(sequence)["Forward_Predictions_per_TSS"]
        _x = [(key, r[key]["dG_total"]) for key in keys(r) |> collect]
        tss = _x[argmin([x[2] for x in _x])][1] + ind - 500
    else
        ind = Int64(maximum(df_no_prom[i, "gene_position"]))
        sequence = wt_sequence[ind:ind+500]
        r = p(sequence)["Reverse_Predictions_per_TSS"]
        _x = [(key, r[key]["dG_total"]) for key in keys(r) |> collect]
        tss = _x[argmin([x[2] for x in _x])][1] + ind
    end
    push!(tss_list, tss)

    push!(name_list, join(df_no_prom[i, "genes"], "_") * "_predicted")
end

if nrow(df_no_prom)  >0
    # Add start sites to genes
    insertcols!(df_no_prom, 2, :tss =>tss_list)
    # Add information that they are predicted
    insertcols!(df_no_prom, 5, :promoter =>name_list)
    insertcols!(df_no_prom, 5, :evidence =>fill(["COMP"], nrow(df_no_prom)))
    # Add promoters to list
    append!(df, df_no_prom)
end

println("Done!")
println()

## For genes with multiple promoters, take the strongest predicted one
println(df)
# Create temporary DataFrame
temp_df = DataFrame()
#=
# Go through each gene
println("Finding best promoter for genes with multiple promoters...")
for gene_list in df.genes |> unique
    # Find promoters for gene
    sub_df = df[map(x -> x == gene_list, df.genes), :]
    println(sub_df)
    # If more than one promoter, find strongest
    if nrow(sub_df) > 1
        inds = map(x -> "EXP" in x, sub_df.evidence)
        if sum(inds) == 0
            append!(temp_df, DataFrame(wgregseq.design.find_best_promoter(sub_df, wt_sequence)))
        elseif sum(inds) == 1
            append!(temp_df, sub_df[inds, :])
        else
            _df = sub_df[inds, :]
            dist_mat = _df.tss' .- _df.tss
            println(dist_mat)
            break
        end
    else
        append!(temp_df, sub_df)
    end
end

#df = temp_df

println("Done!")
println()
=#

df
## Design Sequences for genes with TSS
println("Creating mutated sequences...")
df_sequences = DataFrame()
for row in eachrow(df)
    tss = Int64(row.tss)
    direction = row.direction
    genes = row.genes
    promoter = row.promoter
    seq = wgregseq.design.find_seq(tss, direction, 115, 45, wt_sequence)[1]
    mut_list = wgregseq.design.mutations_rand(seq, 0.1, 1500)
    names = ["$(promoter)_$i" for i in 0:1500]
    _df = DataFrame(sequence=mut_list, genes=fill(genes, 1501), promoter=fill(promoter, 1501), name=names)
    global df_sequences = vcat(df_sequences, _df)
end

if any(length.(df_sequences.sequence) .!= 160)
    throw(ErrorException("Not all sequences are 160bp!"))
else
    println("Done!")
    println()
end
df_sequences
## Check sequences for restriction sites
gdf = groupby(deepcopy(df_sequences), :genes)
df_stack = DataFrame()

enzymes = ["SalI", "SacI", "NheI", "XbaI", "SpeI", "XhoI", "EcoRI", "ApaI", "ScaI", "NcoI", "MluI", "EcoRV", "BbsI", "BamHI", "AgeI", "PstI", "NsiI", "SbfI"]
println("Adding restriction enzymes...")
for enz in enzymes
    if enz ∉ wgregseq.enzyme_list.enzyme
        throw(ErrorException("$enz is not in list of enzymes."))
    end
end

for group in gdf
    group[:, "sequence"] = wgregseq.design.add_primer(convert(Vector{LongSequence{DNAAlphabet{4}}}, (group.sequence)), 100, "both")
    df_restriction = wgregseq.design.find_restriction_sites(enzymes, group[:, "sequence"])
    sort!(df_restriction, "sites")
    dict = Dict{Any, Any}(df_restriction.enzyme .=> df_restriction.sites)
    dict["gene"] = [unique(group.genes)[1]]
    dict["promoter"] = [unique(group.promoter)[1]]
    append!(df_stack, DataFrame(dict))
end
df_stack

dict_enz = Dict{Any, Any}(enzymes .=> sum.(eachcol(df_stack[!, enzymes])))
dict_enz["gene"] = [String31["all"]]
dict_enz["promoter"] = "all"
append!(df_stack, DataFrame(dict_enz))
println("")
println(df_stack)

# Set these enzymes to disable query
enz1 = ""
enz2 = ""

println("Upstream restriction enzyme (default is SpeI by hitting `enter`):")
while enz1 ∉ wgregseq.enzyme_list.enzyme
    global enz1 = readline()
    if enz1 == ""
        global enz1 = "SpeI"
    end
    if enz1 ∉ wgregseq.enzyme_list.enzyme
        println("$enz1 not in list of enzymes")
    end
end
println("Downstream restriction enzyme (default is ApaI by hitting `enter`):")
while enz2 ∉ wgregseq.enzyme_list.enzyme
    global enz2 = readline()
    if enz2 == ""
        global enz2 = "ApaI"
    end
    if enz2 ∉ wgregseq.enzyme_list.enzyme
        println("$enz2 not in list of enzymes")
    end
end

## Add restriction sites
df_sequences.sequence = wgregseq.design.add_re_sites.(df_sequences.sequence, enz1, enz2)
if any(length.(df_sequences.sequence) .!= 172)
    println(length.(df_sequences.sequence) |> unique)
    throw(ErrorException("Not all sequences are 172 after adding primers bp!"))
end

insertcols!(df_sequences, 4, :upstream_re_site => fill(enz1, nrow(df_sequences)))
insertcols!(df_sequences, 5, :downstream_re_site => fill(enz2, nrow(df_sequences)))

println("Done!")
println()

## Adding primers
# First we check that the primer does not contain the restriction site. If that happens
# we try to take the next primer

## Add reverse primers
println("Adding primers for groups of 5 genes...")
n_per_group = 5
cop_df = deepcopy(df_sequences)
insertcols!(cop_df, 4, :rev_primer2 => fill((0, (193, 211)), nrow(cop_df)))
insertcols!(cop_df, 5, :rev_primer3 => fill((0, (212, 230)), nrow(cop_df)))

gdf = groupby(cop_df, :genes)

# Go through groups of n_per_group genes and add primer
primer = wgregseq.design.check_primers_re_sites(enz1, enz2, 100, "both")
primer_0 = deepcopy(primer)
i = 1
while i * n_per_group < length(gdf)
    global primer = wgregseq.design.check_primers_re_sites(enz1, enz2, primer + 1, "rev")
    for _df in gdf[1+(i-1)*n_per_group:i*n_per_group]
        _df[:, "sequence"] = wgregseq.design.add_primer(_df[:, "sequence"], primer, "rev")
        _df[:, "rev_primer2"] =  fill((primer, (193, 211)), nrow(_df))
    end
    global i += 1
end
primer = wgregseq.design.check_primers_re_sites(enz1, enz2, primer+1, "rev")
for _df in gdf[(i - 1)*n_per_group+1:end]
    _df[:, "sequence"] = wgregseq.design.add_primer(_df[:, "sequence"], primer, "rev")
    _df[:, "rev_primer2"] =  fill((primer, (193, 211)), nrow(_df))
end

# Combine sequences to single table
df_final = combine(gdf, :)
df_final.sequence = [ x[1:end-1] for x in df_final.sequence]

# Confirm sequence length
if any(length.(df_final.sequence) .!= 191)
    throw(ErrorException("Not all sequences are 191bp!"))
else
    println("Done!")
    println()
end

# Group sequences again
println("Adding group specific primers...")
gdf = groupby(df_final, :genes)

# Add reverse primer for each group ID that was determined in the beginning
for _df in gdf
    genes = _df[1, "genes"]
    gene_ids = []
    for gene in genes
        if gene in keys(gene_group_ID_dict)
            group = gene_group_ID_dict[gene]
            push!(gene_ids, group)
        end
    end
    gene_ids = gene_ids |> unique

    if length(gene_ids) > 1
        throw(ErrorException("More than one group for genes: $genes."))
    else
        ID = gene_ids[1]
    end
    local primer = wgregseq.design.check_primers_re_sites(enz1, enz2, Int(200 + ID), "rev")
    _df[:, "sequence"] = wgregseq.design.add_primer(_df[:, "sequence"], primer, "rev")
    _df[:, "rev_primer3"] =  fill((primer, (212, 230)), nrow(_df))
end

# Combine groups to single table
df_final = combine(gdf, :)
df_final.sequence = [ x[1:end-1] for x in df_final.sequence]

# Confirm Length
if any(length.(df_final.sequence) .!= 210)
    throw(ErrorException("Not all sequences are 210bp!"))
else
    println("Done!")
    println()
end

println("Adding forward and reverse primers...")

insertcols!(df_final, 4, :fwd_primer => fill((primer_0, (1, 20)), nrow(df_final)))
insertcols!(df_final, 5, :rev_primer1 => fill((primer_0, (231, 250)), nrow(df_final)))
df_final.sequence = wgregseq.design.add_primer(df_final.sequence, primer_0)

# Confirm that the primers are correct
fwd_primer = wgregseq.design.import_primer(primer_0, "fwd")
rev_primer = wgregseq.design.import_primer(primer_0, "rev")
if any([seq[1:20] != fwd_primer for seq in df_final.sequence])
    throw(ErrorException("Not all sequences have the right forward primer!"))
elseif any([seq[end-19:end] != rev_primer for seq in df_final.sequence])
    throw(ErrorException("Not all sequences have the right reverse primer!"))
else
    println("Done!")
    println()
end

##
# Save results
filename = string(Dates.today()) * "_sequence_list.csv"
CSV.write("/$home_dir/data/twist_orders/$filename", df_final)
println("Sequence list saved in `/$home_dir/data/twist_orders/$filename`")
println("Total number of sequences: $(nrow(df_final))")

##
println()
println("Sequences per group.")
println(combine(groupby(df_final, :rev_primer3), nrow))

df_final


text:
# Numercal experiment of Van der Pol oscillator
# Training data sets are generated from the ODE solutions

using DifferentialEquations
using LinearAlgebra
using DiffEqFlux
using LaTeXStrings
using PGFPlotsX
using Statistics

function vdp(du, u, p, t)
    μ = p[1]
    du[1] = u[2]
    return du[2] = 2 * μ * u[2] - u[1]^2 * u[2] - u[1]
end

function get_sol2(sol, ind) # Convert solution of ODE solver to array solution
    lu = length(sol.u)
    u = Vector{Float64}(undef, lu)
    for i in 1:lu
        uv = sol.u[i]
        u[i] = uv[ind]
    end
    return u
end

function zero_measure(u, ind, t) # Measuring the time of zero crossing points from numerical continuation
    l = length(u)
    zero = Array{Float64}(undef, 0)
    T = Array{Float64}(undef, 0)
    low_p = Array{Float64}(undef, 0)
    high_p = Array{Float64}(undef, 0)
    Ti = Array{Int64}(undef, 0)
    for i in 2:(l - 1)
        sign_con2 = u[i][ind + 1] * u[i - 1][ind + 1]
        if sign_con2 < 0
            if (u[i][ind] + u[i - 1][ind]) / 2 < 0
                low_p = vcat(low_p, (u[i][ind] + u[i - 1][ind]) / 2)
            else
                high_p = vcat(high_p, (u[i][ind] + u[i - 1][ind]) / 2)
            end
        end
    end
    h₀ = mean(high_p) + mean(low_p)
    h₀ = h₀ / 2
    for i in 2:(l - 1)
        sign_con = (u[i][ind] - h₀) * (u[i + 1][ind] - h₀)
        if sign_con < 0
            if (u[i][ind + 1] + u[i + 1][ind + 1]) / 2 > 0
                zero = vcat(zero, (u[i][ind] + u[i - 1][ind]) / 2)
                Ti = vcat(Ti, i)
                T = vcat(T, t[i])
            end
        end
    end
    t_l = length(T)
    P = Array{Float64}(undef, t_l - 1)
    for j in 2:t_l
        P[j - 1] = T[j] - T[j - 1]
    end
    return (T=Ti, P=P, hp=high_p, lp=low_p, h₀=h₀)
end

function get_stable_LCO(p, u0, tl, tol, eq, stol, rp, ind1, ind2, u₀, v₀, st) # Get a stable LCO from numerical integration
    u = u0
    dim = length(u0)
    prob = ODEProblem(eq, u, (0, tl * rp), p)
    sol = DifferentialEquations.solve(prob, Tsit5(); reltol=stol, abstol=stol, saveat=st)
    vP = 1
    P = 0
    T = 0
    while vP > tol
        u = sol.u[end]
        prob = ODEProblem(eq, u, (0, tl), p)
        sol = DifferentialEquations.solve(
            prob, Tsit5(); reltol=stol, abstol=stol, saveat=st
        )
        z = zero_measure(sol.u, 1, sol.t)
        vP = Statistics.var(z.hp)
        P = z.P[1]
        T = z.T
    end
    tl = length(sol)
    uu = Array{Float64}(undef, tl, dim)
    for i in 1:dim
        uu[:, i] = get_sol2(sol, i)
    end
    t = Array{Float64}(undef, length(sol))
    r = Array{Float64}(undef, length(sol))
    u = uu[:, ind1]
    v = uu[:, ind2]
    for i in 1:length(u)
        t[i] = atan(v[i] - v₀, u[i] - u₀)
        r[i] = sqrt((u[i] - u₀)^2 + (v[i] - v₀)^2)
    end
    return (u=uu, t=t, r=r, P=P, T=T)
end

function LS_harmonics(r, t, ω, N) # Computing Fourier coefficients of the amplitude in the measued state-variable coordinates
    # Fourier coefficients are computed in least square sence
    c = Array{Float64}(undef, 2 * N + 1)
    M = Array{Float64}(undef, 1, 2 * N + 1)
    tM = Array{Float64}(undef, 0, 2 * N + 1)
    tl = length(t)
    rr = Array{Float64}(undef, tl)
    M[1] = 1
    for j in 1:tl
        for i in 1:N
            M[1 + i] = cos(ω * t[j] * i)
            M[1 + N + i] = sin(ω * t[j] * i)
        end
        tM = vcat(tM, M)
    end
    MM = transpose(tM) * tM
    rN = transpose(tM) * r
    c = inv(MM) * rN
    for j in 1:tl
        rr[j] = c[1]
        for i in 1:N
            rr[j] += c[i + 1] * cos(ω * t[j] * i)
            rr[j] += c[i + 1 + N] * sin(ω * t[j] * i)
        end
    end
    return (coeff=c, rr=rr)
end

function generate_data_vdp(vel_l, Vel, nh)
    #Generate training data
    u0 = Float32[1.0, 0]
    tol = 1e-7
    stol = 1e-8
    eq = vdp
    rp = 5
    ind1 = 1
    ind2 = 2
    AA = zeros(vel_l, Int(nh * 2 + 1))

    p_ = Vel[1]
    g = get_stable_LCO(p_, u0, tl, tol, eq, stol, rp, ind1, ind2, 0.0, 0.0, st)
    u₀ = mean(g.u[:, 1])
    v₀ = mean(g.u[:, 2])
    for i in 1:vel_l
        p_ = Vel[i]
        g = get_stable_LCO(p_, u0, tl, tol, eq, stol, rp, ind1, ind2, 0.0, 0.0, st)
        r = g.r
        t = g.t
        c = LS_harmonics(r, t, 1, nh).coeff
        AA[i, :] = c
    end
    AA = transpose(AA)
    θ = range(0; stop=2π, length=θ_l)
    coθ = cos.(θ)
    siθ = sin.(θ)
    t_series = [
        Transpose(
            get_stable_LCO(Vel[i], u0, tl, tol, eq, stol, rp, ind1, ind2, 0.0, 0.0, st2).u[
                :, [1, 2]
            ],
        ) for i in 1:vel_l
    ]
    θ_series = [
        get_stable_LCO(Vel[i], u0, tl, tol, eq, stol, rp, ind1, ind2, 0.0, 0.0, st2).t for
        i in 1:vel_l
    ]
    return (data=AA, ts=t_series, coθ=coθ, siθ=siθ, theta_s=θ_series)
end

function nf_dis_vdp(U₀, Vel)
    del = Vel - U₀ * ones(length(Vel))
    s_amp = sqrt.(del)
    vl = [s_amp[i] * [coθ'; siθ'] for i in 1:length(Vel)]
    return vl
end

function f_coeff(vlT, Vel, u₀, v₀)
    Pr = zeros(2 * nh + 1, 0)
    for k in 1:length(Vel)
        z1 = vlT[k][1, :] - u₀ * ones(θ_l)
        z2 = vlT[k][2, :] - v₀ * ones(θ_l)
        theta = atan.(z2, z1)
        r = sqrt.(z1 .^ 2 + z2 .^ 2)
        tM = Array{Float64}(undef, 0, 2 * nh + 1)
        rr = Array{Float64}(undef, θ_l)
        for j in 1:θ_l
            tM1 = Array{Float64}(undef, 0, nh + 1)
            tM2 = Array{Float64}(undef, 0, nh)
            tM1_ = [cos(theta[j] * i) for i in 1:nh]
            tM2_ = [sin(theta[j] * i) for i in 1:nh]
            tM1_ = vcat(1, tM1_)
            tM1 = vcat(tM1, Transpose(tM1_))
            tM2 = vcat(tM2, Transpose(tM2_))
            tM_ = hcat(tM1, tM2)
            tM = vcat(tM, tM_)
        end
        MM = Transpose(tM) * tM
        rN = Transpose(tM) * r
        c = inv(MM) * rN
        Pr = hcat(Pr, c)
        Pr
    end
    return Pr
end
# Define the experimental parameter

θ_l = 200;
θ = range(0; stop=2π, length=θ_l)
coθ = cos.(θ)
siθ = sin.(θ)
tl = 30.0;
st = 0.01;
st2 = 0.02;
vel_l = 3
Vel = range(0.2; stop=1.0, length=vel_l)
nh = 20

dat = generate_data_vdp(vel_l, Vel, nh)
AA = dat.data
tl2 = 10.0

tt = Int(tl2 / st2 + 1)
t_series = [dat.ts[i][:, 1:tt] for i in 1:length(Vel)]
θ_series = [dat.theta_s[i][1:tt] for i in 1:length(Vel)]

function predict_lt_vdp(θ_t) #predict the linear transformation
    nf = nf_dis_vdp(0.0, Vel)
    vl = nf
    p1, p2, p3, p4 = θ_t
    T = [p1 p3; p2 p4]
    vlT = [T * (vl[i]) for i in 1:length(Vel)]

    Pr = f_coeff(vlT, Vel, 0, 0)
    return Pr
end

function lt_pp_vdp(θ_t) # This function gives phase portrait of the transformed system from the normal form
    nf = nf_dis_vdp(0.0, Vel)
    vl = nf
    p1, p2, p3, p4 = θ_t
    T = [p1 p3; p2 p4]
    vlT = [T * (vl[i]) for i in 1:length(Vel)]
    return vlT
end

function Array_chain(gu, ann, p) # vectorized input-> vectorized neural net
    al = length(gu[1, :])
    AC = zeros(2, 0)
    for i in 1:al
        AC = hcat(AC, ann(gu[:, i], p))
    end
    return AC
end

function loss_lt_vdp(θ_t)
    pred = predict_lt_vdp(θ_t)
    return sum(abs2, AA .- pred) # + 1e-5*sum(sum.(abs, params(ann)))
end

θ = vec([3.0 0.0; 0.0 3.0])
loss_lt_vdp(θ)

res1 = DiffEqFlux.sciml_train(loss_lt_vdp, θ, ADAM(0.01); maxiters=200)
res1.minimum
θ_ = res1.minimizer

##
function predict_nt_vdp(θ_t)
    p1, p2, p3, p4 = θ_
    np1 = θ_t[1]
    pn = θ_t[2:end]
    nf = nf_dis_vdp(np1, Vel)
    vl = nf
    T = [p1 p3; p2 p4]
    vlT = [
        T * (vl[i]) +
        Array_chain([vl[i]; (Vel[i] - np1) * ones(1, θ_l)], ann, pn) / scale_f_pp for
        i in 1:length(Vel)
    ]

    return Pr = f_coeff(vlT, Vel, 0, 0)
end

function lt_pp_n_vdp(θ_t) # This function gives phase portrait of the transformed system from the normal form (stable LCO)
    p1, p2, p3, p4 = θ_
    np1 = θ_t[1]
    pn = θ_t[2:end]
    nf = nf_dis_vdp(np1, Vel)
    vl = nf
    T = [p1 p3; p2 p4]
    vlT = [
        T * (vl[i]) +
        Array_chain([vl[i]; (Vel[i] - np1) * ones(1, θ_l)], ann, pn) / scale_f_pp for
        i in 1:length(Vel)
    ]
    return vlT
end

function lt_pp_n_vdp(θ_t, Vel) # This function gives phase portrait of the transformed system from the normal form (stable LCO)
    p1, p2, p3, p4 = θ_
    np1 = θ_t[1]
    pn = θ_t[2:end]
    nf = nf_dis_vdp(np1, Vel)
    vl = nf
    T = [p1 p3; p2 p4]
    vlT = [
        T * (vl[i]) +
        Array_chain([vl[i]; (Vel[i] - np1) * ones(1, θ_l)], ann, pn) / scale_f_pp for
        i in 1:length(Vel)
    ]
    return vlT
end

function loss_nt_vdp(θ_t)
    pred = predict_nt_vdp(θ_t)
    return sum(abs2, AA .- pred) # + 1e-5*sum(sum.(abs, params(ann)))
end

hidden = 31
ann = FastChain(
    FastDense(3, hidden, tanh), FastDense(hidden, hidden, tanh), FastDense(hidden, 2)
)
θn = initial_params(ann)
scale_f_pp = 1e1

pp = [0]
θn = vcat(pp, θn)
loss_nt_vdp(θn)

res1 = DiffEqFlux.sciml_train(loss_nt_vdp, θn, ADAM(0.01); maxiters=300)
res1 = DiffEqFlux.sciml_train(
    loss_nt_vdp, res1.minimizer, BFGS(; initial_stepnorm=1e-5); maxiters=10000
)
θ_2 = res1.minimizer

Ap = lt_pp_n_vdp(θ_2)

a = @pgf Axis(
    {
        xlabel = L"$z_1$",
        ylabel = L"$z_2$",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = -3,
        xmax = 3,
        ymin = -5.5,
        ymax = 5.5,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[1][1, :], Ap[1][2, :])),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(t_series[1][1, :], t_series[1][2, :])),
    LegendEntry("Ground truth"),
)
#pgfsave("./Figures/num_vdp/vdp_1.pdf",a)

a = @pgf Axis(
    {
        xlabel = L"$z_1$",
        ylabel = L"$z_2$",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = -3,
        xmax = 3,
        ymin = -5.5,
        ymax = 5.5,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[3][1, :], Ap[3][2, :])),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(t_series[3][1, :], t_series[3][2, :])),
    LegendEntry("Ground truth"),
)

#pgfsave("./Figures/num_vdp/vdp_2.pdf",a)

function Inv_T_u_vdp(th0, vel, tol) # This function gives initial conditions of the model
    p1, p2, p3, p4, = θ_[1:4]
    T = [p1 p3; p2 p4]
    np1 = θ_2[1]
    pn = θ_2[2:end]
    del = vel - np1
    s_amp = sqrt.(del)
    theta = range(-π; stop=π, length=400)

    uu = [[s_amp * cos(theta[i]), s_amp * sin(theta[i])] for i in 1:length(theta)]
    u = [T * uu[i] + ann([uu[i]; vel - np1], pn) / scale_f_pp for i in 1:length(theta)]
    t0 = [abs(atan(u[i][2], u[i][1]) - th0) for i in 1:length(theta)]
    er = minimum(t0)
    while er > tol
        #    global theta,t0
        theta = range(theta[argmin(t0) - 1], theta[argmin(t0) + 1]; length=400)
        uu = [[s_amp * cos(theta[i]), s_amp * sin(theta[i])] for i in 1:length(theta)]
        u = [T * uu[i] + ann([uu[i]; vel - np1], pn) / scale_f_pp for i in 1:length(theta)]
        t0 = [abs(atan(u[i][2], u[i][1]) - th0) for i in 1:length(theta)]
        er = minimum(t0)
    end
    return (th=theta[argmin(t0)], uu=u[argmin(t0)])
end

function dudt_ph2(u, p, t)
    θ = u[1]
    r = u[2]
    c = u[3]
    np1 = θ_2[1]
    δ₀ = np1
    ν = (c - δ₀)
    ω₀ = p[1]
    uu = [r * cos(θ), r * sin(θ), ν]
    du₁ = ω₀ + ann3(uu, p[2:end])[1] / om_scale
    du₂ = 0
    du₃ = 0
    return [du₁, du₂, du₃]
end

function atan2(u)
    θ = atan.(u[2, :], u[1, :])
    for i in 2:length(θ)
        if θ[i] - θ[i - 1] > π
            θ = θ + vcat(zeros(i - 1), -2 * π * ones(length(θ) - i + 1))
        elseif θ[i] - θ[i - 1] < -π
            θ = θ + vcat(zeros(i - 1), 2 * π * ones(length(θ) - i + 1))
        end
    end
    return θ
end

function dudt_nf_vdp(u, p, t)
    c = u[3]
    np1 = θ_2[1]
    a2 = -1
    δ₀ = np1
    ν = (c - δ₀)
    r2 = u[1]^2 + u[2]^2
    r = sqrt(r2)
    θ = atan(u[2], u[1])
    ω₀ = p[1]
    ω₁ = p[2]
    uu = [r * cos(θ), r * sin(θ), ν]
    hc = ann3(ν, p[3:end])
    phs_s = [hc[nh2 + i + 1] * sin(i * θ) for i in 1:nh2]
    phs_s = sum(phs_s)
    phs_c = [hc[i + 1] * cos(i * θ) for i in 1:nh2]
    phs_c = sum(phs_c)
    ph = ω₀ + ν * ω₁ + ann3(uu, p[3:end])[1] / scale_f
    du₁ = ν * u[1] - u[2] * ph + a2 * u[1] * r2
    du₂ = u[1] * ph + ν * u[2] + a2 * u[2] * r2
    du₃ = 0
    return [du₁, du₂, du₃]
end

function predict_time_T_vdp(p) #,uu_t0
    np1 = θ_2[1]
    pn = θ_2[2:end]
    p1, p2, p3, p4 = θ_[1:4]
    T = [p1 p3; p2 p4]
    A1 = [
        Array(
            concrete_solve(
                ODEProblem(dudt_nf_vdp, u_t0[i], (0, tl2), p),
                Tsit5(),
                u_t0[i],
                p;
                saveat=st3,
                abstol=1e-8,
                reltol=1e-8,
                sensealg=InterpolatingAdjoint(; autojacvec=ReverseDiffVJP()),
            ),
        ) for i in 1:length(Vel)
    ]
    uu = [transpose(hcat(A1[i][1, :], A1[i][2, :], A1[i][3, :])) for i in 1:vl]
    delU = zeros(2, spl2)
    delU2 = -np1 * ones(1, spl2)
    delU = vcat(delU, delU2)
    uu = [uu[i] + delU for i in 1:vl]
    vlT = [T * uu[i][1:2, :] + Array_chain(uu[i], ann, pn) / scale_f_pp for i in 1:vl]
    Pr = zeros(0, spl2)
    for i in 1:vl
        theta = transpose(atan2(vlT[i][[1, 2], :]))
        Pr = vcat(Pr, theta)
    end
    return Pr
end

st3 = 0.05

t_s3 = [
    Array(
        concrete_solve(
            ODEProblem(vdp, t_series[i][:, 1], (0, tl2), Vel[i]),
            Tsit5(),
            t_series[i][:, 1],
            Vel[i];
            saveat=st3,
        ),
    ) for i in 1:length(Vel)
]
t_s2 = [atan2(t_s3[i]) for i in 1:length(Vel)]

vl = length(Vel)
spl2 = length(t_s2[1])
t_s = zeros(vl, spl2)

tol = 1e-8
np1 = θ_2[1]
pn = θ_2[2:end]
s_amp = [sqrt(Vel[i] - np1) for i in 1:length(Vel)]
theta0 = [t_s2[i][1] for i in 1:length(Vel)]
θ₀ = [Inv_T_u_vdp(theta0[i], Vel[i], tol).th for i in 1:length(Vel)]
u_t0 = [[s_amp[i] * cos.(θ₀[i]), s_amp[i] * sin.(θ₀[i]), Vel[i]] for i in 1:length(Vel)]

for i in 1:vl
    t_s[i, :] = t_s2[i]
end
A3 = t_s
spl2 = length(A3[1, :])

function loss_time_T_vdp(p)
    pred = predict_time_T_vdp(p)
    return sum(abs2, A3 .- pred) # + 1e-5*sum(sum.(abs, params(ann)))
end

hidden = 31
nh2 = 10
ann3 = FastChain(FastDense(3, hidden, tanh), FastDense(hidden, nh2 * 2 + 1, tanh))
np = initial_params(ann3)
aa = [-1, 0.16]
p = vcat(aa, np)
scale_f = 1
##harmonic representation of phase
res1 = DiffEqFlux.sciml_train(loss_time_T_vdp, p, ADAM(0.01); maxiters=200)
res1 = DiffEqFlux.sciml_train(
    loss_time_T_vdp, p, NADAM(0.002, (0.89, 0.995)); maxiters=4000
)
res1 = DiffEqFlux.sciml_train(
    loss_time_T_vdp, res1.minimizer, BFGS(; initial_stepnorm=1e-5); maxiters=30000
)

res1.minimum
p = res1.minimizer

tv = range(0, tl2; length=length(A3[1, :]))

ind = 1
a = @pgf Axis(
    {
        xlabel = "Time (sec)",
        ylabel = L"$z_1$",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = 0,
        xmax = 10,
        ymax = 4,
        ymin = -4,
    },
    Plot(
        {color = "red", no_marks},
        Coordinates(tv, predict_time_T_vdp2(p)[2 * (ind - 1) + 1, :]),
    ),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(tv, t_s3[ind][1, :])),
    LegendEntry("Ground truth"),
)
#pgfsave("./Figures/num_vdp/vdp_t_020.pdf",a)

ind = 3
a = @pgf Axis(
    {
        xlabel = "Time (sec)",
        ylabel = L"$z_1$",
        legend_pos = "south east",
        height = "9cm",
        width = "9cm",
        xmin = 0,
        xmax = 10,
        ymax = 4,
        ymin = -4,
    },
    Plot(
        {color = "red", no_marks},
        Coordinates(tv, predict_time_T_vdp2(p)[2 * (ind - 1) + 1, :]),
    ),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(tv, t_s3[ind][1, :])),
    LegendEntry("Ground truth"),
)

#pgfsave("./Figures/num_vdp/vdp_t_100.pdf",a)

ext: jl


text:
isconstrained(alg::AbstractMethodOfStepsAlgorithm{constrained}) where constrained =
    constrained

OrdinaryDiffEq.alg_maximum_order(alg::AbstractMethodOfStepsAlgorithm) =
  OrdinaryDiffEq.alg_maximum_order(alg.alg)
OrdinaryDiffEq.alg_extrapolates(alg::AbstractMethodOfStepsAlgorithm) =
  OrdinaryDiffEq.alg_extrapolates(alg.alg)

OrdinaryDiffEq.uses_uprev(alg::AbstractMethodOfStepsAlgorithm, adaptive) = true

DiffEqBase.isadaptive(alg::AbstractMethodOfStepsAlgorithm) =
  DiffEqBase.isadaptive(alg.alg)

"""
    iscomposite(alg)

Return if algorithm `alg` is a composite algorithm.
"""
iscomposite(alg) = false
iscomposite(::OrdinaryDiffEq.OrdinaryDiffEqCompositeAlgorithm) = true
iscomposite(alg::AbstractMethodOfStepsAlgorithm) = iscomposite(alg.alg)


text:
using LinearAlgebra
export BCIPEPS, SquareBCIPEPS
"""
    BCIPEPS{LT<:AbstractLattice, T, N}

Infinite projected entangled pair of states.
`LT` is the type of lattice, `T` and `N` are bulk tensor element type and order.
"""
struct BCIPEPS{LT<:AbstractLattice, T, N, AT<:AbstractArray{<:AbstractArray,2}}
    bulk::AT
end
function BCIPEPS{LT}(bulk::AT) where {LT<:AbstractLattice, AT<:AbstractArray{<:AbstractArray,2}}
    T, N = eltype(bulk[1,1]), ndims(bulk[1,1])
    BCIPEPS{LT,T,N,AT}(bulk)
end

############### BCIPEPS on square lattice ###################
# size of bulk is `d × d × d × d × s`
const SquareBCIPEPS{T} = BCIPEPS{SquareLattice, T, 5}
function SquareBCIPEPS(bulk::AT) where {AT<:AbstractArray{<:AbstractArray,2}}
    # NOTE: from here, wrapping `bulk` with a `BCIPEPS` type can prevent programing from illegal input with incorrect size.
    bulk11 = bulk[1,1]
    size(bulk11,1) == size(bulk11,2) == size(bulk11,3) == size(bulk11,4) || throw(DimensionMismatch(
        "size of tensor error, should be `(d, d, d, d, s)`, got $(size(bulk11))."))
    T = eltype(bulk11)
    BCIPEPS{SquareLattice,T,5,AT}(bulk)
end
getd(bcipeps::SquareBCIPEPS) = size(bcipeps.bulk[1,1], 1)
gets(bcipeps::SquareBCIPEPS) = size(bcipeps.bulk[1,1], 5)

"""
    indexperm_symmetrize(bcipeps::SquareBCIPEPS)

return a `SquareBCIPEPS` based on `bcipeps` that is symmetric under
permutation of its virtual indices, which is ordered below:
```
        4
        │
 1 ── bcipeps ── 3
        │
        2
```
"""
function indexperm_symmetrize(bcipeps::SquareBCIPEPS)
    Ni,Nj = size(bcipeps.bulk)
    bulk = reshape([symmetrize(bcipeps.bulk[i]) for i = 1:Ni*Nj], (Ni, Nj))
    return SquareBCIPEPS(bulk)
end

function symmetrize(x)
    x += permutedims(x, (1,4,3,2,5)) # up-down
    x += permutedims(x, (3,2,1,4,5)) # left-right
    x += permutedims(x, (2,1,4,3,5)) # diagonal
    x += permutedims(x, (4,3,2,1,5)) # rotation
    x / norm(x)
end


text:
import Base.getindex
# using Snopt

struct PrimalVars{V,T}
    Z::V
    X::SubArray{T}
    U::SubArray{T}
end

function PrimalVars(Z::AbstractVector,ind_x::Matrix{Int},ind_u::Matrix{Int})
    n,N = size(ind_x)
    m = size(ind_u,1)
    Nz = N*n + (N-1)*m
    if length(Z) > Nz
        Z = Z[1:Nz]
    end
    X = view(Z,ind_x)
    U = view(Z,ind_u)
    PrimalVars(Z,X,U)
end

function PrimalVars(Z::AbstractVector,n::Int,m::Int,N::Int)
    ind_x, ind_u = get_primal_inds(n,m,N)
    PrimalVars(Z,ind_x,ind_u)
end

function PrimalVars(n::Int,m::Int,N::Int)
    Z = zeros(N*n + (N-1)*m)
    PrimalVars(Z,n,m,N)
end

function PrimalVars(X::Matrix,U::Matrix)
    Z = packZ(X,U)
    n,N = size(X)
    m = size(U,1)
    PrimalVars(Z,n,m,N)
end

function PrimalVars(res::SolverIterResults)
    PrimalVars(to_array(res.X), to_array(res.U))
end

get_sizes(Z::PrimalVars) = size(Z.X,1), size(Z.U,1), size(Z.X,2)
Base.length(Z::PrimalVars) = length(Z.Z)

function getindex(Z::PrimalVars,k::Int)
    n,m,N = get_sizes(Z)
    if k < N
        return Z.Z[(1:n+m) .+ (k-1)*(n+m)]
    elseif k == N
        return Z.Z[end-n+1:end]
    else
        BoundsError("Cannot retrieve time step $k, Z only has $N time steps")
    end
end

struct NewtonVars{T}
    V::Vector{T}
    Z::PrimalVars{V,T} where V
    ν::SubArray{T}
    λ::SubArray{T}
    μ::SubArray{T}
    r::SubArray{T}
    a::Vector{Bool}  # Active set
    part::NamedTuple
    dualinds::UnitRange{Int}
end

function NewtonVars(Z::PrimalVars,ν::T,λ::T,μ::T,r::T) where T <: AbstractVector
    Nz = length(Z)
    Nx = length(ν)
    Nh = length(λ)
    Ng = length(μ)
    n,m,N = get_sizes(Z)
    part = create_partition([Nz,Nx,Nh,Ng,2Nz],(:z,:ν,:λ,:μ,:r))
    V = [Z.Z;ν;λ;μ;r]
    Z = PrimalVars(view(V,part.z),n,m,N)
    ν = view(V,part.ν)
    λ = view(V,part.λ)
    μ = view(V,part.μ)
    r = view(V,part.r)
    a = ones(Bool,length(V))
    dualinds = Nz .+ (1:Nx+Nh+Ng+2Nz)
    NewtonVars(V,Z,ν,λ,μ,r,a,part,dualinds)
end

primals(V::NewtonVars) = V.Z.Z
duals(V::NewtonVars) = view(V.V,V.dualinds)
duals(V::NewtonVars, a::Vector{Bool}) = view(duals(V),a[V.dualinds])
Base.copy(V::NewtonVars) = NewtonVars(V.Z,V.ν,V.λ,V.μ,V.r)
+(V::NewtonVars,A::Vector) = begin V1 = copy(V); V1.V .+= A; return V1 end
Base.length(V::NewtonVars) = length(V.V)

function NewtonVars(solver::Solver, res::SolverIterResults)
    n,m,N = get_sizes(solver)
    p,pI,pE = get_num_constraints(solver)
    p_c,pI_c,pE_c = get_num_constraints(solver,:custom)
    p_N,pI_N,pE_N = get_num_terminal_constraints(solver)
    p_bnd = pI-pI_c  # number of bound constraints

    X = to_array(res.X)
    U = to_array(res.U)
    m = size(U,1)
    ν = zeros(n,N)
    r = vec(get_bound_multipliers(solver,res))

    for k = 1:N
        ν[:,k] = res.s[k]*0
    end
    ν = vec(ν)
    ind = create_partition([p_bnd,pI_c,pE_c],(:bnd,:I,:E))

    if res isa ConstrainedVectorResults
        μ = zeros(pI_c,N-1)
        λ = zeros(pE_c,N-1)
        for k = 1:N-1
            μ[:,k] = res.λ[k][ind.I]
            λ[:,k] = res.λ[k][ind.E]
        end
        μ = vec(μ)
        λ = vec(λ)
        append!(μ, res.λ[N][1:pI_N])
        append!(λ, res.λ[N][pI_N .+ (1:pE_N)])
    else
        λ = Float64[]
        μ = Float64[]
    end
    Z = PrimalVars(X,U)
    NewtonVars(Z,ν,λ,μ,r)
end

function get_primal_inds(n,m,N)
    Nx = N*n
    Nu = (N-1)*m
    Nz = Nx+Nu
    ind_x = zeros(Int,n,N)
    ind_u = zeros(Int,m,N-1)
    ix = 1:n
    iu = n .+ (1:m)
    for k = 1:N-1
        ind_x[:,k] = ix .+ (k-1)*(n+m)
        ind_u[:,k] = iu .+ (k-1)*(n+m)
    end
    ind_x[:,N] = ix .+ (N-1)*(n+m)
    return ind_x, ind_u
end

function get_batch_sizes(solver::Solver)
    p,pI,pE = get_num_constraints(solver,:custom)
    p_N,pI_N,pE_N = get_num_terminal_constraints(solver,:custom)
    n,m,N = get_sizes(solver)
    if solver.obj isa ConstrainedObjective && solver.obj.use_xf_equality_constraint
        pE_N = n
        p_N = pI_N + pE_N
    end

    Nx = N*n
    Nu = (N-1)*m
    Nh = (N-1)*pE + pE_N
    Ng = (N-1)*pI + pI_N
    NN = 2Nx + Nu + Nh + Ng + 2(Nx+Nu)
    return Nx,Nu,Nh,Ng,NN
end

function gen_usrfun_newton(solver::Solver)
    p,pI,pE = get_num_constraints(solver,:custom)
    p_N,pI_N,pE_N = get_num_terminal_constraints(solver,:custom)
    n,m,N = get_sizes(solver)

    # c_function!, c_jacobian!, c_labels, cI!, cE! = generate_constraint_functions(solver.obj)
    cE!,cI! = solver.obj.cE, solver.obj.cI
    cE_N!,cI_N! = solver.obj.cE_N, solver.obj.cI_N
    obj = solver.obj
    costfun = obj.cost
    if obj.use_xf_equality_constraint
        pE_N = n
        p_N = pI_N + pE_N
    end

    Nx,Nu,Nh,Ng,NN = get_batch_sizes(solver)
    Nz = Nx+Nu

    names = (:z,:ν,:λ,:μ,:r)
    ind1 = TrajectoryOptimization.create_partition([Nz,Nx,Nh,Ng,2Nz],names)
    ind2 = TrajectoryOptimization.create_partition2([Nz,Nx,Nh,Ng,2Nz],names)

    ind_x, ind_u = get_primal_inds(n,m,N)
    ind_Z = create_partition([Nx,Nu],(:x,:u))
    ind_z = create_partition([n,m],(:x,:u))
    ind_h = create_partition([(N-1)*pE,pE_N],(:k,:N))
    ind_g = create_partition([(N-1)*pI,pI_N],(:k,:N))
    ind_h = (k=ind_h.k, N=ind_h.N, k_mat=reshape(ind_h.k,pE,N-1))
    ind_g = (k=ind_g.k, N=ind_g.N, k_mat=reshape(ind_g.k,pI,N-1))
    ind_c = create_partition([Nx,Nh,Ng,2Nx],(:d,:E,:I,:b))

    if solver.obj.cost isa QuadraticCost
        Z0 = PrimalVars(n,m,N)
        ∇²J, = taylor_expansion(solver,Z0)
    else
        error("Not yet implemented for non-Quadratic cost functions")
    end

    # Set up jacobians (fast for this application)
    discretizer = eval(solver.integration)
    f!(xdot,x,u) = TrajectoryOptimization.dynamics(solver.model,xdot,x,u)
    fd! = rk4(f!,solver.dt)
    fd_aug!(xdot,z) = fd!(xdot,view(z,ind_z.x),view(z,ind_z.u))
    Fd!(dz,xdot,z) = ForwardDiff.jacobian!(dz,fd_aug!,xdot,z)

    # Bounds
    ind_bnd = create_partition([Nz,Nz],(:U,:L))
    z_L, z_U = get_bounds(solver,false)
    bound_constraint(c,Z) = begin copyto!(view(c,ind_bnd.U), Z - z_U);
                                  copyto!(view(c,ind_bnd.L), z_L - Z); end
    bound_constraint(Z) = begin c = zeros(eltype(Z),2Nz); bound_constraint(c,Z); return c; end
    jac_bnd = [Diagonal(I,Nz); -Diagonal(I,Nz)]

    # Custom Constraint jacobians
    augmenter(f!) = f_aug!(a,z) = f!(a,view(z,ind_z.x),view(z,ind_z.u))
    cE!,cI! = solver.obj.cE, solver.obj.cI
    cE_N!,cI_N! = solver.obj.cE_N, solver.obj.cI_N
    cE_aug = augmenter(cE!)
    cI_aug = augmenter(cI!)
    ∇cE(Cz,c,z) = ForwardDiff.jacobian!(Cz,cE_aug,c,z)
    ∇cI(Cz,c,z) = ForwardDiff.jacobian!(Cz,cI_aug,c,z)
    ∇cE_N(Cx,c,x) = ForwardDiff.jacobian!(Cx,cE_N!,c,x)
    ∇cI_N(Cx,c,x) = ForwardDiff.jacobian!(Cx,cI_N!,c,x)

    ∇cE_blocks = [(pE*(i-1) .+ (1:pE),(n+m)*(i-1) .+ (1:n+m)) for i = 1:N-1]
    push!(∇cE_blocks,(Nh-pE_N+1:Nh, Nz-n+1:Nz))
    ∇cI_blocks = [(pI*(i-1) .+ (1:pI),(n+m)*(i-1) .+ (1:n+m)) for i = 1:N-1]
    push!(∇cI_blocks,(Ng-pI_N+1:Ng, Nz-n+1:Nz))
    if solver.obj.use_xf_equality_constraint
        cE_N!(c,x) = copyto!(c,x-solver.obj.xf)
        ∇cE_N(Cx,c,x) = begin copyto!(Cx,I); cE_N!(c,x); end
    end

    function mycost(Z::PrimalVars)
        X,U = Z.X, Z.U
        J = 0.0
        for k = 1:N-1
            J += stage_cost(costfun,X[:,k],U[:,k])*solver.dt
        end
        J += stage_cost(costfun,X[:,N])
        return J
    end
    mycost(Z::AbstractVector) = mycost(PrimalVars(Z,ind_x,ind_u))
    mycost(V::NewtonVars) = mycost(V.Z)

    function cE(C,Z)
        Z̄ = PrimalVars(Z,ind_x,ind_u)
        X,U = Z̄.X, Z̄.U

        C_ = view(C,ind_h.k_mat)
        for k = 1:N-1
            cE!(view(C_,1:pE,k),X[:,k],U[:,k])
        end
        CN = view(C,ind_h.N)
        cE_N!(CN,X[:,N])
        return nothing
    end
    function cE(Z)
        C = zeros(eltype(Z),Nh)
        cE(C,Z)
        return C
    end

    function cI(C,Z)
        Z̄ = PrimalVars(Z,ind_x,ind_u)
        X,U = Z̄.X, Z̄.U

        C_ = view(C,ind_g.k_mat)
        for k = 1:N-1
            cI!(view(C_,1:pI,k),X[:,k],U[:,k])
        end
        CN = view(C,ind_g.N)
        cI_N!(CN,X[:,N])
        return nothing
    end
    function cI(Z)
        C = zeros(eltype(Z),Ng)
        cI(C,Z)
        return C
    end

    function dynamics(D,Z)
        Z̄ = PrimalVars(Z,ind_x,ind_u)
        X,U = Z̄.X, Z̄.U

        d = reshape(D,n,N)
        d[:,1] = X[:,1] - solver.obj.x0
        for k = 2:N
            solver.fd(view(d,1:n,k),X[:,k-1],U[:,k-1])
            d[:,k] -= X[:,k]
        end
    end
    function dynamics(Z)
        D = zeros(eltype(Z),Nx)
        dynamics(D,Z)
        return D
    end

    grad_J(grad,Z) = gradient!(grad,solver,PrimalVars(Z,ind_x,ind_u))
    grad_J(Z) = gradient(solver,PrimalVars(Z,ind_x,ind_u))

    hess_J(hess,Z) = copyto!(hess,∇²J)
    hess_J(Z) = ∇²J

    function jacob_dynamics(jacob,Z)
        Z̄ = PrimalVars(Z,ind_x,ind_u)
        xdot = zeros(n)
        jacob[1:n,1:n] = Diagonal(I,n)
        for k = 1:N-1
            off1 = k*n
            off2 = (k-1)*(n+m)
            block = view(jacob,off1 .+ (1:n),off2 .+ (1:n+m))
            Fd!(block, xdot, Z̄[k])
            jacob[off1 .+ (1:n),(off2+n+m) .+ (1:n)] = -Diagonal(I,n)
        end
    end
    function jacob_dynamics(Z)
        jacob = spzeros(Nx,Nz)
        jacob_dynamics(jacob,Z)
        return jacob
    end
    # Sparsity structure (1s in all non-zero elements)
    function jacob_dynamics()
        jacob = spzeros(Nx,Nz)
        jacob[1:n,1:n] = Diagonal(I,n)
        for k = 1:N-1
            off1 = k*n
            off2 = (k-1)*(n+m)
            block = view(jacob,off1 .+ (1:n),off2 .+ (1:n+m))
            block .= 1
            jacob[off1 .+ (1:n),(off2+n+m) .+ (1:n)] = Diagonal(I,n)
        end
        return jacob
    end

    function jacob_cE(jacob,Z)
        Z̄ = PrimalVars(Z,ind_x,ind_u)
        X,U = Z̄.X, Z̄.U
        c = zeros(pE)
        for k = 1:N-1
            block = view(jacob,∇cE_blocks[k][1],∇cE_blocks[k][2])
            ∇cE(block,c,Z̄[k])
        end
        block = view(jacob,∇cE_blocks[N][1],∇cE_blocks[N][2])
        c = zeros(pE_N)
        ∇cE_N(block,c,Z̄[N])
        return jacob
    end
    function jacob_cE(Z)
        jacob = spzeros(Nh,Nz)
        jacob_cE(jacob,Z)
        return jacob
    end

    function jacob_cI(jacob,Z)
        Z̄ = PrimalVars(Z,ind_x,ind_u)
        X,U = Z̄.X, Z̄.U
        c = zeros(pI)
        for k = 1:N-1
            block = view(jacob,∇cI_blocks[k][1],∇cI_blocks[k][2])
            ∇cI(block,c,Z̄[k])
        end
        block = view(jacob,∇cI_blocks[N][1],∇cI_blocks[N][2])
        c = zeros(pI_N)
        ∇cI_N(block,c,Z̄[N])
        return jacob
    end
    function jacob_cI(Z)
        jacob = spzeros(Ng,Nz)
        jacob_cI(jacob,Z)
        return jacob
    end

    function constraints(c,Z)
        dynamics(view(c,ind_c.d),Z)
        cE(view(c,ind_c.E),Z)
        cI(view(c,ind_c.I),Z)
        cb(view(c,ind_c.b),Z)
    end

    #function constraint_jacobian(jacob,Z)
    #    jacob_dynamics()

    # Assume only bound constraints right now
    jacob_bnd(jacob,Z) = copyto!(jacob,jac_bnd)
    jacob_bnd(Z) = jac_bnd

    c = (E=cE, I=cI, d=dynamics, b=bound_constraint)
    jacob_c = (E=jacob_cE, I=jacob_cI, d=jacob_dynamics, b=jacob_bnd)

    return mycost, c, grad_J, hess_J, jacob_c
end

"""$(SIGNATURES)
Full 2nd order taylor expansion of a Quadratic Cost with respect to Z = [X;U]
"""
function taylor_expansion(solver::Solver,Z::PrimalVars)
    costfun = solver.obj.cost
    dt = solver.dt
    n,m,N = get_sizes(Z)
    Nx = N*n
    Nu = (N-1)*m
    Nz = Nx + Nu
    hess = spzeros(Nz,Nz)
    grad = spzeros(Nz)

    ind1 = create_partition([n,m],(:x,:u))

    X,U = Z.X, Z.U
    for k = 1:N-1
        off = (k-1)*(n+m)
        Q,R,H,q,r = taylor_expansion(costfun,X[:,k],U[:,k])
        hess[off .+ ind1.x, off .+ ind1.x] = Q
        hess[off .+ ind1.x, off .+ ind1.u] = H'
        hess[off .+ ind1.u, off .+ ind1.x] = H
        hess[off .+ ind1.u, off .+ ind1.u] = R

        grad[off .+ ind1.x] = q
        grad[off .+ ind1.u] = r
    end
    hess .*= dt
    grad .*= dt
    off = (N-1)*(n+m)
    Qf, qf = taylor_expansion(costfun,X[:,N])
    hess[off .+ ind1.x, off .+ ind1.x] = Qf
    grad[off .+ ind1.x] = qf
    return hess, grad
end

"""
$(SIGNATURES)
Full 1st order taylor expansion of a Quadratic Cost with respect to Z = [X;U]
"""
function gradient!(grad::AbstractVector,solver::Solver,Z::PrimalVars)
    costfun = solver.obj.cost
    dt = solver.dt

    n,m,N = get_sizes(Z)
    Nx = N*n
    Nu = (N-1)*m
    Nz = Nx + Nu
    ind1 = create_partition([n,m],(:x,:u))

    X,U = Z.X, Z.U
    for k = 1:N-1
        off = (k-1)*(n+m)
        q,r = gradient(costfun,X[:,k],U[:,k])
        grad[off .+ ind1.x] = q
        grad[off .+ ind1.u] = r
    end
    grad .*= dt
    off = (N-1)*(n+m)
    qf = gradient(costfun,X[:,N])
    grad[off .+ ind1.x] = qf
    return nothing
end
function gradient(solver::Solver,Z::PrimalVars)
    n,m,N = get_sizes(Z)
    Nx = N*n
    Nu = (N-1)*m
    Nz = Nx + Nu
    grad = spzeros(Nx+Nu)
    gradient!(grad,solver,Z)
    return grad
end

function get_bounds(solver::Solver, equalities::Bool=true)
    Nx,Nu,Nh,Ng,NN = get_batch_sizes(solver)
    n,m,N = get_sizes(solver)
    obj = solver.obj

    x_L = repeat(obj.x_min,1,N)
    x_U = repeat(obj.x_max,1,N)
    u_L = repeat(obj.u_min,1,N-1)
    u_U = repeat(obj.u_max,1,N-1)

    if equalities
        # Initial Condition
        x_L[1:n,1] .= obj.x0
        x_U[1:n,1] .= obj.x0

        # Terminal Constraint
        if obj.use_xf_equality_constraint
            x_L[1:n,N] .= obj.xf
            x_U[1:n,N] .= obj.xf
        end
    end

    # Pack them together
    z_L = PrimalVars(x_L,u_L).Z
    z_U = PrimalVars(x_U,u_U).Z

    return z_L, z_U
end

function get_bound_multipliers(solver::Solver, res::ConstrainedVectorResults)
    n,m,N = get_sizes(solver)
    Nz = N*n + (N-1)*m
    λ = zeros(n+m,N-1)
    x_L = isfinite.(solver.obj.x_min)
    x_U = isfinite.(solver.obj.x_max)
    u_L = isfinite.(solver.obj.u_min)
    u_U = isfinite.(solver.obj.u_max)
    labels = get_constraint_labels(solver)
    inds = (x_U=labels .== "state (upper bound)",
            x_L=labels .== "state (lower bound)",
            u_U=labels .== "control (upper bound)",
            u_L=labels .== "control (lower bound)",
            x=1:n, u=n.+(1:m))
    r = zeros(Nz,2)
    for k = 1:N-1
        lambda_ = res.λ[k]
        lambda = zeros(n+m,2)
        lambda[inds.x[x_L],1] = lambda_[inds.x_L]
        lambda[inds.x[x_U],2] = lambda_[inds.x_U]
        lambda[inds.u[u_L],1] = lambda_[inds.u_L]
        lambda[inds.u[u_U],2] = lambda_[inds.u_U]
        r[(k-1)*(n+m) .+ (1:n+m),:] = lambda
    end
    return r
end
function get_bound_multipliers(solver::Solver, res::UnconstrainedVectorResults)
    n,m,N = get_sizes(solver)
    Nz = N*n + (N-1)*m
    return zeros(Nz,2)
end

# Newton Polishing
function newton_snopt(solver::Solver, res::SolverIterResults,tol=1e-10)
    mycost, c, grad_J, hess_J, jacob_c = gen_usrfun_newton(solver)
    Z = PrimalVars(res)

    function usrfun(Z)
        J = mycost(Z)
        gJ = grad_J(Z)

        cineq = Float64[]
        gcineq = Float64[]

        ceq = c.d(Z)
        gceq = jacob_c.d(Z)

        fail = false
        # return J, cineq, ceq, fail
        return J, cineq, ceq, gJ, gcineq, gceq, fail
    end

    lb,ub = get_bounds(solver)
    convertInf!(lb)
    convertInf!(ub)

    options = Dict{String, Any}()
    options["Derivative option"] = 0
    options["Verify level"] = 1
    options["Major feasibility tol"] = tol
    # options["Minor optimality  tol"] = solver.opts.eps_intermediate
    options["Major optimality  tol"] = tol
    sumfile = "logs/snopt-summary.out"
    printfile = "logs/snopt-print.out"

    problem = SnoptProblem(usrfun, Z.Z, lb, ub, options)
    t_eval = @elapsed Zopt,optval,status = Snopt.solve(problem)
    stats = parse_snopt_summary()
    stats["info"] = status
    stats["runtime"] = t_eval

    return Zopt,optval,stats
    Zopt = PrimalVars(Zopt,n,m,N)
    res_snopt = ConstrainedVectorResults(solver,Zopt.X,Zopt.U)
    backwardpass!(res_snopt,solver)
    rollout!(res_snopt,solver,0.0)
    return res_snopt
end

function newton_ipopt(solver::Solver,res::SolverIterResults)
    # Convert to PrimalVar
    Z = PrimalVars(res)

    # Get Constants
    n,m,N = get_sizes(solver)
    Nz = length(Z.Z)   # Number of decision variables
    Nx = N*n           # Number of collocation constraints
    nH = 0             # Number of entries in cost hessian

    # Generate functions
    eval_f, c, grad_J, hess_J, jacob_c = gen_usrfun_newton(solver)

    # Pre-allocate jacobian
    ∇c = jacob_c.d()
    row,col,v = findnz(∇c)
    nG = length(v)  # Number of entries in constraint jacobian

    # Set up functions for Ipopt
    eval_g(Z,g) = c.d(g,Z)
    eval_grad_f(Z, grad_f) = grad_J(grad_f,Z)
    function eval_jac_g(Z, mode, rows, cols, vals)
        if mode == :Structure
            copyto!(rows,row)
            copyto!(cols,col)
        else
            jacob_c.d(∇c,Z)
            copyto!(vals,nonzeros(∇c))
        end
    end

    # Get bounds
    lb,ub = get_bounds(solver)
    convertInf!(lb)
    convertInf!(ub)

    g_L = zeros(Nx)
    g_U = zeros(Nx)
    P = length(g_L)

    # Create Problem
    prob = Ipopt.createProblem(Nz, lb, ub, P,g_L, g_U, nG, nH,
        eval_f, eval_g, eval_grad_f, eval_jac_g)
    prob.x = Z.Z

    # Set options
    dir = root_dir()
    opt_file = joinpath(dir,"ipopt.opt")
    addOption(prob,"option_file_name",opt_file)
    if solver.opts.verbose == false
        addOption(prob,"print_level",0)
    end

    # Solve
    t_eval = @elapsed status = solveProblem(prob)
    stats = parse_ipopt_summary()
    stats["info"] = Ipopt.ApplicationReturnStatus[status]
    stats["runtime"] = t_eval
    return prob.x, stats

    # Convert back to feasible trajectory
    vars = PrimalVars(prob.x,n,m,N)
    res_ipopt = ConstrainedVectorResults(solver,Zopt.X,Zopt.U)
    backwardpass!(res_snopt,solver)
    rollout!(res_snopt,solver,0.0)
    return res_ipopt
end

function gen_newton_functions(solver::Solver)
    p,pI,pE = get_num_constraints(solver)
    p_N,pI_N,pE_N = get_num_terminal_constraints(solver)
    n,m,N = get_sizes(solver)

    obj = solver.obj
    costfun = obj.cost

    Nx,Nu,Nh,Ng,NN = get_batch_sizes(solver)
    Nz = Nx+Nu

    ind_x, ind_u = get_primal_inds(n,m,N)

    names = (:z,:ν,:λ,:μ,:r)
    ind1 = TrajectoryOptimization.create_partition([Nz,Nx,Nh,Ng,2Nz],names)
    ind2 = TrajectoryOptimization.create_partition2([Nz,Nx,Nh,Ng,2Nz],names)
    ind1 = merge(ind1,create_partition([Nx,Nu],(:x,:u)))
    ind1 = merge(ind1,create_partition([Nz,Nx+Nh+Ng+2Nz],(:z,:y)))
    ind2 = merge(ind2,create_partition2([Nz,Nx+Nh+Ng+2Nz],(:z,:y)))

    mycost, c, grad_J, hess_J, jacob_c = gen_usrfun_newton(solver)

    Z0 = PrimalVars(n,m,N)
    ∇²J = hess_J(Z0.Z)

    function active_set(Z,eps=1e-3)
        a = ones(Bool,NN)
        ci = c.I(Z)
        bi = c.b(Z)
        c_inds = ones(Bool,length(ci))
        c_inds = ci .>= -eps
        a[ind1.μ] = c_inds
        b_inds = ones(Bool,length(bi))
        b_inds = ci .>= -eps
        a[ind1.r] = b_inds
        return a
    end

    function active_set(V::NewtonVars,eps=1e-3)
        ci = c.I(primals(V))
        c_inds = ones(Bool,length(ci))
        c_inds = ci .>= -eps
        V.a[ind1.μ] = c_inds
        bi = c.b(primals(V))
        b_inds = ones(Bool,length(bi))
        b_inds = bi .>= -eps
        V.a[ind1.r] = b_inds
    end

    function buildKKT(V::NewtonVars,ρ,type=:kkt)
        Z = primals(V)
        A = spzeros(NN,NN)
        A[ind2.zz...] = ∇²J
        D = view(A,ind2.νz...)
        H = view(A,ind2.λz...)
        G = view(A,ind2.μz...)
        B = view(A,ind2.rz...)
        jacob_c.d(D,Z)
        jacob_c.E(H,Z)
        jacob_c.I(G,Z)
        jacob_c.b(B,Z)
        A = Symmetric(A')

        b = zeros(NN)
        d = view(b,ind1.ν)
        h = view(b,ind1.λ)
        g = view(b,ind1.μ)
        bnd = view(b,ind1.r)
        c.d(d,Z)
        c.E(h,Z)
        c.I(g,Z)
        c.b(bnd,Z)
        return A,b
    end

    #
    # function constraint_projection!(A,b,V::NewtonVars)
    #     Z = primals(V)
    #     D = view(A,ind2.νz)
    #     H = view(A,ind2.λz)
    #     G = view(A,ind2.μz)
    #     jacob_c.d(D,Z)
    #     jacob_c.E(H,Z)
    #     jacob_c.I(G,Z)
    #
    #     d = view(b,ind1.ν)
    #     h = view(b,ind1.λ)
    #     g = view(b,ind1.μ)
    #     c.d(d,Z)
    #     c.E(h,Z)
    #     c.I(g,Z)
    #
    #     y = view(b,ind1.y)
    #     Y = view(A,ind2.yz)
    #     δV̂ = -Y'*((Y*Y')\y)
    #     return δV̂
    # end
    #
    function newton_step(V::NewtonVars,ρ; type=:kkt, eps=1e-2, reg=Diagonal(zeros(length(V))),verbose=solver.opts.verbose)
        max_c2(V) = max(Ng > 0 ? maximum(c.I(V.V)) : 0, norm(c.E(V.V),Inf))
        max_c(V) = max(norm(c.d(V.V),Inf),max_c2(V))

        meritfun = mycost

        V_ = copy(V)
        active_set(V,eps)
        amu = V.a[ind1.μ]
        ar = V.a[ind1.r]
        Z_ = primals(V_)

        d1 = c.d(Z_)
        h1 = c.E(Z_)
        g1 = c.I(Z_)
        b1 = c.b(Z_)
        y = [d1;h1;g1[amu];b1[ar]]
        δz = zero(ind1.z)
        if verbose; println("max y: $(maximum(abs.(y)))") end
        while norm(y,Inf) > 1e-6
            D = jacob_c.d(Z_)
            H = jacob_c.E(Z_)
            G = jacob_c.I(Z_)
            B = jacob_c.b(Z_)
            Y = [D;H;G[amu,:];B[ar,:]]

            #δV̂ = -(∇²J\(Y'))*(Y*(∇²J\(Y'))\y)
            δZ = -Y'*((Y*Y')\y)
            Z_ .+= δZ

            d1 = c.d(Z_)
            h1 = c.E(Z_)
            g1 = c.I(Z_)
            b1 = c.b(Z_)
            y = [d1;h1;g1[amu];b1[ar]]
            if verbose; println("max y: $(maximum(abs.(y)))") end
        end

        J0 = mycost(V_)
        if verbose; println("Initial Cost: $J0") end

        # Build and solve KKT
        A,b = buildKKT(V_,ρ,type)
        active_set(V_,eps)
        a = V_.a
        amu = a[ind1.μ]
        ar = a[ind1.r]
        δV = zero(V.V)
        Ā = A[a,a] + reg[a,a]
        δV[a] = -Ā\b[a]

        # Line Search
        ϕ=0.01
        α = 2
        V1 = copy(V_)
        Z1 = primals(V1)
        δV1 = α.*δV
        J = J0+1e8
        while J > J0 #+ α*ϕ*b'δV1
            α *= 0.5
            δV1 = α.*δV
            V1.V  .= V_.V + δV1

            d1 = c.d(Z1)
            h1 = c.E(Z1)
            g1 = c.I(Z1)
            b1 = c.b(Z1)
            y = [d1;h1;g1[amu];b1[ar]]
            if verbose; println("max y: $(maximum(abs.(y)))") end
            while norm(y,Inf) > 1e-6
                D = jacob_c.d(Z1)
                H = jacob_c.E(Z1)
                G = jacob_c.I(Z1)
                B = jacob_c.b(Z1)
                Y = [D;H;G[amu,:];B[ar,:]]

                #δV̂ = -(∇²J\(Y'))*(Y*(∇²J\(Y'))\y)
                δZ = -Y'*((Y*Y')\y)
                Z1 .+= δZ

                d1 = c.d(Z1)
                h1 = c.E(Z1)
                g1 = c.I(Z1)
                b1 = c.b(Z1)
                y = [d1;h1;g1[amu];b1[ar]]
                if verbose; println("max y: $(maximum(abs.(y)))") end
            end

            J = meritfun(V1)
            if verbose; println("New Cost: $J") end
        end

        # Multiplier projection
        ∇J = grad_J(Z1)
        d1 = c.d(Z1)
        h1 = c.E(Z1)
        g1 = c.I(Z1)
        b1 = c.b(Z1)
        y = [d1;h1;g1[amu];b1[ar]]
        D = jacob_c.d(Z1)
        H = jacob_c.E(Z1)
        G = jacob_c.I(Z1)
        B = jacob_c.b(Z1)
        Y = Array([D;H;G[amu,:];B[ar,:]])

        ν = V1.ν
        μ = V1.μ
        λ = V1.λ
        r = V1.r
        lambda = [ν;λ;μ[amu];r[ar]]
        res = ∇J + Y'lambda
        δlambda = zero(lambda)
        if verbose; println("max residual before: $(norm(res,Inf))") end
        δlambda -= (Y*Y')\(Y*res)
        lambda1 = lambda + δlambda
        r = ∇J + Y'lambda1
        if verbose; println("max residual after: $(norm(res,Inf))") end
        ν1 = lambda1[1:Nx]
        λ1 = lambda1[Nx .+ (1:Nh)]
        μ1 = lambda1[(Nx+Nh) .+ (1:count(amu))]
        r1 = lambda1[(Nx+Nh+count(amu)) .+ (1:count(ar))]
        V1.ν .= ν1
        V1.λ .= λ1
        V1.μ[amu] = μ1
        V1.r[ar] = r1
        J = meritfun(V1)
        if verbose; println("New Cost: $J") end

        # Take KKT Step
        cost = J
        A,b = buildKKT(V1,ρ,type)
        active_set(V1,eps)
        a = V1.a
        grad = norm(b[a])
        c_max = max_c(V1)

        change(x,x0) = format((x0-x)/x0*100,precision=4) * "%"
        if verbose
            println()
            println("  cost: $J $(change(J,J0))")
            println("  step: $(norm(δV1))")
            println("  grad: $(grad)")
            println("  c_max: $(c_max)")
            println("  c_max2: $(max_c2(V1))")
            println("  α: $α")
            # println("  rank: $(rank(Ā))")
            # println("  cond: $(cond(Ā))")
        end
        stats = Dict("cost"=>cost,"grad"=>grad,"c_max"=>c_max)
        return V1
    end

    return newton_step, buildKKT, active_set
end

function newton_projection(solver::Solver,res::SolverIterResults; kwargs...)
    V = NewtonVars(solver,res)
    newton_step, = gen_newton_functions(solver)
    V_ = newton_step(V,0; kwargs...)
end


text:
"""
    validate_date(date::AbstractString)
    validate_date(date::Date)

Returns `month/day/year` string in format accepted by city
calendar. Valid input formats are `"<M>/<D>/<YYYY>"` (`6/1/2017`, `06/01/2017`)
or `Date(yyyy,mm,dd)` (`Date(2017,6,1)`).
"""
function validate_date(date::AbstractString)
    d_split = replace(date, "-" => "/")
    d_split = split(d_split, "/")
    length(d_split) == 3 ||
        throw(ArgumentError("Invalid date: `$date`. Must take form mm-dd-yyyy or mm/dd/yyyy"))
    y = d_split[3]
    length(y) < 4 && throw(ArgumentError("Full 4-digit year required (`$y` invalid)"))
    y = parse(Int, y)
    m = parse(Int, d_split[1])
    d = parse(Int, d_split[2])
    return validate_date(Date(y, m, d))
end

validate_date(date::Date) = string(month(date), "/", day(date), "/", year(date))

agenda_dateformat = dateformat"E, U d, YYYY  H:MM p"

const meeting_version = 1
const meeting_schema = Legolas.Schema("meeting", meeting_version)
const Meeting = Legolas.@row("meeting@1", name::AbstractString, date::DateTime, id::Int,
                             link::AbstractString)

"""
    request_meetings(start_date, stop_date)

Return list of meetings that occurred (or will occur) between `start_date` and `stop_date`.
"""
function request_meetings(start_date, stop_date)
    start = validate_date(start_date)
    stop = validate_date(stop_date)
    url = "http://somervillecityma.iqm2.com/Citizens/calendar.aspx?View=List&From=$start&To=$stop"
    @debug "Requesting $url" start stop
    r = HTTP.get(url)
    html = with_logger(NullLogger()) do
        return root(parsehtml(String(r.body)))
    end
    meeting_nodes = findall("//div[@class='RowLink']", html)
    return DataFrame(map(meeting_nodes) do m
                         deets = split(elements(m)[1]["title"], "\r")
                         name = replace(deets[3], "Board:\t" => "")
                         link = joinpath("http://somervillecityma.iqm2.com",
                                         elements(m)[1]["href"][2:end])
                         date = DateTime(deets[1], agenda_dateformat)
                         id = parse(Int, split(link, "ID=")[2])
                         return Meeting(; name, date, id, link)
                     end)
end


text:
using MAGMA
using Test, Random, LinearAlgebra, CUDA

using Test, LinearAlgebra, CUDA

@testset "dense linear algebra" begin
    include("dense/svd.jl")
    include("dense/linearsystemsolver.jl")
    include("dense/eigen.jl")
end


text:
# Functions to compute the log probability of star parameters and galaxy
# parameters given pixel data
using Distributions
import ..SensitiveFloats: SensitiveFloat
import JLD

const eps_prob_a = 1e-6

###################################
# log likelihood function makers  #
###################################

"""
Creates a vectorized version of the star logpdf as a function of
unconstrained params.
    star_params = [lnr, lnc1, ..., lnc4, ra, dec]
Args:
  - images: Vector of Image types (data for log_likelihood)
  - S: --- from ElboArgs
  - N:
Returns:
  - star_logpdf  : unnormalized logpdf function handle that takes in a flat,
                   unconstrained array as parameter
  - star_logprior: star param log prior function handle that takes in same
                   flat, unconstrained array as parameter
"""
function make_star_logpdf(images::Vector{<:Image},
                          S::Int64,
                          N::Int64,
                          source_params::Vector{Vector{Float64}},
                          patches::Matrix{ImagePatch},
                          active_sources::Vector{Int},
                          psf_K::Int64)
    # define star prior log probability density function
    prior = construct_prior()
    subprior = prior.star

    function star_logprior(state::Vector{Float64})
        brightness, colors, pos = state[1], state[2:5], state[6:end]
        return color_logprior(brightness, colors, prior, true)
    end

    # define star log joint probability density function
    function star_logpdf(state::Vector{Float64})
        ll_prior = star_logprior(state)
        brightness, colors, position = state[1], state[2:5], state[6:end]
        dummy_gal_shape = [.1, .1, .1, .1]
        ll_like  = state_log_likelihood(true, brightness, colors, position,
                                        dummy_gal_shape, images,
                                        patches,
                                        active_sources,
                                        psf_K,
                                        source_params,
                                        S, N)
        return ll_like + ll_prior
    end

    return star_logpdf, star_logprior
end

"""
Creates a vectorized version of a galaxy logpdf as a function of
unconstrained params
    gal_params = [lnr, lnc1, ..., lnc4, ra, dec, shp1, shp2, shp3, shp4]
where the flux and position params are the same as the star case, and the
shape params are unconstrained versions of (gdev, gaxis, gangle, gscale)
Args:
  - images: Vector of Image types (data for log_likelihood)
  - ea: ElboArgs book keeping argument
Returns:
  - gal_logpdf  : unnormalized logpdf function handle that takes in a flat,
                 unconstrained array as parameter
  - gal_logprior: star param log prior function handle that takes in same
                 flat, unconstrained array as parameter
"""
function make_galaxy_logpdf(images::Vector{<:Image},
                            S::Int64,
                            N::Int64,
                            source_params::Vector{Vector{Float64}},
                            patches::Matrix{ImagePatch},
                            active_sources::Vector{Int},
                            psf_K::Int64)
    # define galaxy prior function
    prior    = construct_prior()
    subprior = prior.galaxy

    function galaxy_logprior(state::Vector{Float64})
        brightness, colors, pos, gal_shape =
            state[1], state[2:5], state[6:7], state[8:end]

        # brightness prior
        ll_b = color_logprior(brightness, colors, prior, true)
        ll_s = shape_logprior(constrain_gal_shape(gal_shape), prior)
        return ll_b + ll_s
    end

    # define galaxy log joint probability density function
    function galaxy_logpdf(state::Vector{Float64})
        ll_prior = galaxy_logprior(state)
        brightness, colors, position, gal_shape =
            state[1], state[2:5], state[6:7], state[8:end]
        ll_like  = state_log_likelihood(false, brightness, colors, position,
                                        constrain_gal_shape(gal_shape), images,
                                        patches,
                                        active_sources, psf_K,
                                        source_params,
                                        S, N)
        return ll_like + ll_prior
    end

    return galaxy_logpdf, galaxy_logprior
end

"""
Log likelihood of a single source given source params.
Args:
  - is_star: bool describing the type of source
  - brightness: log r-band value
  - colors: array of colors
  - position: ra/dec of source
  - gal_shape: vector of galaxy shape params (used if is_star=false)
  - images: vector of Image types (data for likelihood)
  - ea: ElboArgs object that maintains params for all sources
Returns:
  - result: a scalar describing the log likelihood of the
            (brightness,colors,position,gal_shape) params conditioned on
            the rest of the args
"""
function state_log_likelihood(is_star::Bool,                # source is star
                              brightness::Float64,          # source log r brightness
                              colors::Vector{Float64},      # source vector of colors
                              position::Vector{Float64},    # source position
                              gal_shape::Vector{Float64},   # source gal shape
                              images::Vector{<:Image},   # list of images with source pixel data
                              patches::Matrix{ImagePatch},    # formerly of ElboArgs
                              active_sources::Vector{Int},  # formerly of ElboArgs
                              psf_K::Int64,                 # number of PSF Comps
                              source_params::Vector{Vector{Float64}}, # list of background sources
                              S::Int64,
                              N::Int64)
    # TODO: cache the background rate image!! --- does not need to be recomputed at each ll eval
    # convert brightness/colors to fluxes for scaling
    fluxes = colors_to_fluxes(brightness, colors)

    # create objects needed to compute the mean poisson value per pixel
    # (similar to ElboDeriv.process_active_pixels!)
    bvn_derivs = BivariateNormalDerivatives{Float64}()
    fs0m = SensitiveFloat{Float64}(length(StarPosParams), 1, false, false)
    fs1m = SensitiveFloat{Float64}(length(GalaxyPosParams), 1, false, false)

    # load star/gal mixture components (make sure these reflect
    gdev, gaxis, gangle, gscale = gal_shape
    source_params[1][lidx.pos]       = position
    source_params[1][lidx.gal_frac_dev]   = gdev
    source_params[1][lidx.gal_axis_ratio]  = gaxis
    source_params[1][lidx.gal_angle] = gangle
    source_params[1][lidx.gal_radius_px] = gscale

    # iterate over the pixels, summing pixel-specific poisson rates
    ll = 0.

    @assert length(active_sources) == 1

    for n in 1:N
        img = images[n]

        star_mcs, gal_mcs = load_bvn_mixtures(S, patches,
                              source_params, active_sources, psf_K, n,
                              calculate_gradient=false,
                              calculate_hessian=false)

        p = patches[active_sources[1], n]
        H2, W2 = size(p.active_pixel_bitmap)
        for w2 in 1:W2, h2 in 1:H2
            # (h2, w2) index the local patch, while (h, w) index the image
            h = p.bitmap_offset[1] + h2
            w = p.bitmap_offset[2] + w2

            if !p.active_pixel_bitmap[h2, w2]
                continue
            end

            # compute the unit-flux pixel values
            pixel_rate = img.sky[h, w]
            for s in 1:S
                params_s  = source_params[s]

                star_light_density!(fs0m, p, h, w, params_s[lidx.pos], false)
                populate_gal_fsm!(fs1m, bvn_derivs, s, h, w, false, p.wcs_jacobian, gal_mcs)

                if s == 1
                    this_rate  = is_star ? fs0m.v[] : fs1m.v[]
                    pixel_rate += fluxes[img.b] * this_rate
                else
                    # determine if background source is star/gal; get fluxes
                    s_is_star = params_s[lidx.is_star[1,1]] > .5
                    type_idx  = s_is_star ? 1 : 2
                    flux_s    = colors_to_fluxes(params_s[lidx.flux[type_idx]],
                                                 params_s[lidx.color[:,type_idx]])
                    rate_s = s_is_star ? fs0m.v[] : fs1m.v[]
                    rate_s *= flux_s[img.b]
                    pixel_rate += rate_s
                end
            end

            # multiply by image's gain for this pixel
            rate = pixel_rate * img.nelec_per_nmgy[h]
            pixel_ll = logpdf(Poisson(rate[1]), round(Int, img.pixels[h, w]))
            ll += pixel_ll
        end
    end

    return ll
end

###########################################################################
# priors - for computing prior log probs of star/gal parameters           #
###########################################################################

struct StarPrior
    brightness::LogNormal
    color_component::Categorical
    colors::Vector{MvNormal}
end

struct GalaxyPrior
    brightness::LogNormal
    color_component::Categorical
    colors::Vector{MvNormal}
    gal_radius_px::LogNormal
    gal_axis_ratio::Beta
    gal_frac_dev::Beta
end

struct Prior
    is_star::Bernoulli
    star::StarPrior
    galaxy::GalaxyPrior
end

"""
Construct a `Prior` object, contains StarPrior and GalaxyPrior objects
"""
function construct_prior()
    pp = load_prior()
    star_prior = StarPrior(
        LogNormal(pp.flux_mean[1], sqrt(pp.flux_var[1])),
        Categorical(pp.k[:,1]),
        [MvNormal(pp.color_mean[:,k,1], pp.color_cov[:,:,k,1]) for k in 1:NUM_COLOR_COMPONENTS])

    gal_prior = GalaxyPrior(
        LogNormal(pp.flux_mean[2], sqrt(pp.flux_var[2])),
        Categorical(pp.k[:,2]),
        [MvNormal(pp.color_mean[:,k,2], pp.color_cov[:,:,k,2]) for k in 1:NUM_COLOR_COMPONENTS],
        LogNormal(pp.gal_radius_px_mean, sqrt(pp.gal_radius_px_var)),
        Beta(1, 1),
        Beta(1, 1))

    return Prior(Bernoulli(.5), star_prior, gal_prior)
end

"""
Return the log prior probability of a (brightness,color) pair based on
prior probability objects constructed by `construct_prior()`
"""
function color_logprior(brightness::Float64,
                        colors::Vector{Float64},
                        prior::Prior,
                        is_star::Bool)
    subprior      = is_star ? prior.star : prior.galaxy
    ll_brightness = logpdf(subprior.brightness, exp(brightness))
    ll_component  = [logpdf(subprior.colors[k], colors) for k in 1:NUM_COLOR_COMPONENTS]
    ll_color      = logsumexp(ll_component + log.(subprior.color_component.p))
    return ll_brightness + ll_color
end

function shape_logprior(gal_shape::Vector{Float64}, prior::Prior)
    # position and gal_angle have uniform priors--we ignore them
    gdev, gaxis, gangle, gscale = gal_shape
    ll_shape = 0.
    ll_shape += logpdf(prior.galaxy.gal_radius_px, gscale)
    ll_shape += logpdf(prior.galaxy.gal_axis_ratio, gaxis)
    ll_shape += logpdf(prior.galaxy.gal_frac_dev, gdev)
    return ll_shape
end

#################################################################
# transformation of fluxes = (u, g, r, i, z) in nanomaggies to  #
# (bright,color) pairs                                          #
#################################################################

"""
Translate between fluxes (nmgy rates) and colors (on which
LogNormal priors are placed)
"""
function fluxes_to_colors(fluxes::Vector{Float64})
    lnr    = log(fluxes[3])
    colors = Vector{Float64}(4)
    colors[1] = log(fluxes[2]) - log(fluxes[1])
    colors[2] = log(fluxes[3]) - log(fluxes[2])
    colors[3] = log(fluxes[4]) - log(fluxes[3])
    colors[4] = log(fluxes[5]) - log(fluxes[4])
    return lnr, colors
end

"""
Translate from the (brightness, color) parameterization to nmgy fluxes
"""
function colors_to_fluxes(brightness::Float64, colors::Vector{Float64})
    # build up log fluxes
    ret    = Vector{Float64}(5)
    lnr    = brightness
    ret[3] = lnr     # r flux
    ret[4] = lnr + colors[3]        # ln(i/r) = c3 => lni = lnr - c3
    ret[5] = ret[4] + colors[4]     # ln(z/i) = c4 => lnz = lni - c4
    ret[2] = -colors[2] + lnr       # ln(r/g) = color_var => lng = color_var + lnr
    ret[1] = -colors[1] + ret[2]    # ln(g/u) = color_mean => lnu = color_mean + lng
    return exp.(ret)
end

#######################################################################
# galaxy shape transformation (unconstrained => constrained and back) #
#######################################################################

function constrain_gal_shape(unc_gal_shape::Vector{Float64})
    gdev, gaxis, gangle, gscale = unc_gal_shape
    constr_shape    = Vector{Float64}(4)
    constr_shape[1] = clamp(sigmoid(gdev), eps_prob_a, 1-eps_prob_a)
    constr_shape[2] = clamp(sigmoid(gaxis), eps_prob_a, 1-eps_prob_a)
    constr_shape[3] = gangle       # TODO put this between [0, 2pi]
    constr_shape[4] = clamp(exp(gscale), eps_prob_a, Inf)
    return constr_shape
end

function unconstrain_gal_shape(con_gal_shape::Vector{Float64})
    gdev, gaxis, gangle, gscale = con_gal_shape
    unc_shape    = Vector{Float64}(4)
    unc_shape[1] = logit(gdev)
    unc_shape[2] = logit(gaxis)
    unc_shape[3] = gangle
    unc_shape[4] = log(gscale)
    return unc_shape
end

######################################################
# initialize from catalog and variational parameters #
######################################################

function init_star_state(entry::CatalogEntry)
    brightness, colors = fluxes_to_colors(entry.star_fluxes)
    param_vec = [brightness; colors; entry.pos]
end

function init_galaxy_state(entry::CatalogEntry)
    brightness, colors = fluxes_to_colors(entry.gal_fluxes)
    #gdev, gaxis, gangle, gscale = gal_shape
    gal_shape = unconstrain_gal_shape([
                    clamp(entry.gal_frac_dev, eps_prob_a, 1.-eps_prob_a),
                    clamp(entry.gal_axis_ratio, eps_prob_a, 1.-eps_prob_a),
                    entry.gal_angle,
                    clamp(entry.gal_radius_px, eps_prob_a, Inf)
                    ])
    param_vec = [brightness; colors; entry.pos; gal_shape]
    return param_vec
end

function catalog_entry_to_latent_state_params(ce::CatalogEntry)
    # create a float array of the appropriate length
    ret = Vector{Float64}(length(lidx))

    # galaxy shape params
    ret[lidx.pos]       = ce.pos
    ret[lidx.gal_frac_dev]   = ce.gal_frac_dev
    ret[lidx.gal_axis_ratio]  = ce.gal_axis_ratio
    ret[lidx.gal_radius_px] = ce.gal_radius_px
    ret[lidx.gal_angle] = ce.gal_angle

    # star, gal r flux
    star_lnr, star_cols = fluxes_to_colors(ce.star_fluxes)
    gal_lnr, gal_cols   = fluxes_to_colors(ce.gal_fluxes)
    ret[lidx.flux]         = [star_lnr, gal_lnr]
    ret[lidx.color]         = hcat([star_cols, gal_cols]...)

    # set the prob star/prob gal
    ret[lidx.is_star] = clamp(ce.is_star, eps_prob_a, 1-eps_prob_a)
    ret
end

function extract_star_state(ls::Array{Float64, 1})
    return [[ls[lidx.flux[1]]]; ls[lidx.color[:, 1]]; ls[lidx.pos]]
end

function extract_galaxy_state(ls::Array{Float64, 1})
    gal_shape = [ls[lidx.gal_frac_dev]  ,
                 ls[lidx.gal_axis_ratio] ,
                 ls[lidx.gal_angle],
                 ls[lidx.gal_radius_px]]
    return [[ls[lidx.flux[2]]]; ls[lidx.color[:, 2]]; ls[lidx.pos];
            unconstrain_gal_shape(gal_shape)]
end

###################
# other util funs #
###################

#TODO is there a better place for this generic function? --- acm
function logsumexp(a::Vector{Float64})
    a_max = maximum(a)
    out = log(sum(exp, a - a_max))
    return out + a_max
end

function logit(a::Float64)
    return log(a) - log(1.0-a)
end

function sigmoid(a::Float64)
    return 1. / (1. + exp(-a))
end


text:
# scopes-of-variables.jl
# Reading through
#
# https://docs.julialang.org/en/v1/manual/variables-and-scoping/
#
# and playing around with the various things I find in there.

------------------------------------------------------------


text:
module Js2x

# Write your package code here.

export Profile, ProfileDataPoint, point_counts, avg_brightness, laser_on_times, read_edger_board

const MAX_LENGTH = 21
const MAX_PROFILES = MAX_LENGTH * 12
const MAX_VERTICAL = 243
mutable struct ProfileDataPoint
    x::Int32
    y::Int32
    brightness::Int32
end
mutable struct Profile
	sequenceNumber::Int32
	location::Int32
    sendLocation::Int32
	laserOnTime::Int32
	timeInHead::Int32
	inputs::Int32
	flags::Int32
	laserIndex::Int32
	reserved2::Int32
	numberPoints::Int32
	data::Vector{ProfileDataPoint}
end

point_counts(profiles) = map(prof->prof.numberPoints,profiles)
laser_on_times(profiles) = map(prof->prof.laserOnTime,profiles)
avg_brightness(profiles) =  map(prof->mean(map(pd->pd.brightness,prof.data)),profiles)

function read_profile(file)
    sequenceNumber = read(file,Int32)
    location = read(file,Int32)
    sendLocation = read(file,Int32)
    laserOnTime = read(file,Int32)
    timeInHead = read(file,Int32)
    inputs = read(file,Int32)
    flags = read(file,Int32)
    laserIndex = read(file,Int32)
    reserved2 = read(file,Int32)
    numberPoints = read(file,Int32)
    data = Vector{ProfileDataPoint}(undef,numberPoints)
    for i = 1:MAX_VERTICAL
        x = read(file,Int32)
        y = read(file,Int32)
        brightness = read(file,Int32)
        if (i <= numberPoints)
            data[i] = ProfileDataPoint(x,y,brightness)
        end
    end
    return Profile(sequenceNumber,location,sendLocation,laserOnTime,timeInHead,inputs,flags,laserIndex,reserved2,numberPoints,data)
end

function read_edger_board(filename)
    file = open(filename)
    seek(file,796184)
    profiles = Array{Profile}(undef,MAX_PROFILES,2)
    for i=1:MAX_PROFILES
        for j = 1:2
            profiles[i,j] = read_profile(file)
        end
    end
    return profiles
end

end


text:
"""
    reinstantiate(old, pairs...)

create a new struct which is a copy of an existing struct with specific fields changed.

source: https://stackoverflow.com/a/58613307
"""
function reinstantiate(old::T, pairs::Pair...) where {T}
    field_values = [getfield(old, field) for field in fieldnames(T)]
    for pair in pairs
        index = findfirst(fieldnames(T) .== pair.first)
        @assert index !== nothing "$(pair.first) is not a field of $(T)"
        field_values[index] = pair.second
    end
    return T(field_values...)
end


text:
module Word2Vec

using Base.Collections      # for priority queue
using Base.Cartesian        # for @nexprs
using Distances
using NumericExtensions
using Blocks
using Compat

if isless(Base.VERSION, v"0.4.0-")
using Iterators
end

export LinearClassifier, train_one, WordEmbedding, train, accuracy
export save, restore
export find_nearest_words

include("utils.jl")
include("tree.jl")
include("word_stream.jl")
include("softmax_classifier.jl")
include("train.jl")
include("query.jl")

end # module


text:
# Standard translational or rotational joint
abstract type Joint{T,N} <: AbstractJoint{T,N} end

Joint0 = Joint{T,0} where T
Joint1 = Joint{T,1} where T
Joint2 = Joint{T,2} where T
Joint3 = Joint{T,3} where T

Base.show(io::IO, joint::Joint) = summary(io, joint)

### Constaint and nullspace matrices
@inline constraintmat(::Joint0{T}) where T = szeros(T,0,3)
@inline nullspacemat(::Joint0{T}) where T = SMatrix{3,3,T,9}(I)
@inline constraintmat(joint::Joint1) = joint.V3
@inline nullspacemat(joint::Joint1) = joint.V12
@inline constraintmat(joint::Joint2) = joint.V12
@inline nullspacemat(joint::Joint2) = joint.V3
@inline constraintmat(::Joint3{T}) where T = SMatrix{3,3,T,9}(I)
@inline nullspacemat(::Joint3{T}) where T = szeros(T,0,3)

### Constraints and derivatives
## Discrete-time position wrappers (for dynamics)
g(joint::Joint, statea::State, stateb::State, Δt) = g(joint, posargsnext(statea, Δt)..., posargsnext(stateb, Δt)...)
g(joint::Joint, stateb::State, Δt) = g(joint, posargsnext(stateb, Δt)...)
g(joint::Joint, statea::State, stateb::State) = g(joint, posargsc(statea)..., posargsc(stateb)...)
g(joint::Joint, stateb::State) = g(joint, posargsc(stateb)...)

@inline g(joint::Joint{T,N}) where {T,N} = szeros(T, N)

### Force derivatives (for linearization)
## Forcing
@inline function setForce!(joint::Joint, Fτ::SVector)
    joint.Fτ = zerodimstaticadjoint(nullspacemat(joint)) * Fτ
    return
end
@inline setForce!(joint::Joint) = return

## Derivative wrappers
@inline function ∂Fτ∂ua(joint::Joint, body1::Body, body2::Body)
    return ∂Fτ∂ua(joint, body1.state, body2.state) * zerodimstaticadjoint(nullspacemat(joint))
end
@inline function ∂Fτ∂ub(joint::Joint, body1::Body, body2::Body)
    if body2.id == joint.childid
        return ∂Fτ∂ub(joint, body1.state, body2.state) * zerodimstaticadjoint(nullspacemat(joint))
    else
        return ∂Fτ∂ub(joint)
    end
end
@inline function ∂Fτ∂ub(joint::Joint, body1::Origin, body2::Body)
    if body2.id == joint.childid
        return return ∂Fτ∂ub(joint, body2.state) * zerodimstaticadjoint(nullspacemat(joint))
    else
        return ∂Fτ∂ub(joint)
    end
end

@inline ∂Fτ∂ub(joint::Joint{T,N}) where {T,N} = szeros(T, 6, 3 - N)

### Minimal coordinates
@inline minimalCoordinates(joint::Joint{T,N}) where {T,N} = szeros(T, 3 - N)


text:
using FastBEAST
using StaticArrays
using LinearAlgebra
using Printf
CM = CompScienceMeshes

## Example1 Evenly distributed sphere
sphere = CM.meshsphere(1,0.05)
spoints = sphere.vertices
spoints = spoints[2:length(spoints)] # First point is [0,0,0]
npoints = length(spoints)

## Example2 Evenly distributed cube
x = 1
y = 1
z = 1
h = 0.1
cuboid = CM.meshcuboid(x, y, z,h)
spoints = cuboid.vertices
npoints = length(spoints)

## Example2 Evenly distributed cylinder
radius = 1.0
height = 2.0
h = 0.1
cylinder = CM.meshcylinder(;radius, height, h)
spoints = cylinder.vertices
spoints = spoints[2:length(spoints)]
npoints = length(spoints)

##
@views OneoverRkernelassembler(matrix, tdata, sdata) = assembler(
    OneoverRkernel,
    matrix,
    spoints[tdata],
    spoints[sdata]
)

stree = create_tree(spoints, BoxTreeOptions(nmin=50))
kmat = assembler(OneoverRkernel, spoints, spoints)
@time hmat = HMatrix(OneoverRkernelassembler, stree, stree, compressor=:aca, T=Float64)

@printf("Accuracy test: %.2e\n", estimate_reldifference(hmat,kmat))
@printf("Compression rate: %.2f %%\n", compressionrate(hmat)*100)

stree = create_tree(spoints, KMeansTreeOptions(iterations=100,nchildren=2,nmin=50))
kmat = assembler(OneoverRkernel, spoints, spoints)
@time hmat = HMatrix(OneoverRkernelassembler, stree, stree, compressor=:aca, T=Float64)

@printf("Accuracy test: %.2e\n", estimate_reldifference(hmat,kmat))
@printf("Compression rate: %.2f %%\n", compressionrate(hmat)*100)

## Plot Figure
scatter(
    [spoints[i][1] for i=1:npoints],
    [spoints[i][2] for i=1:npoints],
    [spoints[i][3] for i=1:npoints]
)

max_stars_repo_name: sbadrian/FastBEAST

max_issues_repo_name: sbadrian/FastBEAST

max_forks_repo_name: sbadrian/FastBEAST


text:
##
using MRIgeneralizedBloch
using DifferentialEquations
using StaticArrays
using Test
R2slT = precompute_R2sl()

## choose random parameters
Npulse = 500
α = π/2 * rand() * sin.(8π * rand() * (1:Npulse) / Npulse)
α[1] = π
TRF = 100e-6 .+ 400e-6 * rand(Npulse)
TRF[1] = 500e-6

TR = 3.5e-3

B1 = 0.7 + 0.6 * rand()
ω0 = 1000 * rand()
m0s = 0.4 * rand()
R1f = 0.15 + 0.2 * rand()
R1s = 2 + rand()
R2f = 1 / (40e-3 + 100e-3 * rand())
T2s = 5e-6 + 10e-6 * rand()
Rx = 10 + 10 * rand()

R1a = 0.5 + rand()

rtol = 1e-5

## m0s
Δm0s = 1e-6
s0 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list=[grad_m0s()])
s1 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s + Δm0s, R1f, R2f, Rx, R1s, T2s, R2slT)

g = s0[:,1,2]
gfd = vec(s1 .- s0[:,1,1]) / Δm0s
@test g ≈ gfd rtol = rtol

## R1f
ΔR1f = 1e-6
s0 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list=[grad_R1f()])
s1 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f + ΔR1f, R2f, Rx, R1s, T2s, R2slT)

g = s0[:,1,2]
gfd = vec(s1 .- s0[:,1,1]) / ΔR1f
@test g ≈ gfd rtol = rtol

## R2f
ΔR2f = 1e-6
s0 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list=[grad_R2f()])
s1 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f + ΔR2f, Rx, R1s, T2s, R2slT)

g = s0[:,1,2]
gfd = vec(s1 .- s0[:,1,1]) / ΔR2f
@test g ≈ gfd rtol = rtol

## Rx
ΔRx = 1e-6
s0 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list=[grad_Rx()])
s1 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx + ΔRx, R1s, T2s, R2slT)

g = s0[:,1,2]
gfd = vec(s1 .- s0[:,1,1]) / ΔRx
@test g ≈ gfd rtol = rtol

## R1s
ΔR1s = 1e-6
s0 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list=[grad_R1s()])
s1 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s + ΔR1s, T2s, R2slT)

g = s0[:,1,2]
gfd = vec(s1 .- s0[:,1,1]) / ΔR1s
@test g ≈ gfd rtol = rtol

## T2s
ΔT2s = 1e-9
s0 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list=[grad_T2s()])
s1 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s + ΔT2s, R2slT)

g = s0[:,1,2]
gfd = vec(s1 .- s0[:,1,1]) / ΔT2s
@test g ≈ gfd rtol = 5e-3 # not as precise because of the Higham's Complex Step Approximation

## ω0
Δω0 = 1e-6
s0 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list=[grad_ω0()])
s1 = calculatesignal_linearapprox(α, TRF, TR, ω0 + Δω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT)

g = s0[:,1,2]
gfd = vec(s1 .- s0[:,1,1]) / Δω0
@test g ≈ gfd rtol = rtol

## B1
ΔB1 = 1e-6
s0 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list=[grad_B1()])
s1 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1 + ΔB1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT)

g = s0[:,1,2]
gfd = vec(s1 .- s0[:,1,1]) / ΔB1
@test g ≈ gfd rtol = 1e-2 # not as precise because of the Higham's Complex Step Approximation

## R1a
ΔR1a = 1e-6
s0 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1a, R2f, Rx, R1a, T2s, R2slT, grad_list=[grad_R1a()])
s1 = calculatesignal_linearapprox(α, TRF, TR, ω0, B1, m0s, R1a + ΔR1a, R2f, Rx, R1a + ΔR1a, T2s, R2slT)

g = s0[:,1,2]
gfd = vec(s1 .- s0[:,1,1]) / ΔR1a
@test g ≈ gfd rtol = rtol