text:
#=
Need plots for:

orbits.#
orbits shaded by property.
orbits shaded by z position.
multiple planets^
overtop images^
phot vs a plots.
histograms of all properties.

astrometry#
pma plots scattered by property.
pma time plots colored by property.

Spectra/photometry

function to generate all.
=#

import Contour as Contourlib
# using CairoMakie
using GLMakie
using DirectOrbits
using StatsBase

function plot_orbits!(gp::GridPosition, args...; kwargs...)
    ax_sky = Axis(
        gp,
        xreversed=true,
        autolimitaspect=1,
        xlabel="ΔRA (mas)",
        ylabel="ΔDEC (mas)"
    )
    plot_orbits!(ax_sky, args...; kwargs...)
end
function plot_orbits!(ax::Axis, chains, planet_key;
    color,
    N=500,
    ii = rand(1:size(chains,1) * size(chains,3), N),
    colorrange=quantile(filter(isfinite, color),(0.05,0.95))
)
    orbits = DirectDetections.construct_elements(chains, planet_key, ii)
    posns = DirectOrbits.kep2cart_ν.(orbits, range(0, 2π, length=90)')
    ras = getproperty.(posns, :x)
    decs = getproperty.(posns, :y)
    color = color[ii]
    for (ras, decs, c) in zip(eachrow(ras), eachrow(decs), color)
        lines!(ax, ras, decs, linewidth=0.5; color=fill(c, size(ras)), colorrange, colormap=:plasma)
    end
    lines!(ax, [0], [0],  color=:black, label=string(planet_key))
    # Makie.scatter!(ax, [0], [0], marker="⋆", color=:white, markersize=30)
    # Makie.scatter!(ax, [0], [0], marker="⋆", color=:black, markersize=20)
end

function plot_image!(ax::Axis, image, platescale;
    colorrange=quantile(filter(isfinite, image),(0.05,0.95)),
    colormap=:magma,
)
    Makie.heatmap!(
        ax,
        (collect(axes(img,1)) .* platescale),
        (collect(axes(img,2)) .* platescale),
        collect(img[end:-1:begin,:]),
        colormap,
        colorrange,
    )
end

function plot_astrom!(ax::Axis, system::System)
    for planet in system.planets
        if !isnothing(planet.astrometry)
            errorbars!(ax,
                collect(planet.astrometry.ra),
                collect(planet.astrometry.dec),
                collect(planet.astrometry.σ_dec),
                direction=:y,
                color=:white,
                linewidth=4,
            )
            errorbars!(ax,
                collect(planet.astrometry.ra),
                collect(planet.astrometry.dec),
                collect(planet.astrometry.σ_ra),
                direction=:x,
                color=:white,
                linewidth=4,
            )
            errorbars!(ax,
                collect(planet.astrometry.ra),
                collect(planet.astrometry.dec),
                collect(planet.astrometry.σ_dec),
                direction=:y,
                color=:red,
            )
            errorbars!(ax,
                collect(planet.astrometry.ra),
                collect(planet.astrometry.dec),
                collect(planet.astrometry.σ_ra),
                direction=:x,
                color=:red,
            )
        end
    end
end

function plot_pma_epoch!(ax::Axis, system::System, ind::Integer)
    system_pma = system.propermotionanom
    Makie.errorbars!(
        ax,
        [system_pma.pm_ra[ind]],
        [system_pma.pm_dec[ind]],
        [system_pma.σ_pm_ra[ind]],
        direction=:x,
        color=:white,
        linewidth=4
    )
    Makie.errorbars!(
        ax,
        [system_pma.pm_ra[ind]],
        [system_pma.pm_dec[ind]],
        [system_pma.σ_pm_dec[ind]],
        direction=:y,
        color=:white,
        linewidth=4
    )
    Makie.errorbars!(
        ax,
        [system_pma.pm_ra[ind]],
        [system_pma.pm_dec[ind]],
        [system_pma.σ_pm_ra[ind]],
        direction=:x,
        color=:red,
    )
    Makie.errorbars!(
        ax,
        [system_pma.pm_ra[ind]],
        [system_pma.pm_dec[ind]],
        [system_pma.σ_pm_dec[ind]],
        direction=:y,
        color=:red,
    )
    Makie.hlines!(ax, 0; color=:black, label="")
    Makie.vlines!(ax, 0; color=:black, label="")
end

function plot_pma_epoch!(ax::Axis, chains, ind::Integer;
    color=:black,
    N=500,
    ii = rand(1:size(chains,1) * size(chains,3), N),
    colorrange=quantile(filter(isfinite, color),(0.05,0.95))
)
    system_pma = chains.info.model.propermotionanom
    vx = zeros(size(ii))
    vy = zeros(size(ii))
    for j in keys(chains.info.model.planets)
        elements = DirectDetections.construct_elements(chains, j, ii)
        mass = vec(chains["$j[mass]"][ii])
        vx .+= getindex.(propmotionanom.(elements, system_pma.ra_epoch[ind], mass.*DirectDetections.mjup2msol),1)
        vy .+= getindex.(propmotionanom.(elements, system_pma.dec_epoch[ind], mass.*DirectDetections.mjup2msol),2)
    end
    if typeof(color) <: AbstractArray
        color = color[ii]
    end
    Makie.scatter!(ax, vx, vy; color, markersize=5, colorrange, colormap=:plasma)
end

function plot_pma_time!(ax::Axis, system::System, direction; orientation=:horizontal)
    system_pma = system.propermotionanom
    if direction == :ra
        y = collect(system_pma.pm_ra)
        σ = collect(system_pma.σ_pm_ra)
        epoch = collect(system_pma.ra_epoch)
    elseif direction == :dec
        y = collect(system_pma.pm_dec)
        σ = collect(system_pma.σ_pm_dec)
        epoch = collect(system_pma.dec_epoch)
    else
        error("Unsupported direction. Must be one of :ra or :dec")
    end
    if orientation == :horizontal
        Makie.errorbars!(ax, epoch, y, σ, direction=:y, color=:white, linewidth=4)
        Makie.errorbars!(ax, epoch, y, σ, direction=:y, color=:red)
    else
        Makie.errorbars!(ax, y, epoch, σ, direction=:x, color=:white, linewidth=4)
        Makie.errorbars!(ax, y, epoch, σ, direction=:x, color=:red)
    end
end
function plot_pma_time!(ax::Axis, chains, direction;
    orientation=:horizontal,
    color,
    N=100,
    ii = rand(1:size(chains,1)*size(chains,3), N),
    colormap=:plasma,
    colorrange=quantile(filter(isfinite, color),(0.05,0.95)),
    linewidth=0.5
)
    system_pma = chains.info.model.propermotionanom
    if direction == :ra
        dir_ind = 1
    elseif direction == :dec
        dir_ind = 2
    else
        error("Unsupported direction. Must be one of :ra or :dec")
    end

    epoch_min = minimum([system_pma.ra_epoch; system_pma.dec_epoch;])
    epoch_max = maximum([system_pma.ra_epoch; system_pma.dec_epoch;])
    epoch_dt = epoch_max - epoch_min
    epoch_min -= 0.2epoch_dt
    epoch_max += 0.2epoch_dt
    epoch = range(epoch_min, epoch_max, length=200)

    vy = zeros(length(ii), length(epoch))
    for j in keys(chains.info.model.planets)
        elements = DirectDetections.construct_elements(chains, j, ii)
        mass = vec(chains["$j[mass]"][ii])
        vy .+= getindex.(propmotionanom.(elements, epoch', mass.*DirectDetections.mjup2msol), dir_ind)
    end
    # if typeof(color) <: AbstractArray
        color = color[ii]
    # end
    if orientation == :horizontal
        for (row,c) in zip(eachrow(vy), color)
            Makie.lines!(ax, epoch, row; color=fill(c, size(epoch)), colormap, colorrange,  linewidth)
        end
    else
        for (row,c) in zip(eachrow(vy), color)
            Makie.lines!(ax, row, epoch; color=fill(c, size(epoch)), colormap, colorrange,  linewidth)
        end
    end
    Makie.xlims!(ax, epoch_min, epoch_max)
    Makie.ylims!(ax, quantile(vec(vy), (0.0001, 0.9999)))

end

function hist_prop!(ax::Axis, prop; color=:blue, label="")
    h = StatsBase.normalize(fit(Histogram, vec(prop)), mode=:pdf)
    Makie.stairs!(ax, h; color, label)
    Makie.ylims!(ax, low=0)
end

function plot_phot!(gp::GridPosition, planet::Planet; axis=(;), kwargs...)
    bands = collect(string.(planet.photometry.band))
    ax = Axis(gp;
        xlabel="band",
        xticks=(collect(eachindex(bands)), bands),
        axis...
    )
    plot_phot!(ax, planet)
    return ax
end
function plot_phot!(ax::Axis, planet::Planet)
    bands = string.(planet.photometry.band)
    Makie.errorbars!(ax,
        eachindex(bands),
        collect(planet.photometry.phot),
        collect(planet.photometry.σ_phot),
        color=:red,
    )
end

# TODO
function plot_phot!(ax::Axis, chains, planet_key;
    N,
    ii = rand(1:size(chains,1)*size(chains,3), N),
    color=:black
)
    planet = chains.info.model.planets[planet_key]
    bands = string.(planet.photometry.band)
    Makie.errorbars!(ax,
        eachindex(bands),
        collect(planet.photometry.phot),
        collect(planet.photometry.σ_phot),
        color=:red,
    )
end

function compute_levels(H)
    # Calculate levels for contours
   levels = 1 .- exp.(-0.5 .* (0.5:0.5:2.1).^2)
   ii = sortperm(reshape(H,:))
   h2flat = H[ii]
   sm = cumsum(h2flat)
   sm /= sm[end]
   if all(isnan, sm) || length(H) <= 1
       @warn "Could not compute valid contours"
       plotcontours = false
       V = [0]
   else
       V = sort(map(v0 -> h2flat[sm .≤ v0][end], levels))
       if any(==(0), diff(V))
           @warn "Too few points to create valid contours"
       end
   end
   return V
end
function margin_hist(
    chains,
    key1,
    key2,
    range1=range(extrema(chains[key1])...,length=30),
    range2=range(extrema(chains[key2])...,length=30);
    color=:black
)

    # TODO: this needs to be generalized somehow
    priors = NamedTuple(chains.info.model.planets[1].priors.priors)

    ha = fit(Histogram, vec(chains[key1]), range1)
    hl = fit(Histogram, vec(chains[key2]), range2)

    ## Main 2D histogram
    hla = fit(
        Histogram,
        (vec(chains[key1]), vec(chains[key2])),
        (range1, range2)
    )

    fig = Figure(
        resolution = (650*2, 400*2),
        figure_padding=1,
        font="Arial",
        fontsize=20,
    )

    ax = Axis(
        fig[2,1],
        xlabel=key1,
        ylabel=key2,

    )
    heatmap!(ax, hla, colormap=cgrad([:white, color]))

    l = compute_levels(hla.weights)[1:1]
    ax1 = hla.edges[1][1:end-1] .+ step(hla.edges[1])/2
    ax2 = hla.edges[2][1:end-1] .+ step(hla.edges[2])/2
    for cl in Contourlib.levels(Contourlib.contours(ax2, ax1, hla.weights', l))
        lvl = Contourlib.level(cl) # the z-value of this contour level
        thefirst=true
        for line in Contourlib.lines(cl)
            xs, ys = coordinates(line) # coordinates of this line segment
#             p = Makie.GeometryBasics.Polygon(
#                 Makie.GeometryBasics.Point2.(zip(ys, xs)),
#             )
#             poly!(ax, p, color=RGBAf0(model_colours[cnum],0.4))
            lines!(ax, ys, xs; color)

            thefirst=false
        end
    end

    # Maginal histograms
    ax_l = Axis(
        fig[2,2],
    )
    hidexdecorations!(ax_l)
    hideydecorations!(ax_l)

    # Also plot the prior pdf scaled to the same height
    # x = range(extrema(range2)..., length=250)
    # prior_pdf = pdf.(priors.mass,x)
    # Makie.lines!(ax_l,  prior_pdf ./ maximum(prior_pdf) .* maximum(hl.weights),x;color=:black, label="prior",linewidth=1)

    stairs!(ax_l, hl.weights, hl.edges[1][1:end-1].+step(hl.edges[1]); color)

    ax_a = Axis(
        fig[1,1],
    )
    hideydecorations!(ax_a)
    hidexdecorations!(ax_a)

    # Also plot the prior pdf scaled to the same height
    # x = range(extrema(range1)..., length=250)
    # prior_pdf = pdf.(priors.a, x)
    # Makie.lines!(ax_a, x, prior_pdf ./ maximum(prior_pdf) .* maximum(ha.weights), color=:black, label="prior",linewidth=1)

    stairs!(ax_a, ha; color)

    linkxaxes!(ax, ax_a)
    linkyaxes!(ax, ax_l)
    xlims!(ax, extrema(range1))
    ylims!(ax, extrema(range2))
    xlims!(ax_l, low=0, high=maximum(hl.weights)*1.1)
    ylims!(ax_a, low=0, high=maximum(ha.weights)*1.1)
    hidedecorations!(ax, label=false, ticks=false, ticklabels=false)

    rowgap!(fig.layout,0)
    colgap!(fig.layout,0)

    # colsize!(fig.layout, 2, )
    # rowsize!(fig.layout, 1, 60)
    rowsize!(fig.layout, 2, Auto(3))
    colsize!(fig.layout, 2, Aspect(1,1.0))

    # Legend(fig[1,2], ax_a, tellheight=false, tellwidth=false, valign=:top, halign=:right, patchsize=(10,10), framevisible=false)
    return fig
end

# N= 1500
# ii = rand(1:size(chain,1)*size(chain,3),N)
# Zs = chain["b[Z]"][ii]
# Js = chain["b[J]"][ii]
# Ls = chain["b[L]"][ii]
# plot!(vcat(Zs', Js', Ls'), color=:black, lw=1,  alpha=0.05, label="")
# # scatter!(Ls', color=:black, lw=1,  alpha=0.05, label="")
# scatter!(vcat(Zs', Js', Ls'), color=:black, lw=1, ms=4, alpha=0.01, label="")
# s

##

function plotall(chains;
    N=500,
    colorkey="b[a]",
    colormap=:plasma,
    colorrange=quantile(filter(isfinite, chains[colorkey]),(0.05,0.95))
    )
    color = chains[colorkey]

    planet_colors = [:black, :blue]

    ii = rand(1:size(chains,1)*size(chains,3), N)

    # fig =  Figure(resolution=(1600,1400))
    fig =  Figure(resolution=(800,700))
    ax_sky = Axis(fig[1,1:2], xreversed=true,autolimitaspect=1, xlabel="ΔRA (mas)", ylabel="ΔDEC (mas)")

    if !isnothing(chains.info.model.images)
        plot_image!(ax_sky, chains.info.model.images.image[end], chains.info.model.images.platescale[end])
    end

    plot_orbits!(ax_sky, chains, :b; color, ii, colorrange)

    if any(!isnothing(planet.astrometry) for planet in chains.info.model.planets)
        plot_astrom!(ax_sky, chains.info.model)
    end

    if !isnothing(chains.info.model.propermotionanom)

        pma_grid = GridLayout()
        fig.layout[2,1:2] = pma_grid

        ax_pma1 = Axis(pma_grid[1:2,1], xlabel="Δv_ra (mas/yr)", ylabel="Δv_dec (mas/yr)", title="HIPPARCOS", autolimitaspect=1)
        plot_pma_epoch!(ax_pma1, chains, 1; color, ii, colorrange)
        plot_pma_epoch!(ax_pma1, chains.info.model, 1)

        ax_pma_time_ra = Axis(pma_grid[1,2], xlabel="epoch (mjd)", ylabel="Δv_ra (mas/yr)")
        plot_pma_time!(ax_pma_time_ra, chains, :ra; color, ii, colorrange)
        plot_pma_time!(ax_pma_time_ra, chains.info.model, :ra)

        ax_pma_time_dec = Axis(pma_grid[2,2], xlabel="epoch (mjd)", ylabel="Δv_dec (mas/yr)")
        plot_pma_time!(ax_pma_time_dec, chains, :dec; color, ii, colorrange)
        plot_pma_time!(ax_pma_time_dec, chains.info.model, :dec)

        ax_pma2 = Axis(pma_grid[1:2,3],xlabel="Δv_ra (mas/yr)", ylabel="Δv_dec (mas/yr)", title="GAIA", autolimitaspect=1)
        plot_pma_epoch!(ax_pma2, chains, 2; color, ii, colorrange)
        plot_pma_epoch!(ax_pma2, chains.info.model, 2)

        Makie.hidexdecorations!(ax_pma_time_ra, grid=false)

        # colsize!(pma_grid, 1, Aspect(1,4))
        # colsize!(pma_grid, 3, Aspect(1,4))

    #     ax_pma1 = Axis(fig[2,1], xlabel="Δv_ra (mas/yr)", ylabel="Δv_dec (mas/yr)", title="HIPPARCOS", autolimitaspect=1)
    #     plot_pma_epoch!(ax_pma1, chains, 1; color, ii, colorrange)
    #     plot_pma_epoch!(ax_pma1, chains.info.model, 1)

    #     ax_pma2 = Axis(fig[2,2],xlabel="Δv_ra (mas/yr)", ylabel="Δv_dec (mas/yr)", title="GAIA", autolimitaspect=1)
    #     plot_pma_epoch!(ax_pma2, chains, 2; color, ii, colorrange)
    #     plot_pma_epoch!(ax_pma2, chains.info.model, 2)

    #     ax_pma_time_ra = Axis(fig[3,1:2], xlabel="epoch (mjd)", ylabel="Δv_ra (mas/yr)")
    #     plot_pma_time!(ax_pma_time_ra, chains, :ra; color, ii, colorrange)
    #     plot_pma_time!(ax_pma_time_ra, chains.info.model, :ra)

    #     ax_pma_time_dec = Axis(fig[4,1:2], xlabel="epoch (mjd)", ylabel="Δv_dec (mas/yr)")
    #     plot_pma_time!(ax_pma_time_dec, chains, :dec; color, ii, colorrange)
    #     plot_pma_time!(ax_pma_time_dec, chains.info.model, :dec)
    end

    # # Histograms
    # ax_hist_a = Axis(
    #     fig[1,3],
    #     ylabel="posterior density",
    #     xlabel="a (au)",
    # )
    # hideydecorations!(ax_hist_a)
    # ax_hist_e = Axis(
    #     fig[2,3],
    #     ylabel="posterior density",
    #     xlabel="e",
    # )
    # hideydecorations!(ax_hist_e)
    # ax_hist_m = Axis(
    #     fig[3,3],
    #     ylabel="posterior density",
    #     xlabel="mass (Mjup)",
    # )
    # hideydecorations!(ax_hist_m)
    # for (i, planet) in enumerate(chains.info.model.planets)
    #     p = planet.name
    #     hist_prop!(ax_hist_m, chains["$p[mass]"], color=:black)
    #     hist_prop!(ax_hist_a, chains["$p[a]"], color=:black)
    #     hist_prop!(ax_hist_e, chains["$p[e]"], color=:black)
    # end

    # # TODO: or images
    # if any(!isnothing(planet.photometry) for planet in chains.info.model.planets)
    #     # TODO: or later planets
    #     plot_phot!(fig[4,3], chains.info.model.planets[1])
    # end

    Colorbar(
        fig[:, end+1];
        limits = colorrange,
        colormap,
        label=colorkey,
        size=34
    )
    linkaxes!(ax_pma1, ax_pma2)
    # linkaxes!(ax_pma_time_ra, ax_pma_time_dec)

    rowsize!(fig.layout, 1, Auto(3))
    # rowsize!(fig.layout, 2, Auto(2))
    # rowsize!(fig.layout, 3, Auto(1))
    # colsize!(fig.layout, 1, Auto(2/3))
    # colsize!(fig.layout, 2, Auto(1/3))

    return fig
end

##
fig = plotall(chains, colorkey="b[mass]", N=10)


text:
# Julia wrapper for header: /usr/include/libavutil/samplefmt.h
# Automatically generated using Clang.jl wrap_c, version 0.0.0

export
    av_get_sample_fmt_name,
    av_get_sample_fmt,
    av_get_alt_sample_fmt,
    av_get_packed_sample_fmt,
    av_get_planar_sample_fmt,
    av_get_sample_fmt_string,
    av_get_bytes_per_sample,
    av_sample_fmt_is_planar,
    av_samples_get_buffer_size,
    av_samples_fill_arrays,
    av_samples_alloc,
    av_samples_alloc_array_and_samples,
    av_samples_copy,
    av_samples_set_silence

function av_get_sample_fmt_name(sample_fmt::AVSampleFormat)
    ccall((:av_get_sample_fmt_name, libavutil), Cstring, (AVSampleFormat,), sample_fmt)
end

function av_get_sample_fmt(name)
    ccall((:av_get_sample_fmt, libavutil), AVSampleFormat, (Cstring,), name)
end

function av_get_alt_sample_fmt(sample_fmt::AVSampleFormat, planar::Integer)
    ccall((:av_get_alt_sample_fmt, libavutil), AVSampleFormat, (AVSampleFormat, Cint), sample_fmt, planar)
end

function av_get_packed_sample_fmt(sample_fmt::AVSampleFormat)
    ccall((:av_get_packed_sample_fmt, libavutil), AVSampleFormat, (AVSampleFormat,), sample_fmt)
end

function av_get_planar_sample_fmt(sample_fmt::AVSampleFormat)
    ccall((:av_get_planar_sample_fmt, libavutil), AVSampleFormat, (AVSampleFormat,), sample_fmt)
end

function av_get_sample_fmt_string(buf, buf_size::Integer, sample_fmt::AVSampleFormat)
    ccall((:av_get_sample_fmt_string, libavutil), Cstring, (Cstring, Cint, AVSampleFormat), buf, buf_size, sample_fmt)
end

function av_get_bytes_per_sample(sample_fmt::AVSampleFormat)
    ccall((:av_get_bytes_per_sample, libavutil), Cint, (AVSampleFormat,), sample_fmt)
end

function av_sample_fmt_is_planar(sample_fmt::AVSampleFormat)
    ccall((:av_sample_fmt_is_planar, libavutil), Cint, (AVSampleFormat,), sample_fmt)
end

function av_samples_get_buffer_size(linesize, nb_channels::Integer, nb_samples::Integer, sample_fmt::AVSampleFormat, align::Integer)
    ccall((:av_samples_get_buffer_size, libavutil), Cint, (Ptr{Cint}, Cint, Cint, AVSampleFormat, Cint), linesize, nb_channels, nb_samples, sample_fmt, align)
end

function av_samples_fill_arrays(audio_data, linesize, buf, nb_channels::Integer, nb_samples::Integer, sample_fmt::AVSampleFormat, align::Integer)
    ccall((:av_samples_fill_arrays, libavutil), Cint, (Ptr{Ptr{UInt8}}, Ptr{Cint}, Ptr{UInt8}, Cint, Cint, AVSampleFormat, Cint), audio_data, linesize, buf, nb_channels, nb_samples, sample_fmt, align)
end

function av_samples_alloc(audio_data, linesize, nb_channels::Integer, nb_samples::Integer, sample_fmt::AVSampleFormat, align::Integer)
    ccall((:av_samples_alloc, libavutil), Cint, (Ptr{Ptr{UInt8}}, Ptr{Cint}, Cint, Cint, AVSampleFormat, Cint), audio_data, linesize, nb_channels, nb_samples, sample_fmt, align)
end

function av_samples_alloc_array_and_samples(audio_data, linesize, nb_channels::Integer, nb_samples::Integer, sample_fmt::AVSampleFormat, align::Integer)
    ccall((:av_samples_alloc_array_and_samples, libavutil), Cint, (Ptr{Ptr{Ptr{UInt8}}}, Ptr{Cint}, Cint, Cint, AVSampleFormat, Cint), audio_data, linesize, nb_channels, nb_samples, sample_fmt, align)
end

function av_samples_copy(dst, src, dst_offset::Integer, src_offset::Integer, nb_samples::Integer, nb_channels::Integer, sample_fmt::AVSampleFormat)
    ccall((:av_samples_copy, libavutil), Cint, (Ptr{Ptr{UInt8}}, Ptr{Ptr{UInt8}}, Cint, Cint, Cint, Cint, AVSampleFormat), dst, src, dst_offset, src_offset, nb_samples, nb_channels, sample_fmt)
end

function av_samples_set_silence(audio_data, offset::Integer, nb_samples::Integer, nb_channels::Integer, sample_fmt::AVSampleFormat)
    ccall((:av_samples_set_silence, libavutil), Cint, (Ptr{Ptr{UInt8}}, Cint, Cint, Cint, AVSampleFormat), audio_data, offset, nb_samples, nb_channels, sample_fmt)
end

ext: jl


text:
@testset "extrema_series" begin
    @testset "test problems" begin
        demo_problems = demo_problem.(eachindex(demo_problems_definitions), true)
        for (f, (a, b), min) in demo_problems
            min1, _, y1 = ArbExtras.extrema_series(f, a, b)
            _, min2, y2 = .-ArbExtras.extrema_series(x -> -f(x), a, b)
            min3, y3 = ArbExtras.minimum_series(f, a, b)
            min4, y4 = .-ArbExtras.maximum_series(x -> -f(x), a, b)

            @test isequal(min1, min2)
            @test isequal(min2, min3)
            @test isequal(min3, min4)

            @test isequal(y1, y2)
            @test isequal(y3, y3)
            @test isequal(y3, y4)

            @test Arblib.overlaps(min1, min)

            @test Arblib.overlaps(y1, f(Arb(a + b) / 2))
        end
    end

    @testset "special cases" begin
        # Thin interval
        @test ArbExtras.extrema_series(x -> -2x, Arf(2), Arf(2)) == (-4, -4, -4)
        @test ArbExtras.minimum_series(x -> -2x, Arf(2), Arf(2)) == (-4, -4)
        @test ArbExtras.maximum_series(x -> -2x, Arf(2), Arf(2)) == (-4, -4)
        @test ArbExtras.extrema_series(x -> -2x, Arf(2), Arf(2), abs_value = true) ==
              (4, 4, 4)
        @test ArbExtras.minimum_series(x -> -2x, Arf(2), Arf(2), abs_value = true) == (4, 4)
        @test ArbExtras.maximum_series(x -> -2x, Arf(2), Arf(2), abs_value = true) == (4, 4)

        # Non-ordered endpoints
        @test_throws ArgumentError ArbExtras.extrema_series(identity, Arf(2), Arf(1))
        @test_throws ArgumentError ArbExtras.minimum_series(identity, Arf(2), Arf(1))
        @test_throws ArgumentError ArbExtras.maximum_series(identity, Arf(2), Arf(1))

        # Infinite endpoints
        @test_throws ArgumentError ArbExtras.extrema_series(identity, Arf(2), Arf(Inf))
        @test_throws ArgumentError ArbExtras.minimum_series(identity, Arf(2), Arf(Inf))
        @test_throws ArgumentError ArbExtras.maximum_series(identity, Arf(2), Arf(Inf))

        # Monotonic on interval
        @test isequal(ArbExtras.extrema_series(identity, Arf(-2), Arf(4)), (-2, 4, NaN))
        @test isequal(ArbExtras.minimum_series(identity, Arf(-2), Arf(4)), (-2, NaN))
        @test isequal(ArbExtras.maximum_series(identity, Arf(-2), Arf(4)), (4, NaN))

        @test isequal(
            ArbExtras.extrema_series(identity, Arf(-2), Arf(4), abs_value = true),
            (0, 4, NaN),
        )
        @test isequal(
            ArbExtras.extrema_series(identity, Arf(-3), Arf(-1), abs_value = true),
            (1, 3, NaN),
        )
        @test isequal(
            ArbExtras.minimum_series(identity, Arf(-2), Arf(4), abs_value = true),
            (0, NaN),
        )
        @test isequal(
            ArbExtras.minimum_series(identity, Arf(-3), Arf(-1), abs_value = true),
            (1, NaN),
        )
        @test isequal(
            ArbExtras.maximum_series(identity, Arf(-6), Arf(4), abs_value = true),
            (6, NaN),
        )

        # Check that it gives a correct result even if the
        # computations with the endpoints cannot be done exactly. I
        # have not been able to come up with an example for which we
        # get a wrong result without adding the endpoints explicitly.
        # If we are working with such large numbers the enclosure we
        # get is typically very large and includes the true result
        # anyway.

        @test all(
            Arblib.overlaps.(
                ArbExtras.extrema_series(x -> x^2, Arf(1), Arf(1e300))[1:2],
                (Arb(-1), Arb(1e300)^2),
            ),
        )
        @test all(
            Arblib.overlaps.(
                ArbExtras.extrema_series(x -> x^2, Arf(-1e300), Arf(-1))[1:2],
                (Arb(-1), Arb(1e300)^2),
            ),
        )
        @test Arblib.overlaps(
            ArbExtras.minimum_series(x -> x^2, Arf(1), Arf(1e300))[1],
            Arb(-1),
        )
        @test Arblib.overlaps(
            ArbExtras.minimum_series(x -> x^2, Arf(-1e300), Arf(-1))[1],
            Arb(-1),
        )
        @test Arblib.overlaps(
            ArbExtras.maximum_series(x -> -x^2, Arf(1), Arf(1e300))[1],
            Arb(1),
        )
        @test Arblib.overlaps(
            ArbExtras.maximum_series(x -> -x^2, Arf(-1e300), Arf(-1))[1],
            Arb(1),
        )

        # Non-finite restterm
        @test all(isnan, ArbExtras.extrema_series(x -> inv(x - sin(x)), Arf(0.1), Arf(1)))
        @test all(isnan, ArbExtras.minimum_series(x -> inv(x - sin(x)), Arf(0.1), Arf(1)))
        @test all(isnan, ArbExtras.maximum_series(x -> inv(x - sin(x)), Arf(0.1), Arf(1)))
    end
end


text:
include("update.jl")

function randomwalk(nsteps::Integer, prob::Real=0.5)
    result = zeros(Int, nsteps+1)
    x = 0
    result[1] = x
    for t in 1:nsteps
        x = update(x,prob)
        result[t+1] = x
    end
    return result
end

max_stars_repo_name: isuke/JuliaBook-Samples

max_issues_repo_name: isuke/JuliaBook-Samples

max_forks_repo_name: isuke/JuliaBook-Samples


text:
using JuLIP, QMMM2, SHIPs, PrettyTables, LinearAlgebra, Plots,
      DataFrames
using JuLIP.MLIPs
include(@__DIR__() * "/swqmmm.jl")

## reference model and MM potential(s)
r0 = rnn(:Si)
sw = StillingerWeber()

## test convergence
ships =  []
degs = [2, 3, 4, 5, 6, 7, 8]
rmse_E = []
rmse_DE = []
rmse_D2E = []
for deg = degs
   🚢, fitinfo = SWqmmm.train_ship(3, deg, wL=1.0, rinfact=0.6, rtol=5e-4)
   push!(ships, 🚢)
   push!(rmse_E, fitinfo["rmse"]["Es"])
   push!(rmse_DE, fitinfo["rmse"]["dEs"])
   push!(rmse_D2E, fitinfo["rmse"]["d2Esh"])
end
println(); println()

println("RMSE-Table")
df_rmse = DataFrame(deg = degs, E = rmse_E, dE = rmse_DE, d2E = rmse_D2E)
display(df_rmse)


text:
using AbstractAlgebra
using Zygote

struct ZeroForm
   rep  :: Function
   name :: String
end

Coefficent = Union{Function, Real}

struct OneForm
	primitive :: ZeroForm
end

struct HigherForm
   components 	:: Array{Union{HigherForm, ZeroForm}}
   coeffs 	:: Array{Coefficent}
   exact 	:: Bool
   degree 	:: Int64
end


text:
__precompile__()

module PolicyGradient

using mjWrap
using LearningStrategies
import LearningStrategies: setup!, update!
using Baseline
using Policy
using Tools
using ExpFunctions

using ReverseDiff

using IterativeSolvers

export PolicyGradModel, NPGStrategy, setup!, update!, PolicyRollout

struct PolicyGradModel
   theta::Vector     # policy parameters
   mjsys::mjWrap.mjSet
   policy::Policy.AbstractPolicy

   # scratch stuff
   samples::mjWrap.TrajSamples
   meansamples::mjWrap.TrajSamples

   trace::Dict{Symbol,Vector{Float64}}

   niter::Int

   function PolicyGradModel(theta,
                            mjsys, policy,
                            T, numT, meanT, niter)
      ns = mjsys.ns
      samples     = mjw.allocateTrajSamples(mjsys, T, numT, ns)
      meansamples = mjw.allocateTrajSamples(mjsys, T, meanT, ns)
      return new(theta, mjsys, policy,
                 samples, meansamples,
                 Dict(:stocR => Vector{Float64}(),
                      :meanR => Vector{Float64}(),
                      :evalscore => Vector{Float64}() ), niter)
                 #zeros(niter), zeros(niter), zeros(niter), niter)
   end
end

struct NPGStrategy{dtype<:AbstractFloat} <: LearningStrategy
   vpg::Vector{dtype}
   npg::Vector{dtype}
   grad::Vector{dtype}

   fullFIM::Bool
   gradll::Matrix{dtype}
   FIM::Matrix{dtype}

   adv::Vector{dtype}
   ret::Vector{dtype}

   baseline::Baseline.AbstractBaseline

   norm_step_size::Float64
   gamma::Float64
   gae::Float64

   cg_iter::Int
   cg_reg::Float64
   cg_tol::Float64

   obsmat::Matrix{dtype}
   ctrlmat::Matrix{dtype}

   function NPGStrategy{dtype}(model,
                               baseline,
                               fullFIM,
                               norm_step_size,
                               gamma, gae,
                               cg_iter, cg_reg, cg_tol) where dtype<:AbstractFloat
      #dtype = eltype(model.policy.theta) #Float64
      T, numT = model.samples.T, model.samples.numT

      nparam  = length(model.theta)
      vpg     = zeros(dtype, nparam)
      npg     = zeros(dtype, nparam)
      grad    = zeros(dtype, nparam)

      if fullFIM
         gradll  = zeros(dtype, nparam, T*numT)
         FIM     = zeros(dtype, nparam, nparam)
      else
         gradll  = zeros(dtype, 1, 1)
         FIM     = zeros(dtype, 1, 1)
      end

      adv     = zeros(dtype, T*numT)
      ret     = zeros(dtype, T*numT)

      tmpobs  = rand(dtype, model.mjsys.ns, T*numT)
      tmpctrl = rand(dtype, model.mjsys.nu, T*numT)

      new(vpg, npg, grad,
          fullFIM, gradll, FIM, adv, ret,
          baseline,
          norm_step_size,
          gamma, gae,
          cg_iter, cg_reg, cg_tol,
          tmpobs, tmpctrl)
   end
end
Base.show(io::IO, s::NPGStrategy) = print(io, "NPG: $(s.fullFIM?"FIM":"HVP"), δ=$(s.norm_step_size)\n")

function grad_builder(f::Function, input)
   ctp = ReverseDiff.GradientTape(f, input)
   if length(ctp.tape) <= 10000
      info("Compiling tape")
      ctp = ReverseDiff.compile(ctp)
   end
   return (out,in)->ReverseDiff.gradient!(out, ctp, in) # taped
end

# feature of learning strats; need to have iter and null to get iteration num
function update!(model::PolicyGradModel, s::NPGStrategy, iter, null)

   tic()

   samples  = model.samples
   mjsys    = model.mjsys
   pol      = model.policy

   b        = s.baseline

   nsamples = samples.T*samples.numT

   # Parallize here......................................................
   rsum = sum(samples.reward, 1) # rollouts done in different strat

   if eltype(pol.theta) == Float64
      ctrlmat = reshape(samples.ctrl, mjsys.nu, nsamples)
      obsmat  = reshape(samples.obs, mjsys.ns, nsamples)
   else
      copy!(s.ctrlmat, convert(Matrix{Float32}, reshape(samples.ctrl, mjsys.nu, nsamples)))
      copy!(s.obsmat, convert(Matrix{Float32}, reshape(samples.obs, mjsys.ns, nsamples)))
      ctrlmat = s.ctrlmat
      obsmat  = s.obsmat
   end

   println("returns")
   @time Tools.compute_returns!(s.ret, samples.reward, s.gamma)

   println("predict")
   @time Baseline.predict!(b, reshape(samples.obs, mjsys.ns, nsamples))

   println("advantages")
   @time Tools.computeGAEadvantages!(s.adv, b.bline,
                                     samples.reward,
                                     reshape(samples.obs, mjsys.ns, nsamples), s.gae, s.gamma)

   println("fit")
   @time Baseline.fit!(b, s.ret) # update fit

   if s.fullFIM
      DT = eltype(s.vpg)

      info("FULL FIM")
      println("gradll")
      @time Policy.gradLL!(s.gradll, pol, obsmat, ctrlmat)

      println("vpg")
      # vpg[:] = gradll*adv / nsamples
      @time BLAS.gemv!('N', 1/DT(nsamples), s.gradll, s.adv, DT(0.0), s.vpg)

      println("fim")
      #FIM[:] = gradll*gradll' / nsamples
      @time begin
         #const α = 0.9
         #FIM = copy(s.FIM)
         A_mul_Bt!(s.FIM, s.gradll, s.gradll)
         scale!(s.FIM, 1/DT(nsamples))
         #@. s.FIM = FIM + α*(s.FIM - FIM) # keep older FIM data around; works for higher α
      end

      println("cg")
      @time s.npg[:] = Tools.cpu_cg_solve(s.FIM, s.vpg, s.cg_iter, s.cg_reg, s.cg_tol)
      #@time IterativeSolvers.cg!(s.npg, s.FIM, s.vpg; tol = s.cg_tol, maxiter = s.cg_iter)
   else
      info("HVP algorithm")
      @time my∇LLxA = grad_builder( (theta, a)->Policy.loglikelihood(theta, pol, obsmat, ctrlmat)*a,
                                   (pol.theta, s.adv) )
      println("vpg")
      @time begin
         my∇LLxA((s.vpg, s.adv), (pol.theta, s.adv))
         scale!(s.vpg, 1/nsamples)
      end

      # TODO don't need in-place Loglikelihood for output; need inplace during comp
      minus = similar(s.adv')
      alpha = similar(s.adv')

      function hvp!(z, vec)
         const ϵ = sqrt(eps(eltype(vec))) * norm(vec)
         const eps_vec = ϵ * vec
         Policy.loglikelihood!(alpha, pol.theta+eps_vec, pol, obsmat, ctrlmat)
         Policy.loglikelihood!(minus, pol.theta-eps_vec, pol, obsmat, ctrlmat)
         alpha .-= minus
         scale!(alpha, 1.0/2ϵ)
         s.adv .= alpha'
         my∇LLxA((z, s.adv), (pol.theta, s.adv))
         #scale!(z, 1/nsamples)
      end
      println("cg")
      @time s.npg[:] = Tools.hvp_cg_solve(hvp!, s.vpg, s.cg_iter, s.cg_reg, s.cg_tol)
      scale!(s.npg, nsamples)
   end

   alpha = sqrt(s.norm_step_size / dot(s.npg, s.vpg))
   println("grad")
   s.grad .= alpha.*s.npg

   println("update")
   Policy.updatetheta(pol, pol.theta + s.grad) # update parameters
   model.theta[:] = pol.theta

   ########################################################## mean rollout
   println("no noise")
   meanreward = sum(model.meansamples.reward, 1)

   ############################################################### cleanup
   iterscore = mean(rsum)
   meanscore = mean(meanreward)
   if iter > 1
      @printf "\n  Stoc: %f, Mean: %f" iterscore meanscore
   else
      @printf "\n  Mean of rewards: %f" iterscore
   end
   @printf "\n         NPG Norm: %f\n\n" norm(s.npg) #normgrad[iter]
   @printf "  Reward of Mean Rollout: %f\n" meanscore
   itertime = toc()
   @printf "  Total time: %f\n" itertime
   push!(model.trace[:stocR], iterscore)
   push!(model.trace[:meanR], meanscore)
end

mutable struct PolicyRollout <: LearningStrategy
   f::FunctionSet
   policyaction::Function
   policymean::Function
   function PolicyRollout(f)
      new(f, ()->nothing, ()->nothing)
   end
end
Base.show(io::IO, s::PolicyRollout) = print(io, "π rollout\n")

function setup!(s::PolicyRollout, model::PolicyGradModel)
   s.policyaction = (c,o)->Policy.getaction!(c, model.policy, o)
   s.policymean   = (c,o)->Policy.getmean!(c, model.policy, o)
end

function stochasticrollout(mjsys::mjWrap.mjSet,
                           samples::mjWrap.TrajSamples,
                           ctrlfunc!::Function, f::FunctionSet, iter::Int)
   # stochastic rollout
   f.setmodel!(mjsys, iter) # change model for this iteration
   f.initstate!(mjsys, samples)
   f.setcontrol!(mjsys, samples.ctrl) # sets policy noise in ctrl
   @time mjWrap.rollout(mjsys, samples,
                        ctrlfunc!, f.observe!, f.reward)
end

function update!(model::PolicyGradModel, s::PolicyRollout, iter, null)
   println("policy rollout, stoc & mean")

   mjsys       = model.mjsys
   samples     = model.samples
   meansamples = model.meansamples

   tic()
   stochasticrollout(mjsys, samples, s.policyaction, s.f, iter)

   # policy's mean rollout
   s.f.initstate!(mjsys, meansamples)
   mjWrap.rollout(mjsys, meansamples,
                  s.policymean,
                  s.f.observe!, s.f.reward)
   toc()
end

end


text:
# Useful byte reprensetations
# - '\t': String([0x09])
# - ':': String([0x3a])
# - '0': String([0x30])
# - '1': String([0x31])
# - '.': String([0x2e])
# - '|': String([0x7c])
# - '/': String([0x2f])

# - 'A': String([0x41])
# - 'T': String([0x54])
# - 'C': String([0x43])
# - 'G': String([0x47])
# - 'N': String([0x4e])
"""
    filter(src, record_index, sample_index; des = "filtered." * src)

Filter vcf files (.vcf or .vcf.gz) according to row indices `record_index`
and column indices `sample_index` and write to a new set of vcf files `des`.

# Input:
- `src`: full vcf file name
- `record_index`: row indices to keep.
- `sample_index`: column indices to keep.

# Optional arguments:
- `des`: output vcf file name; default it `"filtered." * src`.
"""
function filter(
    src::AbstractString,
    record_index::AbstractVector{<:Integer},
    sample_index::AbstractVector{<:Integer};
    des::AbstractString = "filtered." * src
    )
    # create record (row) and sample (column) masks
    records, samples = nrecords(src), nsamples(src)
    if eltype(record_index) == Bool
        record_mask = record_index
    else
        record_mask = falses(records)
        record_mask[record_index] .= true
    end
    if eltype(sample_index) == Bool
        sample_mask = sample_index
    else
        sample_mask = falses(samples)
        sample_mask[sample_index] .= true
    end

    # create input and output VCF files
    reader = VCF.Reader(openvcf(src, "r"))
    writer = VCF.Writer(openvcf(des, "w"), filter_header(reader, sample_mask))

    # write to des
    for (i, record) in enumerate(reader)
        if record_mask[i]
            filter_record!(record, sample_mask)
            VCF.write(writer, record)
        end
    end

    close(reader)
    flush(writer)
    close(VCF.BioCore.IO.stream(writer))
    return nothing
end

"""
    filter_header(reader, sample_mask)

Filters the header info of a VCF reader. Will not keep sample IDs for entry `i` if `sample_mask[i] = false`.

# Inpute:
- `reader`: a VCF reader object.
- `sample_mask`: a BitVector. `sample_mask[i] = true` means keep sample `i`.

TODO: Create VCFTool.jl signature in meta info. If filedate/signature exist already, delete them before writing new ones
"""
function filter_header(
    reader::VCF.Reader,
    sample_mask::BitVector
    )
    # save meta information
    # fileformat = VCF.header(reader).metainfo[1]
    # filedate   = VCF.MetaInfo("##filedate=" * string(Dates.today()))
    # signature  = VCF.MetaInfo("##source=VCFTools.jl")
    # metainfo = vcat(fileformat, filedate, signature, VCF.header(reader).metainfo[2:end])

    # filter sampleID
    metainfo = VCF.header(reader).metainfo
    sampleID = VCF.header(reader).sampleID[sample_mask]
    return VCF.Header(metainfo, sampleID)
end

"""
    filter_record!(record, sample_mask)

Filters a VCF record by samples.

# Inputs:
- `record`: a VCF record
- `sample_mask`: a BitVector. `sample_mask[i] = true` means keep sample `i`.
"""
function filter_record!(
    record::VCF.Record,
    sample_mask::BitVector
    )

    # quick return
    if all(sample_mask)
        return nothing
    end

    p = length(record.genotype)
    new_data = UInt8[]
    new_genotype = Vector{UnitRange{Int64}}[]
    sizehint!(new_data, p)
    sizehint!(new_genotype, p)

    # copy chrom, pos, id, ref, alt, qual, filter, information, format into new_data
    for i in 1:(record.genotype[1][1][1] - 1)
        push!(new_data, record.data[i])
    end

    # copy genotype indices and data
    for i in 1:p
        if sample_mask[i]
            old_geno = record.genotype[i]
            new_geno = UnitRange{Int64}[]
            for g in old_geno
                # save genotypes
                new_start = length(new_data) + 1
                new_end   = new_start + length(g) - 1
                push!(new_geno, new_start:new_end)

                # save strings to data
                old_start = g[1]
                old_end   = g[end]
                [push!(new_data, record.data[i]) for i in old_start:old_end]
                push!(new_data, 0x3a) # 0x3a = byte equivalent of char ':'
            end
            new_data[end] = 0x09 # turn last ':' into '\t'
            push!(new_genotype, new_geno)
        end
    end
    resize!(new_data, length(new_data) - 1) # get rid of last '\t'

    # update pointers to data and genotype indices
    record.data = new_data
    record.genotype = new_genotype
    record.filled = 1:length(new_data)
    return nothing
end

"""
    mask_gt(src, masks; [des = "masked." * src], [separator = '/'])

Creates a new VCF file `des` where genotype entry (i, j) of `src`
is missing if `masks[i, j]` is true. `src` is unchanged.

# Arguments
- `src`: Input VCF file name.
- `masks`: Bit matrix. `masks[i, j] = true` means mask entry (i, j).
- `des`: output VCF file name.
- `separator`: Separator of VCF genotypes. Can be '/' (default) or '|'.
"""
function mask_gt(
    src::AbstractString,
    masks::BitArray{2};
    des::AbstractString = "masked." * src,
    separator::Char = '/'
    )
    records, samples = nrecords(src), nsamples(src)
    p, n = size(masks)
    if !(records == p && samples == n)
        throw(DimensionMismatch("size(src) = ($records, $samples) ≠ size(masks) = ($p, $n)."))
    end

    # define byte representation of separator
    if separator == '/'
        separator = 0x2f
    elseif separator == '|'
        separator = 0x7c
    else
        throw(ArgumentError("separator must be / or | but got $separator."))
    end

    # create input and output VCF files
    reader = VCF.Reader(openvcf(src, "r"))
    writer = VCF.Writer(openvcf(des, "w"), VCF.header(reader))

    # loop over each record
    for (i, record) in enumerate(reader)
        gtkey = VCF.findgenokey(record, "GT")
        if gtkey != nothing
            # loop over genotypes
            for (j, geno) in enumerate(record.genotype)
                if masks[i, j]
                    record.data[geno[gtkey][1]] = 0x2e
                    record.data[geno[gtkey][2]] = separator
                    record.data[geno[gtkey][3]] = 0x2e
                end
            end
        end
        write(writer, record)
    end
    flush(writer); close(reader); close(writer)
    return nothing
end


text:
@GenGlobal g_SharedSimulations
@GenGlobal g_SimulationPrimitives
@GenGlobal g_BaseDataSetofSets
@GenGlobal g_SimulationDrawsMatrix

# ------------------------------------------------------
# sparse state transition matrix
# ------------------------------------------------------

function sparse_state_transition(wp::AbstractUnitProblem)
     nS = length(wp)
     rowval = [ssprime(wp, si, d) for si in 1:nS for d in actionspace(wp, si)]
     colptr = cumsum(vcat(1, [_dmax(wp,si)+1 for si in 1:nS]))
     nzval = zeros(Float64, length(rowval))

     return SparseMatrixCSC(nS, nS, colptr, rowval, nzval)
end

sparse_state_transition(ddm::DynamicDrillModel) = sparse_state_transition(statespace(ddm))

# ------------------------------------------------------
# Temporary variables for simulations
# ------------------------------------------------------

"""
Holds temp variables needed for each simulation of a section `i` given
    draw number `m` from shared Halton arrays `uvs`
"""
struct SimulationTmp <: AbstractTmpVars
     Pprime::SparseMatrixCSC{Float64,Int}
     @addStructFields Vector{Float64} sa sb qm LLJ profit surplus drillcost revenue vdfull Eeps

     function SimulationTmp(Pprime, sa, sb, qm, LLJ, profit, surplus, drillcost, revenue, vdfull, Eeps)
         length(sa) == length(sb) == length(Eeps) == checksquare(Pprime) ||
            throw(DimensionMismatch())

         length(profit) == length(surplus) == length(drillcost) == length(revenue) ==
            length(vdfull) || throw(DimensionMismatch())

        return new(Pprime, sa, sb, qm, LLJ, profit, surplus, drillcost, revenue, vdfull, Eeps)
     end
end

@getFieldFunction SimulationTmp Pprime sa sb qm LLJ profit surplus drillcost revenue vdfull Eeps

function SimulationTmp(data::DataDrill, M::Integer)
    m = _model(data)

    Pprime = sparse_state_transition(statespace(m))

    nS = length(statespace(m))
    sa = Vector{Float64}(undef, nS)
    sb = similar(sa)
    Eeps = similar(sa)
    qm = Vector{Float64}(undef, M)

    LLJ = Vector{Float64}(undef, maxj1length(data))

    profit     = zeros(Float64, _dmax(statespace(m))+1)
    surplus    = similar(profit)
    drillcost  = similar(profit)
    revenue    = similar(profit)
    vdfull     = similar(profit)

     return SimulationTmp(Pprime, sa, sb, qm, LLJ, profit, surplus, drillcost, revenue, vdfull, Eeps)
end

"""
    isodd(t) ? (sa,sb) : (sb,sa)

`tday` holds our current state in `t`
`tmrw` holds the distribution of our state in `t+1` given prices
we switch between `sa` and `sb` to avoid allocating...
"""

function today_tomorrow(x::SimulationTmp, t::Integer)
    if isodd(t)
        return sa(x), sb(x)
    else
        return sb(x), sa(x)
    end
end

SimulationTmp(data::DataSetofSets, M) = SimulationTmp(drill(data), _num_sim(M))

num_states(x::SimulationTmp) = length(sa(x))

function reset!(x::SimulationTmp, s0::Integer)
    fill!(sa(x), 0)
    fill!(sb(x), 0)
    fill!(Eeps(x), 0)
    setindex!(sa(x), 1, s0)
    setindex!(sb(x), 1, s0)
    return nothing
end

function update!(simtmp::SimulationTmp, obs::ObservationDrill, sim::SimulationDraw, θ)

    m = _model(obs)
    wp = statespace(m)
    dograd = false
    newobs = ObservationDrill(m, ichars(obs), zchars(obs), 0, length(wp))

    x = reward(m)
    surp = NoRoyaltyProblem(x)

    # components of payoff. NOTE that these DO NOT depend on the deterministic state... just zₜ
    for d in actionspace(wp)

        r   = flow!(vw_revenue(x, θ), revenue(x), d, newobs, vw_revenue(x, θ), sim, dograd)
        c   = flow!(vw_cost(   x, θ), cost(   x), d, newobs, vw_cost(   x, θ), sim, dograd)
        pft = flow!(θ,                        x,  d, newobs, θ,                sim, dograd)
        s   = flow!(θ,                     surp,  d, newobs, θ,                sim, dograd)

        # assumes extension & eur are fixed
        setindex!( revenue(  simtmp), r  , d+1)  # (1-r)rev - 0
        setindex!( drillcost(simtmp), c  , d+1)  #    0     - drillcost(d)
        setindex!( profit(   simtmp), pft, d+1)  # (1-r)rev - drillcost(d)
        setindex!( surplus(  simtmp), s  , d+1)  #      rev - drillcost(d)
    end

    Q = eur_kernel(revenue(x), 0, newobs, vw_revenue(x,θ), sim)
    ext = extensioncost(extend(x), vw_extend(x,θ))

    return Q, ext
end

# ----------------------
# Holds simulations
# ----------------------

"What kind of simulation to run given datastruct"
struct SimulationPrimitives{D<:DataSetofSets,R}
    data::D
    sim::SimulationDraws{Float64,2,Matrix{Float64}}
    Tstop::Int
    theta::Vector{Float64}
    simtmp::SimulationTmp
    sharedsim::SharedSimulations{R}

    function SimulationPrimitives(
        data::D, sim, Tstop, theta, sharedsim::SharedSimulations{R}
    ) where {D<:DataSetofSets, R}
        _nparm(data) == length(theta) || throw(DimensionMismatch("data & theta"))
        num_i(data) == num_i(sim) || throw(DimensionMismatch("num_i"))
        0 < Tstop || throw(DomainError(Tstop))
        drilldat = drill(data)
        Tstop <= length(zchars(drilldat)) || throw(DomainError(Tstop))
        new{D,R}(data, sim, Tstop, theta, SimulationTmp(data, sim), sharedsim)
    end
end

@noinline function SimulationPrimitives(
    data::DataSetofSets, sim::SimulationDrawsMatrix,
    rwrd::AbstractStaticPayoff, wp::AbstractStateSpace, Tstop, theta, sharedsim
    )
    theta_d = theta_drill(data, theta)
    θρ = theta_drill_ρ(_model(drill(data)), theta_d)
    update!(sim, θρ)
    update!(data, theta)
    datanew = DataSetofSets(data, DataDrill(drill(data), rwrd, wp))
    return SimulationPrimitives(datanew, sim, Tstop, theta, sharedsim)
end

function SimulationPrimitives(rwrd, wp, Tstop, theta)
    data = get_g_BaseDataSetofSets()
    sim = get_g_SimulationDrawsMatrix()
    sharedsim = get_g_SharedSimulations()
    return SimulationPrimitives(data, sim, rwrd, wp, Tstop, theta, sharedsim)
end

function set_g_SimulationPrimitives(rwrd, wp, Tstop, theta)
    simprim = SimulationPrimitives(rwrd, wp, Tstop, theta)
    return set_g_SimulationPrimitives(simprim)
end

SimulationTmp(x::SimulationPrimitives) = x.simtmp
SharedSimulations(x::SimulationPrimitives) = x.sharedsim
SimulationDraws(x::SimulationPrimitives) = x.sim
SimulationDrawsMatrix(x::SimulationPrimitives) = SimulationDraws(x)
SimulationDrawsVector(x::SimulationPrimitives, i) =
    view(SimulationDrawsMatrix(x), i)

@getFieldFunction SimulationPrimitives Tstop theta data

split_thetas(x::SimulationPrimitives) = split_thetas(data(x), theta(x))
theta_drill(x::SimulationPrimitives) = first(split_thetas(x))


text:
using Printf

function is_valid(word)
    for letter in word
        if count(c -> (c == letter), word) > count(c -> (c == letter), letters)
            return false
        end
    end

    return true
end

function load_valid_words()
    return [line for line in readlines("dictionary.txt") if is_valid(line)]
end

const global scores = Dict(
	'a' => 1, 'b' => 3, 'c' => 3, 'd' => 2,
	'e' => 1, 'f' => 4, 'g' => 2, 'h' => 4,
	'i' => 1, 'j' => 8, 'k' => 5, 'l' => 1,
	'm' => 3, 'n' => 1, 'o' => 1, 'p' => 3,
	'q' => 10,'r' => 1, 's' => 1, 't' => 1,
	'u' => 1, 'v' => 4, 'w' => 4, 'x' => 8,
    'y' => 4, 'z' => 10
)

function get_score(word)
    return sum(scores[letter] for letter in word)
end

function display_words(words)
    len = length(words)

    for i = len:-1:1
        index = len - i + 1
        @printf("%i. %s\n", i, words[index])
    end
end

const global letters = "abcdefg"

function main()
    words = load_valid_words()
    sort!(words, by=get_score)
    display_words(words)
end

@time main()


text:
using CSV,DataFrames
import Utils
#

function dsum( df::DataFrame, keys )
    s = zeros( size( df,1 ))
    for (k,v) in keys
        # println( k )
        # println( df[!,k])
        s .+= df[!,k]
    end
    s
end

good_things = (
    :c11111t =>"Rice",
    :c11121t =>"Bread",
    :c11131t =>"Pasta products",
    :c11142t =>"Pastry (savoury)",
    :c11151t =>"Other breads and cereals",
    :c11211t =>"Beef (fresh, chilled or frozen)",
    :c11221t =>"Pork (fresh, chilled or frozen)",
    :c11231t =>"Lamb (fresh, chilled or frozen)",
    :c11241t =>"Poultry (fresh, chilled or frozen)",
    :c11251t =>"Sausages",
    :c11252t =>"Bacon and ham",
    :c11253t =>"Offal, pâté etc.",
    :c11261t =>"Other preserved or processed meat and meat preparations",
    :c11271t =>"Other fresh, chilled or frozen edible meat",
    :c11311t =>"Fish (fresh, chilled or frozen)",
    :c11321t =>"Seafood (fresh, chilled or frozen)",
    :c11331t =>"Dried, smoked or salted fish and seafood",
    :c11341t =>"Other preserved or processed fish and seafood and preparations",
    :c11411t =>"Whole milk",
    :c11421t =>"Low fat milk",
    :c11431t =>"Preserved milk",
    :c11441t =>"Yoghurt",
    :c11451t =>"Cheese and curd",
    :c11461t =>"Other milk products",
    :c11471t =>"Eggs",
    :c11511t =>"Butter",
    :c11521t =>"Margarine and other vegetable fats",
    :c11522t =>"Peanut butter",
    :c11531t =>"Olive oil",
    :c11541t =>"Edible oils",
    :c11551t =>"Other edible animal fats",
    :c11611t =>"Citrus fruits (fresh)",
    :c11621t =>"Bananas (fresh)",
    :c11631t =>"Apples (fresh)",
    :c11641t =>"Pears (fresh)",
    :c11651t =>"Stone fruits (fresh)",
    :c11661t =>"Berries (fresh)",
    :c11671t =>"Other fresh, chilled or frozen fruits",
    :c11681t =>"Dried fruit and nuts",
    :c11691t =>"Preserved fruit and fruit-based products",
    :c11711t =>"Leaf and stem vegetables (fresh or chilled)",
    :c11721t =>"Cabbages (fresh or chilled)",
    :c11731t =>"Vegetables grown for their fruit (fresh, chilled or frozen)",
    :c11741t =>"Root crops, non-starchy bulbs and mushrooms (fresh or frozen)",
    :c11751t =>"Dried vegetables",
    :c11761t =>"Other preserved or processed vegetables",
    :c11771t =>"Potatoes",
    :c11781t =>"Other tubers and products of tuber vegetables",
    :c11911t =>"Sauces, condiments",
    :c11921t =>"Salt, spices and culinary herbs",
    :c11931t =>"Baker's yeast, dessert preparations, soups",
    :c11941t =>"Other food products",
    :c12111t =>"Coffee",
    :c12121t =>"Tea",
    :c12131t =>"Cocoa and powdered chocolate",
    :c12211t =>"Mineral or spring waters",
    :c12241t =>"Vegetable juices"
)

sugars = (
    :c11122t =>"Buns, crispbread and biscuits",
    :c11141t =>"Cakes and puddings",
    :c11811t =>"Sugar",
    :c11821t =>"Jams, marmalades",
    :c11831t =>"Chocolate",
    :c11841t =>"Confectionery products",
    :c11851t =>"Edible ices and ice cream",
    :c11861t =>"Other sugar products",
    :c12131t =>"Cocoa and powdered chocolate",
    :c12221t =>"Soft drinks",
    :c12231t =>"Fruit juices",
    :cb1112t =>"Eat Out Confectionery eaten off premises",
    :cb1113t =>"Eat Out Ice cream eaten off premises",
    :cb1114t =>"Eat Out Takeaway Soft drinks eaten off premises",
    :cb1117t =>"Eat Out Confectionery (child)",
    :cb1118t =>"Eat Out Ice cream (child)",
    :cb1119t =>"Eat Out Soft drinks (child)",
    :cb1122t =>"Eat Out Confectionery",
    :cb1123t =>"Eat Out Ice cream",
    :cb1124t =>"Eat Out Soft drinks"
    )

sugars_at_home = (
    :c11122t =>"Buns, crispbread and biscuits",
    :c11141t =>"Cakes and puddings",
    :c11811t =>"Sugar",
    :c11821t =>"Jams, marmalades",
    :c11831t =>"Chocolate",
    :c11841t =>"Confectionery products",
    :c11851t =>"Edible ices and ice cream",
    :c11861t =>"Other sugar products",
    :c12131t =>"Cocoa and powdered chocolate",
    :c12221t =>"Soft drinks",
    :c12231t =>"Fruit juices"
    )

other_eating_out = (
    :cb1111t =>"Catered food non-alcoholic drink eaten / drunk on premises",
    :cb1115t =>"Hot food eaten off premises",
    :cb1116t =>"Cold food eaten off premises",
    :cb111at =>"Hot food (child)",
    :cb111bt =>"Cold food (child)",
    :cb1121t =>"Food non-alcoholic drinks eaten drunk on premises",
    :cb1125t =>"Hot food",
    :cb1126t =>"Cold food",
    :cb1127t =>"Hot take away meal eaten at home",
    :cb1128t =>"Cold take away meal eaten at home",
    :cb112bt =>"Contract catering (food)",
    :cb1213t =>"Meals bought and eaten at workplace",
    :cb1311t =>"Catered food - eaten on premises",
    :cb1312t =>"Non-alcoholic drink - drunk on premises"
)

sugars_eating_out = (
:cb1112t =>"Eat Out Confectionery eaten off premises",
:cb1113t =>"Eat Out Ice cream eaten off premises",
:cb1114t =>"Eat Out Takeaway Soft drinks eaten off premises",
:cb1117t =>"Eat Out Confectionery (child)",
:cb1118t =>"Eat Out Ice cream (child)",
:cb1119t =>"Eat Out Soft drinks (child)",
:cb1122t =>"Eat Out Confectionery",
:cb1123t =>"Eat Out Ice cream",
:cb1124t =>"Eat Out Soft drinks" )

bad_stuff = (
    :cb111ct =>"Spirits and liqueurs (away from home)",
    :cb111dt =>"Wine from grape or other fruit (away from home)",
    :cb111et =>"Fortified wines (away from home)",
    :cb111ft =>"Ciders and Perry (away from home)",
    :cb111gt =>"Alcopops (away from home)",
    :cb111ht =>"Champagne and sparkling wines (away from home)",
    :cb111it =>"Beer and lager (away from home)",
    :cb111jt =>"Round of drinks (away from home)",
    :c21111t =>"Spirits and liqueurs (brought home)",
    :c21211t =>"Wine from grape or other fruit (brought home)",
    :c21212t =>"Fortified wine (brought home)",
    :c21213t =>"Ciders and Perry (brought home)",
    :c21214t =>"Alcopops (brought home)",
    :c21221t =>"Champagne and sparkling wines (brought home)",
    :c21311t =>"Beer and lager (brought home)",
    :c22111t =>"Cigarettes",
    :c22121t =>"Cigars",
    :c22131t =>"Other tobacco"
 )

  # C21111t +
  # C21211t +
  # C21212t +
  # C21213t +
  # C21214t +
  # C21221t +
  # C21311t +
  # C22111t +
  # C22121t +
  # C22131t +
  # C23111t

allfood = union( good_things, sugars, bad_stuff, other_eating_out )

for (k,v) in allfood
    println(Symbol(Utils.basiccensor(v)))
end

for (k,v) in good_things
    println(Symbol(Utils.basiccensor(v)))
end

lcfraw = Utils.loadtoframe( "/mnt/data/lcf/1718/tab/tab/dvhh_ukanon_2017-18.tab" ) |> DataFrame
lcf = DataFrame()

lcf[!,:age_u_18]=lcfraw[!,:a020]+lcfraw[!,:a021]+lcfraw[!,:a022]+lcfraw[!,:a030]+lcfraw[!,:a031]+lcfraw[!,:a032]
lcf[!,:age_18_plus]=lcfraw[!,:a049]-lcf[!,:age_u_18]
lcf[!,:tenure_type] =  lcfraw[!,:a122]
lcf[!,:region] =  lcfraw[!,:gorx]
lcf[!,:economic_pos] =  lcfraw[!,:a093]
lcf[!,:age_of_oldest] =  lcfraw[!,:a065p]
lcf[!,:healthy_food] = dsum( lcfraw, good_things )
lcf[!,:takeaways] = dsum( lcfraw, other_eating_out )
lcf[!,:alcohol_tobacco] = dsum( lcfraw, bad_stuff )

for (k,v) in sugars
    sk = Symbol(Utils.basiccensor(v))
    lcf[!,sk] = lcfraw[!,k]
end

n = names( lcfraw )

lcf[!,:sumallfood] = dsum( lcfraw, allfood )

lcf[!,:age_u_18]=lcfraw[!,:a020]+lcfraw[!,:a021]+lcfraw[!,:a022]+lcfraw[!,:a030]+lcfraw[!,:a031]+lcfraw[!,:a032]
lcf[!,:age_18_plus]=lcfraw[!,:a049]-lcf[!,:age_u_18]
lcf[!,:tenure_type] =  lcfraw[!,:a122]
lcf[!,:region] =  lcfraw[!,:gorx]
lcf[!,:economic_pos] =  lcfraw[!,:a093]
lcf[!,:age_of_oldest] =  lcfraw[!,:a065p]

lcf[!,:total_consumpt] =  lcfraw[!,:p600t]
lcf[!,:food_and_drink] =  lcfraw[!,:p601t]
lcf[!,:alcohol_tobacco] =  lcfraw[!,:p602t]
lcf[!,:clothing] =  lcfraw[!,:p603t]
lcf[!,:housing] =  lcfraw[!,:p604t]
lcf[!,:household_goods] =  lcfraw[!,:p605t]
lcf[!,:health] =  lcfraw[!,:p606t]
lcf[!,:transport] =  lcfraw[!,:p607t]
lcf[!,:communication] =  lcfraw[!,:p608t]
lcf[!,:recreation] =  lcfraw[!,:p609t]
lcf[!,:education] =  lcfraw[!,:p610t]
lcf[!,:restaurants_etc] =  lcfraw[!,:p611t]
lcf[!,:miscellaneous] =  lcfraw[!,:p612t]
lcf[!,:non_consumption] =  lcfraw[!,:p620tp]
lcf[!,:total_expend] =  lcfraw[!,:p630tp]
lcf[!,:equiv_scale] =  lcfraw[!,:oecdsc]
lcf[!,:weekly_net_inc] =  lcfraw[!,:p389p]

allfood = union( good_things, sugars_at_home )

max_stars_repo_name: grahamstark/ScottishTaxBenefitModel

max_issues_repo_name: grahamstark/ScottishTaxBenefitModel

max_forks_repo_name: grahamstark/ScottishTaxBenefitModel


text:
import GeometricFlux: message, update, propagate

in_channel = 10
out_channel = 5
N = 6
adj = [0. 1. 0. 0. 0. 0.;
       1. 0. 0. 1. 1. 1.;
       0. 0. 0. 0. 0. 1.;
       0. 1. 0. 0. 1. 0.;
       0. 1. 0. 1. 0. 1.;
       0. 1. 1. 0. 1. 0.]
ne = [[2], [1,4,5,6], [6], [2,5], [2,4,6], [2,3,5]]

struct NewLayer <: MessagePassing
    adjlist::AbstractVector{<:AbstractVector}
    weight
end

NewLayer(adjm::AbstractMatrix, m, n) = NewLayer(neighbors(adjm), randn(m,n))

(l::NewLayer)(X) = propagate(l, X=X, aggr=:add)
message(n::NewLayer; x_i=zeros(0), x_j=zeros(0)) = x_j
update(::NewLayer; X=zeros(0), M=zeros(0)) = M

X = Array(reshape(1:N*in_channel, in_channel, N))
l = NewLayer(adj, in_channel, out_channel)

message(n::NewLayer; x_i=zeros(0), x_j=zeros(0)) = n.weight' * x_j

@testset "msgpass" begin
    Y = l(X)
    @test size(Y) == (out_channel, N)
end


text:
#module PowerFlowTests
include("../src/PowerFlowAnalysis.jl")

using Base.Test
using JuMP
using Ipopt
using PowerModels
using PowerFlowAnalysis

case6ww = PowerModels.parse_file("data/case6ww.m")
pm = ACPPowerModel(case6ww)
sol = run_ac_pf(case6ww, IpoptSolver())

net = Network(pm)
sys = PowerFlowSystem(pm, sol)
lf = LurieForm(sys)

known_pairs = [(1,4), (2,5), (3,6)] # lines in the orignal file
@testset "Network topology" for pair in known_pairs
  f, t = pair
  @test (f,t) in net.pairs
  pidx = net.pair_idx[(f,t)]
  @test net.from[pidx] == net.bus_idx[f]
  @test net.to[pidx]   == net.bus_idx[t]
end

@testset "Admittance matrix" begin
  Ymp = Dict(
      (5,5) => 5.293290152818541 - 14.137764902337823im,
      (3,6) => -1.923076923076923 + 9.615384615384615im,
      (1,3) => 0.0im
    ) # Several values extracted from MATPOWER
  for ((f,t), y) in Ymp
    @test y ≈ sys.Y[net.bus_idx[f], net.bus_idx[t]]
  end
end

p, q = calc_pq(sys.Y, sys.v⁰, sys.ϑ⁰)
@testset "Power flow solutions" begin
  @test p[sys.pq] ≈ sys.p⁰[sys.pq]
  @test p[sys.pv] ≈ sys.p⁰[sys.pv]
  @test q[sys.pq] ≈ sys.q⁰[sys.pq]
end

@testset "Lurie form" begin
  v¹ = (1 + 0.1*randn(sys.nb)).*sys.v⁰
  ϑ¹ = (1 + 0.1*randn(sys.nb)).*sys.ϑ⁰
  f¹ = calc_f(v¹, ϑ¹, lf)
  s¹ = lf.S*f¹
  p¹, q¹ = calc_pq(sys.Y, v¹, ϑ¹)

  @test p¹[sys.pv] ≈ s¹[1:sys.npv]
  @test p¹[sys.pq] ≈ s¹[(sys.npv + 1):(sys.npv + sys.npq)]
  @test q¹[sys.pq] ≈ s¹[(sys.npv + sys.npq + 1):(sys.npv + 2*sys.npq)]
end

@testset "Jacobian ∂f/∂x" begin
  dfdx = PowerFlowAnalysis.calc_dfdx(lf)
  dx = 1e-4*randn(size(lf.x⁰))
  v¹,ϑ¹ = PowerFlowAnalysis.calc_vϑ(lf.x⁰ + dx, lf)
  df = PowerFlowAnalysis.calc_f(v¹, ϑ¹, lf) - lf.f⁰
  @test isapprox(df, dfdx*dx; rtol = 1e-3)
end

max_stars_repo_name: hunghtd/PowerFlowAnalysis-tutorial

max_issues_repo_name: hunghtd/PowerFlowAnalysis-tutorial

max_forks_repo_name: hunghtd/PowerFlowAnalysis-tutorial


text:
# ------------ Call sites ------------ #

@inline function (ctx::ScoreContext)(c::typeof(plate),
                                     addr::Address,
                                     call::Function,
                                     args::Vector)
    visit!(ctx, addr => 1)
    ss = get_subselection(ctx, (addr, 1))
    len = length(args)
    ret, w = score(ss, call, args[1]...)
    v_ret = Vector{typeof(ret)}(undef, len)
    v_ret[1] = ret
    increment!(ctx, w)
    for i in 2:len
        visit!(ctx, addr => i)
        ss = get_subselection(ctx, (addr, i))
        ret, w = score(ss, call, args[i]...)
        v_ret[i] = ret
        increment!(ctx, w)
    end
    return v_ret
end


text:
# Desired:
#   using Test
#
# In order that Test need not to be defined in the user environment, it is included via Modia:
import Modia

@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Modia.Test
end

@testset "RunTests" begin

    include("runsymbolictests.jl")
    # include("TestModelsWithError.jl")
    include("runsimulationtests.jl")
    include("../examples/runexamples.jl")
    # include("runlargetests.jl")

end # testset
nothing


text:
using JLD2
using Infiltrator
using LightGraphs
using Random
using JSON
using BenchmarkTools
using MultiAgentPathFinding

using AerialGroundPathFinding

const BASE_SEED = 6789
const MANHATTAN_NODEFILE = "../data/manhattan_node_attribs.json"
const MANHATTAN_WEIGHTFILE = "../data/manhattan_sparse_wts.jld2"

# Only caring about distance for now
const ALPHA_WEIGHT_DISTANCE = 1.0

const NCARS = parse(Int64, ARGS[1])
const NDRONES = parse(Int64, ARGS[2])
const OUTFILEDIR = ARGS[3]
const CAPACITY = parse(Int64, ARGS[4])
const TRIALS = parse(Int64, ARGS[5])
const DECAY = parse(Int64, ARGS[6])

function main()

    @load MANHATTAN_WEIGHTFILE manhattan_sparse_wts

    graph = LightGraphs.SimpleDiGraph(manhattan_sparse_wts)
    location_list = get_location2D_list(MANHATTAN_NODEFILE)

    for t = 1:TRIALS+1

        rng = MersenneTwister(BASE_SEED+t)

        ground_task_list = get_tasks_with_valid_path(graph, manhattan_sparse_wts, location_list, NCARS, rng)
        aerial_task_list = get_tasks_with_valid_path(graph, manhattan_sparse_wts, location_list, NDRONES, rng)

        env = CoordinatedMAPFEnv(ground_task_list=ground_task_list, aerial_task_list=aerial_task_list, road_graph=graph, road_graph_wts=manhattan_sparse_wts, location_list=location_list, alpha_weight_distance=ALPHA_WEIGHT_DISTANCE, car_capacity=CAPACITY)

        time_comps = Float64[]

        t1 = @elapsed dir_drone_paths = compute_independent_paths(env, env.aerial_task_list)
        t2 = @elapsed dir_car_paths = compute_independent_paths(env, env.ground_task_list)
        push!(time_comps, t1+t2)

        # New decay function
        if DECAY == "TANH"
            t3 = @elapsed augment_road_graph_with_aerial_paths!(env, dir_drone_paths, x -> 0.5*(1 + tanh(x)))
        else
            t3 = @elapsed augment_road_graph_with_aerial_paths!(env, dir_drone_paths) # logistic by default
        end
        push!(time_comps, t3)

        initial_gmapf_states = [GroundMAPFState(task.origin, false) for task in env.ground_task_list]
        ground_pp_ordering = collect(1:NCARS)

        t5 = @elapsed ground_pp_soln = plan_prioritized_paths!(env, ground_pp_ordering, initial_gmapf_states, GroundMAPFConstraint, PlanResult{GroundMAPFState,GroundMAPFAction,Float64})
        push!(time_comps, t5)

        t6 = @elapsed update_ground_paths_with_ground_mapf_result!(env, ground_pp_soln)
        push!(time_comps, t6)
        stage1_time = sum(time_comps)

        t7 = @elapsed set_ground_transit_graph!(env, [ddp.total_dist for ddp in dir_drone_paths])
        push!(time_comps, t7)

        initial_amapf_states = [AerialMAPFState(idx=1, ground_transit_idx=sg[1]) for (i, sg) in enumerate(env.gtg_drone_start_goal_idxs)]
        t8 = @elapsed aerial_pp_ordering = get_increasing_dist_ordering(dir_drone_paths)
        push!(time_comps, t8)

        t9 = @elapsed aerial_pp_soln = plan_prioritized_paths!(env, aerial_pp_ordering, initial_amapf_states, GroundTransitConstraint, PlanResult{AerialMAPFState,GroundTransitAction,Float64})
        push!(time_comps, t9)

        t10 = @elapsed update_aerial_ground_paths_cbs_solution!(env, aerial_pp_soln, dir_drone_paths)
        push!(time_comps, t10)

        not_coord_ids = count_cars_not_coordinating(env.ground_paths)
        for ncid in not_coord_ids
            env.ground_paths[ncid] = dir_car_paths[ncid]
        end

        total_dist = compute_total_cost(env, env.aerial_paths, env.ground_paths)

        total_time = sum(time_comps)
        stage2_time = total_time - stage1_time

        res_dict = Dict{String,Float64}("compute_time" => sum(time_comps), "total_path_cost" => total_dist, "stage1_time" => stage1_time, "stage2_time" => stage2_time)

        if t > 1
            fname = string(OUTFILEDIR, "_Cars_", NCARS, "_Drones_", NDRONES, "_Capacity_", CAPACITY, "_Trial_", (t-1), ".json")
            println(fname)
            open(fname, "w") do f
                JSON.print(f, res_dict, 2)
            end
        end
    end
end

main()


text:
function costKproduv(pin, t, u, v)
    p = if typeof(pin) <: Vector
        reshape(pin, length(pin) ÷ 2, 2)
    elseif typeof(pin) <: Matrix
        pin
    end
    np = size(p, 1)
    I, J = indKM_sparse(t)
    SK, SM = FemAdjoint.assembKM_P12D(p, t)
    val = 0.
    for (i, j, s) ∈ zip(I, J, SK)
        val += u[i] * s * v[j]
    end
    return val
end
function ∇costKproduv(pin, t, u, v;
                      jacK::SparseMatrixCSC{Float64, Int64} = spzeros(0, 0),
                      jacM::SparseMatrixCSC{Float64, Int64} = spzeros(0, 0))
    p = if typeof(pin) <: Vector
        reshape(pin, length(pin) ÷ 2, 2)
    elseif typeof(pin) <: Matrix
        pin
    end
    np = size(p, 1)
    I, J = indKM_sparse(t)
    # jacobian structures
    fK(y, x) = FemAdjoint.assembK_P12D_inplace!(x, t, y)
    if length(jacK) == 0
        y = similar(FemAdjoint.assembKM_P12D(p, t)[1])
        sparsity_pattern = Symbolics.jacobian_sparsity(fK, y, p[:])
        jacK = Float64.(sparse(sparsity_pattern))
    end
    # evaluate jacobians
    colorsK = matrix_colors(jacK)
    JK = forwarddiff_color_jacobian!(jacK, fK, p[:], colorvec = colorsK)
    # compute full gradient
    g = zeros(np, 2)
    for (i, j, s) ∈ zip(findnz(JK)...)
        g[j] += u[J[i]] * s * v[I[i]]
    end
    return g
end
function costMproduv(pin, t, u, v)
    p = if typeof(pin) <: Vector
        reshape(pin, length(pin) ÷ 2, 2)
    elseif typeof(pin) <: Matrix
        pin
    end
    np = size(p, 1)
    I, J = indKM_sparse(t)
    SK, SM = FemAdjoint.assembKM_P12D(p, t)
    val = 0.
    for (i, j, s) ∈ zip(I, J, SM)
        val += u[i] * s * v[j]
    end
    return val
end
function ∇costMproduv(pin, t, u, v;
                      jacK::SparseMatrixCSC{Float64, Int64} = spzeros(0, 0),
                      jacM::SparseMatrixCSC{Float64, Int64} = spzeros(0, 0))
    p = if typeof(pin) <: Vector
        reshape(pin, length(pin) ÷ 2, 2)
    elseif typeof(pin) <: Matrix
        pin
    end
    np = size(p, 1)
    I, J = indKM_sparse(t)
    # jacobian structures
    fM(y, x) = FemAdjoint.assembM_P12D_inplace!(x, t, y)
    if length(jacM) == 0
        y = similar(FemAdjoint.assembKM_P12D(p, t)[2])
        sparsity_pattern = Symbolics.jacobian_sparsity(fK, y, p[:])
        jacK = Float64.(sparse(sparsity_pattern))
    end
    # evaluate jacobians
    colorsM = matrix_colors(jacM)
    JM = forwarddiff_color_jacobian!(jacM, fM, p[:], colorvec = colorsM)
    # compute full gradient
    g = zeros(np, 2)
    for (i, j, s) ∈ zip(findnz(JM)...)
        g[j] += u[J[i]] * s * v[I[i]]
    end
    return g
end


text:
using Test
using UltraDark
using NPZ
import MPI
import PencilFFTs

for grids_type in [Grids, PencilGrids]
    @testset "Can initialize $grids_type" begin
        resol = 8

        grids = grids_type(1.0, resol)
        @test typeof(grids) <: grids_type
        @test all(grids.ψx .== 0)
        @test reshape(grids.x, resol) == reshape(grids.y, resol)
        @test reshape(grids.x, resol) == reshape(grids.z, resol)

        grids = grids_type((1.0, 1.0, 1.0), (resol, resol, resol))
        @test typeof(grids) <: grids_type
        @test all(grids.ψx .== 0)
        @test reshape(grids.x, resol) == reshape(grids.y, resol)
        @test reshape(grids.x, resol) == reshape(grids.z, resol)

        grids = grids_type((4.0, 2.0, 1.0), (4resol, 2resol, 1resol))
        @test typeof(grids) <: grids_type
        @test all(grids.ψx .== 0)
        @test size(grids.x) == (4resol, 1, 1)
        @test size(grids.y) == (1, 2resol, 1)
        @test size(grids.z) == (1, 1, 1resol)
        @test reshape(grids.x, 4resol)[1 + 4resol÷2 - 4resol÷8:4resol÷2 + 4resol÷8] == reshape(grids.z, resol)
        @test reshape(grids.y, 2resol)[1 + 2resol÷2 - 2resol÷4:2resol÷2 + 2resol÷4] == reshape(grids.z, resol)
    end
end

@testset "Consistent initialisation" begin
    resol = 16
    dir = mktempdir()

    grids = Grids(1.0, resol)
    pencil_grids = PencilGrids(1.0, resol)

    gathered_pg = PencilFFTs.gather(pencil_grids.k)

    if MPI.Comm_rank(pencil_grids.MPI_COMM) == 0
        @test all(grids.k .== gathered_pg)

        npzwrite(joinpath(dir, "k_grids.npy"), grids.k)

        npzwrite(joinpath(dir, "k_pencilgrids.npy"), gathered_pg)

        @test all(npzread(joinpath(dir, "k_grids.npy")) .== npzread(joinpath(dir, "k_pencilgrids.npy")))
    else
        @test size(gathered_pg) == 0
    end

end


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -4.116064765345723 <= q <= -0.12537054293327043)

                     add_NL_constraint(m, :(gelu(-0.6620590692147021 + -0.9580027962881332*gelu(-0.38090481240177176 + -0.02939531016329644*gelu(0.2067810723024852 + 0.7469716803367321*gelu(-0.2319263638689142 + 0.46965024077813977*$(x[1]) + 0.028636456496167195*$(x[2]) + 0.15542595955334626*$(x[3])) + -0.3209180228281623*gelu(-0.2352871682017641 + -0.6225824541901428*$(x[1]) + -0.2700435280904796*$(x[2]) + -0.7791025613764826*$(x[3]))) + -0.9724733578803808*gelu(-0.4274287128450034 + -0.22403315015562253*gelu(-0.2319263638689142 + 0.46965024077813977*$(x[1]) + 0.028636456496167195*$(x[2]) + 0.15542595955334626*$(x[3])) + -0.6136110647248842*gelu(-0.2352871682017641 + -0.6225824541901428*$(x[1]) + -0.2700435280904796*$(x[2]) + -0.7791025613764826*$(x[3])))) + -0.699289524220764*gelu(0.4498619401955626 + 0.5463276179102481*gelu(0.2067810723024852 + 0.7469716803367321*gelu(-0.2319263638689142 + 0.46965024077813977*$(x[1]) + 0.028636456496167195*$(x[2]) + 0.15542595955334626*$(x[3])) + -0.3209180228281623*gelu(-0.2352871682017641 + -0.6225824541901428*$(x[1]) + -0.2700435280904796*$(x[2]) + -0.7791025613764826*$(x[3]))) + -0.02415027804893688*gelu(-0.4274287128450034 + -0.22403315015562253*gelu(-0.2319263638689142 + 0.46965024077813977*$(x[1]) + 0.028636456496167195*$(x[2]) + 0.15542595955334626*$(x[3])) + -0.6136110647248842*gelu(-0.2352871682017641 + -0.6225824541901428*$(x[1]) + -0.2700435280904796*$(x[2]) + -0.7791025613764826*$(x[3]))))) + gelu(-0.8107501452143855 + 0.008470632510253484*gelu(-0.38090481240177176 + -0.02939531016329644*gelu(0.2067810723024852 + 0.7469716803367321*gelu(-0.2319263638689142 + 0.46965024077813977*$(x[1]) + 0.028636456496167195*$(x[2]) + 0.15542595955334626*$(x[3])) + -0.3209180228281623*gelu(-0.2352871682017641 + -0.6225824541901428*$(x[1]) + -0.2700435280904796*$(x[2]) + -0.7791025613764826*$(x[3]))) + -0.9724733578803808*gelu(-0.4274287128450034 + -0.22403315015562253*gelu(-0.2319263638689142 + 0.46965024077813977*$(x[1]) + 0.028636456496167195*$(x[2]) + 0.15542595955334626*$(x[3])) + -0.6136110647248842*gelu(-0.2352871682017641 + -0.6225824541901428*$(x[1]) + -0.2700435280904796*$(x[2]) + -0.7791025613764826*$(x[3])))) + 0.39739497887500974*gelu(0.4498619401955626 + 0.5463276179102481*gelu(0.2067810723024852 + 0.7469716803367321*gelu(-0.2319263638689142 + 0.46965024077813977*$(x[1]) + 0.028636456496167195*$(x[2]) + 0.15542595955334626*$(x[3])) + -0.3209180228281623*gelu(-0.2352871682017641 + -0.6225824541901428*$(x[1]) + -0.2700435280904796*$(x[2]) + -0.7791025613764826*$(x[3]))) + -0.02415027804893688*gelu(-0.4274287128450034 + -0.22403315015562253*gelu(-0.2319263638689142 + 0.46965024077813977*$(x[1]) + 0.028636456496167195*$(x[2]) + 0.15542595955334626*$(x[3])) + -0.6136110647248842*gelu(-0.2352871682017641 + -0.6225824541901428*$(x[1]) + -0.2700435280904796*$(x[2]) + -0.7791025613764826*$(x[3]))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
export get_nodegraph,
    get_param_types,
    get_new_node_id,
    set_param_types

"""
    get_nodegraph()

Returns the underlying [`NodeGraph`](@ref)
"""
get_nodegraph() = load(NG_FILE, "ng")

"""
    get_param_types()

Returns the types of nodes that can be created, as a list of their type-symbols.
"""
get_param_types() = Symbol.(readlines(PTYPES_FILE))

"""
    get_new_node_id()

Returns the ID of the next node that will be created
"""
get_new_node_id() = read(NNID_FILE, Int)

"""
    set_nodegraph(ng::NodeGraph)

Sets the currently active [`NodeGraph`](@ref). Usage of this function
should not be required. To edit a pre-existing node graph, pass it to
[`run_server`](@ref).
"""
set_nodegraph(ng::NodeGraph) = jldsave(NG_FILE; ng)

"""
    set_param_types(types::Vector{Symbol})

Sets the types of nodes that can be created. Requires a list of valid
type-symbols.
"""
set_param_types(types::Vector{Symbol}) = write(PTYPES_FILE, join(types, "\n"))

"""
    set_new_node_id(nn_id::Int)

Sets the ID of the next node to be created.
"""
set_new_node_id(nn_id::Int) = write(NNID_FILE, nn_id)


text:
using ThingSpeak
using Base.Test

# write your own tests here
using ThingSpeak
function loadsecret(name="THINGSPEAK_API_KEY")
  if name in keys(ENV)
    println("Loading $name from environment variable")
    return ENV[name]
  else
    println("Loading $name from file")
    return open(name) do f
      chomp(readline(f))
    end
  end
end

ThingSpeak.setdefaultchannel(loadsecret()) #this is technically a secret but its a dedicated channel for testing
response = tsupdate(Dict(3=>3.5),status="ran runtests.jl for ThingSpeak.jl")
response = tsupdate(1.1,2.2)
@test tssuccess(response)

badresponse = tsupdate(1,2,channel=ThingSpeak.TSChannel("ABC")) # use invalid API key
@test !tssuccess(badresponse)


text:
using Enigma
using Test

@testset "Basics" begin
@testset "No plugboard" begin
    for rp1 = 1:26, rp2=1:26
        enigma = EnigmaMachine()
        set_rotors!(enigma, 1,2,3)
        set_rotor_positions!(enigma, rp1,rp2,4)
        set_ukw!(enigma, 2)

        message = "This is an Enigma test"
        encoded = encode!(enigma, message)

        # set rotor position again for decoding
        set_rotor_positions!(enigma, rp1,rp2,4)
        decoded = decode!(enigma, encoded)
        @test decoded == enigma_styled_text(message)
    end

    for rp1 = 1:26, rp2=1:26
        enigma = EnigmaMachine()
        set_rotors!(enigma, 4,3,5)
        set_rotor_positions!(enigma, rp1,rp2,4)
        set_ukw!(enigma, 3)

        message = "This is an Enigma test"
        encoded = encode!(enigma, message)

        # set rotor position again for decoding
        set_rotor_positions!(enigma, rp1,rp2,4)
        decoded = decode!(enigma, encoded)
        @test decoded == enigma_styled_text(message)
    end
end

@testset "Plugboard" begin
    alphabet = ['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z']
    for rp1 = 1:26, pb1 in [9,11,15], pb2 in [10,20,12], pb3 in [4,5,6], pb4 in [24,13]
        enigma = EnigmaMachine()
        set_rotors!(enigma, 1,2,3)
        set_rotor_positions!(enigma, rp1,9,10)
        set_ukw!(enigma, 2)
        setting = ""
        setting *= Char(pb1+64)
        setting *= Char(pb2+64)
        setting *= " "
        setting *= Char(pb3+64)
        setting *= Char(pb4+64)
        set_plugboard!(enigma, setting)

        message = "The quick brown fox jumps over the lazy dog"
        encoded = encode!(enigma, message)

        # set rotor position again for decoding
        set_rotor_positions!(enigma, rp1,9,10)
        decoded = decode!(enigma, encoded)
        @test decoded == enigma_styled_text(message)
    end
end

@testset "Plugboard ErrorException" begin
    enigma = EnigmaMachine()
    set_rotors!(enigma, 1,2,3)
    set_rotor_positions!(enigma, 3,9,10)
    set_ukw!(enigma, 2)
    setting = "AB AC"
    @test_throws ErrorException set_plugboard!(enigma, setting)

    enigma = EnigmaMachine()
    set_rotors!(enigma, 1,2,3)
    set_rotor_positions!(enigma, 3,9,10)
    set_ukw!(enigma, 2)
    setting = "AB CB"
    @test_throws ErrorException set_plugboard!(enigma, setting)
end

@testset "https://piotte13.github.io/" begin
    enigma = EnigmaMachine()
    set_rotors!(enigma, 1,2,3)
    set_rotor_positions!(enigma, 1,1,1)
    set_ukw!(enigma, 2)

    message = "Test"
    encoded = encode!(enigma, message)
    @test encoded == "OLPF"

    # =============================== #
    enigma = EnigmaMachine()
    set_rotors!(enigma, 1,2,3)
    set_rotor_positions!(enigma, 4,9,7)
    set_ukw!(enigma, 2)

    message = "TestTestTestTestTest"
    encoded = encode!(enigma, message)
    @test encoded == "UKTCV GCKJF GOJFX POSOH"

    # =============================== #
    enigma = EnigmaMachine()
    set_rotors!(enigma, 3,4,2)
    set_rotor_positions!(enigma, 5,9,5)
    set_ukw!(enigma, 2)
    set_plugboard!(enigma, "TE SA")

    message = "The quick brown fox jumps over the lazy dog The quick brown fox jumps over the lazy dog"
    encoded = encode!(enigma, message)
    @test encoded == "IYAAB ZZPON MPWRT RDKNG VEEHN FQJKD DLWGH JOAHE FPGCM JHGLA OBNJJ CRDJL OCIYW NUSHT"

    # =============================== SECRET MESSAGE ======================= #
    enigma = EnigmaMachine()
    set_rotors!(enigma, 1,2,3)
    set_rotor_positions!(enigma, 1,2,3)
    set_ukw!(enigma, 3)
    set_plugboard!(enigma, "DG XB JL")

    message = "BQGYP VBFTW XOMTP FBMVW NNUMF WDNVX ANCRD TBDZX ZGQGV OMGFB KUPHB ORKZU MSTHH PTMSH UXIDW FVUVJ"
    encoded = encode!(enigma, message)
    @test encoded == "LHJRY PFESX YIPQV UCIRC QVYIV QYUUA KLRVN PQWQJ HOJNF QZBYM XYMOO NDONW IELFM ICEXZ FWBVO DSQFX"
    set_rotor_positions!(enigma, 1,2,3)
    encoded = encode!(enigma, message; output_style=:plain)
    @test encoded == "LHJRYPFESXYIPQVUCIRCQVYIVQYUUAKLRVNPQWQJHOJNFQZBYMXYMOONDONWIELFMICEXZFWBVODSQFX"
end
end

@testset "Bombe" begin
@testset "Cracking given all but plugboard" begin
    crack_message = "VZEKL OYDBD SFFYD LZQVS HAWVP BGAMY ATKPW T"
    hint = "message"
    bombe = BombeMachine("DONT KNOW YET", "HINT?")
    set_secret!(bombe, crack_message)
    set_hint!(bombe, hint)
    set_possible_rotors!(bombe, 1,2,3)
    set_possible_rotor_positions!(bombe, 1,1,1)
    set_possible_ukws!(bombe, 1:1)
    # is at position 12 and 17 and 18 are not possible
    set_possible_hint_positions!(bombe, [12,17,18])
    enigmas = run_cracking(bombe; log=true)
    found = false
    correct_message = enigma_styled_text("This is a test message for testing the Bombe")

    for enigma in enigmas
        encoded = encode!(enigma, crack_message)
        if encoded == correct_message
            found = true
            break
        end
    end
    show(enigmas[1:2])
    show(enigmas[3])
    @test found === true

    # different rotors and positions
    crack_message = "QRTUV RQMHR GTBJU NVXAN OPHTM BZAHZ YFKUL CRVRY MLTVW KRXIH EAWXJ "
    hint = "message"
    bombe = BombeMachine(crack_message, hint)
    set_possible_rotors!(bombe, 3,2,1)
    set_possible_rotor_positions!(bombe, 4,7,8)
    set_possible_ukws!(bombe, 1)
    enigmas = run_cracking(bombe; log=false)
    found = false
    correct_message = enigma_styled_text("A long long message with no real meaning but it is important to test")

    for enigma in enigmas
        encoded = encode!(enigma, crack_message)
        if encoded == correct_message
            found = true
            break
        end
    end
    @test found === true

    # Try a larger set of possibilities
    crack_message = "HGHXI AGYEY NDIFW PRMDD QSMJG DCAKP FMIZL RVQIZ WRLJM "
    hint = "weatherreport"
    bombe = BombeMachine(crack_message, hint)
    set_possible_rotors!(bombe, 3,1:4,1:4)
    set_possible_rotor_positions!(bombe, 1:5,1:26,1:26)
    set_possible_ukws!(bombe, 1)
    set_possible_hint_positions!(bombe, 1:5)
    disable_ambiguous!(bombe)
    enigmas = run_cracking(bombe; log=false)
    found = false
    correct_message = enigma_styled_text("A Weatherreport It is very nice outside so be friendly")

    for enigma in enigmas
        encoded = encode!(enigma, crack_message)
        if encoded == correct_message
            found = true
            break
        end
    end
    @test found === true

    # unknown ukw but rotor position
    crack_message = "BLCCX HTMVX NRYXV REFZO LVXQU YHWFJ KHZDU BYBYK RTGHI GRWH"
    hint = "secret message"
    bombe = BombeMachine(crack_message, hint)
    set_possible_rotors!(bombe, 3,4,1)
    set_possible_hint_positions!(bombe, 10:20)
    enigmas = run_cracking(bombe; log=false)
    found = false
    correct_message = enigma_styled_text("This is another secret message with different ukw setting")
    hint_part = uppercase(replace(hint, " "=>""))
    @test length(enigmas) >= 1

    for enigma in enigmas
        encoded = encode!(enigma, crack_message)
        encoded_str = replace(encoded, " "=>"")
        @test findfirst(hint_part, encoded_str[10:40]) !== nothing
    end

    # need to check all ambiguous possibilities
    crack_message = "NTUPQ KYKLL PZNWY LUMGW DUWVV BLLHM USGVM PKDGV ZXOYV QSGIP XRURK OSYRP BNGVP LALIT PEHVE SODIW"
    hint = "Often it is easy"
    bombe = BombeMachine(crack_message, hint)
    set_possible_rotors!(bombe, 3,4,1)
    set_possible_rotor_positions!(bombe, 7, 11, 12)
    set_possible_hint_positions!(bombe, 1)
    set_possible_ukws!(bombe, 3)
    enable_ambiguous!(bombe)
    enigmas = run_cracking(bombe; log=false)
    found = false
    correct_message = enigma_styled_text("Often it is easy but sometimes it is hard to crack because the first word does not give enough hint")

    for enigma in enigmas
        encoded = encode!(enigma, crack_message)
        if encoded == correct_message
            found = true
            break
        end
    end
    @test found === true
end
end
@testset "plotting" begin
@testset "just let it run" begin
    # Just tests that the methods don't error :D
    enigma = EnigmaMachine()
    plts = get_enigma_decode_plots!(enigma, 'K')
    @test length(plts) == 29
    animate_plots(plts[1:2], "enigma_plot_320948209384023842204804038024"; end_extra=1)
    rm("enigma_plot_320948209384023842204804038024.gif")
end
end


text:
using  LinearAlgebra
using Plots

function newtons_method(∇f, H, x, ϵ, k_max)
    k, Δ = 1, Inf
    while abs(Δ) > ϵ && k ≤ k_max
        Δ =  ∇f(x)/abs(H(x))
        x -= Δ
        k += 1
    end
    return x
end

f = x->sin(10x)+cos(3x)
df(x) = 10*cos(10*x)-3*sin(3*x)
ddf(x) = -100*sin(10x)-9*cos(3x)

#converge para o ponto de minimo
initial_point = 5.1       #ponto inicial
convergence = 1e-8      #convergencia
maximum_iterations = 10 #numero maximo de iteracoes
xs = newtons_method(df,ddf,initial_point,convergence,maximum_iterations)
xs,f(xs),ddf(xs),df(xs)

#converge para o ponto de maximo
#initial_point = 4.45       #ponto inicial
#convergence = 1e-8      #convergencia
#maximum_iterations = 10 #numero maximo de iteracoes
#xs = newtons_method(df,ddf,initial_point,convergence,maximum_iterations)
#xs,f(xs),ddf(xs),df(xs)

plot(f,0.,8.,lw=2,draw_arrow="true")
#scatter!([a,c],[f(a),f(c)])
scatter!([initial_point],[f(initial_point)])
scatter!([xs],[f(xs)])


text:
# Autogenerated wrapper script for Gingko_jll for aarch64-apple-darwin
export libgingko, libgingko_cuda, libgingko_device, libgingko_dpcpp, libgingko_hip, libgingko_omp, libgingko_reference

using CompilerSupportLibraries_jll
JLLWrappers.@generate_wrapper_header("Gingko")
JLLWrappers.@declare_library_product(libgingko, "@rpath/libginkgo.1.4.0.dylib")
JLLWrappers.@declare_library_product(libgingko_cuda, "@rpath/libginkgo_cuda.1.4.0.dylib")
JLLWrappers.@declare_library_product(libgingko_device, "@rpath/libginkgo_device.1.4.0.dylib")
JLLWrappers.@declare_library_product(libgingko_dpcpp, "@rpath/libginkgo_dpcpp.1.4.0.dylib")
JLLWrappers.@declare_library_product(libgingko_hip, "@rpath/libginkgo_hip.1.4.0.dylib")
JLLWrappers.@declare_library_product(libgingko_omp, "@rpath/libginkgo_omp.1.4.0.dylib")
JLLWrappers.@declare_library_product(libgingko_reference, "@rpath/libginkgo_reference.1.4.0.dylib")
function __init__()
    JLLWrappers.@generate_init_header(CompilerSupportLibraries_jll)
    JLLWrappers.@init_library_product(
        libgingko,
        "lib/libginkgo.1.4.0.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libgingko_cuda,
        "lib/libginkgo_cuda.1.4.0.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libgingko_device,
        "lib/libginkgo_device.1.4.0.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libgingko_dpcpp,
        "lib/libginkgo_dpcpp.1.4.0.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libgingko_hip,
        "lib/libginkgo_hip.1.4.0.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libgingko_omp,
        "lib/libginkgo_omp.1.4.0.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libgingko_reference,
        "lib/libginkgo_reference.1.4.0.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
############################################################
## types supporting conversion from local to SA/DA arrays
############################################################

############################################################
## joPAsetup

export joPAsetup, joPAsetupException

# type definition
"""
joPAsetup type

# Atributes

- name::String        # name for identification
- dims::Dims          # dimensions of the array
- procs::Vector{Int}  # ids of workers to use
- chunks::Vector{Int} # number of chunks in each dimension
- idxs::Array{Tuple{Vararg{UnitRange{<:Integer}}}}
- cuts::Vector{Vector{<:Integer}}
- DT::DataType

See help for outer constructors for joPAsetup.

"""
struct joPAsetup
    name::String        # name for identification
    dims::Dims          # dimensions of the array
    procs::Vector{Int}  # ids of workers to use
    chunks::Vector{Int} # number of chunks in each dimension
    idxs::Array{Tuple{Vararg{UnitRange{<:Integer}}}}
    cuts::Vector{Vector{<:Integer}}
    DT::DataType
end

# type exception
struct joPAsetupException <: Exception
    msg :: String
end


text:
import JSON

function read_nodes(fname)
    res = []
    open(fname, "r") do io
        for line in eachline(io)
            line = strip(line)
            if startswith(line, "#")
                continue
            elseif startswith(line, "NumNodes")
                # TODO record
            elseif startswith(line, "NumTerminals")
                # TODO record
            else
                # match
                # m = match(r"(o\d+)\t(\d+)\t(\d+)(\tterminal)?", line)
                m = match(r"(X\d+)\s+(\d+(?:\.\d+)?)\s+(\d+(?:\.\d+)?)(\s+terminal)?", line)

                if m != nothing
                    id = m.captures[1]
                    # seems unimportant
                    width = parse(Float32, m.captures[2])
                    height = parse(Float32, m.captures[3])
                    # seems duplicate with .pl
                    terminal = m.captures[4]
                    # @show id, x, y, terminal
                    push!(res, (id, width, height, terminal))
                end
            end
        end
    end
    return res
end

function read_nets(fname)
    res = []
    cur = []
    open(fname, "r") do io
        for line in eachline(io)
            line = strip(line)
            if startswith(line, "#")
                continue
            elseif startswith(line, "NumNodes")
                # TODO record
            elseif startswith(line, "NumTerminals")
                # TODO record
            elseif startswith(line, "NetDegree")
                m = match(r".*:\s*(\d+)\s*(\w+)", line)
                name = m.captures[2]
                push!(res, cur)
                cur = []
                push!(cur, name)
            else
                # match
                # FIXME what is B? see n112279
                # m = match(r"(o\d+)\s*([IOB]) : (-?\d+\.\d+)\s+(-?\d+\.\d+)", line)
                m = match(r"(X\d+)\s*([IOB]) : (-?(?:\d+\.)?\d+)\s+(-?(?:\d+\.)?\d+)", line)
                if m != nothing
                    id = m.captures[1]
                    io = m.captures[2]
                    xoff = parse(Float64, m.captures[3])
                    yoff = parse(Float64, m.captures[4])
                    push!(cur, (id, io, xoff, yoff))
                end
            end
        end
    end
    push!(res, cur)
    # the first is empty
    return res[2:end]
end

function test()
    read_nets("../out/a.nets")
    read_nodes("../out/a.nodes")
    read_pos("../out/a.pl")
end

function read_pos(fname)
    res = []
    open(fname, "r") do io
        for line in eachline(io)
            line = strip(line)
            # m = match(r"(o\d+)\s+(\d+)\s+(\d+).*: \w+( /FIXED)?", line)
            m = match(r"(X\d+)\s+(\d+)\s+(\d+).*: \w+( /FIXED)?", line)
            if m != nothing
                id = m.captures[1]
                x = parse(Int, m.captures[2])
                y = parse(Int, m.captures[3])
                fixed = m.captures[4]
                # if fixed != nothing
                push!(res, (id, x, y, fixed))
                # end
            end
        end
    end
    return res
end

function test()
    m = match(r"(o\d+)\s+(\d+)\s+(\d+).*: \w+( /FIXED)?", "o211446	4797	5451	: N /FIXED")
    m.captures[4] != nothing
    match(r"(o\d+)\t(\d+)\t(\d+)(\tterminal)?", line)

    # http://www.ispd.cc/contests/05/ispd05-contest/announcement-jan-12.pdf
    nodes = read_nodes("/home/hebi/data/VLSI-benchmarks/ispd-2005/adaptec1/adaptec1.nodes")
    length(nodes)
    nets = read_nets("/home/hebi/data/VLSI-benchmarks/ispd-2005/adaptec1/adaptec1.nets")
    length(nets)
    pos = read_pos("/home/hebi/data/VLSI-benchmarks/ispd-2005/adaptec1/adaptec1.pl")
    ans_pos = read_pos("/home/hebi/data/VLSI-benchmarks/ispd-2005/adaptec1/adaptec1.ntup.pl")
    length(pos)

end

function test()
    x,y,w,h,E,solx,soly,mask,mask_offx,mask_offy = read_bench("/home/hebi/data/VLSI-benchmarks/ispd-2005/adaptec1/", "adaptec1")
    visualize(x, y, w, h)
    visualize(solx, soly, w, h)
end

function read_bench(folder, name)
    # http://www.ispd.cc/contests/05/ispd05-contest/announcement-jan-12.pdf
    # find
    @info "reading nodes .."
    nodes = read_nodes(joinpath(folder, "$name.nodes"))
    @info "reading nets .."
    nets = read_nets(joinpath(folder, "$name.nets"))
    @info "reading pos .."
    pos = read_pos(joinpath(folder, "$name.pl"))
    # # FIXME there may not be solution
    # sol_pos = read_pos(joinpath(folder, "$name.ntup.pl"))
    # # change to x,y,w,h
    # dream_pos = read_pos("/home/hebi/data/VLSI-benchmarks/DREAMPlace/install/results/adaptec1/adaptec1.gp.pl")

    # FIXME sort?
    # FIXME assert names are in order
    @info "manipulating .."
    name = [n[1] for n in nodes]
    name1 = [p[1] for p in pos]
    name == name1 || error("name order mismatch: $name vs. $name1")
    # FIXME name starts from "0", but that doesn't really matter
    name_dict = Dict(Pair.(name, 1:length(name)))
    x = [p[2] for p in pos]
    y = [p[3] for p in pos]
    # fixed macros
    mask = [if isnothing(p[4]) 1 else 0 end for p in pos]
    mask_offx = [if isnothing(p[4]) 0 else p[2] end for p in pos]
    mask_offy = [if isnothing(p[4]) 0 else p[3] end for p in pos]
    # solution
    # solx = [p[2] for p in sol_pos]
    # soly = [p[3] for p in sol_pos]
    # dream_solx = [p[2] for p in dream_pos]
    # dream_soly = [p[3] for p in dream_pos]

    # w, h
    w = [n[2] for n in nodes]
    h = [n[3] for n in nodes]
    # nets
    E = [[name_dict[n[1]] for n in net[2:end]] for net in nets]
    # return x,y,w,h,E,solx,soly,mask,mask_offx,mask_offy, dream_solx, dream_soly
    return x,y,w,h,E,mask
end

function decode_place_spec(jobj)
    xs = jobj["xs"]
    ys = jobj["ys"]
    as = jobj["as"]
    ws = jobj["ws"]
    hs = jobj["hs"]
    Es = jobj["Es"]
    mask = jobj["mask"]
    diearea = jobj["diearea"]
    params = jobj["place-params"]
    xs, ys, as, ws, hs, Es, mask, diearea, params
end

function test()
    # read json directly for debugging
    str = open("../out/a.json") do io
        read(io, String)
    end
    jobj = JSON.parse(str)
    xs, ys, ws, hs, Es, mask = parse_jobj(jobj)
    place(xs, ys, ws, hs, Es, mask)
end

max_stars_repo_name: lihebi/bhdl

max_issues_repo_name: lihebi/bhdl

max_forks_repo_name: lihebi/bhdl


text:
using DelimitedFiles, Plots, Printf

perf_i5 = readdlm("out_diffusion_2D_damp_perf.txt")
perf_loop_i5 = readdlm("out_diffusion_2D_damp_perf_loop.txt")
perf_loop_fun_i5 = readdlm("out_diffusion_2D_damp_perf_loop_fun.txt")
perf_gpu_v100 = readdlm("out_diffusion_2D_damp_perf_gpu.txt")
perf_xpu_v100 = readdlm("out_diffusion_2D_damp_perf_xpu.txt")
xpu_v100 = readdlm("out_diffusion_2D_damp_xpu.txt")

plot_fig = 2
save_fig = false

fontsize = 11
    opts = (yaxis=font(fontsize, "Courier"), xaxis=font(fontsize, "Courier"),
            framestyle=:box, titlefontsize=fontsize, titlefont="Courier", ylabel="T_eff [GB/s]",
            linewidth=:3.0, foreground_color_legend = nothing, legendfontsize=11, legend=:topright,
            size=(600,400) )

if plot_fig==1

x = perf_i5[[1,4,5,6,7],1]
ticks = collect(x)
ticklabels = [ @sprintf("%.f",x) for x in ticks ]

p1 = plot(perf_i5[:,1], perf_i5[:,4]; xlims=(-50,4200), ylims=(-0.5, 12), label="CPU broadcast", opts...)
plot!(perf_loop_i5[:,1], perf_loop_i5[:,4]; label="CPU loop", opts...)
plot!(xticks=(ticks,ticklabels),xrotation=0)
display(plot!(perf_loop_fun_i5[:,1], perf_loop_fun_i5[:,4]; label="CPU loop function", opts...))
if save_fig savefig("perf_cpu.png") end

else

x = perf_gpu_v100[[1,6,7,8,9],1]
ticks = collect(x)
ticklabels = [ @sprintf("%.f",x) for x in ticks ]

p2 = plot(perf_gpu_v100[:,1], perf_gpu_v100[:,4]; xlims=(-50, 17000), ylims=(-10, 1100), label="GPU", xlabel="grid resolution (nx)", opts...)
plot!(perf_xpu_v100[:,1], perf_xpu_v100[:,4]; label="XPU parallel_indices", opts...)
plot!(xticks=(ticks,ticklabels),xrotation=0)
display(plot!(xpu_v100[:,1], xpu_v100[:,4]; label="XPU parallel", opts...))
if save_fig savefig("perf_gpu.png") end

end


text:
import Pkg

Pkg.add(Pkg.PackageSpec(url="https://github.com/JuliaComputing/JuliaAcademyData.jl"))
Pkg.add("Weave")
Pkg.add("IJulia")
Pkg.add("Images")
Pkg.add("ImageMagick")
Pkg.add("ImageIO")
Pkg.add("Interact")
Pkg.add("Statistics")
Pkg.add("Plots")
Pkg.add("GR")
Pkg.add("Plotly")
Pkg.add("Glob")
Pkg.add("WebIO")
Pkg.add("Conda")
Pkg.add("PyCall")
Pkg.add("PyPlot")
Pkg.add("Revise")

# Force PyCall to use Conda.jl
#ENV["PYTHON"] = ""
#Pkg.build("PyCall")

# Interact
#import WebIO
#WebIO.install_jupyter_nbextension()


text:
"""
    diffusion_matrix(domain, n)

Assemble discrete diffusion matrix.
"""
function diffusion_matrix end

function diffusion_matrix(domain::ClosedIntervalDomain, n)
    Δx = 1 // n * (domain.right - domain.left)
    inds = [-1, 0, 1]
    stencil = [1, -2, 1] / Δx^2
    diags = [i => fill(s, n + 1 - abs(i)) for (i, s) ∈ zip(inds, stencil)]
    D = spdiagm(diags...)
    D[1, 1] = stencil[1] + stencil[2]
    D[end, end] = stencil[end-1] + stencil[end]
    D
end

function diffusion_matrix(domain::PeriodicIntervalDomain, n)
    Δx = 1 // n * (domain.right - domain.left)
    inds = [-1, 0, 1]
    stencil = [1, -2, 1] / Δx^2
    diags = [i => fill(s, n - abs(i)) for (i, s) ∈ zip(inds, stencil)]
    D = spdiagm(diags...)
    D[1, end] = stencil[1]
    D[end, 1] = stencil[end]
    D
end


text:
ENV["JULIA_DEBUG"] = "HierarchicalTemporalMemory"
using Documenter, Test, HierarchicalTemporalMemory

@enum GROUPS AllTests SpatialPoolerTest TemporalMemoryTest Basic RegionTest NetworkTest

const GROUP= try
  get(ENV, "GROUP", "AllTests")|> Symbol|> eval
catch e
  error("ENV[\"GROUP\"]: Invalid value. Supported values: $(GROUPS|>instances)\n Given: $e")
end
global plot_enabled= false

#using Aqua
#Aqua.test_all(HierarchicalTemporalMemory)

@testset "Doctests" begin
  DocMeta.setdocmeta!(HierarchicalTemporalMemory, :DocTestSetup, :(
      using HierarchicalTemporalMemory;
      const HTM=HierarchicalTemporalMemory
    ); recursive=true)
  doctest(HierarchicalTemporalMemory)
end

@testset "Topology" begin include("topology_test.jl") end
if GROUP == AllTests || GROUP == SpatialPoolerTest
  @testset "Spatial Pooler" begin
    include("spatial_pooler_test.jl")
  end
end
if GROUP == AllTests || GROUP == TemporalMemoryTest
  @testset "Temporal Memory" begin
    include("temporal_memory_test.jl")
  end
end
if GROUP == AllTests || GROUP == RegionTest
  @testset "Region" begin
    include("test_region_projection.jl")
  end
end
if GROUP == AllTests || GROUP == NetworkTest
  @testset "Network" begin
    include("test_network.jl")
  end
end